plug-sdk/pkg/plug/scrape.go

package plug

import (
	"context"
	"slices"

	"git.anthrove.art/Anthrove/otter-space-sdk/v4/pkg/database"
	"git.anthrove.art/Anthrove/otter-space-sdk/v4/pkg/models"
	log "github.com/sirupsen/logrus"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"gorm.io/gorm"
)

var BatchSize = 50
var BasicLoggingFields log.Fields

type BatchSummery struct {
	AddedPosts     int64
	AddedFavorites int64
}

func BatchPostProcessing(ctx context.Context, userSource models.UserSource, posts []models.Post) error {
	_, err := BatchPostProcessingWithSummery(ctx, userSource, posts)
	return err
}

func BatchPostProcessingWithSummery(ctx context.Context, userSource models.UserSource, posts []models.Post) (BatchSummery, error) {
	ctx, span := tracer.Start(ctx, "BatchPostProcessing")
	defer span.End()

	span.SetAttributes(
		attribute.String("user_source_id", string(userSource.ID)),
		attribute.String("user_source_user_id", string(userSource.UserID)),
		attribute.String("user_source_source_id", string(userSource.SourceID)),
	)

	BasicLoggingFields = log.Fields{
		"user_source_id":        userSource.ID,
		"user_source_user_id":   userSource.UserID,
		"user_source_source_id": userSource.SourceID,
	}

	log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Starting BatchPostProcessing")

	db, err := database.GetGorm(ctx)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to get Gorm DB")
		return BatchSummery{}, err
	}

	postIDs := make([]string, 0, len(posts))
	for _, post := range posts {
		postIDs = append(postIDs, post.References[0].SourcePostID)
	}
	span.AddEvent("Collected post IDs", trace.WithAttributes(attribute.Int("post_count", len(postIDs))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).WithField("post_count", len(postIDs)).Info("Collected post IDs")

	existingPosts, err := getAnthrovePost(ctx, db, userSource.SourceID, postIDs)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to fetch existing posts")
		return BatchSummery{}, err
	}
	span.AddEvent("Fetched existing posts", trace.WithAttributes(attribute.Int("existing_post_count", len(existingPosts))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).WithField("existing_post_count", len(existingPosts)).Info("Fetched existing posts")

	var existingPostIDs []models.PostID
	for _, post := range existingPosts {
		existingPostIDs = append(existingPostIDs, models.PostID(post.PostID))
	}

	var existingFavPostIDs []models.PostID
	existingFavPostIDs, err = getAlreadyFavoritesPostIDs(ctx, db, existingPostIDs, userSource.ID)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to fetch existing favorite posts")
		return BatchSummery{}, err
	}
	span.AddEvent("Fetched existing favorite posts", trace.WithAttributes(attribute.Int("existing_fav_post_count", len(existingFavPostIDs))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).WithField("existing_fav_post_count", len(existingFavPostIDs)).Info("Fetched existing favorite posts")

	anthroveFaves := make([]models.UserFavorite, 0, len(existingPosts))
	newPosts := make([]models.Post, 0, len(existingPosts))
	for _, post := range posts {
		if !slices.ContainsFunc(existingPosts, func(reference models.PostReference) bool {
			found := reference.SourcePostID == post.References[0].SourcePostID
			if found {
				if !slices.Contains(existingFavPostIDs, models.PostID(reference.PostID)) {
					anthroveFaves = append(anthroveFaves, models.UserFavorite{
						UserID:       userSource.UserID,
						PostID:       models.PostID(reference.PostID),
						UserSourceID: userSource.ID,
					})
				}
			}
			return found
		}) {
			anthroveFaves = append(anthroveFaves, models.UserFavorite{
				UserID:       userSource.UserID,
				PostID:       post.ID,
				UserSourceID: userSource.ID,
			})
			newPosts = append(newPosts, post)
		}
	}
	span.AddEvent("Processed posts for favorites and new posts", trace.WithAttributes(attribute.Int("new_post_count", len(newPosts)), attribute.Int("new_fav_count", len(anthroveFaves))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Processed posts for favorites and new posts")

	if len(newPosts) > 0 {
		err = database.CreatePostInBatch(ctx, newPosts, BatchSize)
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
			log.WithContext(ctx).WithError(err).Error("Failed to create new posts in batch")
			return BatchSummery{}, err
		}
		span.AddEvent("Created new posts in batch", trace.WithAttributes(attribute.Int("batch_size", BatchSize)))
		log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Created new posts in batch")
	}

	if len(anthroveFaves) > 0 {
		err = database.CreateUserFavoriteInBatch(ctx, anthroveFaves, BatchSize)
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
			log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to create user favorites in batch")
			return BatchSummery{}, err
		}
		span.AddEvent("Created user favorites in batch", trace.WithAttributes(attribute.Int("batch_size", BatchSize)))
		log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Created user favorites in batch")
	}

	return BatchSummery{
		AddedPosts:     int64(len(newPosts)),
		AddedFavorites: int64(len(anthroveFaves)),
	}, nil
}

func getAnthrovePost(ctx context.Context, gorm *gorm.DB, id models.SourceID, postIDs []string) ([]models.PostReference, error) {
	ctx, span := tracer.Start(ctx, "getAnthrovePost")
	defer span.End()

	log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Starting getAnthrovePost")

	var existingPosts []models.PostReference
	err := gorm.WithContext(ctx).Model(models.PostReference{}).Find(&existingPosts, "source_id = ? AND source_post_id IN ?", id, postIDs).Error
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to fetch Anthrove posts")
		return existingPosts, err
	}
	span.AddEvent("Fetched Anthrove posts", trace.WithAttributes(attribute.Int("post_count", len(existingPosts))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).WithField("post_count", len(existingPosts)).Info("Fetched Anthrove posts")

	return existingPosts, nil
}

func getAlreadyFavoritesPostIDs(ctx context.Context, gorm *gorm.DB, existingPostIDs []models.PostID, userSourceID models.UserSourceID) ([]models.PostID, error) {
	ctx, span := tracer.Start(ctx, "getAlreadyFavoritesPostIDs")
	defer span.End()

	log.WithContext(ctx).WithFields(BasicLoggingFields).Info("Starting getAlreadyFavoritesPostIDs")

	var existingFavPostIDS []models.PostID
	err := gorm.WithContext(ctx).Model(&models.UserFavorite{}).Select("post_id").Find(&existingFavPostIDS, "user_source_id = ? AND post_id IN ?", userSourceID, existingPostIDs).Error
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		log.WithContext(ctx).WithError(err).WithFields(BasicLoggingFields).Error("Failed to fetch already favorite post IDs")
		return existingFavPostIDS, err
	}
	span.AddEvent("Fetched already favorite post IDs", trace.WithAttributes(attribute.Int("fav_post_count", len(existingFavPostIDS))))
	log.WithContext(ctx).WithFields(BasicLoggingFields).WithField("fav_post_count", len(existingFavPostIDS)).Info("Fetched already favorite post IDs")

	return existingFavPostIDS, nil
}