e621-sdk-go/main.go

package main

import (
	"context"
	"e621_to_neo4j/e621/api"
	"e621_to_neo4j/e621/api/models"
	"e621_to_neo4j/neo4jAPI"
	"e621_to_neo4j/utils"
	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
	"log"
	"time"
)

func main() {
	config, err := utils.LoadConfig()
	if err != nil {
		log.Println(err)
	}

	driver, err := neo4jAPI.NewConnection(config.Neo4jURL, config.Neo4jUsername, config.Neo4jPassword)
	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)

	e621Client := api.NewClient(config.E621APIKey, config.E621Username)
	user, err := e621Client.GetUserInfo("selloo")
	if err != nil {
		log.Fatal(err)
	}

	err = neo4jAPI.CreateUserNode(ctx, driver, user)
	if err != nil {
		log.Fatal(err)
	}

	favorites, err := e621Client.GetFavorites(user.Name)
	if err != nil {
		log.Fatal(err)
	}

	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)

	start = time.Now()
	uploadSources(favorites, ctx, driver)
	elapsed = time.Since(start)
	log.Printf("upload of Sources took %d", 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
		}

	}
	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
}

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
}