Files
relata/internal/storage/mongo/entities.go
T
2026-08-11 10:40:10 +01:00

606 lines
17 KiB
Go

package mongo
import (
"context"
"strings"
"time"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
"github.com/andydixon/relata/internal/domain"
"github.com/andydixon/relata/internal/ids"
)
// EntityRepo stores extracted entities.
type EntityRepo struct{ base }
var _ domain.EntityRepository = (*EntityRepo)(nil)
// maxSourceRefsPerEntity bounds how many mention references one entity keeps.
// A name appearing thousands of times does not need thousands of stored refs;
// the count is tracked separately in Mentions.
const maxSourceRefsPerEntity = 50
// UpsertByKey merges an observation into the entity with the same
// (workspace, normalised_key), or creates it.
//
// Merging happens ONLY on an exact key match. Similar-looking names create
// separate entities and, elsewhere, a candidate alias for a human to review.
// The unique index on (workspace_id, normalised_key) makes this safe under
// concurrent workers: a racing insert loses and is retried as an update.
func (r *EntityRepo) UpsertByKey(ctx context.Context, e *domain.Entity) (*domain.Entity, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(e.WorkspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "normalised_key", Value: e.NormalisedKey})
now := time.Now().UTC()
if e.ID == "" {
e.ID = ids.New()
}
setOnInsert := bson.D{
{Key: "_id", Value: e.ID},
{Key: "workspace_id", Value: e.WorkspaceID},
{Key: "normalised_key", Value: e.NormalisedKey},
{Key: "type", Value: string(e.Type)},
{Key: "canonical_name", Value: e.CanonicalName},
{Key: "method", Value: string(e.Method)},
{Key: "created_at", Value: now},
{Key: "deleted", Value: false},
}
addToSet := bson.D{}
if len(e.ItemIDs) > 0 {
addToSet = append(addToSet, bson.E{Key: "item_ids", Value: bson.D{{Key: "$each", Value: stringsToAny(e.ItemIDs)}}})
}
if len(e.Aliases) > 0 {
addToSet = append(addToSet, bson.E{Key: "aliases", Value: bson.D{{Key: "$each", Value: stringsToAny(e.Aliases)}}})
}
set := bson.D{
{Key: "updated_at", Value: now},
{Key: "evidence_score", Value: e.EvidenceScore},
}
if e.ModelVersion != "" {
set = append(set, bson.E{Key: "model_version", Value: e.ModelVersion})
}
update := bson.D{
{Key: "$setOnInsert", Value: setOnInsert},
{Key: "$set", Value: set},
{Key: "$inc", Value: bson.D{{Key: "mentions", Value: maxIntVal(e.Mentions, 1)}}},
}
if len(addToSet) > 0 {
update = append(update, bson.E{Key: "$addToSet", Value: addToSet})
}
if len(e.SourceRefs) > 0 {
refs := make([]any, 0, len(e.SourceRefs))
for _, ref := range e.SourceRefs {
refs = append(refs, ref)
}
update = append(update, bson.E{Key: "$push", Value: bson.D{
{Key: "source_refs", Value: bson.D{
{Key: "$each", Value: refs},
// Keep the most recent references only, so a heavily-mentioned
// entity cannot grow without bound.
{Key: "$slice", Value: -maxSourceRefsPerEntity},
}},
}})
}
opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
var out domain.Entity
err = r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&out)
r.observe("upsert_by_key", started, err)
if err != nil {
return nil, translate(err)
}
return &out, nil
}
func maxIntVal(a, b int) int {
if a > b {
return a
}
return b
}
// Get fetches one entity within a workspace.
func (r *EntityRepo) Get(ctx context.Context, workspaceID, id string) (*domain.Entity, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "_id", Value: id})
var e domain.Entity
err = r.coll.FindOne(ctx, f).Decode(&e)
r.observe("get", started, err)
if err != nil {
return nil, translate(err)
}
return &e, nil
}
// GetMany fetches entities by ID within a workspace.
func (r *EntityRepo) GetMany(ctx context.Context, workspaceID string, entityIDs []string) ([]*domain.Entity, error) {
if len(entityIDs) == 0 {
return nil, nil
}
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(entityIDs)}}})
cur, err := r.coll.Find(ctx, f)
r.observe("get_many", started, err)
if err != nil {
return nil, translate(err)
}
var out []*domain.Entity
if err := cur.All(ctx, &out); err != nil {
return nil, err
}
return out, nil
}
// FindByNormalised looks up entities by their exact normalised keys.
func (r *EntityRepo) FindByNormalised(ctx context.Context, workspaceID string, keys []string) ([]*domain.Entity, error) {
if len(keys) == 0 {
return nil, nil
}
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "normalised_key", Value: bson.D{{Key: "$in", Value: stringsToAny(keys)}}})
cur, err := r.coll.Find(ctx, f)
r.observe("find_by_normalised", started, err)
if err != nil {
return nil, translate(err)
}
var out []*domain.Entity
if err := cur.All(ctx, &out); err != nil {
return nil, err
}
return out, nil
}
// List returns a page of entities matching a typed filter.
func (r *EntityRepo) List(ctx context.Context, workspaceID string, ef domain.EntityFilter, p domain.Page) ([]*domain.Entity, int64, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, 0, err
}
f = notDeleted(f)
if len(ef.Types) > 0 {
vals := make([]any, 0, len(ef.Types))
for _, t := range ef.Types {
if t.Valid() {
vals = append(vals, string(t))
}
}
if len(vals) > 0 {
f = append(f, bson.E{Key: "type", Value: bson.D{{Key: "$in", Value: vals}}})
}
}
if ef.NormalisedKey != "" {
f = append(f, bson.E{Key: "normalised_key", Value: ef.NormalisedKey})
}
if ef.ItemID != "" {
f = append(f, bson.E{Key: "item_ids", Value: ef.ItemID})
}
if ef.MinMentions > 0 {
f = append(f, bson.E{Key: "mentions", Value: bson.D{{Key: "$gte", Value: ef.MinMentions}}})
}
if n := strings.TrimSpace(ef.NameContains); n != "" {
f = append(f, bson.E{Key: "canonical_name", Value: bson.D{
{Key: "$regex", Value: regexpQuote(n)},
{Key: "$options", Value: "i"},
}})
}
p = p.Normalise(50, 500)
total, err := r.coll.CountDocuments(ctx, f)
if err != nil {
r.observe("list", started, err)
return nil, 0, translate(err)
}
cur, err := r.coll.Find(ctx, f, options.Find().
SetSort(bson.D{{Key: "mentions", Value: -1}, {Key: "canonical_name", Value: 1}}).
SetSkip(int64(p.Offset)).SetLimit(int64(p.Limit)))
r.observe("list", started, err)
if err != nil {
return nil, 0, translate(err)
}
var out []*domain.Entity
if err := cur.All(ctx, &out); err != nil {
return nil, 0, err
}
return out, total, nil
}
// ListByItem returns the entities referencing an item.
func (r *EntityRepo) ListByItem(ctx context.Context, workspaceID, itemID string) ([]*domain.Entity, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = notDeleted(append(f, bson.E{Key: "item_ids", Value: itemID}))
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{{Key: "mentions", Value: -1}}))
r.observe("list_by_item", started, err)
if err != nil {
return nil, translate(err)
}
var out []*domain.Entity
if err := cur.All(ctx, &out); err != nil {
return nil, err
}
return out, nil
}
// ItemsSharingEntities returns item IDs referencing any of the given entities.
//
// This is the indexed path that keeps candidate discovery bounded: instead of
// comparing a new item with every stored item, it asks only for items that
// already share an entity.
func (r *EntityRepo) ItemsSharingEntities(ctx context.Context, workspaceID string, entityIDs []string, excludeItem string, limit int) ([]string, error) {
if len(entityIDs) == 0 {
return nil, nil
}
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = notDeleted(append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(entityIDs)}}}))
if limit <= 0 || limit > 500 {
limit = 100
}
cur, err := r.coll.Aggregate(ctx, mongo.Pipeline{
{{Key: "$match", Value: f}},
{{Key: "$unwind", Value: "$item_ids"}},
{{Key: "$match", Value: bson.D{{Key: "item_ids", Value: bson.D{{Key: "$ne", Value: excludeItem}}}}}},
{{Key: "$group", Value: bson.D{
{Key: "_id", Value: "$item_ids"},
{Key: "shared", Value: bson.D{{Key: "$sum", Value: 1}}},
}}},
// Items sharing the most entities are the best candidates.
{{Key: "$sort", Value: bson.D{{Key: "shared", Value: -1}}}},
{{Key: "$limit", Value: limit}},
})
r.observe("items_sharing_entities", started, err)
if err != nil {
return nil, translate(err)
}
defer cur.Close(ctx)
var out []string
for cur.Next(ctx) {
var row struct {
ID string `bson:"_id"`
}
if err := cur.Decode(&row); err == nil && row.ID != "" {
out = append(out, row.ID)
}
}
return out, cur.Err()
}
// RemoveItemRefs strips an item from every entity that references it and
// returns the IDs of entities left with no references at all.
//
// Those orphans are what item deletion removes: an entity that exists only
// because of a deleted item has no evidence left to justify it.
func (r *EntityRepo) RemoveItemRefs(ctx context.Context, workspaceID, itemID string) ([]string, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "item_ids", Value: itemID})
update := bson.D{
{Key: "$pull", Value: bson.D{
{Key: "item_ids", Value: itemID},
{Key: "source_refs", Value: bson.D{{Key: "item_id", Value: itemID}}},
}},
{Key: "$set", Value: bson.D{{Key: "updated_at", Value: time.Now().UTC()}}},
}
if _, err := r.coll.UpdateMany(ctx, f, update); err != nil {
r.observe("remove_item_refs", started, err)
return nil, translate(err)
}
// Find entities that now reference nothing.
orphanFilter, _ := wsFilter(workspaceID)
orphanFilter = append(orphanFilter, bson.E{Key: "$or", Value: bson.A{
bson.D{{Key: "item_ids", Value: bson.D{{Key: "$size", Value: 0}}}},
bson.D{{Key: "item_ids", Value: bson.D{{Key: "$exists", Value: false}}}},
}})
cur, err := r.coll.Find(ctx, orphanFilter, options.Find().SetProjection(bson.D{{Key: "_id", Value: 1}}))
r.observe("remove_item_refs", started, err)
if err != nil {
return nil, translate(err)
}
defer cur.Close(ctx)
var orphans []string
for cur.Next(ctx) {
var row struct {
ID string `bson:"_id"`
}
if err := cur.Decode(&row); err == nil {
orphans = append(orphans, row.ID)
}
}
return orphans, cur.Err()
}
// DeleteMany removes entities by ID within a workspace.
func (r *EntityRepo) DeleteMany(ctx context.Context, workspaceID string, entityIDs []string) error {
if len(entityIDs) == 0 {
return nil
}
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return err
}
f = append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(entityIDs)}}})
_, err = r.coll.DeleteMany(ctx, f)
r.observe("delete_many", started, err)
return translate(err)
}
// Count returns the number of live entities in a workspace.
func (r *EntityRepo) Count(ctx context.Context, workspaceID string) (int64, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return 0, err
}
n, err := r.coll.CountDocuments(ctx, notDeleted(f))
r.observe("count", started, err)
return n, translate(err)
}
// CorpusFrequency reports how many items reference each of the given entities,
// which the scorer uses to discount ubiquitous names.
func (r *EntityRepo) CorpusFrequency(ctx context.Context, workspaceID string, entityIDs []string) (map[string]int, error) {
if len(entityIDs) == 0 {
return nil, nil
}
ctx, cancel := r.ctx(ctx)
defer cancel()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(entityIDs)}}})
cur, err := r.coll.Find(ctx, f, options.Find().SetProjection(bson.D{{Key: "item_ids", Value: 1}}))
if err != nil {
return nil, translate(err)
}
defer cur.Close(ctx)
out := map[string]int{}
for cur.Next(ctx) {
var row struct {
ID string `bson:"_id"`
ItemIDs []string `bson:"item_ids"`
}
if err := cur.Decode(&row); err == nil {
out[row.ID] = len(row.ItemIDs)
}
}
return out, cur.Err()
}
// ---------------------------------------------------------------------------
// AliasRepo stores candidate and confirmed entity aliases.
type AliasRepo struct{ base }
var _ domain.AliasRepository = (*AliasRepo)(nil)
// Upsert records an alias proposal, leaving an existing human decision intact.
func (r *AliasRepo) Upsert(ctx context.Context, a *domain.EntityAlias) error {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(a.WorkspaceID)
if err != nil {
return err
}
f = append(f,
bson.E{Key: "entity_id", Value: a.EntityID},
bson.E{Key: "normalised", Value: a.Normalised},
)
now := time.Now().UTC()
if a.ID == "" {
a.ID = ids.New()
}
// Status is set only on insert: re-running discovery must not reset an
// alias a person has already confirmed or rejected.
update := bson.D{
{Key: "$setOnInsert", Value: bson.D{
{Key: "_id", Value: a.ID},
{Key: "workspace_id", Value: a.WorkspaceID},
{Key: "entity_id", Value: a.EntityID},
{Key: "normalised", Value: a.Normalised},
{Key: "status", Value: string(domain.AliasCandidate)},
{Key: "created_at", Value: now},
}},
{Key: "$set", Value: bson.D{
{Key: "alias", Value: a.Alias},
{Key: "candidate_entity_id", Value: a.CandidateEntityID},
{Key: "method", Value: string(a.Method)},
{Key: "similarity", Value: a.Similarity},
{Key: "reason", Value: a.Reason},
{Key: "updated_at", Value: now},
}},
}
_, err = r.coll.UpdateOne(ctx, f, update, options.UpdateOne().SetUpsert(true))
r.observe("upsert", started, err)
if isDup(err) {
return nil // a concurrent worker inserted the same proposal
}
return translate(err)
}
// ListByEntity returns aliases proposed for an entity.
func (r *AliasRepo) ListByEntity(ctx context.Context, workspaceID, entityID string) ([]*domain.EntityAlias, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, err
}
f = append(f, bson.E{Key: "entity_id", Value: entityID})
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{{Key: "similarity", Value: -1}}))
r.observe("list_by_entity", started, err)
if err != nil {
return nil, translate(err)
}
var out []*domain.EntityAlias
if err := cur.All(ctx, &out); err != nil {
return nil, err
}
return out, nil
}
// List returns aliases with a given status.
func (r *AliasRepo) List(ctx context.Context, workspaceID string, status domain.AliasStatus, p domain.Page) ([]*domain.EntityAlias, int64, error) {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return nil, 0, err
}
if status != "" {
f = append(f, bson.E{Key: "status", Value: string(status)})
}
p = p.Normalise(50, 200)
total, err := r.coll.CountDocuments(ctx, f)
if err != nil {
return nil, 0, translate(err)
}
cur, err := r.coll.Find(ctx, f, options.Find().
SetSort(bson.D{{Key: "similarity", Value: -1}}).
SetSkip(int64(p.Offset)).SetLimit(int64(p.Limit)))
r.observe("list", started, err)
if err != nil {
return nil, 0, translate(err)
}
var out []*domain.EntityAlias
if err := cur.All(ctx, &out); err != nil {
return nil, 0, err
}
return out, total, nil
}
// SetStatus records a human decision about an alias.
func (r *AliasRepo) SetStatus(ctx context.Context, workspaceID, id string, st domain.AliasStatus) error {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return err
}
f = append(f, bson.E{Key: "_id", Value: id})
res, err := r.coll.UpdateOne(ctx, f, bson.D{{Key: "$set", Value: bson.D{
{Key: "status", Value: string(st)},
{Key: "updated_at", Value: time.Now().UTC()},
}}})
r.observe("set_status", started, err)
if err != nil {
return translate(err)
}
if res.MatchedCount == 0 {
return domain.ErrNotFound
}
return nil
}
// DeleteByEntity removes an entity's aliases.
func (r *AliasRepo) DeleteByEntity(ctx context.Context, workspaceID, entityID string) error {
ctx, cancel := r.ctx(ctx)
defer cancel()
started := time.Now()
f, err := wsFilter(workspaceID)
if err != nil {
return err
}
f = append(f, bson.E{Key: "entity_id", Value: entityID})
_, err = r.coll.DeleteMany(ctx, f)
r.observe("delete_by_entity", started, err)
return translate(err)
}