e621-sdk-go/services/user.go

package services

import (
	"context"
	"e621_to_neo4j/database"
	"e621_to_neo4j/e621"
	"e621_to_neo4j/e621/models"
	"e621_to_neo4j/utils"
	log "github.com/sirupsen/logrus"
	"time"
)

func ScrapeUser(ctx context.Context, graphConnection database.GraphConnection, e621Client e621.Client, username string) error {
	var err error

	e621User, err := e621Client.GetUserInfo(username)
	if err != nil {
		return err
	}

	if e621User.IsBanned {
		log.WithFields(log.Fields{
			"user": e621User.Name,
			"id":   e621User.ID,
			"bann": e621User.IsBanned,
		}).Info("User is Banned")
		return nil
	}

	log.WithFields(log.Fields{
		"user": e621User.Name,
		"id":   e621User.ID,
	}).Info("Processing user")

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

	log.WithFields(log.Fields{
		"user": e621User.Name,
		"id":   e621User.ID,
	}).Info("Getting favorites for user")
	start := time.Now()
	userFavorites, err := e621Client.GetFavorites(e621User)
	if err != nil {
		log.Fatal(err)
	}
	log.WithFields(log.Fields{
		"user":        e621User.Name,
		"id":          e621User.ID,
		"post_amount": len(userFavorites),
		"scrape_time": time.Since(start),
	}).Info("Getting favorites for user")

	startUploadPosts := time.Now()

	// Uploads all Tags, Posts as Nodes to Neo4j
	for i, post := range userFavorites {
		if exists, err := graphConnection.CheckUserToPostLink(ctx, post.ID, e621User.ID); err == nil && exists {
			log.WithFields(log.Fields{
				"user":         e621User.Name,
				"id":           e621User.ID,
				"last_post_id": post.ID,
			}).Info("No new favorites found")
			break
		} else if err != nil {
			return err
		}

		start = time.Now()
		err = uploadNodes(ctx, graphConnection, post)
		if err != nil {
			return err
		}
		log.WithFields(log.Fields{
			"user":        e621User.Name,
			"id":          e621User.ID,
			"post_number": i,
			"post_amount": len(userFavorites),
			"post_id":     post.ID,
			"upload_time": time.Since(start),
		}).Debug("Uploading post")

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

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

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

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

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

		err = uploadArtistTagRelationship(ctx, graphConnection, post)
		if err != nil {
			log.Fatal(err)
			return err
		}
		log.WithFields(log.Fields{
			"user":        e621User.Name,
			"id":          e621User.ID,
			"post_number": i,
			"post_amount": len(userFavorites),
			"post_id":     post.ID,
			"upload_time": time.Since(start),
		}).Debug("Making relationship")
	}
	log.WithFields(log.Fields{
		"user":        e621User.Name,
		"id":          e621User.ID,
		"upload_time": time.Since(startUploadPosts),
	}).Info("Upload to Database finished")

	return nil
}

// uploadNodes uploads the post to the database and creates the nodes
func uploadNodes(ctx context.Context, graphConnection database.GraphConnection, post models.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 = utils.UniqueNonEmptyElementsOf(allGeneralTags)
	uniqueCharacterTags = utils.UniqueNonEmptyElementsOf(allCharacterTags)
	uniqueCopyrightTags = utils.UniqueNonEmptyElementsOf(allCopyrightTags)
	uniqueArtistTags = utils.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 database.GraphConnection, post models.Post, e621User models.E621User) 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 database.GraphConnection, post models.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 database.GraphConnection, post models.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 database.GraphConnection, post models.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 database.GraphConnection, post models.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 copyrigh tag: %s", post.ID, copyrightTag)

	}
	return nil
}

// uploadArtistTagRelationship creates a relationship between the post and the artist tag
func uploadArtistTagRelationship(ctx context.Context, graphConnection database.GraphConnection, post models.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
}