e621-sdk-go/internal/service/manager.go

package service

import (
	"context"
	"git.dragse.it/anthrove/e621-sdk-go/pkg/e621"
	"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/model"
	"git.dragse.it/anthrove/e621-to-graph/pkg/logic"
	"git.dragse.it/anthrove/e621-to-graph/pkg/util"
	log "github.com/sirupsen/logrus"
	"time"
)

func ScrapeUser(ctx context.Context, graphConnection logic.GraphConnection, client *e621.Client, username string) error {
	var err error

	scrapeTime := time.Now()

	e621User, err := client.GetUserByName(username).Execute()
	if err != nil {
		return err
	}

	if e621User.IsBanned {
		log.WithFields(log.Fields{
			"e621_username":  e621User.Name,
			"e621_user_id":   e621User.ID,
			"e621_user_bann": e621User.IsBanned,
		}).Info("service: user is Banned")
		return nil
	}

	err = graphConnection.UploadUser(ctx, e621User)
	if err != nil {
		return err
	}

	currentDBFavCount, err := graphConnection.GetUserFavoriteCount(ctx, e621User.ID)
	if err != nil {
		return err
	}

	favoriteBuilder, err := client.GetFavoritesForUser(e621User.Name)
	if err != nil {
		return err
	}

	if currentDBFavCount > e621User.FavoriteCount {
		//TODO: IMPLEMENT USER MARKED FOR DELETED FAVS
		log.WithFields(log.Fields{
			"e621_username":                  e621User.Name,
			"e621_user_id":                   e621User.ID,
			"e621_current_db_favorite_count": currentDBFavCount,
			"e621_user_favorite_count":       e621User.FavoriteCount,
		}).Debug("service: user has favorites deleted")
	}

	var pageIndex = 1
	for currentDBFavCount < e621User.FavoriteCount {

		favorites, err := favoriteBuilder.Page(pageIndex).Execute()
		if err != nil {
			return err
		}

		if len(favorites) <= 0 {
			return nil
		}

		for _, favorite := range favorites {

			if currentDBFavCount == e621User.FavoriteCount {
				break
			}

			isFaved, err := graphConnection.CheckUserToPostLink(ctx, favorite.ID, e621User.ID)
			if err != nil {
				return err
			}

			if !isFaved {
				err = uploadDataToDB(ctx, graphConnection, favorite, e621User)
				if err != nil {
					return err
				}
				currentDBFavCount++
			}

		}
		pageIndex++
	}

	log.WithFields(log.Fields{
		"e621_username": e621User.Name,
		"e621_user_id":  e621User.ID,
		"post_amount":   e621User.FavoriteCount,
		"scrape_time":   time.Since(scrapeTime),
	}).Info("service: finished processing favorites")

	return nil
}

func uploadDataToDB(ctx context.Context, graphConnection logic.GraphConnection, favorite model.Post, e621User model.User) error {
	start := time.Now()
	err := uploadNodes(ctx, graphConnection, favorite)
	if err != nil {
		return err
	}
	log.WithFields(log.Fields{
		"e621_username": e621User.Name,
		"e621_user_id":  e621User.ID,
		"post_id":       favorite.ID,
		"upload_time":   time.Since(start),
	}).Debug("service: uploaded post")

	start = time.Now()
	err = uploadPostToUserRelationship(ctx, graphConnection, favorite, e621User)
	if err != nil {
		log.Fatal(err)
		return err
	}

	err = uploadSourceTagRelationship(ctx, graphConnection, favorite)
	if err != nil {
		log.Fatal(err)
		return err
	}

	err = uploadGeneralTagRelationship(ctx, graphConnection, favorite)
	if err != nil {
		log.Fatal(err)
		return err
	}

	err = uploadCharacterTagtRelationship(ctx, graphConnection, favorite)
	if err != nil {
		log.Fatal(err)
		return err
	}

	err = uploadCopyrightTagRelationship(ctx, graphConnection, favorite)
	if err != nil {
		log.Fatal(err)
		return err
	}

	err = uploadArtistTagRelationship(ctx, graphConnection, favorite)
	if err != nil {
		log.Fatal(err)
		return err
	}
	log.WithFields(log.Fields{
		"e621_username": e621User.Name,
		"e621_user_id":  e621User.ID,
		"post_id":       favorite.ID,
		"upload_time":   time.Since(start),
	}).Debug("service: made relationship")
	return nil
}

