427 lines
19 KiB
Go
427 lines
19 KiB
Go
// 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
|
|
}
|