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)

	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
		}

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

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
}