// uploadNodes uploads the post to the database and creates the nodes
func uploadNodes(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) 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)

	allGeneralTags = append(allGeneralTags, post.Tags.General...)
	allCharacterTags = append(allCharacterTags, post.Tags.Character...)
	allCopyrightTags = append(allCopyrightTags, post.Tags.Copyright...)
	allArtistTags = append(allArtistTags, post.Tags.Artist...)

	uniqueGeneralTags = util.UniqueNonEmptyElementsOf(allGeneralTags)
	uniqueCharacterTags = util.UniqueNonEmptyElementsOf(allCharacterTags)
	uniqueCopyrightTags = util.UniqueNonEmptyElementsOf(allCopyrightTags)
	uniqueArtistTags = util.UniqueNonEmptyElementsOf(allArtistTags)

	err := graphConnection.UploadPost(ctx, post.ID)
	if err != nil {
		return err
	}

	// Uploads the source to the database
	for _, source := range post.Sources {
		err := graphConnection.UploadSource(ctx, source)
		if err != nil {
			return err
		}
	}

	for _, uniqueGeneralTag := range uniqueGeneralTags {
		err := graphConnection.UploadTag(ctx, uniqueGeneralTag, "general")
		if err != nil {
			return err
		}
	}
	for _, uniqueCharacterTag := range uniqueCharacterTags {
		err := graphConnection.UploadTag(ctx, uniqueCharacterTag, "character")
		if err != nil {
			return err
		}
	}

	for _, uniqueCopyrightTag := range uniqueCopyrightTags {
		err := graphConnection.UploadTag(ctx, uniqueCopyrightTag, "copyright")
		if err != nil {
			return err
		}
	}

	for _, uniqueArtistTag := range uniqueArtistTags {
		err := graphConnection.UploadTag(ctx, uniqueArtistTag, "artist")
		if err != nil {
			return err
		}
	}

	return nil
}

// uploadPostToUserRelationship creates a relationship between the user and the post
func uploadPostToUserRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post, e621User model.User) error {
	err := graphConnection.EstablishUserToPostLink(ctx, post.ID, e621User.ID)
	if err != nil {
		return err
	}
	// log.Printf("Created UserToPostRelationship for user: %s to post: %d", e621User.Name, post.ID)
	return nil
}

// uploadSourceTagRelationship creates a relationship between the post and the source
func uploadSourceTagRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) error {
	for _, source := range post.Sources {
		err := graphConnection.EstablishPostToSourceLink(ctx, post.ID, source)
		if err != nil {
			return err
		}
		// log.Printf("Created PostToSourceRelationship for Post: %d to source: %s", post.ID, source)

	}
	return nil
}

// uploadGeneralTagRelationship creates a relationship between the post and the general tag
func uploadGeneralTagRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) error {
	for _, generalTag := range post.Tags.General {
		err := graphConnection.EstablishPostToTagLink(ctx, post.ID, generalTag)
		if err != nil {
			return err
		}
		// log.Printf("Created PostToTagRelationship for post: %d to general tag: %s", post.ID, generalTag)

	}
	return nil
}

// uploadCharacterTagtRelationship creates a relationship between the post and the character tag
func uploadCharacterTagtRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) error {
	for _, characterTag := range post.Tags.Character {
		err := graphConnection.EstablishPostToTagLink(ctx, post.ID, characterTag)
		if err != nil {
			return err
		}
		// log.Printf("Created PostToTagRelationship for post: %d to character tag: %s", post.ID, characterTag)

	}
	return nil
}

// uploadCopyrightTagRelationship creates a relationship between the post and the copyright tag
func uploadCopyrightTagRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) error {
	for _, copyrightTag := range post.Tags.Copyright {
		err := graphConnection.EstablishPostToTagLink(ctx, post.ID, copyrightTag)
		if err != nil {
			return err
		}
		// log.Printf("Created PostToTagRelationship for post: %d to copyright tag: %s", post.ID, copyrightTag)

	}
	return nil
}

// uploadArtistTagRelationship creates a relationship between the post and the artist tag
func uploadArtistTagRelationship(ctx context.Context, graphConnection logic.GraphConnection, post model.Post) error {
	for _, artistTag := range post.Tags.Artist {
		err := graphConnection.EstablishPostToTagLink(ctx, post.ID, artistTag)
		if err != nil {
			return err
		}
		// log.Printf("Created PostToTagRelationship for post: %d to artist tag: %s", post.ID, artistTag)

	}
	return nil
}