This repository has been archived on 2024-07-22. You can view files and clone it, but cannot push or open issues or pull requests.
e621-sdk-go/main.go

245 lines
6.5 KiB
Go
Raw Normal View History

2023-05-22 11:08:08 +00:00
package main
import (
"context"
"e621_to_neo4j/e621/api"
"e621_to_neo4j/e621/api/models"
"e621_to_neo4j/neo4jAPI"
2023-05-22 11:08:08 +00:00
"e621_to_neo4j/utils"
"github.com/neo4j/neo4j-go-driver/v5/neo4j"
2023-05-22 11:08:08 +00:00
"log"
"time"
)
func main() {
config, err := utils.LoadConfig()
if err != nil {
log.Println(err)
}
driver, err := neo4jAPI.NewConnection(config.Neo4jURL, config.Neo4jUsername, config.Neo4jPassword)
2023-05-22 11:08:08 +00:00
if err != nil {
log.Fatal(err)
}
ctx := context.Background()
defer func(driver neo4j.DriverWithContext, ctx context.Context) {
err := driver.Close(ctx)
if err != nil {
log.Fatal(err)
}
}(driver, ctx)
2023-05-22 11:08:08 +00:00
e621Client := api.NewClient(config.E621APIKey, config.E621Username)
user, err := e621Client.GetUserInfo("selloo")
if err != nil {
log.Fatal(err)
}
err = neo4jAPI.CreateUserNode(ctx, driver, user)
2023-05-22 11:08:08 +00:00
if err != nil {
log.Fatal(err)
}
favorites, err := e621Client.GetFavorites(user.Name)
2023-05-22 11:08:08 +00:00
if err != nil {
log.Fatal(err)
}
2023-05-22 11:08:08 +00:00
start := time.Now()
//uploadTags(favorites, ctx, driver)
elapsed := time.Since(start)
log.Printf("upload of Tags took %d", elapsed)
start = time.Now()
//uploadPosts(favorites, ctx, driver)
elapsed = time.Since(start)
log.Printf("upload of Posts took %d", elapsed)
2023-05-22 11:08:08 +00:00
start = time.Now()
//uploadSources(favorites, ctx, driver)
elapsed = time.Since(start)
log.Printf("upload of Sources took %d", elapsed)
2023-05-22 11:08:08 +00:00
start = time.Now()
// err = makePostToTagRelationship(favorites, ctx, driver)
elapsed = time.Since(start)
log.Printf("upload of Sources took %v", elapsed)
start = time.Now()
// makeUserToPostRelationship(favorites, user, ctx, driver)
elapsed = time.Since(start)
log.Printf("upload of Sources took %v", elapsed)
start = time.Now()
err = makePostToSourceRelationship(favorites, ctx, driver)
if err != nil {
log.Fatal(err)
}
elapsed = time.Since(start)
log.Printf("upload of Sources took %v", elapsed)
}
func uploadPosts(posts []models.Post, ctx context.Context, driver neo4j.DriverWithContext) error {
for _, fav := range posts {
log.Printf("e621 PostID: %d", fav.ID)
err := neo4jAPI.CreatePostNode(ctx, driver, fav.ID)
if err != nil {
return err
2023-05-22 11:08:08 +00:00
}
}
return nil
}
func uploadTags(posts []models.Post, ctx context.Context, driver neo4j.DriverWithContext) error {
uniqueGeneralTags := make([]string, 0)
uniqueCharacterTags := make([]string, 0)
uniqueCopyrightTags := make([]string, 0)
uniqueArtistTags := make([]string, 0)
allGeneralTags := make([]string, 0)
allCharacterTags := make([]string, 0)
allCopyrightTags := make([]string, 0)
allArtistTags := make([]string, 0)
// add all tags together
for _, fav := range posts {
allGeneralTags = append(allGeneralTags, fav.Tags.General...)
allCharacterTags = append(allCharacterTags, fav.Tags.Character...)
allCopyrightTags = append(allCopyrightTags, fav.Tags.Character...)
allArtistTags = append(allArtistTags, fav.Tags.Artist...)
}
// Process General PostTags
uniqueGeneralTags = utils.UniqueNonEmptyElementsOf(allGeneralTags)
// Process Character PostTags
uniqueCharacterTags = utils.UniqueNonEmptyElementsOf(allCharacterTags)
// Process Copyright PostTags
uniqueCopyrightTags = utils.UniqueNonEmptyElementsOf(allCopyrightTags)
// Process Artist PostTags
uniqueArtistTags = utils.UniqueNonEmptyElementsOf(allArtistTags)
log.Printf("uniqueGeneralTags length: %d", len(uniqueGeneralTags))
log.Printf("uniqueCopyrightTags length: %d", len(uniqueCopyrightTags))
log.Printf("uniqueArtistTags length: %d", len(uniqueArtistTags))
for _, uniqueGeneralTag := range uniqueGeneralTags {
log.Printf("TagType: General - Tag: %s", uniqueGeneralTag)
err := neo4jAPI.CreateTagNode(ctx, driver, uniqueGeneralTag, "general")
if err != nil {
return err
}
}
for _, uniqueCharacterTag := range uniqueCharacterTags {
log.Printf("TagType: Character - Tag: %s", uniqueCharacterTag)
err := neo4jAPI.CreateTagNode(ctx, driver, uniqueCharacterTag, "character")
if err != nil {
return err
}
}
for _, uniqueCopyrightTag := range uniqueCopyrightTags {
log.Printf("TagType: Copyright - Tag: %s", uniqueCopyrightTag)
err := neo4jAPI.CreateTagNode(ctx, driver, uniqueCopyrightTag, "copyright")
if err != nil {
return err
}
}
for _, uniqueArtistTag := range uniqueArtistTags {
log.Printf("TagType: Artist - Tag: %s", uniqueArtistTag)
err := neo4jAPI.CreateTagNode(ctx, driver, uniqueArtistTag, "artist")
if err != nil {
return err
}
}
return nil
2023-05-22 11:08:08 +00:00
}
func uploadSources(posts []models.Post, ctx context.Context, driver neo4j.DriverWithContext) error {
for _, fav := range posts {
for _, source := range fav.Sources {
log.Printf("e621 Source: %s", source)
err := neo4jAPI.CreateSourceNode(ctx, driver, source)
if err != nil {
return err
}
}
}
return nil
}
func makePostToTagRelationship(posts []models.Post, ctx context.Context, driver neo4j.DriverWithContext) error {
for _, fav := range posts {
id := fav.ID
for _, generalTags := range fav.Tags.General {
log.Printf("PostID: %d has General Tag: %s", id, generalTags)
err := neo4jAPI.PostToTagRelationship(ctx, driver, id, generalTags)
if err != nil {
return err
}
}
for _, characterTags := range fav.Tags.Character {
log.Printf("PostID: %d has Character Tag: %s", id, characterTags)
err := neo4jAPI.PostToTagRelationship(ctx, driver, id, characterTags)
if err != nil {
return err
}
}
for _, copyrightTags := range fav.Tags.Copyright {
log.Printf("PostID: %d has Copyright Tag: %s", id, copyrightTags)
err := neo4jAPI.PostToTagRelationship(ctx, driver, id, copyrightTags)
if err != nil {
return err
}
}
for _, artistTags := range fav.Tags.Artist {
log.Printf("PostID: %d has Artist Tag: %s", id, artistTags)
err := neo4jAPI.PostToTagRelationship(ctx, driver, id, artistTags)
if err != nil {
return err
}
}
}
return nil
}
func makeUserToPostRelationship(posts []models.Post, user models.E621User, ctx context.Context, driver neo4j.DriverWithContext) error {
for _, post := range posts {
log.Printf("UserID: %d has PostID: %d in his favorites", user.ID, post.ID)
err := neo4jAPI.UserToPostRelationship(ctx, driver, post.ID, user.ID)
if err != nil {
return err
}
}
return nil
}
func makePostToSourceRelationship(posts []models.Post, ctx context.Context, driver neo4j.DriverWithContext) error {
for _, post := range posts {
postID := post.ID
for _, source := range post.Sources {
log.Printf("PostID: %d has Source URL: %s", postID, source)
err := neo4jAPI.PostToSourceRelationship(ctx, driver, postID, source)
if err != nil {
return err
}
}
}
return nil
}