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 }