From c4d13ca4a7f2a1d04349dc7f792089daa85bf5db Mon Sep 17 00:00:00 2001 From: viettranx Date: Fri, 10 Apr 2026 19:33:03 +0700 Subject: [PATCH] perf: eliminate N+1 query patterns across import, vault, and usage pipelines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Convert importCron/importUserProfiles/importUserOverrides to multi-row UPSERT with chunking (200→2 queries per import) - Add GetTenantsByIDs batch method to TenantStore (PG+SQLite), refactor handleMine() to single-query tenant fetch (N→1) - Refactor buildAgentKeyMap() to use existing GetByKeys with tenant-scoped context instead of per-key SELECT (N→1) - Add GetDocumentsByIDs, GetDocumentByBasename, CreateLinks batch methods to VaultStore (PG+SQLite) - Batch-fetch docs in enrichment Phase 0, carry title through pipeline to avoid per-doc refetch in classify phase - Replace ListDocuments(limit=500) fallback in wikilink resolution with targeted GetDocumentByBasename DB query - Batch CreateLink calls in SyncDocLinks and classifyLinks - Refactor usage handleGet/handleSummary to use ListPagedRich instead of List+GetOrCreate N+1 loop (100→1 queries) - Batch INSERT for team import members/comments/events/links (250→5) - Add migration 47: UNIQUE constraint on cron_jobs(agent_id,tenant_id,name) with dedup, SQLite schema v15 --- internal/gateway/methods/tenants.go | 20 +- internal/gateway/methods/usage.go | 44 ++-- internal/http/agents_import_sections.go | 138 +++++----- internal/http/agents_import_team.go | 235 +++++++++++++----- internal/http/auth_test.go | 3 + internal/store/pg/tenant_store.go | 22 ++ internal/store/pg/vault_documents.go | 51 ++++ internal/store/pg/vault_links.go | 87 +++++++ internal/store/sqlitestore/schema.go | 10 +- internal/store/sqlitestore/schema.sql | 1 + internal/store/sqlitestore/tenants.go | 28 +++ internal/store/sqlitestore/vault_documents.go | 54 ++++ internal/store/sqlitestore/vault_links.go | 86 +++++++ internal/store/tenant_store.go | 3 + internal/store/vault_store.go | 6 + internal/upgrade/version.go | 2 +- internal/vault/enrich_classify.go | 14 +- internal/vault/enrich_worker.go | 49 +++- internal/vault/links.go | 40 ++- ...00047_cron_jobs_unique_constraint.down.sql | 3 + .../000047_cron_jobs_unique_constraint.up.sql | 19 ++ 21 files changed, 724 insertions(+), 191 deletions(-) create mode 100644 migrations/000047_cron_jobs_unique_constraint.down.sql create mode 100644 migrations/000047_cron_jobs_unique_constraint.up.sql diff --git a/internal/gateway/methods/tenants.go b/internal/gateway/methods/tenants.go index 7aff4966..fd57465b 100644 --- a/internal/gateway/methods/tenants.go +++ b/internal/gateway/methods/tenants.go @@ -381,10 +381,26 @@ func (m *TenantsMethods) handleMine(ctx context.Context, client *gateway.Client, return } + // Batch-fetch all tenant data in a single query instead of per-membership. + ids := make([]uuid.UUID, 0, len(memberships)) + for _, mem := range memberships { + ids = append(ids, mem.TenantID) + } + tenants, tErr := m.tenantStore.GetTenantsByIDs(ctx, ids) + if tErr != nil { + slog.Error("tenants.mine: batch fetch failed", "error", tErr, "user_id", userID) + client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, i18n.T(locale, i18n.MsgFailedToList, "tenants"))) + return + } + tenantMap := make(map[uuid.UUID]*store.TenantData, len(tenants)) + for i := range tenants { + tenantMap[tenants[i].ID] = &tenants[i] + } + entries := make([]tenantEntry, 0, len(memberships)) for _, mem := range memberships { - t, err := m.tenantStore.GetTenant(ctx, mem.TenantID) - if err != nil || t == nil || t.Status != store.TenantStatusActive { + t := tenantMap[mem.TenantID] + if t == nil || t.Status != store.TenantStatusActive { continue } entries = append(entries, tenantEntry{ID: t.ID.String(), Name: t.Name, Slug: t.Slug, Role: mem.Role, Status: t.Status}) diff --git a/internal/gateway/methods/usage.go b/internal/gateway/methods/usage.go index e779a79c..2e817683 100644 --- a/internal/gateway/methods/usage.go +++ b/internal/gateway/methods/usage.go @@ -50,28 +50,28 @@ func (m *UsageMethods) handleGet(ctx context.Context, client *gateway.Client, re params.Limit = 20 } - sessions := m.sessions.List(ctx, params.AgentID) + // Use ListPagedRich: single query returns model, provider, tokens — no N+1 GetOrCreate loop. + // Fetch large batch to filter non-zero tokens, then paginate in-memory. + result := m.sessions.ListPagedRich(ctx, store.SessionListOpts{ + AgentID: params.AgentID, + Limit: 10000, + }) - records := make([]UsageRecord, 0, len(sessions)) - for _, s := range sessions { - // Get full session data for token info - data := m.sessions.GetOrCreate(ctx, s.Key) - if data.InputTokens == 0 && data.OutputTokens == 0 { + records := make([]UsageRecord, 0, len(result.Sessions)) + for _, s := range result.Sessions { + if s.InputTokens == 0 && s.OutputTokens == 0 { continue } - - // Extract agentID from session key (format: "agent::") agentID := extractAgentIDFromKey(s.Key) - records = append(records, UsageRecord{ AgentID: agentID, SessionKey: s.Key, - Model: data.Model, - Provider: data.Provider, - InputTokens: data.InputTokens, - OutputTokens: data.OutputTokens, - TotalTokens: data.InputTokens + data.OutputTokens, - Timestamp: data.Updated.UnixMilli(), + Model: s.Model, + Provider: s.Provider, + InputTokens: s.InputTokens, + OutputTokens: s.OutputTokens, + TotalTokens: s.InputTokens + s.OutputTokens, + Timestamp: s.Updated.UnixMilli(), }) } @@ -96,7 +96,8 @@ func (m *UsageMethods) handleGet(ctx context.Context, client *gateway.Client, re } func (m *UsageMethods) handleSummary(ctx context.Context, client *gateway.Client, req *protocol.RequestFrame) { - sessions := m.sessions.List(ctx, "") // all agents + // Use ListPagedRich: single query returns all token data — no N+1 GetOrCreate loop. + result := m.sessions.ListPagedRich(ctx, store.SessionListOpts{Limit: 10000}) type agentSummary struct { InputTokens int64 `json:"inputTokens"` @@ -108,9 +109,8 @@ func (m *UsageMethods) handleSummary(ctx context.Context, client *gateway.Client byAgent := make(map[string]*agentSummary) var totalRecords int - for _, s := range sessions { - data := m.sessions.GetOrCreate(ctx, s.Key) - if data.InputTokens == 0 && data.OutputTokens == 0 { + for _, s := range result.Sessions { + if s.InputTokens == 0 && s.OutputTokens == 0 { continue } @@ -119,9 +119,9 @@ func (m *UsageMethods) handleSummary(ctx context.Context, client *gateway.Client byAgent[agentID] = &agentSummary{} } - byAgent[agentID].InputTokens += data.InputTokens - byAgent[agentID].OutputTokens += data.OutputTokens - byAgent[agentID].TotalTokens += data.InputTokens + data.OutputTokens + byAgent[agentID].InputTokens += s.InputTokens + byAgent[agentID].OutputTokens += s.OutputTokens + byAgent[agentID].TotalTokens += s.InputTokens + s.OutputTokens byAgent[agentID].Sessions++ totalRecords++ } diff --git a/internal/http/agents_import_sections.go b/internal/http/agents_import_sections.go index 9f5d8a34..bfc5cbe9 100644 --- a/internal/http/agents_import_sections.go +++ b/internal/http/agents_import_sections.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log/slog" + "strings" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/config" @@ -151,40 +152,42 @@ func (h *AgentsHandler) importKG(ctx context.Context, ag *store.AgentData, arc * func (h *AgentsHandler) importCron(ctx context.Context, ag *store.AgentData, arc *importArchive, summary *ImportSummary, progressFn func(ProgressEvent)) { tid := importTenantID(ctx) - for _, j := range arc.cronJobs { - // cron_jobs has no UNIQUE constraint on (agent_id, name), so use SELECT+UPDATE/INSERT - var exists bool - _ = h.db.QueryRowContext(ctx, - `SELECT EXISTS(SELECT 1 FROM cron_jobs WHERE agent_id = $1 AND name = $2 AND tenant_id = $3)`, - ag.ID, j.Name, tid, - ).Scan(&exists) - if exists { - _, err := h.db.ExecContext(ctx, - `UPDATE cron_jobs - SET schedule_kind=$1, cron_expression=$2, interval_ms=$3, - run_at=$4, timezone=$5, payload=$6, delete_after_run=$7 - WHERE agent_id=$8 AND name=$9 AND tenant_id=$10`, - j.ScheduleKind, j.CronExpression, j.IntervalMS, - nullStr(j.RunAt), j.Timezone, j.Payload, j.DeleteAfterRun, - ag.ID, j.Name, tid, - ) - if err != nil { - slog.Warn("agents.import.cron_job.update", "agent_id", ag.ID, "name", j.Name, "error", err) - } - } else { - _, err := h.db.ExecContext(ctx, - `INSERT INTO cron_jobs - (agent_id, name, enabled, schedule_kind, cron_expression, interval_ms, run_at, timezone, payload, delete_after_run, tenant_id) - VALUES ($1, $2, false, $3, $4, $5, $6, $7, $8, $9, $10)`, - ag.ID, j.Name, j.ScheduleKind, + const paramsPerRow = 10 // agent_id, name, schedule_kind, cron_expression, interval_ms, run_at, timezone, payload, delete_after_run, tenant_id (enabled is literal false) + const chunkSize = 5000 + + for start := 0; start < len(arc.cronJobs); start += chunkSize { + end := min(start+chunkSize, len(arc.cronJobs)) + chunk := arc.cronJobs[start:end] + + args := make([]any, 0, len(chunk)*paramsPerRow) + placeholders := make([]string, 0, len(chunk)) + for i, j := range chunk { + base := i * paramsPerRow + placeholders = append(placeholders, fmt.Sprintf( + "($%d,$%d,false,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d)", + base+1, base+2, base+3, base+4, base+5, base+6, base+7, base+8, base+9, base+10, + )) + args = append(args, ag.ID, j.Name, j.ScheduleKind, j.CronExpression, j.IntervalMS, nullStr(j.RunAt), j.Timezone, j.Payload, j.DeleteAfterRun, tid, ) - if err != nil { - slog.Warn("agents.import.cron_job.insert", "agent_id", ag.ID, "name", j.Name, "error", err) - } } - summary.CronJobs++ + + q := `INSERT INTO cron_jobs + (agent_id, name, enabled, schedule_kind, cron_expression, interval_ms, run_at, timezone, payload, delete_after_run, tenant_id) + VALUES ` + strings.Join(placeholders, ",") + ` + ON CONFLICT (agent_id, tenant_id, name) DO UPDATE SET + schedule_kind = EXCLUDED.schedule_kind, + cron_expression = EXCLUDED.cron_expression, + interval_ms = EXCLUDED.interval_ms, + run_at = EXCLUDED.run_at, + timezone = EXCLUDED.timezone, + payload = EXCLUDED.payload, + delete_after_run = EXCLUDED.delete_after_run` + if _, err := h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("agents.import.cron_jobs.batch", "agent_id", ag.ID, "count", len(chunk), "error", err) + } + summary.CronJobs += len(chunk) } if progressFn != nil { progressFn(ProgressEvent{Phase: "cron", Status: "done", Current: summary.CronJobs, Total: len(arc.cronJobs)}) @@ -194,18 +197,28 @@ func (h *AgentsHandler) importCron(ctx context.Context, ag *store.AgentData, arc func (h *AgentsHandler) importUserProfiles(ctx context.Context, ag *store.AgentData, arc *importArchive, summary *ImportSummary, progressFn func(ProgressEvent)) { // workspace=NULL for portability (auto-created via GetOrCreateUserProfile on first user access) tid := importTenantID(ctx) - for _, p := range arc.userProfiles { - _, err := h.db.ExecContext(ctx, - `INSERT INTO user_agent_profiles (agent_id, user_id, workspace, tenant_id) - VALUES ($1, $2, NULL, $3) - ON CONFLICT DO NOTHING`, - ag.ID, p.UserID, tid, - ) - if err != nil { - slog.Warn("agents.import.user_profile", "agent_id", ag.ID, "user_id", p.UserID, "error", err) - continue + const colsPerRow = 3 // agent_id, user_id, tenant_id + const chunkSize = 5000 + + for start := 0; start < len(arc.userProfiles); start += chunkSize { + end := min(start+chunkSize, len(arc.userProfiles)) + chunk := arc.userProfiles[start:end] + + args := make([]any, 0, len(chunk)*colsPerRow) + placeholders := make([]string, 0, len(chunk)) + for i, p := range chunk { + base := i * colsPerRow + placeholders = append(placeholders, fmt.Sprintf("($%d,$%d,NULL,$%d)", base+1, base+2, base+3)) + args = append(args, ag.ID, p.UserID, tid) } - summary.UserProfiles++ + + q := `INSERT INTO user_agent_profiles (agent_id, user_id, workspace, tenant_id) + VALUES ` + strings.Join(placeholders, ",") + ` + ON CONFLICT (agent_id, user_id) DO NOTHING` + if _, err := h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("agents.import.user_profiles.batch", "agent_id", ag.ID, "count", len(chunk), "error", err) + } + summary.UserProfiles += len(chunk) } if progressFn != nil { progressFn(ProgressEvent{Phase: "user_profiles", Status: "done", Current: summary.UserProfiles, Total: len(arc.userProfiles)}) @@ -214,22 +227,35 @@ func (h *AgentsHandler) importUserProfiles(ctx context.Context, ag *store.AgentD func (h *AgentsHandler) importUserOverrides(ctx context.Context, ag *store.AgentData, arc *importArchive, summary *ImportSummary, progressFn func(ProgressEvent)) { tid := importTenantID(ctx) - for _, o := range arc.userOverrides { - _, err := h.db.ExecContext(ctx, - `INSERT INTO user_agent_overrides (agent_id, user_id, provider, model, settings, tenant_id) - VALUES ($1, $2, $3, $4, $5, $6) - ON CONFLICT (agent_id, user_id) DO UPDATE SET - provider = EXCLUDED.provider, - model = EXCLUDED.model, - settings = EXCLUDED.settings, - updated_at = NOW()`, - ag.ID, o.UserID, o.Provider, o.Model, coalesceJSON(o.Settings), tid, - ) - if err != nil { - slog.Warn("agents.import.user_override", "agent_id", ag.ID, "user_id", o.UserID, "error", err) - continue + const colsPerRow = 6 // agent_id, user_id, provider, model, settings, tenant_id + const chunkSize = 5000 + + for start := 0; start < len(arc.userOverrides); start += chunkSize { + end := min(start+chunkSize, len(arc.userOverrides)) + chunk := arc.userOverrides[start:end] + + args := make([]any, 0, len(chunk)*colsPerRow) + placeholders := make([]string, 0, len(chunk)) + for i, o := range chunk { + base := i * colsPerRow + placeholders = append(placeholders, fmt.Sprintf( + "($%d,$%d,$%d,$%d,$%d,$%d)", + base+1, base+2, base+3, base+4, base+5, base+6, + )) + args = append(args, ag.ID, o.UserID, o.Provider, o.Model, coalesceJSON(o.Settings), tid) } - summary.UserOverrides++ + + q := `INSERT INTO user_agent_overrides (agent_id, user_id, provider, model, settings, tenant_id) + VALUES ` + strings.Join(placeholders, ",") + ` + ON CONFLICT (agent_id, user_id) DO UPDATE SET + provider = EXCLUDED.provider, + model = EXCLUDED.model, + settings = EXCLUDED.settings, + updated_at = NOW()` + if _, err := h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("agents.import.user_overrides.batch", "agent_id", ag.ID, "count", len(chunk), "error", err) + } + summary.UserOverrides += len(chunk) } if progressFn != nil { progressFn(ProgressEvent{Phase: "user_overrides", Status: "done", Current: summary.UserOverrides, Total: len(arc.userOverrides)}) diff --git a/internal/http/agents_import_team.go b/internal/http/agents_import_team.go index 0e246261..9d0a3199 100644 --- a/internal/http/agents_import_team.go +++ b/internal/http/agents_import_team.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "path/filepath" + "strings" "github.com/google/uuid" @@ -60,20 +61,37 @@ func (h *AgentsHandler) importTeamSection(ctx context.Context, ag *store.AgentDa // Always include the importing agent itself agentKeyToID[ag.AgentKey] = ag.ID - // Insert non-lead members - for _, m := range arc.teamMembers { - memberID, ok := agentKeyToID[m.AgentKey] - if !ok { - slog.Info("import.team: member not found, skipping", "agent_key", m.AgentKey) - continue + // Insert non-lead members (batch) + { + type memberRow struct { + agentID uuid.UUID + role string } - if _, err = h.db.ExecContext(ctx, - `INSERT INTO agent_team_members (team_id, agent_id, role, tenant_id, joined_at) - VALUES ($1, $2, $3, $4, NOW()) - ON CONFLICT (team_id, agent_id) DO NOTHING`, - teamID, memberID, m.Role, tid, - ); err != nil { - slog.Warn("import.team: insert member", "agent_key", m.AgentKey, "error", err) + var rows []memberRow + for _, m := range arc.teamMembers { + memberID, ok := agentKeyToID[m.AgentKey] + if !ok { + slog.Info("import.team: member not found, skipping", "agent_key", m.AgentKey) + continue + } + rows = append(rows, memberRow{agentID: memberID, role: m.Role}) + } + const cols = 4 // team_id, agent_id, role, tenant_id + for start := 0; start < len(rows); start += 1000 { + end := min(start+1000, len(rows)) + chunk := rows[start:end] + args := make([]any, 0, len(chunk)*cols) + ph := make([]string, 0, len(chunk)) + for i, r := range chunk { + b := i * cols + ph = append(ph, fmt.Sprintf("($%d,$%d,$%d,$%d,NOW())", b+1, b+2, b+3, b+4)) + args = append(args, teamID, r.agentID, r.role, tid) + } + q := `INSERT INTO agent_team_members (team_id, agent_id, role, tenant_id, joined_at) + VALUES ` + strings.Join(ph, ",") + ` ON CONFLICT (team_id, agent_id) DO NOTHING` + if _, err = h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("import.team: batch insert members", "count", len(chunk), "error", err) + } } } @@ -130,62 +148,132 @@ func (h *AgentsHandler) importTeamSection(ctx context.Context, ag *store.AgentDa } } - // Insert comments - for _, c := range arc.teamComments { - if c.TaskIdx < 0 || c.TaskIdx >= len(taskIDByIdx) { - continue + // Insert comments (batch) + { + type commentRow struct { + id uuid.UUID + taskID uuid.UUID + agentID *uuid.UUID + userID *string + content string + commentType string + metadata any } - taskID := taskIDByIdx[c.TaskIdx] - - var agentID *uuid.UUID - if c.AgentKey != nil { - if id, ok := agentKeyToID[*c.AgentKey]; ok { - agentID = &id + var cRows []commentRow + for _, c := range arc.teamComments { + if c.TaskIdx < 0 || c.TaskIdx >= len(taskIDByIdx) { + continue + } + var agentID *uuid.UUID + if c.AgentKey != nil { + if id, ok := agentKeyToID[*c.AgentKey]; ok { + agentID = &id + } + } + cRows = append(cRows, commentRow{ + id: uuid.Must(uuid.NewV7()), taskID: taskIDByIdx[c.TaskIdx], + agentID: agentID, userID: c.UserID, content: c.Content, + commentType: c.CommentType, metadata: nullJSON(c.Metadata), + }) + } + const cols = 8 + for start := 0; start < len(cRows); start += 1000 { + end := min(start+1000, len(cRows)) + chunk := cRows[start:end] + args := make([]any, 0, len(chunk)*cols) + ph := make([]string, 0, len(chunk)) + for i, r := range chunk { + b := i * cols + ph = append(ph, fmt.Sprintf("($%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,NOW())", b+1, b+2, b+3, b+4, b+5, b+6, b+7, b+8)) + args = append(args, r.id, r.taskID, r.agentID, r.userID, r.content, r.commentType, r.metadata, tid) + } + q := `INSERT INTO team_task_comments (id, task_id, agent_id, user_id, content, comment_type, metadata, tenant_id, created_at) + VALUES ` + strings.Join(ph, ",") + if _, err = h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("import.team: batch insert comments", "count", len(chunk), "error", err) } } - if _, err = h.db.ExecContext(ctx, - `INSERT INTO team_task_comments (id, task_id, agent_id, user_id, content, comment_type, metadata, tenant_id, created_at) - VALUES ($1,$2,$3,$4,$5,$6,$7,$8,NOW())`, - uuid.Must(uuid.NewV7()), taskID, agentID, c.UserID, - c.Content, c.CommentType, nullJSON(c.Metadata), tid, - ); err != nil { - slog.Warn("import.team: insert comment", "task_id", taskID, "error", err) + } + + // Insert events (batch) + { + type eventRow struct { + id uuid.UUID + taskID uuid.UUID + eventType string + actorType string + actorID string + data any + } + var eRows []eventRow + for _, ev := range arc.teamEvents { + if ev.TaskIdx < 0 || ev.TaskIdx >= len(taskIDByIdx) { + continue + } + eRows = append(eRows, eventRow{ + id: uuid.Must(uuid.NewV7()), taskID: taskIDByIdx[ev.TaskIdx], + eventType: ev.EventType, actorType: ev.ActorType, + actorID: ev.ActorID, data: nullJSON(ev.Data), + }) + } + const cols = 7 + for start := 0; start < len(eRows); start += 1000 { + end := min(start+1000, len(eRows)) + chunk := eRows[start:end] + args := make([]any, 0, len(chunk)*cols) + ph := make([]string, 0, len(chunk)) + for i, r := range chunk { + b := i * cols + ph = append(ph, fmt.Sprintf("($%d,$%d,$%d,$%d,$%d,$%d,$%d,NOW())", b+1, b+2, b+3, b+4, b+5, b+6, b+7)) + args = append(args, r.id, r.taskID, r.eventType, r.actorType, r.actorID, r.data, tid) + } + q := `INSERT INTO team_task_events (id, task_id, event_type, actor_type, actor_id, data, tenant_id, created_at) + VALUES ` + strings.Join(ph, ",") + if _, err = h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("import.team: batch insert events", "count", len(chunk), "error", err) + } } } - // Insert events - for _, ev := range arc.teamEvents { - if ev.TaskIdx < 0 || ev.TaskIdx >= len(taskIDByIdx) { - continue + // Insert agent_links (batch) + { + type linkRow struct { + id uuid.UUID + srcID uuid.UUID + tgtID uuid.UUID + dir string + desc string } - taskID := taskIDByIdx[ev.TaskIdx] - if _, err = h.db.ExecContext(ctx, - `INSERT INTO team_task_events (id, task_id, event_type, actor_type, actor_id, data, tenant_id, created_at) - VALUES ($1,$2,$3,$4,$5,$6,$7,NOW())`, - uuid.Must(uuid.NewV7()), taskID, - ev.EventType, ev.ActorType, ev.ActorID, nullJSON(ev.Data), tid, - ); err != nil { - slog.Warn("import.team: insert event", "task_id", taskID, "error", err) + var lRows []linkRow + for _, l := range arc.teamLinks { + srcID, srcOK := agentKeyToID[l.SourceAgentKey] + tgtID, tgtOK := agentKeyToID[l.TargetAgentKey] + if !srcOK || !tgtOK { + slog.Info("import.team: agent_link skipped — agent not found", + "source", l.SourceAgentKey, "target", l.TargetAgentKey) + continue + } + lRows = append(lRows, linkRow{ + id: uuid.Must(uuid.NewV7()), srcID: srcID, tgtID: tgtID, + dir: l.Direction, desc: l.Description, + }) } - } - - // Insert agent_links - for _, l := range arc.teamLinks { - srcID, srcOK := agentKeyToID[l.SourceAgentKey] - tgtID, tgtOK := agentKeyToID[l.TargetAgentKey] - if !srcOK || !tgtOK { - slog.Info("import.team: agent_link skipped — agent not found", - "source", l.SourceAgentKey, "target", l.TargetAgentKey) - continue - } - if _, err = h.db.ExecContext(ctx, - `INSERT INTO agent_links (id, source_agent_id, target_agent_id, direction, description, created_by, tenant_id) - VALUES ($1,$2,$3,$4,$5,$6,$7) - ON CONFLICT DO NOTHING`, - uuid.Must(uuid.NewV7()), srcID, tgtID, - l.Direction, l.Description, userID, tid, - ); err != nil { - slog.Warn("import.team: insert link", "error", err) + const cols = 7 + for start := 0; start < len(lRows); start += 1000 { + end := min(start+1000, len(lRows)) + chunk := lRows[start:end] + args := make([]any, 0, len(chunk)*cols) + ph := make([]string, 0, len(chunk)) + for i, r := range chunk { + b := i * cols + ph = append(ph, fmt.Sprintf("($%d,$%d,$%d,$%d,$%d,$%d,$%d)", b+1, b+2, b+3, b+4, b+5, b+6, b+7)) + args = append(args, r.id, r.srcID, r.tgtID, r.dir, r.desc, userID, tid) + } + q := `INSERT INTO agent_links (id, source_agent_id, target_agent_id, direction, description, created_by, tenant_id) + VALUES ` + strings.Join(ph, ",") + ` ON CONFLICT DO NOTHING` + if _, err = h.db.ExecContext(ctx, q, args...); err != nil { + slog.Warn("import.team: batch insert links", "count", len(chunk), "error", err) + } } } @@ -208,6 +296,7 @@ func (h *AgentsHandler) importTeamSection(ctx context.Context, ag *store.AgentDa } // buildAgentKeyMap resolves all agent_keys referenced in the team section to UUIDs. +// Uses batch GetByKeys with tenant-scoped context instead of per-key SELECT. func (h *AgentsHandler) buildAgentKeyMap(ctx context.Context, tid uuid.UUID, arc *importArchive) map[string]uuid.UUID { keys := make(map[string]struct{}) for _, m := range arc.teamMembers { @@ -231,14 +320,22 @@ func (h *AgentsHandler) buildAgentKeyMap(ctx context.Context, tid uuid.UUID, arc keys[l.TargetAgentKey] = struct{}{} } - out := make(map[string]uuid.UUID, len(keys)) - for key := range keys { - var id uuid.UUID - if err := h.db.QueryRowContext(ctx, - `SELECT id FROM agents WHERE agent_key = $1 AND tenant_id = $2`, key, tid, - ).Scan(&id); err == nil { - out[key] = id - } + keyList := make([]string, 0, len(keys)) + for k := range keys { + keyList = append(keyList, k) + } + + // Ensure tenant-scoped context for GetByKeys (Red Team Fix #2). + scopedCtx := store.WithTenantID(ctx, tid) + agents, err := h.agents.GetByKeys(scopedCtx, keyList) + if err != nil { + slog.Warn("import.team: batch agent key lookup failed", "error", err) + return make(map[string]uuid.UUID) + } + + out := make(map[string]uuid.UUID, len(agents)) + for _, a := range agents { + out[a.AgentKey] = a.ID } return out } diff --git a/internal/http/auth_test.go b/internal/http/auth_test.go index c5c07552..bc2869e8 100644 --- a/internal/http/auth_test.go +++ b/internal/http/auth_test.go @@ -117,6 +117,9 @@ func (m *mockTenantStore) GetTenantUser(context.Context, uuid.UUID) (*store.Tena func (m *mockTenantStore) CreateTenantUserReturning(context.Context, uuid.UUID, string, string, string) (*store.TenantUserData, error) { return nil, fmt.Errorf("not implemented") } +func (m *mockTenantStore) GetTenantsByIDs(context.Context, []uuid.UUID) ([]store.TenantData, error) { + return nil, nil +} type mockPairingStore struct { paired map[string]bool diff --git a/internal/store/pg/tenant_store.go b/internal/store/pg/tenant_store.go index e837166a..e5b72906 100644 --- a/internal/store/pg/tenant_store.go +++ b/internal/store/pg/tenant_store.go @@ -8,6 +8,7 @@ import ( "time" "github.com/google/uuid" + "github.com/lib/pq" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -80,6 +81,27 @@ func (s *PGTenantStore) ListTenants(ctx context.Context) ([]store.TenantData, er return tenants, nil } +func (s *PGTenantStore) GetTenantsByIDs(ctx context.Context, ids []uuid.UUID) ([]store.TenantData, error) { + if len(ids) == 0 { + return nil, nil + } + // Chunk by 500 to stay well within PG param limits. + const chunkSize = 500 + var all []store.TenantData + for start := 0; start < len(ids); start += chunkSize { + end := min(start+chunkSize, len(ids)) + var chunk []store.TenantData + if err := pkgSqlxDB.SelectContext(ctx, &chunk, + `SELECT id, name, slug, status, settings, created_at, updated_at + FROM tenants WHERE id = ANY($1)`, + pq.Array(ids[start:end])); err != nil { + return nil, err + } + all = append(all, chunk...) + } + return all, nil +} + func (s *PGTenantStore) UpdateTenant(ctx context.Context, id uuid.UUID, updates map[string]any) error { return execMapUpdate(ctx, s.db, "tenants", id, updates) } diff --git a/internal/store/pg/vault_documents.go b/internal/store/pg/vault_documents.go index 86457608..66f7bb92 100644 --- a/internal/store/pg/vault_documents.go +++ b/internal/store/pg/vault_documents.go @@ -176,6 +176,57 @@ func (s *PGVaultStore) GetDocumentByID(ctx context.Context, tenantID, id string) return &doc, nil } +// GetDocumentsByIDs returns documents matching the given IDs with tenant isolation. +// Chunks by 500 to stay within PG param limits. +func (s *PGVaultStore) GetDocumentsByIDs(ctx context.Context, tenantID string, docIDs []string) ([]store.VaultDocument, error) { + if len(docIDs) == 0 { + return nil, nil + } + tid := mustParseUUID(tenantID) + const chunkSize = 500 + var all []store.VaultDocument + for start := 0; start < len(docIDs); start += chunkSize { + end := min(start+chunkSize, len(docIDs)) + var scanned []vaultDocRow + if err := pkgSqlxDB.SelectContext(ctx, &scanned, + `SELECT id, tenant_id, agent_id, team_id, scope, custom_scope, path, title, doc_type, content_hash, summary, metadata, created_at, updated_at + FROM vault_documents WHERE id = ANY($1) AND tenant_id = $2`, + pqStringArray(docIDs[start:end]), tid); err != nil { + return nil, err + } + for i := range scanned { + all = append(all, scanned[i].toVaultDocument()) + } + } + return all, nil +} + +// GetDocumentByBasename finds a document by path basename (case-insensitive). +// Uses the stored generated column path_basename + index for fast lookup. +func (s *PGVaultStore) GetDocumentByBasename(ctx context.Context, tenantID, agentID, basename string) (*store.VaultDocument, error) { + tid := mustParseUUID(tenantID) + q := `SELECT id, tenant_id, agent_id, team_id, scope, custom_scope, path, title, doc_type, content_hash, summary, metadata, created_at, updated_at + FROM vault_documents + WHERE tenant_id = $1 AND path_basename = lower($2)` + args := []any{tid, basename} + p := 3 + if agentID != "" { + q += fmt.Sprintf(" AND agent_id = $%d", p) + args = append(args, mustParseUUID(agentID)) + } + q += " LIMIT 1" + var row vaultDocRow + err := s.db.QueryRowContext(ctx, q, args...).Scan( + &row.ID, &row.TenantID, &row.AgentID, &row.TeamID, &row.Scope, &row.CustomScope, + &row.Path, &row.Title, &row.DocType, &row.ContentHash, &row.Summary, + &row.MetaJSON, &row.CreatedAt, &row.UpdatedAt) + if err != nil { + return nil, err + } + doc := row.toVaultDocument() + return &doc, nil +} + // DeleteDocument removes a vault document by tenant, agent, and path. // Empty agentID means no agent filter. // Team scoping via RunContext (same rules as GetDocument). diff --git a/internal/store/pg/vault_links.go b/internal/store/pg/vault_links.go index 66ef4fb8..54903073 100644 --- a/internal/store/pg/vault_links.go +++ b/internal/store/pg/vault_links.go @@ -11,6 +11,93 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" ) +// CreateLinks batch-inserts vault links with tenant validation. +// Validates same-tenant boundary for all links, then multi-row INSERT. +func (s *PGVaultStore) CreateLinks(ctx context.Context, links []store.VaultLink) error { + if len(links) == 0 { + return nil + } + + // Collect all unique doc IDs for tenant validation. + docIDSet := make(map[string]struct{}) + for _, l := range links { + docIDSet[l.FromDocID] = struct{}{} + docIDSet[l.ToDocID] = struct{}{} + } + docIDs := make([]string, 0, len(docIDSet)) + for id := range docIDSet { + docIDs = append(docIDs, id) + } + + // Batch-fetch tenant_id for all referenced docs. + type docMeta struct { + tenantID uuid.UUID + teamID *uuid.UUID + } + docMap := make(map[string]docMeta, len(docIDs)) + rows, err := s.db.QueryContext(ctx, + `SELECT id, tenant_id, team_id FROM vault_documents WHERE id = ANY($1)`, + pqStringArray(docIDs)) + if err != nil { + return fmt.Errorf("vault batch link: fetch docs: %w", err) + } + defer rows.Close() + for rows.Next() { + var id, tid uuid.UUID + var teamID *uuid.UUID + if err := rows.Scan(&id, &tid, &teamID); err != nil { + return fmt.Errorf("vault batch link: scan doc: %w", err) + } + docMap[id.String()] = docMeta{tenantID: tid, teamID: teamID} + } + if err := rows.Err(); err != nil { + return fmt.Errorf("vault batch link: rows err: %w", err) + } + + // Validate and build multi-row INSERT (chunk by 500). + const chunkSize = 500 + var valid []store.VaultLink + for _, l := range links { + from, fromOK := docMap[l.FromDocID] + to, toOK := docMap[l.ToDocID] + if !fromOK || !toOK { + continue // skip if doc not found + } + if from.tenantID != to.tenantID { + continue // cross-tenant + } + if from.teamID != nil && to.teamID != nil && *from.teamID != *to.teamID { + continue // cross-team + } + valid = append(valid, l) + } + + for start := 0; start < len(valid); start += chunkSize { + end := min(start+chunkSize, len(valid)) + chunk := valid[start:end] + const paramsPerRow = 6 // id, from_doc_id, to_doc_id, link_type, context, created_at + args := make([]any, 0, len(chunk)*paramsPerRow) + ph := make([]string, 0, len(chunk)) + now := time.Now().UTC() + for i, l := range chunk { + b := i * paramsPerRow + ph = append(ph, fmt.Sprintf("($%d,$%d,$%d,$%d,$%d,$%d)", b+1, b+2, b+3, b+4, b+5, b+6)) + args = append(args, + uuid.Must(uuid.NewV7()), + mustParseUUID(l.FromDocID), mustParseUUID(l.ToDocID), + l.LinkType, l.Context, now, + ) + } + q := `INSERT INTO vault_links (id, from_doc_id, to_doc_id, link_type, context, created_at) + VALUES ` + strings.Join(ph, ",") + ` + ON CONFLICT (from_doc_id, to_doc_id, link_type) DO UPDATE SET context = EXCLUDED.context` + if _, err := s.db.ExecContext(ctx, q, args...); err != nil { + return fmt.Errorf("vault batch create links: %w", err) + } + } + return nil +} + // CreateLink inserts a vault link, updating context on conflict. // Validates same-tenant + same-team boundary before insert. func (s *PGVaultStore) CreateLink(ctx context.Context, link *store.VaultLink) error { diff --git a/internal/store/sqlitestore/schema.go b/internal/store/sqlitestore/schema.go index 5ae63f90..2b53e969 100644 --- a/internal/store/sqlitestore/schema.go +++ b/internal/store/sqlitestore/schema.go @@ -15,7 +15,7 @@ var schemaSQL string // SchemaVersion is the current SQLite schema version. // Bump this when adding new migration steps below. -const SchemaVersion = 14 +const SchemaVersion = 15 // migrations maps version → SQL to apply when upgrading FROM that version. // schema.sql always represents the LATEST full schema (for fresh DBs). @@ -404,6 +404,14 @@ CREATE INDEX IF NOT EXISTS idx_vault_docs_agent_scope ON vault_documents(agent_i CREATE INDEX IF NOT EXISTS idx_vault_docs_type ON vault_documents(agent_id, doc_type); CREATE INDEX IF NOT EXISTS idx_vault_docs_hash ON vault_documents(content_hash); CREATE INDEX IF NOT EXISTS idx_vault_docs_team ON vault_documents(team_id);`, + + // Version 14 → 15: cron_jobs UNIQUE constraint on (agent_id, tenant_id, name). + // Dedup first: keep one row per combo (SQLite has no DISTINCT ON — use GROUP BY + MIN). + 14: `DELETE FROM cron_jobs WHERE id NOT IN ( + SELECT MIN(id) FROM cron_jobs GROUP BY agent_id, tenant_id, name +); +CREATE UNIQUE INDEX IF NOT EXISTS uq_cron_jobs_agent_tenant_name + ON cron_jobs(agent_id, tenant_id, name);`, } // EnsureSchema creates tables if they don't exist and applies incremental migrations. diff --git a/internal/store/sqlitestore/schema.sql b/internal/store/sqlitestore/schema.sql index f71e2bae..7efb3241 100644 --- a/internal/store/sqlitestore/schema.sql +++ b/internal/store/sqlitestore/schema.sql @@ -482,6 +482,7 @@ CREATE INDEX IF NOT EXISTS idx_cron_jobs_user_id ON cron_jobs(user_id); CREATE INDEX IF NOT EXISTS idx_cron_jobs_agent_user ON cron_jobs(agent_id, user_id); CREATE INDEX IF NOT EXISTS idx_cron_jobs_team ON cron_jobs(team_id) WHERE team_id IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_cron_jobs_tenant ON cron_jobs(tenant_id); +CREATE UNIQUE INDEX IF NOT EXISTS uq_cron_jobs_agent_tenant_name ON cron_jobs(agent_id, tenant_id, name); -- ============================================================ -- Table: cron_run_logs diff --git a/internal/store/sqlitestore/tenants.go b/internal/store/sqlitestore/tenants.go index f5bfba49..133f3b6d 100644 --- a/internal/store/sqlitestore/tenants.go +++ b/internal/store/sqlitestore/tenants.go @@ -7,6 +7,7 @@ import ( "database/sql" "encoding/json" "errors" + "strings" "time" "github.com/google/uuid" @@ -86,6 +87,33 @@ func (s *SQLiteTenantStore) ListTenants(ctx context.Context) ([]store.TenantData return result, nil } +func (s *SQLiteTenantStore) GetTenantsByIDs(ctx context.Context, ids []uuid.UUID) ([]store.TenantData, error) { + if len(ids) == 0 { + return nil, nil + } + const chunkSize = 500 + var all []store.TenantData + for start := 0; start < len(ids); start += chunkSize { + end := min(start+chunkSize, len(ids)) + chunk := ids[start:end] + ph := make([]string, len(chunk)) + args := make([]any, len(chunk)) + for i, id := range chunk { + ph[i] = "?" + args[i] = id.String() + } + q := `SELECT ` + tenantSelectCols + ` FROM tenants WHERE id IN (` + strings.Join(ph, ",") + `)` + var rows []tenantRow + if err := pkgSqlxDB.SelectContext(ctx, &rows, q, args...); err != nil { + return nil, err + } + for _, r := range rows { + all = append(all, r.toTenantData()) + } + } + return all, nil +} + func (s *SQLiteTenantStore) UpdateTenant(ctx context.Context, id uuid.UUID, updates map[string]any) error { return execMapUpdate(ctx, s.db, "tenants", id, updates) } diff --git a/internal/store/sqlitestore/vault_documents.go b/internal/store/sqlitestore/vault_documents.go index 86ddc45a..7852340d 100644 --- a/internal/store/sqlitestore/vault_documents.go +++ b/internal/store/sqlitestore/vault_documents.go @@ -119,6 +119,60 @@ func (s *SQLiteVaultStore) GetDocumentByID(ctx context.Context, tenantID, id str return scanVaultDoc(row) } +// GetDocumentsByIDs returns documents matching the given IDs with tenant isolation. +func (s *SQLiteVaultStore) GetDocumentsByIDs(ctx context.Context, tenantID string, docIDs []string) ([]store.VaultDocument, error) { + if len(docIDs) == 0 { + return nil, nil + } + const chunkSize = 500 + var all []store.VaultDocument + for start := 0; start < len(docIDs); start += chunkSize { + end := min(start+chunkSize, len(docIDs)) + chunk := docIDs[start:end] + ph := make([]string, len(chunk)) + args := make([]any, len(chunk)) + for i, id := range chunk { + ph[i] = "?" + args[i] = id + } + args = append(args, tenantID) + q := `SELECT id, tenant_id, agent_id, team_id, scope, custom_scope, path, title, doc_type, content_hash, summary, metadata, created_at, updated_at + FROM vault_documents WHERE id IN (` + strings.Join(ph, ",") + `) AND tenant_id = ?` + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + for rows.Next() { + doc, scanErr := scanVaultDocRow(rows) + if scanErr != nil { + rows.Close() + return nil, scanErr + } + all = append(all, *doc) + } + rows.Close() + if err := rows.Err(); err != nil { + return nil, err + } + } + return all, nil +} + +// GetDocumentByBasename finds a document by path basename (case-insensitive). +func (s *SQLiteVaultStore) GetDocumentByBasename(ctx context.Context, tenantID, agentID, basename string) (*store.VaultDocument, error) { + q := `SELECT id, tenant_id, agent_id, team_id, scope, custom_scope, path, title, doc_type, content_hash, summary, metadata, created_at, updated_at + FROM vault_documents + WHERE tenant_id = ? AND lower(replace(path, rtrim(path, replace(path, '/', '')), '')) = lower(?)` + args := []any{tenantID, basename} + if agentID != "" { + q += " AND agent_id = ?" + args = append(args, agentID) + } + q += " LIMIT 1" + row := s.db.QueryRowContext(ctx, q, args...) + return scanVaultDoc(row) +} + // DeleteDocument removes a vault document (FK cascades delete vault_links). // Empty agentID means no agent filter. // Team scoping via RunContext (same rules as GetDocument). diff --git a/internal/store/sqlitestore/vault_links.go b/internal/store/sqlitestore/vault_links.go index 65d3f186..ff78953c 100644 --- a/internal/store/sqlitestore/vault_links.go +++ b/internal/store/sqlitestore/vault_links.go @@ -13,6 +13,92 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" ) +// CreateLinks batch-inserts vault links with tenant validation. +func (s *SQLiteVaultStore) CreateLinks(ctx context.Context, links []store.VaultLink) error { + if len(links) == 0 { + return nil + } + + // Collect all unique doc IDs for tenant validation. + docIDSet := make(map[string]struct{}) + for _, l := range links { + docIDSet[l.FromDocID] = struct{}{} + docIDSet[l.ToDocID] = struct{}{} + } + docIDs := make([]string, 0, len(docIDSet)) + for id := range docIDSet { + docIDs = append(docIDs, id) + } + + // Batch-fetch tenant_id for all referenced docs. + type docMeta struct { + tenantID string + teamID *string + } + docMap := make(map[string]docMeta, len(docIDs)) + ph := make([]string, len(docIDs)) + args := make([]any, len(docIDs)) + for i, id := range docIDs { + ph[i] = "?" + args[i] = id + } + rows, err := s.db.QueryContext(ctx, + `SELECT id, tenant_id, team_id FROM vault_documents WHERE id IN (`+strings.Join(ph, ",")+`)`, args...) + if err != nil { + return fmt.Errorf("vault batch link: fetch docs: %w", err) + } + defer rows.Close() + for rows.Next() { + var id, tid string + var teamID *string + if err := rows.Scan(&id, &tid, &teamID); err != nil { + return fmt.Errorf("vault batch link: scan doc: %w", err) + } + docMap[id] = docMeta{tenantID: tid, teamID: teamID} + } + if err := rows.Err(); err != nil { + return fmt.Errorf("vault batch link: rows err: %w", err) + } + + // Validate and insert (chunk by 500). + const chunkSize = 500 + var valid []store.VaultLink + for _, l := range links { + from, fromOK := docMap[l.FromDocID] + to, toOK := docMap[l.ToDocID] + if !fromOK || !toOK { + continue + } + if from.tenantID != to.tenantID { + continue + } + if from.teamID != nil && to.teamID != nil && *from.teamID != *to.teamID { + continue + } + valid = append(valid, l) + } + + now := time.Now().UTC().Format(time.RFC3339Nano) + for start := 0; start < len(valid); start += chunkSize { + end := min(start+chunkSize, len(valid)) + chunk := valid[start:end] + const cols = 6 + iArgs := make([]any, 0, len(chunk)*cols) + iPh := make([]string, 0, len(chunk)) + for _, l := range chunk { + iPh = append(iPh, "(?,?,?,?,?,?)") + iArgs = append(iArgs, uuid.Must(uuid.NewV7()).String(), l.FromDocID, l.ToDocID, l.LinkType, l.Context, now) + } + q := `INSERT INTO vault_links (id, from_doc_id, to_doc_id, link_type, context, created_at) + VALUES ` + strings.Join(iPh, ",") + ` + ON CONFLICT (from_doc_id, to_doc_id, link_type) DO UPDATE SET context = excluded.context` + if _, err := s.db.ExecContext(ctx, q, iArgs...); err != nil { + return fmt.Errorf("vault batch create links: %w", err) + } + } + return nil +} + // CreateLink inserts a vault link, updating context on conflict. // Validates same-tenant + same-team boundary before insert. func (s *SQLiteVaultStore) CreateLink(ctx context.Context, link *store.VaultLink) error { diff --git a/internal/store/tenant_store.go b/internal/store/tenant_store.go index 08d4b68f..e592fdac 100644 --- a/internal/store/tenant_store.go +++ b/internal/store/tenant_store.go @@ -67,6 +67,9 @@ type TenantStore interface { ListUsers(ctx context.Context, tenantID uuid.UUID) ([]TenantUserData, error) ListUserTenants(ctx context.Context, userID string) ([]TenantUserData, error) + // GetTenantsByIDs returns tenants matching the given UUIDs in a single query. + GetTenantsByIDs(ctx context.Context, ids []uuid.UUID) ([]TenantData, error) + // ResolveUserTenant returns the tenant_id for a user. // If user belongs to multiple tenants, returns the first (by created_at). // If no membership, returns MasterTenantID (backward compat). diff --git a/internal/store/vault_store.go b/internal/store/vault_store.go index 4db4ba12..7c09a9bc 100644 --- a/internal/store/vault_store.go +++ b/internal/store/vault_store.go @@ -83,10 +83,16 @@ type VaultStore interface { CountDocuments(ctx context.Context, tenantID, agentID string, opts VaultListOptions) (int, error) UpdateHash(ctx context.Context, tenantID, id, newHash string) error + // GetDocumentsByIDs returns documents matching the given IDs with tenant isolation. + GetDocumentsByIDs(ctx context.Context, tenantID string, docIDs []string) ([]VaultDocument, error) + // GetDocumentByBasename finds a document by path basename (case-insensitive). + GetDocumentByBasename(ctx context.Context, tenantID, agentID, basename string) (*VaultDocument, error) + // Search (FTS + vector hybrid) Search(ctx context.Context, opts VaultSearchOptions) ([]VaultSearchResult, error) // Links + CreateLinks(ctx context.Context, links []VaultLink) error CreateLink(ctx context.Context, link *VaultLink) error DeleteLink(ctx context.Context, tenantID, id string) error GetOutLinks(ctx context.Context, tenantID, docID string) ([]VaultLink, error) diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index c20a3461..7ba0c007 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -2,4 +2,4 @@ package upgrade // RequiredSchemaVersion is the schema migration version this binary requires. // Bump this whenever adding a new SQL migration file. -const RequiredSchemaVersion uint = 46 +const RequiredSchemaVersion uint = 47 diff --git a/internal/vault/enrich_classify.go b/internal/vault/enrich_classify.go index 19bfa9a7..83252c15 100644 --- a/internal/vault/enrich_classify.go +++ b/internal/vault/enrich_classify.go @@ -106,10 +106,8 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri if err := w.vault.DeleteDocLinksByTypes(ctx, tenantID, sourceDocID, allTypes); err != nil { slog.Warn("vault.classify: delete_old", "doc", sourceDocID, "err", err) } - for i := range newLinks { - if err := w.vault.CreateLink(ctx, &newLinks[i]); err != nil { - slog.Debug("vault.classify: create_link", "from", sourceDocID, "to", newLinks[i].ToDocID, "err", err) - } + if err := w.vault.CreateLinks(ctx, newLinks); err != nil { + slog.Debug("vault.classify: batch_create_links", "from", sourceDocID, "count", len(newLinks), "err", err) } } } @@ -125,10 +123,10 @@ func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, agentID s slog.Warn("vault.classify: find_similar", "doc", r.payload.DocID, "err", err) continue } - // Derive title from path (payload has no Title field; neighbor docs have it from DB). - title := r.payload.Path - if doc, err := w.vault.GetDocumentByID(ctx, tenantID, r.payload.DocID); err == nil && doc != nil { - title = doc.Title + // Use title carried from Phase 0 batch-fetch (avoids per-doc refetch). + title := r.title + if title == "" { + title = r.payload.Path } src := classifyDoc{ DocID: r.payload.DocID, diff --git a/internal/vault/enrich_worker.go b/internal/vault/enrich_worker.go index 9a3f2065..fe2c36a3 100644 --- a/internal/vault/enrich_worker.go +++ b/internal/vault/enrich_worker.go @@ -111,6 +111,7 @@ func (w *enrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) e type enriched struct { payload eventbus.VaultDocUpsertedPayload summary string + title string // carried from DB for classify phase (avoids refetch) } // processBatch drains and processes queued vault doc events in a loop. @@ -165,13 +166,16 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) { // processChunk runs the 4-phase enrichment pipeline for a single chunk of docs. // Phase 1 batches all files into a single LLM call for summarization. func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultDocUpsertedPayload) { - // Phase 0 — Prepare: dedup check, fetch existing doc, read file content. + // Phase 0 — Prepare: dedup check, batch-fetch existing docs, read file content. type prepared struct { payload eventbus.VaultDocUpsertedPayload content string // non-empty = needs LLM summarization summary string // non-empty = already summarized, skip LLM + title string // carried from DB for classify phase } - var all []prepared + + // Filter dedup'd items first. + var pending []eventbus.VaultDocUpsertedPayload for _, item := range items { w.dedupMu.Lock() if prev, exists := w.dedup[item.DocID]; exists && prev == item.ContentHash { @@ -179,18 +183,37 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD continue } w.dedupMu.Unlock() + pending = append(pending, item) + } + if len(pending) == 0 { + return + } - existing, err := w.vault.GetDocumentByID(ctx, item.TenantID, item.DocID) - if err != nil { - slog.Warn("vault.enrich: get_doc", "doc", item.DocID, "err", err) - continue - } + // Batch-fetch all existing docs in a single query. + tenantID := pending[0].TenantID + docIDs := make([]string, len(pending)) + for i, item := range pending { + docIDs[i] = item.DocID + } + existingDocs, err := w.vault.GetDocumentsByIDs(ctx, tenantID, docIDs) + if err != nil { + slog.Warn("vault.enrich: batch_fetch_docs", "count", len(docIDs), "err", err) + return + } + docMap := make(map[string]*store.VaultDocument, len(existingDocs)) + for i := range existingDocs { + docMap[existingDocs[i].ID] = &existingDocs[i] + } + + var all []prepared + for _, item := range pending { + existing := docMap[item.DocID] if existing != nil && existing.Summary != "" { - all = append(all, prepared{payload: item, summary: existing.Summary}) + all = append(all, prepared{payload: item, summary: existing.Summary, title: existing.Title}) continue } if existing != nil && existing.DocType == "media" { - all = append(all, prepared{payload: item}) + all = append(all, prepared{payload: item, title: existing.Title}) continue } @@ -204,7 +227,11 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD if len(runes) > enrichBatchItemMaxRunes { runes = runes[:enrichBatchItemMaxRunes] } - all = append(all, prepared{payload: item, content: string(runes)}) + title := "" + if existing != nil { + title = existing.Title + } + all = append(all, prepared{payload: item, content: string(runes), title: title}) } if len(all) == 0 { return @@ -235,7 +262,7 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD // Build enriched results. var results []enriched for _, p := range all { - results = append(results, enriched{payload: p.payload, summary: p.summary}) + results = append(results, enriched{payload: p.payload, summary: p.summary, title: p.title}) } // Phase 2 — Embed: update summary + embed per doc. diff --git a/internal/vault/links.go b/internal/vault/links.go index a890be79..ce14c62a 100644 --- a/internal/vault/links.go +++ b/internal/vault/links.go @@ -48,7 +48,7 @@ func ExtractWikilinks(content string) []WikilinkMatch { } // ResolveWikilinkTarget finds a vault_document matching the wikilink target. -// Strategy: exact path -> path+.md -> basename match -> nil. +// Strategy: exact path -> path+.md -> basename DB query -> nil. func ResolveWikilinkTarget(ctx context.Context, vs store.VaultStore, target, tenantID, agentID string) (*store.VaultDocument, error) { // 1. Exact path match doc, err := vs.GetDocument(ctx, tenantID, agentID, target) @@ -64,20 +64,17 @@ func ResolveWikilinkTarget(ctx context.Context, vs store.VaultStore, target, ten } } - // 3. Basename search: list all docs and find by basename - docs, err := vs.ListDocuments(ctx, tenantID, agentID, store.VaultListOptions{Limit: 500}) - if err != nil { - return nil, err + // 3. Basename search via targeted DB query (replaces ListDocuments N+1). + targetBase := filepath.Base(target) + doc, err = vs.GetDocumentByBasename(ctx, tenantID, agentID, targetBase) + if err == nil && doc != nil { + return doc, nil } - targetBase := strings.ToLower(filepath.Base(target)) - targetBaseMD := targetBase - if !strings.HasSuffix(targetBaseMD, ".md") { - targetBaseMD = targetBase + ".md" - } - for i := range docs { - base := strings.ToLower(filepath.Base(docs[i].Path)) - if base == targetBase || base == targetBaseMD { - return &docs[i], nil + // Try with .md suffix. + if !strings.HasSuffix(targetBase, ".md") { + doc, err = vs.GetDocumentByBasename(ctx, tenantID, agentID, targetBase+".md") + if err == nil && doc != nil { + return doc, nil } } @@ -99,7 +96,8 @@ func SyncDocLinks(ctx context.Context, vs store.VaultStore, doc *store.VaultDocu return err } - // Resolve and create new links + // Resolve all wikilinks, then batch-create links in a single call. + var links []store.VaultLink for _, m := range matches { target, err := ResolveWikilinkTarget(ctx, vs, m.Target, tenantID, agentID) if err != nil { @@ -110,15 +108,15 @@ func SyncDocLinks(ctx context.Context, vs store.VaultStore, doc *store.VaultDocu slog.Debug("vault.link_unresolved", "target", m.Target) continue } - link := &store.VaultLink{ + links = append(links, store.VaultLink{ FromDocID: doc.ID, ToDocID: target.ID, LinkType: "wikilink", Context: m.Context, - } - if err := vs.CreateLink(ctx, link); err != nil { - slog.Warn("vault.create_link", "from", doc.Path, "to", target.Path, "err", err) - } + }) } - return nil + if len(links) == 0 { + return nil + } + return vs.CreateLinks(ctx, links) } diff --git a/migrations/000047_cron_jobs_unique_constraint.down.sql b/migrations/000047_cron_jobs_unique_constraint.down.sql new file mode 100644 index 00000000..307a5510 --- /dev/null +++ b/migrations/000047_cron_jobs_unique_constraint.down.sql @@ -0,0 +1,3 @@ +DROP INDEX IF EXISTS idx_vault_docs_basename; +ALTER TABLE vault_documents DROP COLUMN IF EXISTS path_basename; +DROP INDEX IF EXISTS uq_cron_jobs_agent_tenant_name; diff --git a/migrations/000047_cron_jobs_unique_constraint.up.sql b/migrations/000047_cron_jobs_unique_constraint.up.sql new file mode 100644 index 00000000..c18168d8 --- /dev/null +++ b/migrations/000047_cron_jobs_unique_constraint.up.sql @@ -0,0 +1,19 @@ +-- Deduplicate cron_jobs: keep most recently updated row per (agent_id, tenant_id, name). +DELETE FROM cron_jobs +WHERE id NOT IN ( + SELECT DISTINCT ON (agent_id, tenant_id, name) id + FROM cron_jobs + ORDER BY agent_id, tenant_id, name, updated_at DESC +); + +-- Create unique index (will succeed after dedup). +CREATE UNIQUE INDEX IF NOT EXISTS uq_cron_jobs_agent_tenant_name + ON cron_jobs (agent_id, tenant_id, name); + +-- Generated column for path basename (case-insensitive wikilink resolution). +ALTER TABLE vault_documents ADD COLUMN IF NOT EXISTS path_basename TEXT + GENERATED ALWAYS AS (lower(regexp_replace(path, '.+/', ''))) STORED; + +-- Index for fast basename lookup by tenant. +CREATE INDEX IF NOT EXISTS idx_vault_docs_basename + ON vault_documents(tenant_id, path_basename);