Code sync, security fixes
This commit is contained in:
Andy Dixon
@@ -0,0 +1,191 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
)
|
||||
|
||||
// APITokenRepo stores API credentials.
|
||||
//
|
||||
// Every method here works in hashes. There is no query by prefix, no partial
|
||||
// match and no method that returns a secret, because the secret is not stored:
|
||||
// a database read, a backup or a dump yields nothing that can authenticate.
|
||||
type APITokenRepo struct{ base }
|
||||
|
||||
var _ domain.APITokenRepository = (*APITokenRepo)(nil)
|
||||
|
||||
// Create inserts a token record.
|
||||
func (r *APITokenRepo) Create(ctx context.Context, t *domain.APIToken) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if t.WorkspaceID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
if t.TokenHash == "" {
|
||||
return fmt.Errorf("%w: a token record must carry a hash", domain.ErrInvalidInput)
|
||||
}
|
||||
_, err := r.coll.InsertOne(ctx, t)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// GetByHash finds a token by the SHA-256 of its secret.
|
||||
//
|
||||
// The lookup is on the whole hash against a unique index. There is no
|
||||
// comparison in application code for an attacker to time, and no way to probe
|
||||
// for a prefix.
|
||||
func (r *APITokenRepo) GetByHash(ctx context.Context, hash string) (*domain.APIToken, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if hash == "" {
|
||||
return nil, domain.ErrNotFound
|
||||
}
|
||||
var t domain.APIToken
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "token_hash", Value: hash}}).Decode(&t)
|
||||
r.observe("get_by_hash", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// ListForUser returns a person's tokens, newest first.
|
||||
//
|
||||
// Revoked and expired tokens are included: the point of the list is to show
|
||||
// what exists and what happened to it, and hiding a revocation would make the
|
||||
// page a worse audit surface than the log.
|
||||
func (r *APITokenRepo) ListForUser(ctx context.Context, workspaceID, userID string) ([]*domain.APIToken, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if userID == "" {
|
||||
return nil, fmt.Errorf("%w: user id must not be empty", domain.ErrInvalidInput)
|
||||
}
|
||||
f = append(f, bson.E{Key: "user_id", Value: userID})
|
||||
|
||||
opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}}).SetLimit(500)
|
||||
cur, err := r.coll.Find(ctx, f, opts)
|
||||
if err != nil {
|
||||
r.observe("list_for_user", started, err)
|
||||
return nil, translate(err)
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
var out []*domain.APIToken
|
||||
for cur.Next(ctx) {
|
||||
var t domain.APIToken
|
||||
if err := cur.Decode(&t); err != nil {
|
||||
r.observe("list_for_user", started, err)
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, &t)
|
||||
}
|
||||
r.observe("list_for_user", started, cur.Err())
|
||||
return out, cur.Err()
|
||||
}
|
||||
|
||||
// Get fetches one token within a workspace.
|
||||
func (r *APITokenRepo) Get(ctx context.Context, workspaceID, id string) (*domain.APIToken, 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 t domain.APIToken
|
||||
err = r.coll.FindOne(ctx, f).Decode(&t)
|
||||
r.observe("get", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// Revoke marks a token unusable.
|
||||
//
|
||||
// The record is kept rather than deleted, so the token's existence and its end
|
||||
// stay visible. A revoked token is refused on the next request; there is no
|
||||
// cached grant to wait out.
|
||||
func (r *APITokenRepo) Revoke(ctx context.Context, workspaceID, id, revokedBy string, at time.Time) 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})
|
||||
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "revoked_at", Value: at},
|
||||
{Key: "revoked_by", Value: revokedBy},
|
||||
}}}
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("revoke", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TouchLastUsed records that a token authenticated a request.
|
||||
//
|
||||
// Callers throttle this; the repository does not, because it has no view of how
|
||||
// often it is being asked.
|
||||
func (r *APITokenRepo) TouchLastUsed(ctx context.Context, id string, at time.Time) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.UpdateOne(ctx,
|
||||
bson.D{{Key: "_id", Value: id}},
|
||||
bson.D{{Key: "$set", Value: bson.D{{Key: "last_used_at", Value: at}}}})
|
||||
r.observe("touch_last_used", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// DeleteForUser removes every token belonging to a user.
|
||||
func (r *APITokenRepo) DeleteForUser(ctx context.Context, workspaceID, userID string) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if userID == "" {
|
||||
return 0, fmt.Errorf("%w: user id must not be empty", domain.ErrInvalidInput)
|
||||
}
|
||||
f = append(f, bson.E{Key: "user_id", Value: userID})
|
||||
|
||||
res, err := r.coll.DeleteMany(ctx, f)
|
||||
r.observe("delete_for_user", started, err)
|
||||
if err != nil {
|
||||
return 0, translate(err)
|
||||
}
|
||||
return res.DeletedCount, nil
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
"github.com/andydixon/relata/internal/ids"
|
||||
)
|
||||
|
||||
// AuditRepo is an append-only record of consequential actions.
|
||||
type AuditRepo struct{ base }
|
||||
|
||||
var _ domain.AuditRepository = (*AuditRepo)(nil)
|
||||
|
||||
// Record appends an audit event. There is deliberately no update or delete.
|
||||
func (r *AuditRepo) Record(ctx context.Context, ev *domain.AuditEvent) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if ev.ID == "" {
|
||||
ev.ID = ids.New()
|
||||
}
|
||||
if ev.At.IsZero() {
|
||||
ev.At = time.Now().UTC()
|
||||
}
|
||||
_, err := r.coll.InsertOne(ctx, ev)
|
||||
r.observe("record", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// List returns a page of audit events, newest first.
|
||||
func (r *AuditRepo) List(ctx context.Context, workspaceID string, p domain.Page) ([]*domain.AuditEvent, int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
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: "at", 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.AuditEvent
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return out, total, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// RunRepo stores processing run history.
|
||||
type RunRepo struct{ base }
|
||||
|
||||
var _ domain.RunRepository = (*RunRepo)(nil)
|
||||
|
||||
// Create opens a processing run.
|
||||
func (r *RunRepo) Create(ctx context.Context, run *domain.ProcessingRun) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if run.ID == "" {
|
||||
run.ID = ids.New()
|
||||
}
|
||||
if run.StartedAt.IsZero() {
|
||||
run.StartedAt = time.Now().UTC()
|
||||
}
|
||||
if run.Outcome == "" {
|
||||
run.Outcome = "running"
|
||||
}
|
||||
_, err := r.coll.InsertOne(ctx, run)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Finish closes a processing run with its outcome and statistics.
|
||||
//
|
||||
// Runs are never deleted: a reprocess adds a new run, so the record of what a
|
||||
// previous parser or model produced survives for audit.
|
||||
func (r *RunRepo) Finish(ctx context.Context, id, outcome, errMsg string, stats map[string]int) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
now := time.Now().UTC()
|
||||
_, err := r.coll.UpdateOne(ctx,
|
||||
bson.D{{Key: "_id", Value: id}},
|
||||
bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "ended_at", Value: now},
|
||||
{Key: "outcome", Value: outcome},
|
||||
{Key: "error", Value: errMsg},
|
||||
{Key: "stats", Value: stats},
|
||||
}}})
|
||||
r.observe("finish", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// ListByItem returns an item's processing history, newest first.
|
||||
func (r *RunRepo) ListByItem(ctx context.Context, workspaceID, itemID string, limit int) ([]*domain.ProcessingRun, 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_id", Value: itemID})
|
||||
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 20
|
||||
}
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().
|
||||
SetSort(bson.D{{Key: "started_at", Value: -1}}).SetLimit(int64(limit)))
|
||||
r.observe("list_by_item", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.ProcessingRun
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// DeleteByItem removes an item's run history, used by hard deletion.
|
||||
func (r *RunRepo) DeleteByItem(ctx context.Context, workspaceID, itemID string) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f = append(f, bson.E{Key: "item_id", Value: itemID})
|
||||
_, err = r.coll.DeleteMany(ctx, f)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// ModelVersionRepo tracks which models have produced stored data.
|
||||
type ModelVersionRepo struct{ base }
|
||||
|
||||
var _ domain.ModelVersionRepository = (*ModelVersionRepo)(nil)
|
||||
|
||||
// Observe records that a model was used, so a configuration change can be
|
||||
// detected and the affected records identified for reprocessing.
|
||||
func (r *ModelVersionRepo) Observe(ctx context.Context, workspaceID, kind, name string, dims int) error {
|
||||
if name == "" {
|
||||
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: "kind", Value: kind}, bson.E{Key: "name", Value: name})
|
||||
|
||||
now := time.Now().UTC()
|
||||
update := bson.D{
|
||||
{Key: "$setOnInsert", Value: bson.D{
|
||||
{Key: "_id", Value: ids.New()},
|
||||
{Key: "workspace_id", Value: workspaceID},
|
||||
{Key: "kind", Value: kind},
|
||||
{Key: "name", Value: name},
|
||||
{Key: "first_seen_at", Value: now},
|
||||
}},
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "last_seen_at", Value: now},
|
||||
{Key: "dimensions", Value: dims},
|
||||
}},
|
||||
}
|
||||
_, err = r.coll.UpdateOne(ctx, f, update, options.UpdateOne().SetUpsert(true))
|
||||
if isDup(err) {
|
||||
err = nil
|
||||
}
|
||||
r.observe("observe", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// List returns the models seen in a workspace.
|
||||
func (r *ModelVersionRepo) List(ctx context.Context, workspaceID string) ([]*domain.ModelVersion, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{
|
||||
{Key: "kind", Value: 1}, {Key: "name", Value: 1},
|
||||
}))
|
||||
r.observe("list", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.ModelVersion
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
)
|
||||
|
||||
// UserRepo stores local accounts.
|
||||
type UserRepo struct{ base }
|
||||
|
||||
var _ domain.UserRepository = (*UserRepo)(nil)
|
||||
|
||||
// Create inserts a user. Usernames are stored lowercased and a unique index
|
||||
// enforces that two accounts cannot differ only by case.
|
||||
func (r *UserRepo) Create(ctx context.Context, u *domain.User) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
u.Username = strings.ToLower(strings.TrimSpace(u.Username))
|
||||
_, err := r.coll.InsertOne(ctx, u)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// GetByUsername looks up an account for authentication.
|
||||
func (r *UserRepo) GetByUsername(ctx context.Context, username string) (*domain.User, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f := bson.D{{Key: "username", Value: strings.ToLower(strings.TrimSpace(username))}}
|
||||
var u domain.User
|
||||
err := r.coll.FindOne(ctx, f).Decode(&u)
|
||||
r.observe("get_by_username", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
// GetByOAuthSubject looks up a federated account.
|
||||
//
|
||||
// The filter is the (provider, subject) pair and nothing else. Falling back to
|
||||
// the email address would let a provider that recycled one hand over the
|
||||
// previous holder's workspace.
|
||||
func (r *UserRepo) GetByOAuthSubject(ctx context.Context, provider, subject string) (*domain.User, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
provider, subject = strings.TrimSpace(provider), strings.TrimSpace(subject)
|
||||
if provider == "" || subject == "" {
|
||||
return nil, domain.ErrNotFound
|
||||
}
|
||||
|
||||
f := bson.D{
|
||||
{Key: "oauth_provider", Value: provider},
|
||||
{Key: "oauth_subject", Value: subject},
|
||||
}
|
||||
var u domain.User
|
||||
err := r.coll.FindOne(ctx, f).Decode(&u)
|
||||
r.observe("get_by_oauth_subject", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
// GetByID fetches a user by ID.
|
||||
func (r *UserRepo) GetByID(ctx context.Context, id string) (*domain.User, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
var u domain.User
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "_id", Value: id}}).Decode(&u)
|
||||
r.observe("get_by_id", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
// Update replaces a user document.
|
||||
func (r *UserRepo) Update(ctx context.Context, u *domain.User) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
u.UpdatedAt = time.Now().UTC()
|
||||
res, err := r.coll.ReplaceOne(ctx, bson.D{{Key: "_id", Value: u.ID}}, u)
|
||||
r.observe("update", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Count returns the number of accounts, used to detect first-run setup.
|
||||
func (r *UserRepo) Count(ctx context.Context) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
n, err := r.coll.CountDocuments(ctx, bson.D{})
|
||||
r.observe("count", started, err)
|
||||
return n, translate(err)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// SessionRepo stores browser sessions.
|
||||
type SessionRepo struct{ base }
|
||||
|
||||
var _ domain.SessionRepository = (*SessionRepo)(nil)
|
||||
|
||||
// Create inserts a session record.
|
||||
func (r *SessionRepo) Create(ctx context.Context, s *domain.Session) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.InsertOne(ctx, s)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// GetByTokenHash looks up a session by the hash of its cookie token.
|
||||
//
|
||||
// The lookup is by hash, never by the token itself, so the stored record cannot
|
||||
// be replayed as a credential by anyone who can read the database.
|
||||
func (r *SessionRepo) GetByTokenHash(ctx context.Context, hash string) (*domain.Session, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
var s domain.Session
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "token_hash", Value: hash}}).Decode(&s)
|
||||
r.observe("get_by_token", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &s, nil
|
||||
}
|
||||
|
||||
// Touch extends a session's idle expiry.
|
||||
func (r *SessionRepo) Touch(ctx context.Context, id string, lastSeen, expires time.Time) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.UpdateOne(ctx,
|
||||
bson.D{{Key: "_id", Value: id}},
|
||||
bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "last_seen_at", Value: lastSeen},
|
||||
{Key: "expires_at", Value: expires},
|
||||
}}})
|
||||
r.observe("touch", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Revoke invalidates one session.
|
||||
func (r *SessionRepo) Revoke(ctx context.Context, id string) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.UpdateOne(ctx,
|
||||
bson.D{{Key: "_id", Value: id}},
|
||||
bson.D{{Key: "$set", Value: bson.D{{Key: "revoked", Value: true}}}})
|
||||
r.observe("revoke", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// RevokeAllForUser invalidates every session belonging to a user.
|
||||
func (r *SessionRepo) RevokeAllForUser(ctx context.Context, userID string) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.UpdateMany(ctx,
|
||||
bson.D{{Key: "user_id", Value: userID}},
|
||||
bson.D{{Key: "$set", Value: bson.D{{Key: "revoked", Value: true}}}})
|
||||
r.observe("revoke_all", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// DeleteExpired removes sessions past their absolute expiry. A TTL index does
|
||||
// this too; the explicit sweep makes the behaviour testable and prompt.
|
||||
func (r *SessionRepo) DeleteExpired(ctx context.Context, before time.Time) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
res, err := r.coll.DeleteMany(ctx, bson.D{
|
||||
{Key: "absolute_expiry", Value: bson.D{{Key: "$lt", Value: before}}},
|
||||
})
|
||||
r.observe("delete_expired", started, err)
|
||||
if err != nil {
|
||||
return 0, translate(err)
|
||||
}
|
||||
return res.DeletedCount, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// WorkspaceRepo stores workspaces.
|
||||
type WorkspaceRepo struct{ base }
|
||||
|
||||
var _ domain.WorkspaceRepository = (*WorkspaceRepo)(nil)
|
||||
|
||||
// Create inserts a workspace.
|
||||
func (r *WorkspaceRepo) Create(ctx context.Context, w *domain.Workspace) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
_, err := r.coll.InsertOne(ctx, w)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Get fetches a workspace by ID.
|
||||
func (r *WorkspaceRepo) Get(ctx context.Context, id string) (*domain.Workspace, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
var w domain.Workspace
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "_id", Value: id}}).Decode(&w)
|
||||
r.observe("get", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &w, nil
|
||||
}
|
||||
|
||||
// GetBySlug fetches a workspace by its URL-safe name.
|
||||
func (r *WorkspaceRepo) GetBySlug(ctx context.Context, slug string) (*domain.Workspace, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
var w domain.Workspace
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "slug", Value: slug}}).Decode(&w)
|
||||
r.observe("get_by_slug", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &w, nil
|
||||
}
|
||||
|
||||
// List returns every workspace.
|
||||
func (r *WorkspaceRepo) List(ctx context.Context) ([]*domain.Workspace, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
cur, err := r.coll.Find(ctx, bson.D{}, options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}}))
|
||||
r.observe("list", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Workspace
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
)
|
||||
|
||||
// ChunkRepo stores item chunks.
|
||||
type ChunkRepo struct{ base }
|
||||
|
||||
var _ domain.ChunkRepository = (*ChunkRepo)(nil)
|
||||
|
||||
// ReplaceForItem swaps an item's chunks for a new set.
|
||||
//
|
||||
// Delete-then-insert is what makes re-chunking idempotent: reprocessing an item
|
||||
// leaves exactly one generation of chunks, never two overlapping ones. The
|
||||
// caller is responsible for removing the corresponding vectors, which it does
|
||||
// by deleting the item's vectors before re-embedding.
|
||||
func (r *ChunkRepo) ReplaceForItem(ctx context.Context, workspaceID, itemID string, chunks []*domain.Chunk) 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: "item_id", Value: itemID})
|
||||
|
||||
if _, err := r.coll.DeleteMany(ctx, f); err != nil {
|
||||
r.observe("replace_for_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
if len(chunks) == 0 {
|
||||
r.observe("replace_for_item", started, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
docs := make([]any, 0, len(chunks))
|
||||
for _, c := range chunks {
|
||||
if c.WorkspaceID != workspaceID || c.ItemID != itemID {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
docs = append(docs, c)
|
||||
}
|
||||
_, err = r.coll.InsertMany(ctx, docs)
|
||||
r.observe("replace_for_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// ListByItem returns an item's chunks in sequence order.
|
||||
func (r *ChunkRepo) ListByItem(ctx context.Context, workspaceID, itemID string) ([]*domain.Chunk, 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_id", Value: itemID})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{{Key: "sequence", Value: 1}}))
|
||||
r.observe("list_by_item", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Chunk
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// GetMany fetches chunks by ID within a workspace.
|
||||
func (r *ChunkRepo) GetMany(ctx context.Context, workspaceID string, ids []string) ([]*domain.Chunk, error) {
|
||||
if len(ids) == 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(ids)}}})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f)
|
||||
r.observe("get_many", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Chunk
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// MarkEmbedded records that chunks now have vectors under a given model.
|
||||
func (r *ChunkRepo) MarkEmbedded(ctx context.Context, workspaceID string, ids []string, model string) error {
|
||||
if len(ids) == 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(ids)}}})
|
||||
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "embedded", Value: true},
|
||||
{Key: "embedding_model", Value: model},
|
||||
}}}
|
||||
_, err = r.coll.UpdateMany(ctx, f, update)
|
||||
r.observe("mark_embedded", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// ListUnembedded finds chunks lacking a vector for the given model.
|
||||
//
|
||||
// A chunk embedded under a different model counts as unembedded, which is what
|
||||
// makes re-embedding after a model change straightforward.
|
||||
func (r *ChunkRepo) ListUnembedded(ctx context.Context, workspaceID, model string, limit int) ([]*domain.Chunk, 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: "$or", Value: bson.A{
|
||||
bson.D{{Key: "embedded", Value: bson.D{{Key: "$ne", Value: true}}}},
|
||||
bson.D{{Key: "embedding_model", Value: bson.D{{Key: "$ne", Value: model}}}},
|
||||
}})
|
||||
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 200
|
||||
}
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().SetLimit(int64(limit)).
|
||||
SetSort(bson.D{{Key: "item_id", Value: 1}, {Key: "sequence", Value: 1}}))
|
||||
r.observe("list_unembedded", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Chunk
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// DeleteByItem removes an item's chunks.
|
||||
func (r *ChunkRepo) DeleteByItem(ctx context.Context, workspaceID, itemID 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: "item_id", Value: itemID})
|
||||
|
||||
_, err = r.coll.DeleteMany(ctx, f)
|
||||
r.observe("delete_by_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Count returns the number of chunks in a workspace.
|
||||
func (r *ChunkRepo) 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, f)
|
||||
r.observe("count", started, err)
|
||||
return n, translate(err)
|
||||
}
|
||||
|
||||
// TextSearch finds chunks lexically, used to build result excerpts.
|
||||
func (r *ChunkRepo) TextSearch(ctx context.Context, workspaceID, query string, limit int) ([]domain.ChunkHit, error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
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: "$text", Value: bson.D{{Key: "$search", Value: query}}})
|
||||
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 60
|
||||
}
|
||||
opts := options.Find().
|
||||
SetProjection(bson.D{
|
||||
{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}},
|
||||
{Key: "item_id", Value: 1},
|
||||
{Key: "text", Value: 1},
|
||||
}).
|
||||
SetSort(bson.D{{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}}}).
|
||||
SetLimit(int64(limit))
|
||||
|
||||
cur, err := r.coll.Find(ctx, f, opts)
|
||||
r.observe("text_search", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
var out []domain.ChunkHit
|
||||
for cur.Next(ctx) {
|
||||
var row struct {
|
||||
ID string `bson:"_id"`
|
||||
ItemID string `bson:"item_id"`
|
||||
Text string `bson:"text"`
|
||||
Score float64 `bson:"score"`
|
||||
}
|
||||
if err := cur.Decode(&row); err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, domain.ChunkHit{
|
||||
ChunkID: row.ID, ItemID: row.ItemID, Text: row.Text, Score: row.Score,
|
||||
})
|
||||
}
|
||||
return out, cur.Err()
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
)
|
||||
|
||||
// ClaimRepo stores extracted claims.
|
||||
type ClaimRepo struct{ base }
|
||||
|
||||
var _ domain.ClaimRepository = (*ClaimRepo)(nil)
|
||||
|
||||
// ReplaceForItem swaps an item's claims for a new set, making claim extraction
|
||||
// idempotent across reprocessing runs.
|
||||
func (r *ClaimRepo) ReplaceForItem(ctx context.Context, workspaceID, itemID string, claims []*domain.Claim) 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: "item_id", Value: itemID})
|
||||
|
||||
if _, err := r.coll.DeleteMany(ctx, f); err != nil {
|
||||
r.observe("replace_for_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
if len(claims) == 0 {
|
||||
r.observe("replace_for_item", started, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Deduplicate within the batch: two chunks can yield the same claim, and
|
||||
// the unique dedupe_key index would otherwise reject the whole insert.
|
||||
seen := map[string]bool{}
|
||||
docs := make([]any, 0, len(claims))
|
||||
for _, c := range claims {
|
||||
if c.WorkspaceID != workspaceID || c.ItemID != itemID {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
if seen[c.DedupeKey] {
|
||||
continue
|
||||
}
|
||||
seen[c.DedupeKey] = true
|
||||
docs = append(docs, c)
|
||||
}
|
||||
|
||||
_, err = r.coll.InsertMany(ctx, docs, options.InsertMany().SetOrdered(false))
|
||||
if isDup(err) {
|
||||
// A concurrent run inserted an identical claim; that is the intended
|
||||
// outcome, not a failure.
|
||||
err = nil
|
||||
}
|
||||
r.observe("replace_for_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// ListByItem returns an item's claims.
|
||||
func (r *ClaimRepo) ListByItem(ctx context.Context, workspaceID, itemID string) ([]*domain.Claim, 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_id", Value: itemID})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}}))
|
||||
r.observe("list_by_item", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Claim
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// FindBySubjects locates claims about the same coarse subject in other items.
|
||||
//
|
||||
// This is how support and contradiction candidates are found without comparing
|
||||
// every claim to every other claim: the subject key is indexed, so the search
|
||||
// is a bounded lookup rather than a scan.
|
||||
func (r *ClaimRepo) FindBySubjects(ctx context.Context, workspaceID string, subjects []string, excludeItem string, limit int) ([]*domain.Claim, error) {
|
||||
if len(subjects) == 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: "subject", Value: bson.D{{Key: "$in", Value: stringsToAny(subjects)}}})
|
||||
if excludeItem != "" {
|
||||
f = append(f, bson.E{Key: "item_id", Value: bson.D{{Key: "$ne", Value: excludeItem}}})
|
||||
}
|
||||
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 100
|
||||
}
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().SetLimit(int64(limit)))
|
||||
r.observe("find_by_subjects", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Claim
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// GetMany fetches claims by ID within a workspace.
|
||||
func (r *ClaimRepo) GetMany(ctx context.Context, workspaceID string, claimIDs []string) ([]*domain.Claim, error) {
|
||||
if len(claimIDs) == 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(claimIDs)}}})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f)
|
||||
r.observe("get_many", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Claim
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// DeleteByItem removes an item's claims.
|
||||
func (r *ClaimRepo) DeleteByItem(ctx context.Context, workspaceID, itemID 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: "item_id", Value: itemID})
|
||||
|
||||
_, err = r.coll.DeleteMany(ctx, f)
|
||||
r.observe("delete_by_item", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Count returns the number of claims in a workspace.
|
||||
func (r *ClaimRepo) 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, f)
|
||||
r.observe("count", started, err)
|
||||
return n, translate(err)
|
||||
}
|
||||
@@ -0,0 +1,605 @@
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,446 @@
|
||||
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"
|
||||
)
|
||||
|
||||
// ItemRepo stores source items.
|
||||
type ItemRepo struct{ base }
|
||||
|
||||
var _ domain.ItemRepository = (*ItemRepo)(nil)
|
||||
|
||||
// Create inserts a new item.
|
||||
func (r *ItemRepo) Create(ctx context.Context, it *domain.Item) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if it.WorkspaceID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
_, err := r.coll.InsertOne(ctx, it)
|
||||
r.observe("create", started, err)
|
||||
// translate maps a duplicate-key rejection to domain.ErrDuplicate. That is
|
||||
// how the partial unique index on content_hash reaches the caller when a
|
||||
// concurrent request stored identical content first.
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Get fetches one item within a workspace.
|
||||
func (r *ItemRepo) Get(ctx context.Context, workspaceID, id string) (*domain.Item, 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 it domain.Item
|
||||
err = r.coll.FindOne(ctx, f).Decode(&it)
|
||||
r.observe("get", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &it, nil
|
||||
}
|
||||
|
||||
// GetMany fetches items by ID within a workspace, preserving nothing about the
|
||||
// requested order — callers that need order re-map by ID.
|
||||
func (r *ItemRepo) GetMany(ctx context.Context, workspaceID string, ids []string) ([]*domain.Item, error) {
|
||||
if len(ids) == 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(ids)}}})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f)
|
||||
r.observe("get_many", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var out []*domain.Item
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Update replaces an item's mutable fields.
|
||||
func (r *ItemRepo) Update(ctx context.Context, it *domain.Item) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(it.WorkspaceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f = append(f, bson.E{Key: "_id", Value: it.ID})
|
||||
it.UpdatedAt = time.Now().UTC()
|
||||
|
||||
res, err := r.coll.ReplaceOne(ctx, f, it)
|
||||
r.observe("update", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetStatus updates processing status and error text.
|
||||
func (r *ItemRepo) SetStatus(ctx context.Context, workspaceID, id string, st domain.ProcessingStatus, errMsg 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: "_id", Value: id})
|
||||
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(st)},
|
||||
{Key: "error", Value: errMsg},
|
||||
{Key: "updated_at", Value: time.Now().UTC()},
|
||||
}}}
|
||||
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("set_status", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddWarning appends a parse warning without disturbing existing ones.
|
||||
//
|
||||
// $addToSet rather than $push: reprocessing an item that still has the same
|
||||
// problem must not accumulate duplicate warnings.
|
||||
func (r *ItemRepo) AddWarning(ctx context.Context, workspaceID, id string, w domain.ParseWarning) 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})
|
||||
|
||||
update := bson.D{{Key: "$addToSet", Value: bson.D{{Key: "warnings", Value: w}}}}
|
||||
_, err = r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("add_warning", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// buildItemFilter converts a typed filter into a BSON document.
|
||||
//
|
||||
// Every value here originates as a Go typed value. Client strings become
|
||||
// operands, never keys or operators.
|
||||
func buildItemFilter(workspaceID string, f domain.ItemFilter) (bson.D, error) {
|
||||
q, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !f.IncludeDeleted {
|
||||
q = notDeleted(q)
|
||||
}
|
||||
if len(f.SourceTypes) > 0 {
|
||||
vals := make([]any, 0, len(f.SourceTypes))
|
||||
for _, s := range f.SourceTypes {
|
||||
if s.Valid() {
|
||||
vals = append(vals, string(s))
|
||||
}
|
||||
}
|
||||
if len(vals) > 0 {
|
||||
q = append(q, bson.E{Key: "source_type", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
}
|
||||
if len(f.MIMETypes) > 0 {
|
||||
q = append(q, bson.E{Key: "mime_type", Value: bson.D{{Key: "$in", Value: stringsToAny(f.MIMETypes)}}})
|
||||
}
|
||||
if len(f.Statuses) > 0 {
|
||||
vals := make([]any, 0, len(f.Statuses))
|
||||
for _, s := range f.Statuses {
|
||||
vals = append(vals, string(s))
|
||||
}
|
||||
q = append(q, bson.E{Key: "status", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
if f.ContentHash != "" {
|
||||
q = append(q, bson.E{Key: "content_hash", Value: f.ContentHash})
|
||||
}
|
||||
if f.CreatedFrom != nil || f.CreatedTo != nil {
|
||||
rng := bson.D{}
|
||||
if f.CreatedFrom != nil {
|
||||
rng = append(rng, bson.E{Key: "$gte", Value: *f.CreatedFrom})
|
||||
}
|
||||
if f.CreatedTo != nil {
|
||||
rng = append(rng, bson.E{Key: "$lte", Value: *f.CreatedTo})
|
||||
}
|
||||
q = append(q, bson.E{Key: "created_at", Value: rng})
|
||||
}
|
||||
if t := strings.TrimSpace(f.TitleContains); t != "" {
|
||||
// The user's text is escaped and used as a literal substring: it can
|
||||
// never become a regular expression with its own semantics.
|
||||
q = append(q, bson.E{Key: "title", Value: bson.D{
|
||||
{Key: "$regex", Value: regexpQuote(t)},
|
||||
{Key: "$options", Value: "i"},
|
||||
}})
|
||||
}
|
||||
return q, nil
|
||||
}
|
||||
|
||||
// regexpQuote escapes every regular-expression metacharacter so that a search
|
||||
// string is matched literally. Without this, a query of ".*" would match
|
||||
// everything and "(((" would be a denial-of-service.
|
||||
func regexpQuote(s string) string {
|
||||
const special = `\.+*?()|[]{}^$`
|
||||
var b strings.Builder
|
||||
b.Grow(len(s) * 2)
|
||||
for _, r := range s {
|
||||
if r < 128 && strings.ContainsRune(special, r) {
|
||||
b.WriteByte('\\')
|
||||
}
|
||||
b.WriteRune(r)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// List returns a page of items and the total matching count.
|
||||
func (r *ItemRepo) List(ctx context.Context, workspaceID string, f domain.ItemFilter, p domain.Page) ([]*domain.Item, int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
q, err := buildItemFilter(workspaceID, f)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
p = p.Normalise(20, 200)
|
||||
|
||||
total, err := r.coll.CountDocuments(ctx, q)
|
||||
if err != nil {
|
||||
r.observe("list", started, err)
|
||||
return nil, 0, translate(err)
|
||||
}
|
||||
|
||||
opts := options.Find().
|
||||
SetSort(bson.D{{Key: "ingested_at", Value: -1}, {Key: "_id", Value: 1}}).
|
||||
SetSkip(int64(p.Offset)).
|
||||
SetLimit(int64(p.Limit))
|
||||
|
||||
cur, err := r.coll.Find(ctx, q, opts)
|
||||
r.observe("list", started, err)
|
||||
if err != nil {
|
||||
return nil, 0, translate(err)
|
||||
}
|
||||
var out []*domain.Item
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return out, total, nil
|
||||
}
|
||||
|
||||
// FindByContentHash locates an existing item with identical content.
|
||||
func (r *ItemRepo) FindByContentHash(ctx context.Context, workspaceID, hash string) (*domain.Item, 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: "content_hash", Value: hash}))
|
||||
|
||||
var it domain.Item
|
||||
err = r.coll.FindOne(ctx, f, options.FindOne().SetSort(bson.D{{Key: "ingested_at", Value: 1}})).Decode(&it)
|
||||
r.observe("find_by_hash", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &it, nil
|
||||
}
|
||||
|
||||
// TextSearch runs MongoDB's indexed lexical search within a workspace.
|
||||
//
|
||||
// The query string is passed as a value to $text, which treats it as search
|
||||
// terms rather than as a document — it cannot introduce operators.
|
||||
func (r *ItemRepo) TextSearch(ctx context.Context, workspaceID, query string, f domain.ItemFilter, limit int) ([]domain.LexicalHit, error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return nil, nil
|
||||
}
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
q, err := buildItemFilter(workspaceID, f)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
q = append(q, bson.E{Key: "$text", Value: bson.D{{Key: "$search", Value: query}}})
|
||||
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 60
|
||||
}
|
||||
opts := options.Find().
|
||||
SetProjection(bson.D{{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}}}).
|
||||
SetSort(bson.D{{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}}}).
|
||||
SetLimit(int64(limit))
|
||||
|
||||
cur, err := r.coll.Find(ctx, q, opts)
|
||||
r.observe("text_search", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
var out []domain.LexicalHit
|
||||
for cur.Next(ctx) {
|
||||
var row struct {
|
||||
ID string `bson:"_id"`
|
||||
Score float64 `bson:"score"`
|
||||
}
|
||||
if err := cur.Decode(&row); err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, domain.LexicalHit{ItemID: row.ID, Score: row.Score})
|
||||
}
|
||||
return out, cur.Err()
|
||||
}
|
||||
|
||||
// Tombstone marks an item deleted while retaining the record, so that audit
|
||||
// history and any references to it remain resolvable.
|
||||
func (r *ItemRepo) Tombstone(ctx context.Context, workspaceID, id 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: "_id", Value: id})
|
||||
|
||||
// The original content is cleared on tombstone: a deletion request should
|
||||
// remove the user's data, while the shell of the record documents that
|
||||
// something existed and was removed.
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "deleted", Value: true},
|
||||
{Key: "status", Value: string(domain.StatusTombstoned)},
|
||||
{Key: "original_text", Value: ""},
|
||||
{Key: "extracted_text", Value: ""},
|
||||
{Key: "updated_at", Value: time.Now().UTC()},
|
||||
}}}
|
||||
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("tombstone", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// HardDelete removes the item document entirely.
|
||||
func (r *ItemRepo) HardDelete(ctx context.Context, workspaceID, id 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: "_id", Value: id})
|
||||
|
||||
res, err := r.coll.DeleteOne(ctx, f)
|
||||
r.observe("hard_delete", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.DeletedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Count returns the number of live items in a workspace.
|
||||
func (r *ItemRepo) 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)
|
||||
}
|
||||
|
||||
// CountByStatus groups live items by processing status, for the dashboard.
|
||||
func (r *ItemRepo) CountByStatus(ctx context.Context, workspaceID string) (map[domain.ProcessingStatus]int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pipeline := mongo.Pipeline{
|
||||
{{Key: "$match", Value: notDeleted(f)}},
|
||||
{{Key: "$group", Value: bson.D{
|
||||
{Key: "_id", Value: "$status"},
|
||||
{Key: "n", Value: bson.D{{Key: "$sum", Value: 1}}},
|
||||
}}},
|
||||
}
|
||||
cur, err := r.coll.Aggregate(ctx, pipeline)
|
||||
r.observe("count_by_status", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
out := map[domain.ProcessingStatus]int64{}
|
||||
for cur.Next(ctx) {
|
||||
var row struct {
|
||||
ID string `bson:"_id"`
|
||||
N int64 `bson:"n"`
|
||||
}
|
||||
if err := cur.Decode(&row); err != nil {
|
||||
continue
|
||||
}
|
||||
out[domain.ProcessingStatus(row.ID)] = row.N
|
||||
}
|
||||
return out, cur.Err()
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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"
|
||||
)
|
||||
|
||||
// JobRepo is the persistent, leased job queue.
|
||||
type JobRepo struct{ base }
|
||||
|
||||
var _ domain.JobRepository = (*JobRepo)(nil)
|
||||
|
||||
// maxJobHistory bounds how many attempt records one job retains.
|
||||
const maxJobHistory = 10
|
||||
|
||||
// Enqueue adds a job to the queue.
|
||||
//
|
||||
// When the job carries an idempotency key, a unique partial index on pending
|
||||
// jobs makes a duplicate enqueue a no-op rather than a second unit of work.
|
||||
// That is what lets the pipeline re-enqueue freely after a crash.
|
||||
func (r *JobRepo) Enqueue(ctx context.Context, j *domain.Job) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if j.WorkspaceID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
if !j.Type.Valid() {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
if j.CreatedAt.IsZero() {
|
||||
j.CreatedAt = now
|
||||
}
|
||||
j.UpdatedAt = now
|
||||
if j.RunAfter.IsZero() {
|
||||
j.RunAfter = now
|
||||
}
|
||||
if j.Status == "" {
|
||||
j.Status = domain.JobPending
|
||||
}
|
||||
|
||||
_, err := r.coll.InsertOne(ctx, j)
|
||||
if isDup(err) {
|
||||
// Equivalent work is already queued.
|
||||
r.observe("enqueue", started, nil)
|
||||
return nil
|
||||
}
|
||||
r.observe("enqueue", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Claim atomically leases the next runnable job.
|
||||
//
|
||||
// A lease, not a lock: the document records who holds it and until when. A
|
||||
// worker that dies stops renewing, the lease expires, and ReleaseExpired
|
||||
// returns the job to the pool. Nothing external is needed to detect the crash.
|
||||
func (r *JobRepo) Claim(ctx context.Context, owner string, types []domain.JobType, lease time.Duration) (*domain.Job, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
now := time.Now().UTC()
|
||||
f := bson.D{
|
||||
{Key: "status", Value: string(domain.JobPending)},
|
||||
{Key: "run_after", Value: bson.D{{Key: "$lte", Value: now}}},
|
||||
}
|
||||
if len(types) > 0 {
|
||||
vals := make([]any, 0, len(types))
|
||||
for _, t := range types {
|
||||
vals = append(vals, string(t))
|
||||
}
|
||||
f = append(f, bson.E{Key: "type", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
|
||||
expires := now.Add(lease)
|
||||
update := bson.D{
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(domain.JobRunning)},
|
||||
{Key: "lease_owner", Value: owner},
|
||||
{Key: "lease_expires", Value: expires},
|
||||
{Key: "updated_at", Value: now},
|
||||
}},
|
||||
{Key: "$inc", Value: bson.D{{Key: "attempts", Value: 1}}},
|
||||
}
|
||||
|
||||
opts := options.FindOneAndUpdate().
|
||||
SetSort(bson.D{{Key: "priority", Value: -1}, {Key: "run_after", Value: 1}}).
|
||||
SetReturnDocument(options.After)
|
||||
|
||||
var j domain.Job
|
||||
err := r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&j)
|
||||
if err == mongo.ErrNoDocuments {
|
||||
r.observe("claim", started, nil)
|
||||
return nil, nil // an empty queue is not an error
|
||||
}
|
||||
r.observe("claim", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &j, nil
|
||||
}
|
||||
|
||||
// Heartbeat extends a lease. It fails if the lease was lost — for instance
|
||||
// because it expired and another worker claimed the job — which is the signal
|
||||
// for the original worker to abandon its work rather than write a stale result.
|
||||
func (r *JobRepo) Heartbeat(ctx context.Context, id, owner string, lease time.Duration) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
now := time.Now().UTC()
|
||||
f := bson.D{
|
||||
{Key: "_id", Value: id},
|
||||
{Key: "lease_owner", Value: owner},
|
||||
{Key: "status", Value: string(domain.JobRunning)},
|
||||
}
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "lease_expires", Value: now.Add(lease)},
|
||||
{Key: "updated_at", Value: now},
|
||||
}}}
|
||||
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("heartbeat", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Complete marks a job succeeded. It is scoped to the lease owner so a worker
|
||||
// that lost its lease cannot mark someone else's attempt as done.
|
||||
func (r *JobRepo) Complete(ctx context.Context, id, owner string, att domain.JobAttempt) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
now := time.Now().UTC()
|
||||
f := bson.D{{Key: "_id", Value: id}, {Key: "lease_owner", Value: owner}}
|
||||
update := bson.D{
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(domain.JobSucceeded)},
|
||||
{Key: "completed_at", Value: now},
|
||||
{Key: "updated_at", Value: now},
|
||||
{Key: "last_error", Value: ""},
|
||||
{Key: "lease_owner", Value: ""},
|
||||
}},
|
||||
{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
|
||||
{Key: "$push", Value: bson.D{{Key: "history", Value: bson.D{
|
||||
{Key: "$each", Value: bson.A{att}},
|
||||
{Key: "$slice", Value: -maxJobHistory},
|
||||
}}}},
|
||||
}
|
||||
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("complete", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fail records a failed attempt.
|
||||
//
|
||||
// When retryAfter is non-nil the job returns to pending for another attempt;
|
||||
// otherwise it is terminal. The decision belongs to the worker, which knows
|
||||
// whether the error was transient — the queue does not guess.
|
||||
func (r *JobRepo) Fail(ctx context.Context, id, owner string, att domain.JobAttempt, retryAfter *time.Time) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
now := time.Now().UTC()
|
||||
set := bson.D{
|
||||
{Key: "updated_at", Value: now},
|
||||
{Key: "last_error", Value: att.Error},
|
||||
{Key: "lease_owner", Value: ""},
|
||||
}
|
||||
if retryAfter != nil {
|
||||
set = append(set,
|
||||
bson.E{Key: "status", Value: string(domain.JobPending)},
|
||||
bson.E{Key: "run_after", Value: *retryAfter},
|
||||
)
|
||||
} else {
|
||||
set = append(set,
|
||||
bson.E{Key: "status", Value: string(domain.JobFailed)},
|
||||
bson.E{Key: "completed_at", Value: now},
|
||||
)
|
||||
}
|
||||
|
||||
f := bson.D{{Key: "_id", Value: id}, {Key: "lease_owner", Value: owner}}
|
||||
update := bson.D{
|
||||
{Key: "$set", Value: set},
|
||||
{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
|
||||
{Key: "$push", Value: bson.D{{Key: "history", Value: bson.D{
|
||||
{Key: "$each", Value: bson.A{att}},
|
||||
{Key: "$slice", Value: -maxJobHistory},
|
||||
}}}},
|
||||
}
|
||||
|
||||
res, err := r.coll.UpdateOne(ctx, f, update)
|
||||
r.observe("fail", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get fetches a job by ID.
|
||||
//
|
||||
// Jobs are infrastructure records rather than user content, so this is not
|
||||
// workspace-scoped; callers that expose a job to a user check its workspace.
|
||||
func (r *JobRepo) Get(ctx context.Context, id string) (*domain.Job, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
var j domain.Job
|
||||
err := r.coll.FindOne(ctx, bson.D{{Key: "_id", Value: id}}).Decode(&j)
|
||||
r.observe("get", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &j, nil
|
||||
}
|
||||
|
||||
// List returns a page of jobs in a workspace.
|
||||
func (r *JobRepo) List(ctx context.Context, workspaceID string, statuses []domain.JobStatus, types []domain.JobType, p domain.Page) ([]*domain.Job, int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if len(statuses) > 0 {
|
||||
vals := make([]any, 0, len(statuses))
|
||||
for _, s := range statuses {
|
||||
vals = append(vals, string(s))
|
||||
}
|
||||
f = append(f, bson.E{Key: "status", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
if len(types) > 0 {
|
||||
vals := make([]any, 0, len(types))
|
||||
for _, t := range 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}}})
|
||||
}
|
||||
}
|
||||
p = p.Normalise(50, 200)
|
||||
|
||||
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: "created_at", 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.Job
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return out, total, nil
|
||||
}
|
||||
|
||||
// Retry returns a failed job to the queue and resets its attempt budget.
|
||||
func (r *JobRepo) Retry(ctx context.Context, workspaceID, id string) (*domain.Job, 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},
|
||||
bson.E{Key: "status", Value: bson.D{{Key: "$in", Value: bson.A{
|
||||
string(domain.JobFailed), string(domain.JobCancelled),
|
||||
}}}},
|
||||
)
|
||||
|
||||
now := time.Now().UTC()
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(domain.JobPending)},
|
||||
{Key: "run_after", Value: now},
|
||||
{Key: "attempts", Value: 0},
|
||||
{Key: "lease_owner", Value: ""},
|
||||
{Key: "updated_at", Value: now},
|
||||
}}}
|
||||
|
||||
opts := options.FindOneAndUpdate().SetReturnDocument(options.After)
|
||||
var j domain.Job
|
||||
err = r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&j)
|
||||
r.observe("retry", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &j, nil
|
||||
}
|
||||
|
||||
// CountByStatus groups jobs by status, for the dashboard and metrics.
|
||||
func (r *JobRepo) CountByStatus(ctx context.Context, workspaceID string) (map[domain.JobStatus]int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cur, err := r.coll.Aggregate(ctx, mongo.Pipeline{
|
||||
{{Key: "$match", Value: f}},
|
||||
{{Key: "$group", Value: bson.D{
|
||||
{Key: "_id", Value: "$status"},
|
||||
{Key: "n", Value: bson.D{{Key: "$sum", Value: 1}}},
|
||||
}}},
|
||||
})
|
||||
r.observe("count_by_status", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
out := map[domain.JobStatus]int64{}
|
||||
for cur.Next(ctx) {
|
||||
var row struct {
|
||||
ID string `bson:"_id"`
|
||||
N int64 `bson:"n"`
|
||||
}
|
||||
if err := cur.Decode(&row); err == nil {
|
||||
out[domain.JobStatus(row.ID)] = row.N
|
||||
}
|
||||
}
|
||||
return out, cur.Err()
|
||||
}
|
||||
|
||||
// ReleaseExpired returns jobs whose lease elapsed to the pending pool.
|
||||
//
|
||||
// This is the crash-recovery path. A job that has exhausted its attempts is
|
||||
// marked failed instead, so a job that reliably kills its worker cannot loop
|
||||
// forever.
|
||||
func (r *JobRepo) ReleaseExpired(ctx context.Context, now time.Time) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f := bson.D{
|
||||
{Key: "status", Value: string(domain.JobRunning)},
|
||||
{Key: "lease_expires", Value: bson.D{{Key: "$lt", Value: now}}},
|
||||
}
|
||||
|
||||
// Jobs still within their attempt budget go back to pending.
|
||||
retryable := append(bson.D{}, f...)
|
||||
retryable = append(retryable, bson.E{Key: "$expr", Value: bson.D{
|
||||
{Key: "$lt", Value: bson.A{"$attempts", "$max_attempts"}},
|
||||
}})
|
||||
|
||||
res, err := r.coll.UpdateMany(ctx, retryable, bson.D{
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(domain.JobPending)},
|
||||
{Key: "run_after", Value: now},
|
||||
{Key: "lease_owner", Value: ""},
|
||||
{Key: "last_error", Value: "worker lease expired; the job was returned to the queue"},
|
||||
{Key: "updated_at", Value: now},
|
||||
}},
|
||||
{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
|
||||
})
|
||||
if err != nil {
|
||||
r.observe("release_expired", started, err)
|
||||
return 0, translate(err)
|
||||
}
|
||||
released := res.ModifiedCount
|
||||
|
||||
// Jobs out of attempts are terminal.
|
||||
exhausted := append(bson.D{}, f...)
|
||||
exhausted = append(exhausted, bson.E{Key: "$expr", Value: bson.D{
|
||||
{Key: "$gte", Value: bson.A{"$attempts", "$max_attempts"}},
|
||||
}})
|
||||
_, err = r.coll.UpdateMany(ctx, exhausted, bson.D{
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "status", Value: string(domain.JobFailed)},
|
||||
{Key: "lease_owner", Value: ""},
|
||||
{Key: "completed_at", Value: now},
|
||||
{Key: "last_error", Value: "worker lease expired and the job has no attempts remaining"},
|
||||
{Key: "updated_at", Value: now},
|
||||
}},
|
||||
{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
|
||||
})
|
||||
r.observe("release_expired", started, err)
|
||||
if err != nil {
|
||||
return released, translate(err)
|
||||
}
|
||||
return released, nil
|
||||
}
|
||||
|
||||
// DeleteForItem removes an item's queued work, used during deletion.
|
||||
func (r *JobRepo) DeleteForItem(ctx context.Context, workspaceID, itemID string) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
f = append(f,
|
||||
bson.E{Key: "item_id", Value: itemID},
|
||||
// A running job is left alone; it will notice the item is gone.
|
||||
bson.E{Key: "status", Value: bson.D{{Key: "$ne", Value: string(domain.JobRunning)}}},
|
||||
// The deletion job itself must survive.
|
||||
bson.E{Key: "type", Value: bson.D{{Key: "$ne", Value: string(domain.JobDeleteItemData)}}},
|
||||
)
|
||||
|
||||
res, err := r.coll.DeleteMany(ctx, f)
|
||||
r.observe("delete_for_item", started, err)
|
||||
if err != nil {
|
||||
return 0, translate(err)
|
||||
}
|
||||
return res.DeletedCount, nil
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
"github.com/andydixon/relata/internal/ids"
|
||||
)
|
||||
|
||||
// RelationshipRepo stores relationships and their review history.
|
||||
type RelationshipRepo struct{ base }
|
||||
|
||||
var _ domain.RelationshipRepository = (*RelationshipRepo)(nil)
|
||||
|
||||
// FindByDedupeKey returns the relationship with the given dedupe key.
|
||||
func (r *RelationshipRepo) FindByDedupeKey(ctx context.Context, workspaceID, dedupeKey string) (*domain.Relationship, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if dedupeKey == "" {
|
||||
return nil, fmt.Errorf("%w: dedupe key must not be empty", domain.ErrInvalidInput)
|
||||
}
|
||||
f = append(f, bson.E{Key: "dedupe_key", Value: dedupeKey})
|
||||
|
||||
out, err := r.findOne(ctx, notDeleted(f))
|
||||
r.observe("find_by_dedupe_key", started, err)
|
||||
return out, err
|
||||
}
|
||||
|
||||
// UpsertByDedupeKey inserts or refreshes a discovered relationship.
|
||||
//
|
||||
// Two properties matter here and are the reason this is not a plain upsert:
|
||||
//
|
||||
// - Reprocessing must not create duplicates. The dedupe key is derived from
|
||||
// the endpoints and type, and a unique index enforces one document per key.
|
||||
// - A human decision is never overwritten by a later analysis run. When the
|
||||
// relationship has been confirmed, rejected or edited, only the evidence is
|
||||
// refreshed; the type, provenance and review status stay as the person left
|
||||
// them, and the score keeps reflecting their decision.
|
||||
func (r *RelationshipRepo) UpsertByDedupeKey(ctx context.Context, rel *domain.Relationship) (*domain.Relationship, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(rel.WorkspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f = append(f, bson.E{Key: "dedupe_key", Value: rel.DedupeKey})
|
||||
|
||||
// If a person has already ruled on this relationship, refresh only the
|
||||
// evidence and leave their decision intact.
|
||||
existing, err := r.findOne(ctx, f)
|
||||
if err != nil && err != domain.ErrNotFound {
|
||||
return nil, err
|
||||
}
|
||||
if existing != nil && existing.ReviewStatus != domain.ReviewPending {
|
||||
update := bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "evidence", Value: rel.Evidence},
|
||||
{Key: "updated_at", Value: time.Now().UTC()},
|
||||
}}}
|
||||
if _, err := r.coll.UpdateOne(ctx, f, update); err != nil {
|
||||
r.observe("upsert", started, err)
|
||||
return nil, translate(err)
|
||||
}
|
||||
refreshed, err := r.findOne(ctx, f)
|
||||
r.observe("upsert", started, err)
|
||||
return refreshed, err
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
if rel.ID == "" {
|
||||
rel.ID = ids.New()
|
||||
}
|
||||
|
||||
update := bson.D{
|
||||
{Key: "$setOnInsert", Value: bson.D{
|
||||
{Key: "_id", Value: rel.ID},
|
||||
{Key: "workspace_id", Value: rel.WorkspaceID},
|
||||
{Key: "dedupe_key", Value: rel.DedupeKey},
|
||||
{Key: "source_kind", Value: string(rel.SourceKind)},
|
||||
{Key: "source_id", Value: rel.SourceID},
|
||||
{Key: "target_kind", Value: string(rel.TargetKind)},
|
||||
{Key: "target_id", Value: rel.TargetID},
|
||||
{Key: "original_type", Value: string(rel.Type)},
|
||||
{Key: "original_provenance", Value: string(rel.Provenance)},
|
||||
{Key: "review_status", Value: string(domain.ReviewPending)},
|
||||
{Key: "created_at", Value: now},
|
||||
{Key: "deleted", Value: false},
|
||||
}},
|
||||
{Key: "$set", Value: bson.D{
|
||||
{Key: "type", Value: string(rel.Type)},
|
||||
{Key: "provenance", Value: string(rel.Provenance)},
|
||||
{Key: "score", Value: rel.Score},
|
||||
{Key: "score_version", Value: rel.ScoreVersion},
|
||||
{Key: "evidence", Value: rel.Evidence},
|
||||
{Key: "updated_at", Value: now},
|
||||
}},
|
||||
}
|
||||
|
||||
opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
|
||||
var out domain.Relationship
|
||||
err = r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&out)
|
||||
r.observe("upsert", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
func (r *RelationshipRepo) findOne(ctx context.Context, f bson.D) (*domain.Relationship, error) {
|
||||
var rel domain.Relationship
|
||||
if err := r.coll.FindOne(ctx, f).Decode(&rel); err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &rel, nil
|
||||
}
|
||||
|
||||
// Create inserts a relationship, typically one a user added by hand.
|
||||
func (r *RelationshipRepo) Create(ctx context.Context, rel *domain.Relationship) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
if rel.WorkspaceID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
_, err := r.coll.InsertOne(ctx, rel)
|
||||
r.observe("create", started, err)
|
||||
return translate(err)
|
||||
}
|
||||
|
||||
// Get fetches one relationship within a workspace.
|
||||
func (r *RelationshipRepo) Get(ctx context.Context, workspaceID, id string) (*domain.Relationship, 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})
|
||||
|
||||
rel, err := r.findOne(ctx, f)
|
||||
r.observe("get", started, err)
|
||||
return rel, err
|
||||
}
|
||||
|
||||
// buildRelFilter converts a typed filter into BSON.
|
||||
func buildRelFilter(workspaceID string, rf domain.RelationshipFilter) (bson.D, error) {
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f = notDeleted(f)
|
||||
|
||||
if len(rf.Types) > 0 {
|
||||
vals := make([]any, 0, len(rf.Types))
|
||||
for _, t := range rf.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 len(rf.Provenances) > 0 {
|
||||
vals := make([]any, 0, len(rf.Provenances))
|
||||
for _, p := range rf.Provenances {
|
||||
if p.Valid() {
|
||||
vals = append(vals, string(p))
|
||||
}
|
||||
}
|
||||
if len(vals) > 0 {
|
||||
f = append(f, bson.E{Key: "provenance", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
}
|
||||
if rf.ConfirmedOnly {
|
||||
f = append(f, bson.E{Key: "review_status", Value: string(domain.ReviewConfirmed)})
|
||||
} else if len(rf.ReviewStatus) > 0 {
|
||||
vals := make([]any, 0, len(rf.ReviewStatus))
|
||||
for _, s := range rf.ReviewStatus {
|
||||
vals = append(vals, string(s))
|
||||
}
|
||||
f = append(f, bson.E{Key: "review_status", Value: bson.D{{Key: "$in", Value: vals}}})
|
||||
}
|
||||
if rf.MinScore > 0 {
|
||||
f = append(f, bson.E{Key: "score", Value: bson.D{{Key: "$gte", Value: rf.MinScore}}})
|
||||
}
|
||||
if rf.NodeID != "" {
|
||||
switch rf.Direction {
|
||||
case "out":
|
||||
f = append(f, bson.E{Key: "source_id", Value: rf.NodeID})
|
||||
case "in":
|
||||
f = append(f, bson.E{Key: "target_id", Value: rf.NodeID})
|
||||
default:
|
||||
f = append(f, bson.E{Key: "$or", Value: bson.A{
|
||||
bson.D{{Key: "source_id", Value: rf.NodeID}},
|
||||
bson.D{{Key: "target_id", Value: rf.NodeID}},
|
||||
}})
|
||||
}
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
// List returns a page of relationships.
|
||||
func (r *RelationshipRepo) List(ctx context.Context, workspaceID string, rf domain.RelationshipFilter, p domain.Page) ([]*domain.Relationship, int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := buildRelFilter(workspaceID, rf)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
p = p.Normalise(25, 200)
|
||||
|
||||
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: "score", Value: -1}, {Key: "_id", 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.Relationship
|
||||
if err := cur.All(ctx, &out); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return out, total, nil
|
||||
}
|
||||
|
||||
// ListForNode returns a node's outgoing and incoming relationships separately,
|
||||
// which is what the item detail page shows.
|
||||
func (r *RelationshipRepo) ListForNode(ctx context.Context, workspaceID string, kind domain.NodeKind, id string, rf domain.RelationshipFilter) (out, in []*domain.Relationship, err error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
fetch := func(direction string) ([]*domain.Relationship, error) {
|
||||
sub := rf
|
||||
sub.NodeID = id
|
||||
sub.NodeKind = kind
|
||||
sub.Direction = direction
|
||||
f, err := buildRelFilter(workspaceID, sub)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().
|
||||
SetSort(bson.D{{Key: "score", Value: -1}}).SetLimit(200))
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var rels []*domain.Relationship
|
||||
if err := cur.All(ctx, &rels); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rels, nil
|
||||
}
|
||||
|
||||
out, err = fetch("out")
|
||||
if err != nil {
|
||||
r.observe("list_for_node", started, err)
|
||||
return nil, nil, err
|
||||
}
|
||||
in, err = fetch("in")
|
||||
r.observe("list_for_node", started, err)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return out, in, nil
|
||||
}
|
||||
|
||||
// ListForItems fetches relationships touching any of the given items, used for
|
||||
// relationship expansion during search.
|
||||
func (r *RelationshipRepo) ListForItems(ctx context.Context, workspaceID string, itemIDs []string, minScore float64) ([]*domain.Relationship, error) {
|
||||
if len(itemIDs) == 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(f)
|
||||
vals := stringsToAny(itemIDs)
|
||||
f = append(f, bson.E{Key: "$or", Value: bson.A{
|
||||
bson.D{{Key: "source_id", Value: bson.D{{Key: "$in", Value: vals}}}},
|
||||
bson.D{{Key: "target_id", Value: bson.D{{Key: "$in", Value: vals}}}},
|
||||
}})
|
||||
if minScore > 0 {
|
||||
f = append(f, bson.E{Key: "score", Value: bson.D{{Key: "$gte", Value: minScore}}})
|
||||
}
|
||||
// Rejected relationships are excluded from expansion: a person has said
|
||||
// this link is wrong, so it must not pull results in.
|
||||
f = append(f, bson.E{Key: "review_status", Value: bson.D{{Key: "$ne", Value: string(domain.ReviewRejected)}}})
|
||||
|
||||
cur, err := r.coll.Find(ctx, f, options.Find().
|
||||
SetSort(bson.D{{Key: "score", Value: -1}}).SetLimit(500))
|
||||
r.observe("list_for_items", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
var rels []*domain.Relationship
|
||||
if err := cur.All(ctx, &rels); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rels, nil
|
||||
}
|
||||
|
||||
// Update replaces a relationship document.
|
||||
func (r *RelationshipRepo) Update(ctx context.Context, rel *domain.Relationship) error {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := wsFilter(rel.WorkspaceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f = append(f, bson.E{Key: "_id", Value: rel.ID})
|
||||
rel.UpdatedAt = time.Now().UTC()
|
||||
|
||||
res, err := r.coll.ReplaceOne(ctx, f, rel)
|
||||
r.observe("update", started, err)
|
||||
if err != nil {
|
||||
return translate(err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AppendReview records a human decision.
|
||||
//
|
||||
// The review history is append-only and the original proposal fields are never
|
||||
// touched, so rejecting a relationship documents the disagreement rather than
|
||||
// erasing what analysis suggested.
|
||||
func (r *RelationshipRepo) AppendReview(ctx context.Context, workspaceID, id string, ev domain.ReviewEvent, newStatus domain.ReviewStatus, newType domain.RelationshipType, newProv domain.Provenance) (*domain.Relationship, 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})
|
||||
|
||||
set := bson.D{
|
||||
{Key: "review_status", Value: string(newStatus)},
|
||||
{Key: "updated_at", Value: time.Now().UTC()},
|
||||
}
|
||||
if newType != "" {
|
||||
set = append(set, bson.E{Key: "type", Value: string(newType)})
|
||||
}
|
||||
if newProv != "" {
|
||||
set = append(set, bson.E{Key: "provenance", Value: string(newProv)})
|
||||
}
|
||||
|
||||
update := bson.D{
|
||||
{Key: "$set", Value: set},
|
||||
{Key: "$push", Value: bson.D{{Key: "review_history", Value: ev}}},
|
||||
}
|
||||
|
||||
opts := options.FindOneAndUpdate().SetReturnDocument(options.After)
|
||||
var out domain.Relationship
|
||||
err = r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&out)
|
||||
r.observe("append_review", started, err)
|
||||
if err != nil {
|
||||
return nil, translate(err)
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// TombstoneForItem marks an item's relationships deleted while retaining them.
|
||||
func (r *RelationshipRepo) TombstoneForItem(ctx context.Context, workspaceID, itemID string) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := r.itemFilter(workspaceID, itemID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
res, err := r.coll.UpdateMany(ctx, f, bson.D{{Key: "$set", Value: bson.D{
|
||||
{Key: "deleted", Value: true},
|
||||
{Key: "updated_at", Value: time.Now().UTC()},
|
||||
}}})
|
||||
r.observe("tombstone_for_item", started, err)
|
||||
if err != nil {
|
||||
return 0, translate(err)
|
||||
}
|
||||
return res.ModifiedCount, nil
|
||||
}
|
||||
|
||||
// DeleteForItem removes an item's relationships outright.
|
||||
func (r *RelationshipRepo) DeleteForItem(ctx context.Context, workspaceID, itemID string) (int64, error) {
|
||||
ctx, cancel := r.ctx(ctx)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
|
||||
f, err := r.itemFilter(workspaceID, itemID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
res, err := r.coll.DeleteMany(ctx, f)
|
||||
r.observe("delete_for_item", started, err)
|
||||
if err != nil {
|
||||
return 0, translate(err)
|
||||
}
|
||||
return res.DeletedCount, nil
|
||||
}
|
||||
|
||||
func (r *RelationshipRepo) itemFilter(workspaceID, itemID string) (bson.D, error) {
|
||||
f, err := wsFilter(workspaceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return append(f, bson.E{Key: "$or", Value: bson.A{
|
||||
bson.D{{Key: "source_id", Value: itemID}},
|
||||
bson.D{{Key: "target_id", Value: itemID}},
|
||||
}}), nil
|
||||
}
|
||||
|
||||
// Count returns the number of live relationships in a workspace.
|
||||
func (r *RelationshipRepo) 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)
|
||||
}
|
||||
|
||||
// CountPendingReview returns how many relationships await human review.
|
||||
func (r *RelationshipRepo) CountPendingReview(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
|
||||
}
|
||||
f = notDeleted(append(f, bson.E{Key: "review_status", Value: string(domain.ReviewPending)}))
|
||||
|
||||
n, err := r.coll.CountDocuments(ctx, f)
|
||||
r.observe("count_pending", started, err)
|
||||
return n, translate(err)
|
||||
}
|
||||
@@ -0,0 +1,426 @@
|
||||
// Package mongo implements Relata's repository ports against MongoDB.
|
||||
//
|
||||
// Two rules apply to every function in this package:
|
||||
//
|
||||
// 1. Workspace enforcement. Every query touching user content includes
|
||||
// workspace_id in its filter. The helper wsFilter is the only sanctioned
|
||||
// way to start a filter, and no exported method accepts a filter document
|
||||
// from a caller.
|
||||
// 2. No client-controlled operators. Filters are assembled here from typed Go
|
||||
// values. A string from a request body becomes a value in a bson.D, never a
|
||||
// key, so a payload like {"$ne": null} is matched as the literal string it
|
||||
// is rather than interpreted as an operator.
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/readpref"
|
||||
|
||||
"github.com/andydixon/relata/internal/config"
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
"github.com/andydixon/relata/internal/observability"
|
||||
)
|
||||
|
||||
// Collection names.
|
||||
const (
|
||||
CollUsers = "users"
|
||||
CollSessions = "sessions"
|
||||
CollWorkspaces = "workspaces"
|
||||
CollItems = "items"
|
||||
CollChunks = "chunks"
|
||||
CollEntities = "entities"
|
||||
CollEntityAliases = "entity_aliases"
|
||||
CollClaims = "claims"
|
||||
CollRelationships = "relationships"
|
||||
CollJobs = "jobs"
|
||||
CollRuns = "processing_runs"
|
||||
CollModelVersions = "model_versions"
|
||||
CollAudit = "audit_events"
|
||||
CollAPITokens = "api_tokens"
|
||||
)
|
||||
|
||||
// apiTokenRetention is how long a token record is kept past its expiry, so the
|
||||
// history remains readable before it is swept.
|
||||
const apiTokenRetention = 90 * 24 * time.Hour
|
||||
|
||||
// Store owns the MongoDB connection and exposes the repositories.
|
||||
type Store struct {
|
||||
client *mongo.Client
|
||||
db *mongo.Database
|
||||
timeout time.Duration
|
||||
metrics *observability.Metrics
|
||||
|
||||
Items *ItemRepo
|
||||
Chunks *ChunkRepo
|
||||
Entities *EntityRepo
|
||||
Aliases *AliasRepo
|
||||
Claims *ClaimRepo
|
||||
Relationships *RelationshipRepo
|
||||
Jobs *JobRepo
|
||||
Users *UserRepo
|
||||
APITokens *APITokenRepo
|
||||
Sessions *SessionRepo
|
||||
Workspaces *WorkspaceRepo
|
||||
Audit *AuditRepo
|
||||
Runs *RunRepo
|
||||
ModelVersions *ModelVersionRepo
|
||||
}
|
||||
|
||||
// Connect dials MongoDB and verifies the connection.
|
||||
func Connect(ctx context.Context, cfg config.MongoConfig, m *observability.Metrics) (*Store, error) {
|
||||
opts := options.Client().
|
||||
ApplyURI(cfg.URI).
|
||||
SetConnectTimeout(cfg.Timeout).
|
||||
SetServerSelectionTimeout(cfg.Timeout).
|
||||
SetTimeout(cfg.Timeout).
|
||||
SetAppName("relata")
|
||||
|
||||
client, err := mongo.Connect(opts)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("connect to mongodb: %w", err)
|
||||
}
|
||||
|
||||
pingCtx, cancel := context.WithTimeout(ctx, cfg.Timeout)
|
||||
defer cancel()
|
||||
if err := client.Ping(pingCtx, readpref.Primary()); err != nil {
|
||||
_ = client.Disconnect(context.Background())
|
||||
return nil, fmt.Errorf("ping mongodb at %s: %w", redactURI(cfg.URI), err)
|
||||
}
|
||||
|
||||
db := client.Database(cfg.Database)
|
||||
s := &Store{client: client, db: db, timeout: cfg.Timeout, metrics: m}
|
||||
|
||||
s.Items = &ItemRepo{base: s.base(CollItems)}
|
||||
s.Chunks = &ChunkRepo{base: s.base(CollChunks)}
|
||||
s.Entities = &EntityRepo{base: s.base(CollEntities)}
|
||||
s.Aliases = &AliasRepo{base: s.base(CollEntityAliases)}
|
||||
s.Claims = &ClaimRepo{base: s.base(CollClaims)}
|
||||
s.Relationships = &RelationshipRepo{base: s.base(CollRelationships)}
|
||||
s.Jobs = &JobRepo{base: s.base(CollJobs)}
|
||||
s.Users = &UserRepo{base: s.base(CollUsers)}
|
||||
s.APITokens = &APITokenRepo{base: s.base(CollAPITokens)}
|
||||
s.Sessions = &SessionRepo{base: s.base(CollSessions)}
|
||||
s.Workspaces = &WorkspaceRepo{base: s.base(CollWorkspaces)}
|
||||
s.Audit = &AuditRepo{base: s.base(CollAudit)}
|
||||
s.Runs = &RunRepo{base: s.base(CollRuns)}
|
||||
s.ModelVersions = &ModelVersionRepo{base: s.base(CollModelVersions)}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Close disconnects from MongoDB.
|
||||
func (s *Store) Close(ctx context.Context) error { return s.client.Disconnect(ctx) }
|
||||
|
||||
// DB exposes the database handle, used by migrations and tests.
|
||||
func (s *Store) DB() *mongo.Database { return s.db }
|
||||
|
||||
// Ping checks connectivity for the readiness probe.
|
||||
func (s *Store) Ping(ctx context.Context) error {
|
||||
return s.client.Ping(ctx, readpref.Primary())
|
||||
}
|
||||
|
||||
// Drop removes the database. Used only by integration tests.
|
||||
func (s *Store) Drop(ctx context.Context) error { return s.db.Drop(ctx) }
|
||||
|
||||
func (s *Store) base(name string) base {
|
||||
return base{coll: s.db.Collection(name), name: name, metrics: s.metrics, timeout: s.timeout}
|
||||
}
|
||||
|
||||
// base carries what every repository needs: its collection, its name for
|
||||
// metrics, and the default operation timeout.
|
||||
type base struct {
|
||||
coll *mongo.Collection
|
||||
name string
|
||||
metrics *observability.Metrics
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
// ctx applies the storage timeout unless the caller's context is tighter.
|
||||
// Every MongoDB call is bounded, so a stalled database cannot pin a worker.
|
||||
func (b base) ctx(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
if b.timeout <= 0 {
|
||||
return context.WithCancel(parent)
|
||||
}
|
||||
if dl, ok := parent.Deadline(); ok && time.Until(dl) < b.timeout {
|
||||
return context.WithCancel(parent)
|
||||
}
|
||||
return context.WithTimeout(parent, b.timeout)
|
||||
}
|
||||
|
||||
// observe records the outcome of a storage operation.
|
||||
func (b base) observe(op string, started time.Time, err error) {
|
||||
if b.metrics == nil {
|
||||
return
|
||||
}
|
||||
outcome := "success"
|
||||
if err != nil && !errors.Is(err, domain.ErrNotFound) {
|
||||
outcome = "error"
|
||||
}
|
||||
b.metrics.StorageOps.WithLabelValues("mongo", b.name+"."+op, outcome).Inc()
|
||||
b.metrics.StorageDuration.WithLabelValues("mongo", b.name+"."+op).Observe(time.Since(started).Seconds())
|
||||
}
|
||||
|
||||
// wsFilter starts a filter scoped to a workspace.
|
||||
//
|
||||
// This is the single entry point for building a query over user content. It
|
||||
// rejects an empty workspace so that a missing principal can never widen a
|
||||
// query to every workspace.
|
||||
func wsFilter(workspaceID string) (bson.D, error) {
|
||||
if workspaceID == "" {
|
||||
return nil, fmt.Errorf("%w: workspace is required for this query", domain.ErrWorkspaceMiss)
|
||||
}
|
||||
return bson.D{{Key: "workspace_id", Value: workspaceID}}, nil
|
||||
}
|
||||
|
||||
// notDeleted appends the tombstone exclusion used by most reads.
|
||||
func notDeleted(f bson.D) bson.D {
|
||||
return append(f, bson.E{Key: "deleted", Value: bson.D{{Key: "$ne", Value: true}}})
|
||||
}
|
||||
|
||||
// isDup reports whether err is a duplicate-key error.
|
||||
func isDup(err error) bool {
|
||||
var we mongo.WriteException
|
||||
if errors.As(err, &we) {
|
||||
for _, e := range we.WriteErrors {
|
||||
if e.Code == 11000 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
var be mongo.BulkWriteException
|
||||
if errors.As(err, &be) {
|
||||
for _, e := range be.WriteErrors {
|
||||
if e.Code == 11000 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// translate converts driver errors into domain errors.
|
||||
func translate(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return domain.ErrNotFound
|
||||
}
|
||||
if isDup(err) {
|
||||
return domain.ErrDuplicate
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// redactURI removes credentials from a connection string before logging it.
|
||||
func redactURI(uri string) string {
|
||||
at := -1
|
||||
for i := len(uri) - 1; i >= 0; i-- {
|
||||
if uri[i] == '@' {
|
||||
at = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if at < 0 {
|
||||
return uri
|
||||
}
|
||||
scheme := "mongodb://"
|
||||
if len(uri) > 14 && uri[:14] == "mongodb+srv://" {
|
||||
scheme = "mongodb+srv://"
|
||||
}
|
||||
return scheme + "[redacted]" + uri[at:]
|
||||
}
|
||||
|
||||
// stringsToAny converts a typed slice for use as an $in value. Values only ever
|
||||
// appear on the right-hand side of a comparison, never as document keys.
|
||||
func stringsToAny(in []string) []any {
|
||||
out := make([]any, 0, len(in))
|
||||
for _, s := range in {
|
||||
if s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// EnsureIndexes creates every index Relata relies on.
|
||||
//
|
||||
// Index creation is idempotent, so this runs on every startup and after any
|
||||
// schema change. Uniqueness constraints here are load-bearing: they are what
|
||||
// make ingestion and reprocessing idempotent under concurrency rather than
|
||||
// merely usually-idempotent.
|
||||
func (s *Store) EnsureIndexes(ctx context.Context) error {
|
||||
type spec struct {
|
||||
coll string
|
||||
models []mongo.IndexModel
|
||||
}
|
||||
|
||||
specs := []spec{
|
||||
{CollWorkspaces, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "slug", Value: 1}}, Options: options.Index().SetUnique(true).SetName("workspace_slug_unique")},
|
||||
}},
|
||||
{CollUsers, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "username", Value: 1}}, Options: options.Index().SetUnique(true).SetName("username_unique")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}}, Options: options.Index().SetName("user_workspace")},
|
||||
// One account per federated identity. This is what makes first
|
||||
// sign-in idempotent: two concurrent callbacks race here, and the
|
||||
// loser re-reads the winner rather than provisioning a second
|
||||
// workspace for the same person. Partial, so local accounts —
|
||||
// which have no subject — are not all collapsed onto one another.
|
||||
{
|
||||
Keys: bson.D{{Key: "oauth_provider", Value: 1}, {Key: "oauth_subject", Value: 1}},
|
||||
Options: options.Index().SetUnique(true).SetName("user_oauth_subject_unique").
|
||||
SetPartialFilterExpression(bson.D{
|
||||
{Key: "oauth_subject", Value: bson.D{{Key: "$type", Value: "string"}}},
|
||||
}),
|
||||
},
|
||||
}},
|
||||
{CollAPITokens, []mongo.IndexModel{
|
||||
// Authentication is a lookup on this index and nothing else.
|
||||
{Keys: bson.D{{Key: "token_hash", Value: 1}}, Options: options.Index().SetUnique(true).SetName("api_token_hash_unique")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "user_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("api_token_ws_user")},
|
||||
// Retention, not enforcement. A token past its expiry is refused on
|
||||
// the next request regardless of this index; the index removes the
|
||||
// record ninety days *after* that expiry, so a token's existence
|
||||
// and its end stay visible for a while and then stop accumulating.
|
||||
{Keys: bson.D{{Key: "expires_at", Value: 1}}, Options: options.Index().SetExpireAfterSeconds(int32(apiTokenRetention.Seconds())).SetName("api_token_ttl")},
|
||||
}},
|
||||
{CollSessions, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "token_hash", Value: 1}}, Options: options.Index().SetUnique(true).SetName("session_token_unique")},
|
||||
{Keys: bson.D{{Key: "user_id", Value: 1}}, Options: options.Index().SetName("session_user")},
|
||||
// A TTL index removes expired sessions even if nothing sweeps them.
|
||||
{Keys: bson.D{{Key: "absolute_expiry", Value: 1}}, Options: options.Index().SetExpireAfterSeconds(0).SetName("session_ttl")},
|
||||
}},
|
||||
{CollItems, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "ingested_at", Value: -1}}, Options: options.Index().SetName("item_ws_ingested")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("item_ws_created")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "content_hash", Value: 1}}, Options: options.Index().SetName("item_ws_hash")},
|
||||
// Duplicate detection is a read followed by a write, so concurrent
|
||||
// requests carrying identical content can both pass the read. This index
|
||||
// makes the database the authority: at most one live, non-deliberate item
|
||||
// per content hash. Deliberate duplicates carry a non-empty duplicate_of
|
||||
// and tombstoned items carry deleted=true, so both fall outside it.
|
||||
{
|
||||
Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "content_hash", Value: 1}},
|
||||
Options: options.Index().SetUnique(true).SetName("item_ws_hash_unique").
|
||||
SetPartialFilterExpression(bson.D{
|
||||
{Key: "duplicate_of", Value: ""},
|
||||
{Key: "deleted", Value: false},
|
||||
}),
|
||||
},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "status", Value: 1}}, Options: options.Index().SetName("item_ws_status")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "source_type", Value: 1}}, Options: options.Index().SetName("item_ws_source")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "mime_type", Value: 1}}, Options: options.Index().SetName("item_ws_mime")},
|
||||
{
|
||||
Keys: bson.D{{Key: "title", Value: "text"}, {Key: "extracted_text", Value: "text"}},
|
||||
Options: options.Index().SetName("item_text").SetWeights(bson.D{{Key: "title", Value: 5}, {Key: "extracted_text", Value: 1}}),
|
||||
},
|
||||
}},
|
||||
{CollChunks, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "item_id", Value: 1}, {Key: "sequence", Value: 1}}, Options: options.Index().SetName("chunk_ws_item_seq")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "embedded", Value: 1}, {Key: "embedding_model", Value: 1}}, Options: options.Index().SetName("chunk_ws_embedded")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "content_hash", Value: 1}}, Options: options.Index().SetName("chunk_ws_hash")},
|
||||
{Keys: bson.D{{Key: "text", Value: "text"}}, Options: options.Index().SetName("chunk_text")},
|
||||
}},
|
||||
{CollEntities, []mongo.IndexModel{
|
||||
// Identity: one entity per (workspace, type, normalised name).
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "normalised_key", Value: 1}}, Options: options.Index().SetUnique(true).SetName("entity_ws_key_unique")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "item_ids", Value: 1}}, Options: options.Index().SetName("entity_ws_items")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "type", Value: 1}}, Options: options.Index().SetName("entity_ws_type")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "mentions", Value: -1}}, Options: options.Index().SetName("entity_ws_mentions")},
|
||||
}},
|
||||
{CollEntityAliases, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "entity_id", Value: 1}, {Key: "normalised", Value: 1}}, Options: options.Index().SetUnique(true).SetName("alias_unique")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "status", Value: 1}}, Options: options.Index().SetName("alias_ws_status")},
|
||||
}},
|
||||
{CollClaims, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "item_id", Value: 1}}, Options: options.Index().SetName("claim_ws_item")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "subject", Value: 1}}, Options: options.Index().SetName("claim_ws_subject")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "dedupe_key", Value: 1}}, Options: options.Index().SetUnique(true).SetName("claim_dedupe_unique")},
|
||||
}},
|
||||
{CollRelationships, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "dedupe_key", Value: 1}}, Options: options.Index().SetUnique(true).SetName("rel_dedupe_unique")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "source_id", Value: 1}, {Key: "score", Value: -1}}, Options: options.Index().SetName("rel_ws_source")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "target_id", Value: 1}, {Key: "score", Value: -1}}, Options: options.Index().SetName("rel_ws_target")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "type", Value: 1}}, Options: options.Index().SetName("rel_ws_type")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "review_status", Value: 1}, {Key: "score", Value: -1}}, Options: options.Index().SetName("rel_ws_review")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "provenance", Value: 1}}, Options: options.Index().SetName("rel_ws_prov")},
|
||||
}},
|
||||
{CollJobs, []mongo.IndexModel{
|
||||
// The claim query: pending jobs whose lease has expired, by priority.
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "run_after", Value: 1}, {Key: "priority", Value: -1}}, Options: options.Index().SetName("job_claim")},
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "lease_expires", Value: 1}}, Options: options.Index().SetName("job_lease")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "status", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("job_ws_status")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "item_id", Value: 1}}, Options: options.Index().SetName("job_ws_item")},
|
||||
// Partial unique index: at most one pending job per idempotency key,
|
||||
// so a retried enqueue does not duplicate work, while completed jobs
|
||||
// with the same key remain as history.
|
||||
{
|
||||
Keys: bson.D{{Key: "idempotency_key", Value: 1}},
|
||||
Options: options.Index().SetName("job_idempotency_pending").SetUnique(true).
|
||||
SetPartialFilterExpression(bson.D{{Key: "status", Value: "pending"}}),
|
||||
},
|
||||
}},
|
||||
{CollRuns, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "item_id", Value: 1}, {Key: "started_at", Value: -1}}, Options: options.Index().SetName("run_ws_item")},
|
||||
}},
|
||||
{CollModelVersions, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "kind", Value: 1}, {Key: "name", Value: 1}}, Options: options.Index().SetUnique(true).SetName("model_version_unique")},
|
||||
}},
|
||||
{CollAudit, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "at", Value: -1}}, Options: options.Index().SetName("audit_ws_at")},
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "action", Value: 1}, {Key: "at", Value: -1}}, Options: options.Index().SetName("audit_ws_action")},
|
||||
{Keys: bson.D{{Key: "target_id", Value: 1}}, Options: options.Index().SetName("audit_target")},
|
||||
}},
|
||||
}
|
||||
|
||||
if err := s.backfillDuplicateOf(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, sp := range specs {
|
||||
if _, err := s.db.Collection(sp.coll).Indexes().CreateMany(ctx, sp.models); err != nil {
|
||||
if sp.coll == CollItems && isDup(err) {
|
||||
// The unique content-hash index cannot be built because the data
|
||||
// already violates it — duplicates admitted before this index
|
||||
// existed. Say so plainly and name the fix; a bare index error
|
||||
// here would look like a configuration fault.
|
||||
return fmt.Errorf("cannot create the unique content-hash index on %s: the "+
|
||||
"collection already contains more than one live item with the same "+
|
||||
"content hash, admitted before this constraint existed. Remove or mark "+
|
||||
"the extras (set duplicate_of to the item being kept) and restart. "+
|
||||
"See migrations/README.md: %w", sp.coll, err)
|
||||
}
|
||||
return fmt.Errorf("create indexes on %s: %w", sp.coll, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// backfillDuplicateOf gives older items the empty duplicate_of the partial
|
||||
// unique index keys on. Items written before that field was stored
|
||||
// unconditionally omit it, and an absent field does not match the index's
|
||||
// filter — so without this they would sit outside the constraint and duplicates
|
||||
// could still slip past. Idempotent, and a no-op once applied.
|
||||
func (s *Store) backfillDuplicateOf(ctx context.Context) error {
|
||||
res, err := s.db.Collection(CollItems).UpdateMany(ctx,
|
||||
bson.D{{Key: "duplicate_of", Value: bson.D{{Key: "$exists", Value: false}}}},
|
||||
bson.D{{Key: "$set", Value: bson.D{{Key: "duplicate_of", Value: ""}}}},
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("backfill duplicate_of on %s: %w", CollItems, err)
|
||||
}
|
||||
if res.ModifiedCount > 0 {
|
||||
slog.Info("migrated items for the unique content-hash index",
|
||||
"collection", CollItems, "items", res.ModifiedCount)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,492 @@
|
||||
// Package qdrant implements the vector store port against Qdrant's REST API.
|
||||
//
|
||||
// Payloads deliberately hold only what retrieval needs — workspace, item,
|
||||
// chunk, source type, creation date and embedding model. Document text and
|
||||
// relationship data stay in MongoDB, which remains the authoritative store.
|
||||
// Keeping the payload thin is what stops Qdrant drifting into a second,
|
||||
// unversioned copy of the corpus.
|
||||
package qdrant
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/andydixon/relata/internal/config"
|
||||
"github.com/andydixon/relata/internal/domain"
|
||||
"github.com/andydixon/relata/internal/observability"
|
||||
)
|
||||
|
||||
// Client talks to Qdrant over REST.
|
||||
type Client struct {
|
||||
base string
|
||||
collection string
|
||||
apiKey string
|
||||
http *http.Client
|
||||
metrics *observability.Metrics
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
var _ domain.VectorStore = (*Client)(nil)
|
||||
|
||||
// New builds a Qdrant client.
|
||||
func New(cfg config.QdrantConfig, m *observability.Metrics) *Client {
|
||||
timeout := cfg.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 15 * time.Second
|
||||
}
|
||||
return &Client{
|
||||
base: strings.TrimRight(cfg.URL, "/"),
|
||||
collection: cfg.Collection,
|
||||
apiKey: cfg.APIKey,
|
||||
timeout: timeout,
|
||||
metrics: m,
|
||||
http: &http.Client{
|
||||
Timeout: timeout,
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Collection returns the configured collection name.
|
||||
func (c *Client) Collection() string { return c.collection }
|
||||
|
||||
// do performs a bounded JSON request against Qdrant.
|
||||
func (c *Client) do(ctx context.Context, method, path string, body, out any) error {
|
||||
started := time.Now()
|
||||
op := method + " " + strings.SplitN(strings.TrimPrefix(path, "/"), "/", 2)[0]
|
||||
|
||||
var reader io.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode qdrant request: %w", err)
|
||||
}
|
||||
reader = bytes.NewReader(b)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.apiKey != "" {
|
||||
req.Header.Set("api-key", c.apiKey)
|
||||
}
|
||||
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
c.observe(op, started, err)
|
||||
return fmt.Errorf("qdrant %s %s: %w", method, path, err)
|
||||
}
|
||||
defer func() {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
|
||||
_ = resp.Body.Close()
|
||||
}()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
err := fmt.Errorf("qdrant %s %s returned HTTP %d: %s",
|
||||
method, path, resp.StatusCode, strings.TrimSpace(string(msg)))
|
||||
c.observe(op, started, err)
|
||||
return err
|
||||
}
|
||||
|
||||
if out != nil {
|
||||
// Bound the response so a huge reply cannot exhaust memory.
|
||||
if err := json.NewDecoder(io.LimitReader(resp.Body, 256<<20)).Decode(out); err != nil {
|
||||
c.observe(op, started, err)
|
||||
return fmt.Errorf("decode qdrant response: %w", err)
|
||||
}
|
||||
}
|
||||
c.observe(op, started, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) observe(op string, started time.Time, err error) {
|
||||
if c.metrics == nil {
|
||||
return
|
||||
}
|
||||
outcome := "success"
|
||||
if err != nil {
|
||||
outcome = "error"
|
||||
}
|
||||
c.metrics.StorageOps.WithLabelValues("qdrant", op, outcome).Inc()
|
||||
c.metrics.StorageDuration.WithLabelValues("qdrant", op).Observe(time.Since(started).Seconds())
|
||||
}
|
||||
|
||||
// Payload field names. These are also the names indexed for filtering.
|
||||
const (
|
||||
fWorkspace = "workspace_id"
|
||||
fItem = "item_id"
|
||||
fChunk = "chunk_id"
|
||||
fSource = "source_type"
|
||||
fCreated = "created_at"
|
||||
fModel = "embedding_model"
|
||||
)
|
||||
|
||||
type collectionInfo struct {
|
||||
Result struct {
|
||||
Config struct {
|
||||
Params struct {
|
||||
Vectors struct {
|
||||
Size int `json:"size"`
|
||||
Distance string `json:"distance"`
|
||||
} `json:"vectors"`
|
||||
} `json:"params"`
|
||||
} `json:"config"`
|
||||
PointsCount int64 `json:"points_count"`
|
||||
} `json:"result"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// EnsureCollection creates the collection if absent and verifies its geometry.
|
||||
//
|
||||
// A dimension mismatch is fatal rather than silently tolerated: writing 768-wide
|
||||
// vectors into a 1024-wide collection would either error per-write or, worse,
|
||||
// produce a collection whose contents cannot all be compared.
|
||||
func (c *Client) EnsureCollection(ctx context.Context, dim int) error {
|
||||
if dim <= 0 {
|
||||
return fmt.Errorf("vector dimension must be positive, got %d", dim)
|
||||
}
|
||||
|
||||
var info collectionInfo
|
||||
err := c.do(ctx, http.MethodGet, "/collections/"+c.collection, nil, &info)
|
||||
if err == nil {
|
||||
existing := info.Result.Config.Params.Vectors.Size
|
||||
if existing != 0 && existing != dim {
|
||||
return fmt.Errorf(
|
||||
"qdrant collection %q stores %d-dimensional vectors but the embedding model produces %d; "+
|
||||
"either restore the previous embedding model or drop the collection and re-embed the workspace",
|
||||
c.collection, existing, dim)
|
||||
}
|
||||
return c.ensurePayloadIndexes(ctx)
|
||||
}
|
||||
|
||||
// Create it. Cosine distance matches the normalised embeddings Relata uses.
|
||||
body := map[string]any{
|
||||
"vectors": map[string]any{
|
||||
"size": dim,
|
||||
"distance": "Cosine",
|
||||
},
|
||||
}
|
||||
if err := c.do(ctx, http.MethodPut, "/collections/"+c.collection, body, nil); err != nil {
|
||||
return fmt.Errorf("create qdrant collection %q: %w", c.collection, err)
|
||||
}
|
||||
return c.ensurePayloadIndexes(ctx)
|
||||
}
|
||||
|
||||
// ensurePayloadIndexes creates the payload indexes the workspace filter needs.
|
||||
//
|
||||
// Without an index on workspace_id every filtered search degrades to a full
|
||||
// scan, so this is a correctness-adjacent performance requirement rather than
|
||||
// an optimisation.
|
||||
func (c *Client) ensurePayloadIndexes(ctx context.Context) error {
|
||||
fields := map[string]string{
|
||||
fWorkspace: "keyword",
|
||||
fItem: "keyword",
|
||||
fChunk: "keyword",
|
||||
fSource: "keyword",
|
||||
fModel: "keyword",
|
||||
fCreated: "integer",
|
||||
}
|
||||
for field, schema := range fields {
|
||||
body := map[string]any{"field_name": field, "field_schema": schema}
|
||||
// Qdrant returns an error for an index that already exists; that is
|
||||
// expected on every start after the first.
|
||||
if err := c.do(ctx, http.MethodPut,
|
||||
"/collections/"+c.collection+"/index?wait=true", body, nil); err != nil {
|
||||
if !strings.Contains(err.Error(), "already exists") {
|
||||
observability.LoggerFrom(ctx).Debug("qdrant payload index not created",
|
||||
"field", field, "error", err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type upsertPoint struct {
|
||||
ID string `json:"id"`
|
||||
Vector []float32 `json:"vector"`
|
||||
Payload map[string]any `json:"payload"`
|
||||
}
|
||||
|
||||
// maxUpsertBatch bounds one write request.
|
||||
const maxUpsertBatch = 128
|
||||
|
||||
// Upsert writes vectors, replacing any point with the same chunk ID.
|
||||
//
|
||||
// Chunk IDs are UUIDs and are used directly as point IDs, so re-embedding a
|
||||
// chunk overwrites its vector rather than adding a second one. That is what
|
||||
// keeps reprocessing free of orphaned vectors.
|
||||
func (c *Client) Upsert(ctx context.Context, points []domain.VectorPoint) error {
|
||||
if len(points) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
for start := 0; start < len(points); start += maxUpsertBatch {
|
||||
end := start + maxUpsertBatch
|
||||
if end > len(points) {
|
||||
end = len(points)
|
||||
}
|
||||
|
||||
batch := make([]upsertPoint, 0, end-start)
|
||||
for _, p := range points[start:end] {
|
||||
if p.WorkspaceID == "" || p.ChunkID == "" || p.ItemID == "" {
|
||||
return fmt.Errorf("vector point is missing workspace, item or chunk identity")
|
||||
}
|
||||
if len(p.Vector) == 0 {
|
||||
return fmt.Errorf("vector point %s has an empty vector", p.ChunkID)
|
||||
}
|
||||
batch = append(batch, upsertPoint{
|
||||
ID: p.ChunkID,
|
||||
Vector: p.Vector,
|
||||
Payload: map[string]any{
|
||||
fWorkspace: p.WorkspaceID,
|
||||
fItem: p.ItemID,
|
||||
fChunk: p.ChunkID,
|
||||
fSource: string(p.SourceType),
|
||||
fCreated: p.CreatedAt.Unix(),
|
||||
fModel: p.EmbeddingModel,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
body := map[string]any{"points": batch}
|
||||
if err := c.do(ctx, http.MethodPut,
|
||||
"/collections/"+c.collection+"/points?wait=true", body, nil); err != nil {
|
||||
return fmt.Errorf("upsert %d vectors: %w", len(batch), err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type searchRequest struct {
|
||||
Vector []float32 `json:"vector"`
|
||||
Limit int `json:"limit"`
|
||||
WithPayload bool `json:"with_payload"`
|
||||
Filter map[string]any `json:"filter,omitempty"`
|
||||
ScoreThresh *float64 `json:"score_threshold,omitempty"`
|
||||
}
|
||||
|
||||
type searchResponse struct {
|
||||
Result []struct {
|
||||
ID string `json:"id"`
|
||||
Score float64 `json:"score"`
|
||||
Payload map[string]any `json:"payload"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
// buildFilter converts a typed filter into Qdrant's filter document.
|
||||
//
|
||||
// The workspace clause is mandatory and is always added first: a search without
|
||||
// it would cross the workspace boundary, so the function refuses to build one.
|
||||
func buildFilter(f domain.VectorFilter) (map[string]any, error) {
|
||||
if f.WorkspaceID == "" {
|
||||
return nil, fmt.Errorf("%w: vector search requires a workspace", domain.ErrWorkspaceMiss)
|
||||
}
|
||||
|
||||
must := []any{
|
||||
map[string]any{"key": fWorkspace, "match": map[string]any{"value": f.WorkspaceID}},
|
||||
}
|
||||
var mustNot []any
|
||||
|
||||
if len(f.ItemIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"key": fItem, "match": map[string]any{"any": f.ItemIDs},
|
||||
})
|
||||
}
|
||||
if len(f.ExcludeItemIDs) > 0 {
|
||||
mustNot = append(mustNot, map[string]any{
|
||||
"key": fItem, "match": map[string]any{"any": f.ExcludeItemIDs},
|
||||
})
|
||||
}
|
||||
if len(f.SourceTypes) > 0 {
|
||||
vals := make([]string, 0, len(f.SourceTypes))
|
||||
for _, s := range f.SourceTypes {
|
||||
vals = append(vals, string(s))
|
||||
}
|
||||
must = append(must, map[string]any{
|
||||
"key": fSource, "match": map[string]any{"any": vals},
|
||||
})
|
||||
}
|
||||
if f.EmbeddingModel != "" {
|
||||
must = append(must, map[string]any{
|
||||
"key": fModel, "match": map[string]any{"value": f.EmbeddingModel},
|
||||
})
|
||||
}
|
||||
if f.CreatedFrom != nil || f.CreatedTo != nil {
|
||||
rng := map[string]any{}
|
||||
if f.CreatedFrom != nil {
|
||||
rng["gte"] = f.CreatedFrom.Unix()
|
||||
}
|
||||
if f.CreatedTo != nil {
|
||||
rng["lte"] = f.CreatedTo.Unix()
|
||||
}
|
||||
must = append(must, map[string]any{"key": fCreated, "range": rng})
|
||||
}
|
||||
|
||||
out := map[string]any{"must": must}
|
||||
if len(mustNot) > 0 {
|
||||
out["must_not"] = mustNot
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Search finds the nearest vectors within a workspace.
|
||||
func (c *Client) Search(ctx context.Context, vector []float32, limit int, f domain.VectorFilter) ([]domain.VectorMatch, error) {
|
||||
if len(vector) == 0 {
|
||||
return nil, fmt.Errorf("search vector is empty")
|
||||
}
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
filter, err := buildFilter(f)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var resp searchResponse
|
||||
req := searchRequest{Vector: vector, Limit: limit, WithPayload: true, Filter: filter}
|
||||
if err := c.do(ctx, http.MethodPost,
|
||||
"/collections/"+c.collection+"/points/search", req, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := make([]domain.VectorMatch, 0, len(resp.Result))
|
||||
for _, r := range resp.Result {
|
||||
ws := payloadString(r.Payload, fWorkspace)
|
||||
// Defence in depth: even though the filter scoped the query, a result
|
||||
// from another workspace is dropped rather than returned.
|
||||
if ws != f.WorkspaceID {
|
||||
continue
|
||||
}
|
||||
m := domain.VectorMatch{
|
||||
ChunkID: payloadString(r.Payload, fChunk),
|
||||
ItemID: payloadString(r.Payload, fItem),
|
||||
WorkspaceID: ws,
|
||||
Score: r.Score,
|
||||
SourceType: domain.SourceType(payloadString(r.Payload, fSource)),
|
||||
EmbeddingModel: payloadString(r.Payload, fModel),
|
||||
}
|
||||
if m.ChunkID == "" {
|
||||
m.ChunkID = r.ID
|
||||
}
|
||||
if ts, ok := payloadInt(r.Payload, fCreated); ok {
|
||||
m.CreatedAt = time.Unix(ts, 0).UTC()
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// DeleteByItem removes every vector belonging to an item.
|
||||
//
|
||||
// Deletion is by filter rather than by enumerated IDs, so it removes vectors
|
||||
// even for chunks MongoDB no longer knows about — which is exactly the orphan
|
||||
// case that would otherwise accumulate.
|
||||
func (c *Client) DeleteByItem(ctx context.Context, workspaceID, itemID string) error {
|
||||
if workspaceID == "" || itemID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
body := map[string]any{
|
||||
"filter": map[string]any{
|
||||
"must": []any{
|
||||
map[string]any{"key": fWorkspace, "match": map[string]any{"value": workspaceID}},
|
||||
map[string]any{"key": fItem, "match": map[string]any{"value": itemID}},
|
||||
},
|
||||
},
|
||||
}
|
||||
return c.do(ctx, http.MethodPost,
|
||||
"/collections/"+c.collection+"/points/delete?wait=true", body, nil)
|
||||
}
|
||||
|
||||
// DeleteByChunks removes specific vectors, scoped to a workspace.
|
||||
func (c *Client) DeleteByChunks(ctx context.Context, workspaceID string, chunkIDs []string) error {
|
||||
if workspaceID == "" {
|
||||
return domain.ErrWorkspaceMiss
|
||||
}
|
||||
if len(chunkIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
body := map[string]any{
|
||||
"filter": map[string]any{
|
||||
"must": []any{
|
||||
map[string]any{"key": fWorkspace, "match": map[string]any{"value": workspaceID}},
|
||||
map[string]any{"key": fChunk, "match": map[string]any{"any": chunkIDs}},
|
||||
},
|
||||
},
|
||||
}
|
||||
return c.do(ctx, http.MethodPost,
|
||||
"/collections/"+c.collection+"/points/delete?wait=true", body, nil)
|
||||
}
|
||||
|
||||
type countResponse struct {
|
||||
Result struct {
|
||||
Count int64 `json:"count"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
// CountByWorkspace reports how many vectors a workspace holds, used to
|
||||
// reconcile the vector store against MongoDB.
|
||||
func (c *Client) CountByWorkspace(ctx context.Context, workspaceID string) (int64, error) {
|
||||
if workspaceID == "" {
|
||||
return 0, domain.ErrWorkspaceMiss
|
||||
}
|
||||
body := map[string]any{
|
||||
"exact": true,
|
||||
"filter": map[string]any{
|
||||
"must": []any{
|
||||
map[string]any{"key": fWorkspace, "match": map[string]any{"value": workspaceID}},
|
||||
},
|
||||
},
|
||||
}
|
||||
var resp countResponse
|
||||
if err := c.do(ctx, http.MethodPost,
|
||||
"/collections/"+c.collection+"/points/count", body, &resp); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return resp.Result.Count, nil
|
||||
}
|
||||
|
||||
// Health checks that Qdrant is reachable.
|
||||
func (c *Client) Health(ctx context.Context) error {
|
||||
return c.do(ctx, http.MethodGet, "/healthz", nil, nil)
|
||||
}
|
||||
|
||||
// DropCollection removes the collection. Used by tests and by the documented
|
||||
// re-embedding procedure after a dimension change.
|
||||
func (c *Client) DropCollection(ctx context.Context) error {
|
||||
return c.do(ctx, http.MethodDelete, "/collections/"+c.collection, nil, nil)
|
||||
}
|
||||
|
||||
func payloadString(p map[string]any, key string) string {
|
||||
if v, ok := p[key].(string); ok {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func payloadInt(p map[string]any, key string) (int64, bool) {
|
||||
switch v := p[key].(type) {
|
||||
case float64:
|
||||
return int64(v), true
|
||||
case int64:
|
||||
return v, true
|
||||
case json.Number:
|
||||
n, err := v.Int64()
|
||||
return n, err == nil
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
Reference in New Issue
Block a user