haixunMaster/haixun-backend/internal/model/knowledge_graph/repository/mongo.go

package repository

import (
	"context"
	"strings"

	"haixun-backend/internal/library/clock"
	app "haixun-backend/internal/library/errors"
	"haixun-backend/internal/library/errors/code"
	libkg "haixun-backend/internal/library/knowledge"
	libmongo "haixun-backend/internal/library/mongo"
	"haixun-backend/internal/model/knowledge_graph/domain/entity"
	domrepo "haixun-backend/internal/model/knowledge_graph/domain/repository"

	"github.com/google/uuid"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type mongoRepository struct {
	collection *mongo.Collection
}

func NewMongoRepository(db *mongo.Database) domrepo.Repository {
	if db == nil {
		return &mongoRepository{}
	}
	return &mongoRepository{collection: db.Collection(entity.CollectionName)}
}

func (r *mongoRepository) EnsureIndexes(ctx context.Context) error {
	if r.collection == nil {
		return nil
	}
	return libmongo.EnsureIndexes(ctx, r.collection, []mongo.IndexModel{
		{Keys: bson.D{{Key: "tenant_id", Value: 1}, {Key: "owner_uid", Value: 1}, {Key: "brand_id", Value: 1}}, Options: options.Index().SetUnique(true).SetPartialFilterExpression(bson.M{"brand_id": bson.M{"$gt": ""}})},
		{Keys: bson.D{{Key: "tenant_id", Value: 1}, {Key: "owner_uid", Value: 1}, {Key: "persona_id", Value: 1}}, Options: options.Index().SetUnique(true).SetPartialFilterExpression(bson.M{"persona_id": bson.M{"$gt": ""}})},
		{Keys: bson.D{{Key: "brand_id", Value: 1}, {Key: "update_at", Value: -1}}},
		{Keys: bson.D{{Key: "persona_id", Value: 1}, {Key: "update_at", Value: -1}}},
	})
}

func brandOwnerFilter(tenantID, ownerUID, brandID string) bson.M {
	filter := bson.M{
		"tenant_id": tenantID,
		"owner_uid": ownerUID,
	}
	for k, v := range libmongo.BrandScopeFilter(brandID) {
		filter[k] = v
	}
	return filter
}

func (r *mongoRepository) UpsertByBrand(ctx context.Context, graph *entity.Graph) (*entity.Graph, error) {
	if r.collection == nil {
		return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured")
	}
	if graph == nil {
		return nil, app.For(code.Brand).InputMissingRequired("graph is required")
	}
	now := clock.NowUnixNano()
	graph.UpdateAt = now
	if graph.CreateAt == 0 {
		graph.CreateAt = now
	}
	if strings.TrimSpace(graph.ID) == "" {
		graph.ID = uuid.NewString()
	}

	filter := brandOwnerFilter(graph.TenantID, graph.OwnerUID, graph.BrandID)
	update := bson.M{
		"$set": bson.M{
			"seed":           graph.Seed,
			"nodes":          graph.Nodes,
			"edges":          graph.Edges,
			"brave_sources":  graph.BraveSources,
			"pain_tag_count": graph.PainTagCount,
			"generated_at":   graph.GeneratedAt,
			"update_at":      graph.UpdateAt,
			"brand_id":       graph.BrandID,
		},
		"$setOnInsert": bson.M{
			"_id":       graph.ID,
			"tenant_id": graph.TenantID,
			"owner_uid": graph.OwnerUID,
			"create_at": graph.CreateAt,
		},
	}
	opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
	var out entity.Graph
	err := r.collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(&out)
	if err != nil {
		return nil, err
	}
	return &out, nil
}

func (r *mongoRepository) FindByBrand(ctx context.Context, tenantID, ownerUID, brandID string) (*entity.Graph, error) {
	if r.collection == nil {
		return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured")
	}
	var out entity.Graph
	err := r.collection.FindOne(ctx, brandOwnerFilter(tenantID, ownerUID, brandID)).Decode(&out)
	if err == mongo.ErrNoDocuments {
		return nil, app.For(code.Brand).ResNotFound("knowledge graph not found")
	}
	if err != nil {
		return nil, err
	}
	return &out, nil
}

func (r *mongoRepository) UpdateNodes(
	ctx context.Context,
	tenantID, ownerUID, brandID string,
	nodes []libkg.Node,
	painTagCount int,
) (*entity.Graph, error) {
	if r.collection == nil {
		return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured")
	}
	now := clock.NowUnixNano()
	var out entity.Graph
	err := r.collection.FindOneAndUpdate(
		ctx,
		brandOwnerFilter(tenantID, ownerUID, brandID),
		bson.M{"$set": bson.M{
			"nodes":          nodes,
			"pain_tag_count": painTagCount,
			"update_at":      now,
		}},
		options.FindOneAndUpdate().SetReturnDocument(options.After),
	).Decode(&out)
	if err == mongo.ErrNoDocuments {
		return nil, app.For(code.Brand).ResNotFound("knowledge graph not found")
	}
	if err != nil {
		return nil, err
	}
	return &out, nil
}