diff --git a/cmd/spinifex/cmd/admin_ochre_appliance.go b/cmd/spinifex/cmd/admin_ochre_appliance.go new file mode 100644 index 000000000..302997dda --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_appliance.go @@ -0,0 +1,87 @@ +package cmd + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/spf13/cobra" +) + +// applianceTeardownTimeout bounds the teardown NATS round trip: it covers an +// RDS DeleteDBInstance call plus a KV delete, not a VM boot, so it is far +// shorter than the appliance's own launch timeout. +const applianceTeardownTimeout = 5 * time.Minute + +var ochreApplianceCmd = &cobra.Command{ + Use: "appliance", + Short: "Manage the shared platform Postgres appliance backing every Ochre vector index", +} + +var ochreApplianceTeardownCmd = &cobra.Command{ + Use: "teardown", + Short: "Destroy the shared platform Postgres appliance", + Long: `teardown destroys the platform appliance singleton: its backing RDS DB +instance, this node's VPC host port, and its KV record. + +This is a destructive, cluster-wide operator action -- every account's vector +index data stored in the appliance is lost with it, and it is not +automatically rebuilt. A future daemon startup re-provisions a fresh +appliance (a new instance, a new password) only if Ochre remains enabled; +existing index registry records are not restored. + +Requires --confirm.`, + Run: runOchreApplianceTeardown, +} + +func init() { + ochreCmd.AddCommand(ochreApplianceCmd) + ochreApplianceCmd.AddCommand(ochreApplianceTeardownCmd) + + ochreApplianceTeardownCmd.Flags().Bool("confirm", false, + "Required: confirms you intend to destroy the shared platform appliance and every account's vector index data with it") +} + +// applianceTeardownFn indirects the NATS call so runApplianceTeardownGuarded +// is testable without a live daemon, mirroring vectorServiceFn/ +// endpointServiceFn. +var applianceTeardownFn = func(ctx context.Context) error { + _, nc, err := loadConfigAndConnectFn() + if err != nil { + return err + } + defer nc.Close() + _, err = utils.NATSRequest[handlers_ochrevector.TeardownApplianceResponse](ctx, nc, + handlers_ochrevector.SubjectTeardownAppliance, &handlers_ochrevector.TeardownApplianceRequest{}, applianceTeardownTimeout, utils.GlobalAccountID) + return err +} + +// runApplianceTeardownGuarded is the testable core of 'ochre appliance +// teardown': it refuses without --confirm before making any NATS call, so a +// missing flag never even opens a connection to the cluster. +func runApplianceTeardownGuarded(ctx context.Context, confirmed bool, teardown func(context.Context) error) (string, error) { + if !confirmed { + return "", errors.New("refusing to tear down the shared platform appliance without --confirm: " + + "this destroys every account's vector index data and is not automatically rebuilt") + } + if err := teardown(ctx); err != nil { + return "", err + } + return "✅ Platform appliance torn down.", nil +} + +func runOchreApplianceTeardown(cmd *cobra.Command, _ []string) { + confirmed, _ := cmd.Flags().GetBool("confirm") + + msg, err := runApplianceTeardownGuarded(context.Background(), confirmed, applianceTeardownFn) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(msg) +} diff --git a/cmd/spinifex/cmd/admin_ochre_appliance_test.go b/cmd/spinifex/cmd/admin_ochre_appliance_test.go new file mode 100644 index 000000000..6eaa9ec0b --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_appliance_test.go @@ -0,0 +1,104 @@ +// Exercises unexported ochre appliance CLI command internals with no +// exported surface to drive them through. +// +//test:in-package +package cmd + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestRunApplianceTeardownGuarded_RefusesWithoutConfirm proves the guard +// refuses before ever invoking teardown -- a missing --confirm must never +// reach a NATS call, let alone destroy the appliance. +func TestRunApplianceTeardownGuarded_RefusesWithoutConfirm(t *testing.T) { + called := false + teardown := func(context.Context) error { + called = true + return nil + } + + _, err := runApplianceTeardownGuarded(context.Background(), false, teardown) + require.Error(t, err) + require.ErrorContains(t, err, "--confirm") + require.False(t, called, "teardown must not run without --confirm") +} + +// TestRunApplianceTeardownGuarded_ConfirmedInvokesTeardown proves --confirm +// drives the teardown call through and reports success. +func TestRunApplianceTeardownGuarded_ConfirmedInvokesTeardown(t *testing.T) { + called := false + teardown := func(context.Context) error { + called = true + return nil + } + + msg, err := runApplianceTeardownGuarded(context.Background(), true, teardown) + require.NoError(t, err) + require.True(t, called) + require.Contains(t, msg, "torn down") +} + +// TestRunApplianceTeardownGuarded_TeardownErrorSurfaces proves a teardown +// failure (e.g. the daemon reporting no appliance up) surfaces to the caller +// rather than being swallowed behind a generic success message. +func TestRunApplianceTeardownGuarded_TeardownErrorSurfaces(t *testing.T) { + teardown := func(context.Context) error { + return errors.New("ochrevector: platform appliance is not enabled or not up on this node") + } + + _, err := runApplianceTeardownGuarded(context.Background(), true, teardown) + require.Error(t, err) + require.ErrorContains(t, err, "not enabled or not up") +} + +// withApplianceTeardown swaps applianceTeardownFn for one that never dials +// NATS, so the Run wrapper's real flag/exit control flow is exercised +// without a live daemon. +func withApplianceTeardown(t *testing.T, fn func(context.Context) error) { + t.Helper() + orig := applianceTeardownFn + t.Cleanup(func() { applianceTeardownFn = orig }) + applianceTeardownFn = fn +} + +func TestRunOchreApplianceTeardown_MissingConfirmExits1(t *testing.T) { + withApplianceTeardown(t, func(context.Context) error { + t.Fatal("teardown must not be called without --confirm") + return nil + }) + + cmd := *ochreApplianceTeardownCmd + code := withOchreExitCapture(t, func() { runOchreApplianceTeardown(&cmd, nil) }) + require.Equal(t, 1, code) +} + +func TestRunOchreApplianceTeardown_ConfirmedPrintsSuccess(t *testing.T) { + withApplianceTeardown(t, func(context.Context) error { return nil }) + + cmd := *ochreApplianceTeardownCmd + require.NoError(t, cmd.Flags().Set("confirm", "true")) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreApplianceTeardown(&cmd, nil) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, "torn down") +} + +func TestRunOchreApplianceTeardown_TeardownErrorExits1(t *testing.T) { + withApplianceTeardown(t, func(context.Context) error { + return errors.New("ochrevector: platform appliance is not enabled or not up on this node") + }) + + cmd := *ochreApplianceTeardownCmd + require.NoError(t, cmd.Flags().Set("confirm", "true")) + + code := withOchreExitCapture(t, func() { runOchreApplianceTeardown(&cmd, nil) }) + require.Equal(t, 1, code) +} diff --git a/cmd/spinifex/cmd/admin_ochre_vector.go b/cmd/spinifex/cmd/admin_ochre_vector.go new file mode 100644 index 000000000..2d8f6879c --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_vector.go @@ -0,0 +1,450 @@ +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strconv" + "time" + + handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/pterm/pterm" + "github.com/spf13/cobra" +) + +// defaultVectorEmbeddingModel is the platform's default embedding model for +// a new index when --model is not given. +const defaultVectorEmbeddingModel = "nomic-embed-text-v1.5" + +// cliChunkPreviewChars bounds a chunk's preview in the query results table +// (D10 payload-size discipline extended to terminal readability): the full, +// server-truncated chunk text is always available via --json. +const cliChunkPreviewChars = 120 + +var ochreVectorCmd = &cobra.Command{ + Use: "vector", + Short: "Manage Ochre tenant vector-store indexes, ingestion and queries", + Long: `Operator surface over the daemon's tenant vector store: create/delete/list +indexes, start an ingestion job from a predastore bucket, check a job's +status, and run a similarity query -- the live-verify surface for the +internal ochre.vector.* NATS API.`, +} + +var ochreVectorIndexCmd = &cobra.Command{ + Use: "index", + Short: "Manage vector index lifecycle", +} + +var ochreVectorIndexCreateCmd = &cobra.Command{ + Use: "create", + Short: "Create a new vector index", + Long: `create reserves a new index under the platform account, provisions its +account schema and backing table, and returns once it is READY. The index's +embedding model and dimension are fixed for its lifetime (D6): every ingest +and query against it uses exactly these values, regardless of what a caller +sends.`, + Run: runOchreVectorIndexCreate, +} + +var ochreVectorIndexDeleteCmd = &cobra.Command{ + Use: "delete ", + Short: "Delete a vector index", + Long: `delete drops the index's backing table and its registry record. Idempotent: +deleting an already-absent index reports success.`, + Args: cobra.ExactArgs(1), + Run: runOchreVectorIndexDelete, +} + +var ochreVectorIndexListCmd = &cobra.Command{ + Use: "list", + Short: "List every vector index", + Run: runOchreVectorIndexList, +} + +var ochreVectorIngestCmd = &cobra.Command{ + Use: "ingest", + Short: "Start an ingestion job for a vector index from a predastore bucket", + Long: `ingest lists --bucket/--prefix, chunks and embeds each object, and upserts +its rows into --index. The embedding model and dimension used are always the +index's own pinned values, not any --model flag -- an index's embedding +model is a property of the index, not of one ingest call. Returns as soon as +the job is claimed (PENDING); use 'job describe' to follow its progress.`, + Run: runOchreVectorIngest, +} + +var ochreVectorJobCmd = &cobra.Command{ + Use: "job", + Short: "Inspect vector-store ingestion jobs", +} + +var ochreVectorJobDescribeCmd = &cobra.Command{ + Use: "describe ", + Short: "Show an ingestion job's current record", + Args: cobra.ExactArgs(1), + Run: runOchreVectorJobDescribe, +} + +var ochreVectorQueryCmd = &cobra.Command{ + Use: "query", + Short: "Run a similarity query against a vector index", + Long: `query embeds --text server-side against --index's own pinned embedding model +(D8) and returns the nearest chunks, nearest first. --filter takes a +JSON-encoded metadata Filter AST (Op/Key/Value/Children). Long chunk text is +truncated in the default table; pass --json for the full response.`, + Run: runOchreVectorQuery, +} + +func init() { + ochreCmd.AddCommand(ochreVectorCmd) + ochreVectorCmd.AddCommand(ochreVectorIndexCmd) + ochreVectorIndexCmd.AddCommand(ochreVectorIndexCreateCmd) + ochreVectorIndexCmd.AddCommand(ochreVectorIndexDeleteCmd) + ochreVectorIndexCmd.AddCommand(ochreVectorIndexListCmd) + ochreVectorCmd.AddCommand(ochreVectorIngestCmd) + ochreVectorCmd.AddCommand(ochreVectorJobCmd) + ochreVectorJobCmd.AddCommand(ochreVectorJobDescribeCmd) + ochreVectorCmd.AddCommand(ochreVectorQueryCmd) + + ochreVectorIndexCreateCmd.Flags().String("name", "", "Human-readable index name (required)") + ochreVectorIndexCreateCmd.Flags().Int("dimension", 0, "Embedding vector dimension (required)") + ochreVectorIndexCreateCmd.Flags().String("model", defaultVectorEmbeddingModel, "Embedding model this index is pinned to") + _ = ochreVectorIndexCreateCmd.MarkFlagRequired("name") + _ = ochreVectorIndexCreateCmd.MarkFlagRequired("dimension") + + ochreVectorIngestCmd.Flags().String("index", "", "Target index ID (required)") + ochreVectorIngestCmd.Flags().String("bucket", "", "Source predastore bucket (required)") + ochreVectorIngestCmd.Flags().String("prefix", "", "Source object key prefix") + ochreVectorIngestCmd.Flags().Int("chunk-size", 0, "Chunk size in runes (0 takes the package default)") + ochreVectorIngestCmd.Flags().Int("chunk-overlap", 0, "Chunk overlap in runes (0 takes the package default)") + ochreVectorIngestCmd.Flags().StringToString("meta", nil, "Static metadata tag stamped on every ingested row, k=v (repeatable)") + _ = ochreVectorIngestCmd.MarkFlagRequired("index") + _ = ochreVectorIngestCmd.MarkFlagRequired("bucket") + + ochreVectorQueryCmd.Flags().String("index", "", "Index ID to query (required)") + ochreVectorQueryCmd.Flags().String("text", "", "Query text, embedded server-side against the index's pinned model (required)") + ochreVectorQueryCmd.Flags().Int("k", 0, "Number of results (0 takes the service default; hard-capped at 100)") + ochreVectorQueryCmd.Flags().String("filter", "", "JSON-encoded metadata Filter AST, e.g. {\"Op\":\"equals\",\"Key\":\"category\",\"Value\":\"faq\"}") + ochreVectorQueryCmd.Flags().Bool("json", false, "Print full, untruncated results as JSON") + _ = ochreVectorQueryCmd.MarkFlagRequired("index") + _ = ochreVectorQueryCmd.MarkFlagRequired("text") +} + +// vectorServiceFn indirects the NATS-backed client so the Run functions' +// connect/exit control flow can be tested without a live daemon, mirroring +// endpointServiceFn. +var vectorServiceFn = func() (handlers_ochrevector.VectorService, func(), error) { + _, nc, err := loadConfigAndConnectFn() + if err != nil { + return nil, nil, err + } + return handlers_ochrevector.NewNATSVectorService(nc), nc.Close, nil +} + +// formatIndexRecord renders one index record as aligned key/value lines. +func formatIndexRecord(rec handlers_ochrevector.Record) string { + rows := [][2]string{ + {"Index ID", rec.ID}, + {"Name", rec.Name}, + {"State", rec.State}, + {"Dimension", strconv.Itoa(rec.Dimension)}, + {"Embedding model", rec.EmbeddingModel}, + } + if !rec.CreatedAt.IsZero() { + rows = append(rows, [2]string{"Created at", rec.CreatedAt.Format(time.RFC3339)}) + } + out := "" + for _, row := range rows { + out += fmt.Sprintf("%-16s %s\n", row[0]+":", row[1]) + } + return out +} + +// formatJobRecord renders one ingestion job record as aligned key/value lines. +func formatJobRecord(job handlers_ochrevector.JobRecord) string { + rows := [][2]string{ + {"Job ID", job.ID}, + {"Index ID", job.IndexID}, + {"State", job.State}, + {"Documents", fmt.Sprintf("%d/%d", job.DocumentsDone, job.DocumentsTotal)}, + } + if len(job.FailedDocuments) > 0 { + rows = append(rows, [2]string{"Failed docs", strconv.Itoa(len(job.FailedDocuments))}) + } + if job.Error != "" { + rows = append(rows, [2]string{"Error", job.Error}) + } + if !job.CreatedAt.IsZero() { + rows = append(rows, [2]string{"Created at", job.CreatedAt.Format(time.RFC3339)}) + } + if !job.UpdatedAt.IsZero() { + rows = append(rows, [2]string{"Updated at", job.UpdatedAt.Format(time.RFC3339)}) + } + out := "" + for _, row := range rows { + out += fmt.Sprintf("%-14s %s\n", row[0]+":", row[1]) + } + return out +} + +// truncateForTable bounds s to maxRunes runes for a table cell, so one +// oversized chunk cannot blow out terminal width; the untruncated text is +// always available via --json. +func truncateForTable(s string, maxRunes int) string { + runes := []rune(s) + if len(runes) <= maxRunes { + return s + } + return string(runes[:maxRunes]) + "…" +} + +// formatQueryResults renders 'ochre vector query' output: a truncated table +// by default, or the full response as indented JSON when asJSON is set +// (D10's "full text via source ref" -- here, via --json). +func formatQueryResults(results []handlers_ochrevector.QueryResult, asJSON bool) (string, error) { + if asJSON { + //nolint:musttag // QueryResult is an internal, non-wire Go type (predates this CLI); it round-trips fine on its default field names. + data, err := json.MarshalIndent(results, "", " ") + if err != nil { + return "", fmt.Errorf("encode results as JSON: %w", err) + } + return string(data), nil + } + if len(results) == 0 { + return "No results.", nil + } + tableData := pterm.TableData{{"SCORE", "SOURCE", "OFFSET", "CHUNK"}} + for _, r := range results { + tableData = append(tableData, []string{ + strconv.FormatFloat(float64(r.Score), 'f', 4, 32), + r.SourceKey, + strconv.Itoa(r.SourceOffset), + truncateForTable(r.Chunk, cliChunkPreviewChars), + }) + } + return pterm.DefaultTable.WithHasHeader().WithLeftAlignment().WithData(tableData).Srender() +} + +// parseFilterFlag decodes --filter's JSON-encoded Filter AST. An empty raw +// string is not an error: it means no filter, matching QueryRequest.Filter's +// nil-is-unfiltered contract. +func parseFilterFlag(raw string) (*handlers_ochrevector.Filter, error) { + if raw == "" { + return nil, nil + } + var f handlers_ochrevector.Filter + //nolint:musttag // Filter is the existing filter-AST type (D9); it round-trips fine on its default field names. + if err := json.Unmarshal([]byte(raw), &f); err != nil { + return nil, fmt.Errorf("decode --filter: %w", err) + } + return &f, nil +} + +// runIndexCreate is the testable core of 'ochre vector index create': it +// mints a fresh index ID (indexes are named, not caller-numbered) and +// creates it. Returns the message to print. +func runIndexCreate(ctx context.Context, svc handlers_ochrevector.VectorService, name string, dimension int, model string) (string, error) { + out, err := svc.CreateIndex(ctx, &handlers_ochrevector.CreateIndexRequest{ + IndexID: utils.GenerateResourceID("idx"), + Name: name, + Dimension: dimension, + EmbeddingModel: model, + }, utils.GlobalAccountID) + if err != nil { + return "", err + } + return fmt.Sprintf("✅ Created index %s.\n\n%s", out.Index.ID, formatIndexRecord(out.Index)), nil +} + +// runIndexDelete is the testable core of 'ochre vector index delete'. +func runIndexDelete(ctx context.Context, svc handlers_ochrevector.VectorService, indexID string) (string, error) { + if _, err := svc.DeleteIndex(ctx, &handlers_ochrevector.DeleteIndexRequest{IndexID: indexID}, utils.GlobalAccountID); err != nil { + return "", err + } + return fmt.Sprintf("✅ Index %s deleted.", indexID), nil +} + +// listIndexesOutput renders 'ochre vector index list'. Split from its Run +// function so it is testable against a fake service with no NATS connection. +func listIndexesOutput(ctx context.Context, svc handlers_ochrevector.VectorService) (string, error) { + out, err := svc.ListIndexes(ctx, &handlers_ochrevector.ListIndexesRequest{}, utils.GlobalAccountID) + if err != nil { + return "", err + } + if len(out.Indexes) == 0 { + return "No vector indexes.", nil + } + tableData := pterm.TableData{{"INDEX ID", "NAME", "STATE", "DIMENSION", "MODEL"}} + for _, r := range out.Indexes { + tableData = append(tableData, []string{r.ID, r.Name, r.State, strconv.Itoa(r.Dimension), r.EmbeddingModel}) + } + return pterm.DefaultTable.WithHasHeader().WithLeftAlignment().WithData(tableData).Srender() +} + +// runIngest is the testable core of 'ochre vector ingest'. The daemon stamps +// the index's own embedding model/dimension onto the source spec server-side +// (see handlers_ochrevector.vectorService.Ingest), so none is sent here. +func runIngest(ctx context.Context, svc handlers_ochrevector.VectorService, indexID, bucket, prefix string, chunkSize, chunkOverlap int, meta map[string]string) (string, error) { + out, err := svc.Ingest(ctx, &handlers_ochrevector.IngestRequest{ + IndexID: indexID, + Source: handlers_ochrevector.SourceSpec{ + Bucket: bucket, + Prefix: prefix, + ChunkSize: chunkSize, + ChunkOverlap: chunkOverlap, + Metadata: meta, + }, + }, utils.GlobalAccountID) + if err != nil { + return "", err + } + return fmt.Sprintf("✅ Started ingestion job %s for index %s (state %s).", out.Job.ID, indexID, out.Job.State), nil +} + +// runJobDescribe is the testable core of 'ochre vector job describe'. +func runJobDescribe(ctx context.Context, svc handlers_ochrevector.VectorService, jobID string) (string, error) { + out, err := svc.DescribeJob(ctx, &handlers_ochrevector.DescribeJobRequest{JobID: jobID}, utils.GlobalAccountID) + if err != nil { + return "", err + } + return formatJobRecord(out.Job), nil +} + +// runQuery is the testable core of 'ochre vector query'. +func runQuery(ctx context.Context, svc handlers_ochrevector.VectorService, indexID, text string, k int, filterJSON string, asJSON bool) (string, error) { + filter, err := parseFilterFlag(filterJSON) + if err != nil { + return "", err + } + out, err := svc.Query(ctx, &handlers_ochrevector.QueryRequest{IndexID: indexID, Text: text, K: k, Filter: filter}, utils.GlobalAccountID) + if err != nil { + return "", err + } + return formatQueryResults(out.Results, asJSON) +} + +func runOchreVectorIndexCreate(cmd *cobra.Command, _ []string) { + name, _ := cmd.Flags().GetString("name") + dimension, _ := cmd.Flags().GetInt("dimension") + model, _ := cmd.Flags().GetString("model") + + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + msg, err := runIndexCreate(context.Background(), svc, name, dimension, model) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(msg) +} + +func runOchreVectorIndexDelete(_ *cobra.Command, args []string) { + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + msg, err := runIndexDelete(context.Background(), svc, args[0]) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(msg) +} + +func runOchreVectorIndexList(_ *cobra.Command, _ []string) { + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + out, err := listIndexesOutput(context.Background(), svc) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(out) +} + +func runOchreVectorIngest(cmd *cobra.Command, _ []string) { + indexID, _ := cmd.Flags().GetString("index") + bucket, _ := cmd.Flags().GetString("bucket") + prefix, _ := cmd.Flags().GetString("prefix") + chunkSize, _ := cmd.Flags().GetInt("chunk-size") + chunkOverlap, _ := cmd.Flags().GetInt("chunk-overlap") + meta, _ := cmd.Flags().GetStringToString("meta") + + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + msg, err := runIngest(context.Background(), svc, indexID, bucket, prefix, chunkSize, chunkOverlap, meta) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(msg) +} + +func runOchreVectorJobDescribe(_ *cobra.Command, args []string) { + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + out, err := runJobDescribe(context.Background(), svc, args[0]) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Print(out) +} + +func runOchreVectorQuery(cmd *cobra.Command, _ []string) { + indexID, _ := cmd.Flags().GetString("index") + text, _ := cmd.Flags().GetString("text") + k, _ := cmd.Flags().GetInt("k") + filterJSON, _ := cmd.Flags().GetString("filter") + asJSON, _ := cmd.Flags().GetBool("json") + + svc, closeFn, err := vectorServiceFn() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer closeFn() + + out, err := runQuery(context.Background(), svc, indexID, text, k, filterJSON, asJSON) + if err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + ochreExit(1) + return + } + fmt.Println(out) +} diff --git a/cmd/spinifex/cmd/admin_ochre_vector_test.go b/cmd/spinifex/cmd/admin_ochre_vector_test.go new file mode 100644 index 000000000..8d2a44cee --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_vector_test.go @@ -0,0 +1,375 @@ +// Exercises unexported ochre vector CLI command internals with no +// exported surface to drive them through. +// +//test:in-package +package cmd + +import ( + "context" + "errors" + "strings" + "testing" + + handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" + "github.com/stretchr/testify/require" +) + +// fakeVectorService serves scripted responses so the CLI's testable cores +// can be exercised without a daemon or a live NATS connection. +type fakeVectorService struct { + createIndex handlers_ochrevector.CreateIndexResponse + createIndexErr error + deleteIndexErr error + list []handlers_ochrevector.Record + listErr error + ingest handlers_ochrevector.IngestResponse + ingestErr error + describeJob handlers_ochrevector.JobRecord + describeErr error + query []handlers_ochrevector.QueryResult + queryErr error + + createIndexInputs []*handlers_ochrevector.CreateIndexRequest + deleteIndexInputs []*handlers_ochrevector.DeleteIndexRequest + ingestInputs []*handlers_ochrevector.IngestRequest + queryInputs []*handlers_ochrevector.QueryRequest +} + +func (f *fakeVectorService) CreateIndex(_ context.Context, in *handlers_ochrevector.CreateIndexRequest, _ string) (*handlers_ochrevector.CreateIndexResponse, error) { + f.createIndexInputs = append(f.createIndexInputs, in) + if f.createIndexErr != nil { + return nil, f.createIndexErr + } + out := f.createIndex + if out.Index.ID == "" { + out.Index.ID = in.IndexID + } + return &out, nil +} + +func (f *fakeVectorService) DeleteIndex(_ context.Context, in *handlers_ochrevector.DeleteIndexRequest, _ string) (*handlers_ochrevector.DeleteIndexResponse, error) { + f.deleteIndexInputs = append(f.deleteIndexInputs, in) + if f.deleteIndexErr != nil { + return nil, f.deleteIndexErr + } + return &handlers_ochrevector.DeleteIndexResponse{}, nil +} + +func (f *fakeVectorService) ListIndexes(_ context.Context, _ *handlers_ochrevector.ListIndexesRequest, _ string) (*handlers_ochrevector.ListIndexesResponse, error) { + if f.listErr != nil { + return nil, f.listErr + } + return &handlers_ochrevector.ListIndexesResponse{Indexes: f.list}, nil +} + +func (f *fakeVectorService) Ingest(_ context.Context, in *handlers_ochrevector.IngestRequest, _ string) (*handlers_ochrevector.IngestResponse, error) { + f.ingestInputs = append(f.ingestInputs, in) + if f.ingestErr != nil { + return nil, f.ingestErr + } + return &f.ingest, nil +} + +func (f *fakeVectorService) DescribeJob(_ context.Context, _ *handlers_ochrevector.DescribeJobRequest, _ string) (*handlers_ochrevector.DescribeJobResponse, error) { + if f.describeErr != nil { + return nil, f.describeErr + } + return &handlers_ochrevector.DescribeJobResponse{Job: f.describeJob}, nil +} + +func (f *fakeVectorService) Query(_ context.Context, in *handlers_ochrevector.QueryRequest, _ string) (*handlers_ochrevector.QueryResponse, error) { + f.queryInputs = append(f.queryInputs, in) + if f.queryErr != nil { + return nil, f.queryErr + } + return &handlers_ochrevector.QueryResponse{Results: f.query}, nil +} + +var _ handlers_ochrevector.VectorService = (*fakeVectorService)(nil) + +const testIndexID = "idx-0123456789abcdef0" + +func TestRunIndexCreate_MintsIDAndPrintsRecord(t *testing.T) { + svc := &fakeVectorService{createIndex: handlers_ochrevector.CreateIndexResponse{ + Index: handlers_ochrevector.Record{Name: "kb1", State: handlers_ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"}, + }} + + msg, err := runIndexCreate(context.Background(), svc, "kb1", 768, "nomic-embed-text-v1.5") + require.NoError(t, err) + require.Contains(t, msg, "READY") + require.Contains(t, msg, "nomic-embed-text-v1.5") + + require.Len(t, svc.createIndexInputs, 1) + require.NotEmpty(t, svc.createIndexInputs[0].IndexID, "the CLI must mint an index ID, not leave it empty") + require.Equal(t, 768, svc.createIndexInputs[0].Dimension) +} + +func TestRunIndexCreate_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{createIndexErr: errors.New("ValidationException")} + _, err := runIndexCreate(context.Background(), svc, "kb1", 768, "m") + require.ErrorContains(t, err, "ValidationException") +} + +func TestRunIndexDelete_ReportsDeletion(t *testing.T) { + svc := &fakeVectorService{} + msg, err := runIndexDelete(context.Background(), svc, testIndexID) + require.NoError(t, err) + require.Contains(t, msg, testIndexID) + require.Len(t, svc.deleteIndexInputs, 1) + require.Equal(t, testIndexID, svc.deleteIndexInputs[0].IndexID) +} + +func TestRunIndexDelete_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{deleteIndexErr: errors.New("nats: timeout")} + _, err := runIndexDelete(context.Background(), svc, testIndexID) + require.ErrorContains(t, err, "timeout") +} + +func TestListIndexesOutput_NoIndexes(t *testing.T) { + msg, err := listIndexesOutput(context.Background(), &fakeVectorService{}) + require.NoError(t, err) + require.Equal(t, "No vector indexes.", msg) +} + +func TestListIndexesOutput_ListsIndexes(t *testing.T) { + svc := &fakeVectorService{list: []handlers_ochrevector.Record{ + {ID: testIndexID, Name: "kb1", State: handlers_ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"}, + }} + msg, err := listIndexesOutput(context.Background(), svc) + require.NoError(t, err) + require.Contains(t, msg, testIndexID) + require.Contains(t, msg, "READY") +} + +func TestListIndexesOutput_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{listErr: errors.New("nats: no responders")} + _, err := listIndexesOutput(context.Background(), svc) + require.ErrorContains(t, err, "no responders") +} + +func TestRunIngest_NoModelFlagSentByCLI(t *testing.T) { + svc := &fakeVectorService{ingest: handlers_ochrevector.IngestResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStatePending}, + }} + + msg, err := runIngest(context.Background(), svc, testIndexID, "docs", "kb/", 100, 10, map[string]string{"team": "eng"}) + require.NoError(t, err) + require.Contains(t, msg, "job-abc") + require.Contains(t, msg, "PENDING") + + require.Len(t, svc.ingestInputs, 1) + in := svc.ingestInputs[0] + require.Equal(t, testIndexID, in.IndexID) + require.Equal(t, "docs", in.Source.Bucket) + require.Equal(t, "kb/", in.Source.Prefix) + require.Equal(t, 100, in.Source.ChunkSize) + require.Equal(t, "eng", in.Source.Metadata["team"]) + // The CLI has no --model flag: the daemon stamps the index's pinned + // model server-side, so the request's EmbeddingModel is left zero. + require.Empty(t, in.Source.EmbeddingModel) +} + +func TestRunIngest_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{ingestErr: errors.New("ResourceNotFoundException")} + _, err := runIngest(context.Background(), svc, "no-such-index", "docs", "", 0, 0, nil) + require.ErrorContains(t, err, "ResourceNotFoundException") +} + +func TestRunJobDescribe_PrintsRecord(t *testing.T) { + svc := &fakeVectorService{describeJob: handlers_ochrevector.JobRecord{ + ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStateReady, DocumentsDone: 3, DocumentsTotal: 3, + }} + out, err := runJobDescribe(context.Background(), svc, "job-abc") + require.NoError(t, err) + require.Contains(t, out, "job-abc") + require.Contains(t, out, "3/3") +} + +func TestRunJobDescribe_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{describeErr: errors.New("ochrevector: job not found")} + _, err := runJobDescribe(context.Background(), svc, "no-such-job") + require.ErrorContains(t, err, "job not found") +} + +func TestRunQuery_TableTruncatesLongChunk(t *testing.T) { + longChunk := strings.Repeat("x", cliChunkPreviewChars+50) + svc := &fakeVectorService{query: []handlers_ochrevector.QueryResult{ + {Chunk: longChunk, SourceKey: "docs/a.txt", SourceOffset: 0, Score: 0.9}, + }} + + out, err := runQuery(context.Background(), svc, testIndexID, "hello", 5, "", false) + require.NoError(t, err) + require.Contains(t, out, "docs/a.txt") + require.NotContains(t, out, longChunk, "the table must not print the full untruncated chunk") + + require.Len(t, svc.queryInputs, 1) + require.Equal(t, testIndexID, svc.queryInputs[0].IndexID) + require.Equal(t, "hello", svc.queryInputs[0].Text) + require.Equal(t, 5, svc.queryInputs[0].K) +} + +func TestRunQuery_JSONFlagPrintsFullChunk(t *testing.T) { + longChunk := strings.Repeat("x", cliChunkPreviewChars+50) + svc := &fakeVectorService{query: []handlers_ochrevector.QueryResult{ + {Chunk: longChunk, SourceKey: "docs/a.txt", Score: 0.9}, + }} + + out, err := runQuery(context.Background(), svc, testIndexID, "hello", 0, "", true) + require.NoError(t, err) + require.Contains(t, out, longChunk, "--json must print the full, untruncated chunk") +} + +func TestRunQuery_NoResults(t *testing.T) { + out, err := runQuery(context.Background(), &fakeVectorService{}, testIndexID, "hello", 0, "", false) + require.NoError(t, err) + require.Equal(t, "No results.", out) +} + +func TestRunQuery_FilterFlagReachesRequest(t *testing.T) { + svc := &fakeVectorService{} + filterJSON := `{"Op":"equals","Key":"category","Value":"faq"}` + + _, err := runQuery(context.Background(), svc, testIndexID, "hello", 0, filterJSON, false) + require.NoError(t, err) + require.Len(t, svc.queryInputs, 1) + require.NotNil(t, svc.queryInputs[0].Filter) + require.Equal(t, handlers_ochrevector.FilterEquals, svc.queryInputs[0].Filter.Op) + require.Equal(t, "category", svc.queryInputs[0].Filter.Key) +} + +func TestRunQuery_InvalidFilterJSONErrors(t *testing.T) { + _, err := runQuery(context.Background(), &fakeVectorService{}, testIndexID, "hello", 0, "{not json", false) + require.Error(t, err) +} + +func TestRunQuery_ErrorSurfaces(t *testing.T) { + svc := &fakeVectorService{queryErr: errors.New("nats: timeout")} + _, err := runQuery(context.Background(), svc, testIndexID, "hello", 0, "", false) + require.ErrorContains(t, err, "timeout") +} + +// withVectorService swaps vectorServiceFn for one returning svc, so a Run +// wrapper's real connect/validate/exit control flow runs without a daemon. +func withVectorService(t *testing.T, svc handlers_ochrevector.VectorService, connErr error) { + t.Helper() + orig := vectorServiceFn + t.Cleanup(func() { vectorServiceFn = orig }) + vectorServiceFn = func() (handlers_ochrevector.VectorService, func(), error) { + if connErr != nil { + return nil, nil, connErr + } + return svc, func() {}, nil + } +} + +func TestRunOchreVectorIndexCreate_PrintsRecord(t *testing.T) { + withVectorService(t, &fakeVectorService{createIndex: handlers_ochrevector.CreateIndexResponse{ + Index: handlers_ochrevector.Record{ID: testIndexID, State: handlers_ochrevector.StateReady}, + }}, nil) + + cmd := *ochreVectorIndexCreateCmd + require.NoError(t, cmd.Flags().Set("name", "kb1")) + require.NoError(t, cmd.Flags().Set("dimension", "768")) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorIndexCreate(&cmd, nil) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, testIndexID) +} + +func TestRunOchreVectorIndexCreate_ConnectFailureExits1(t *testing.T) { + withVectorService(t, nil, errors.New("dial nats: connection refused")) + + cmd := *ochreVectorIndexCreateCmd + require.NoError(t, cmd.Flags().Set("name", "kb1")) + require.NoError(t, cmd.Flags().Set("dimension", "768")) + + code := withOchreExitCapture(t, func() { runOchreVectorIndexCreate(&cmd, nil) }) + require.Equal(t, 1, code) +} + +func TestRunOchreVectorIndexDelete_ReportsDeletion(t *testing.T) { + svc := &fakeVectorService{} + withVectorService(t, svc, nil) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorIndexDelete(ochreVectorIndexDeleteCmd, []string{testIndexID}) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, testIndexID) + require.Len(t, svc.deleteIndexInputs, 1) +} + +func TestRunOchreVectorIndexList_PrintsTable(t *testing.T) { + withVectorService(t, &fakeVectorService{list: []handlers_ochrevector.Record{ + {ID: testIndexID, Name: "kb1", State: handlers_ochrevector.StateReady}, + }}, nil) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorIndexList(ochreVectorIndexListCmd, nil) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, testIndexID) +} + +func TestRunOchreVectorIngest_PrintsJobID(t *testing.T) { + withVectorService(t, &fakeVectorService{ingest: handlers_ochrevector.IngestResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStatePending}, + }}, nil) + + cmd := *ochreVectorIngestCmd + require.NoError(t, cmd.Flags().Set("index", testIndexID)) + require.NoError(t, cmd.Flags().Set("bucket", "docs")) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorIngest(&cmd, nil) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, "job-abc") +} + +func TestRunOchreVectorJobDescribe_PrintsRecord(t *testing.T) { + withVectorService(t, &fakeVectorService{describeJob: handlers_ochrevector.JobRecord{ + ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStateReady, + }}, nil) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorJobDescribe(ochreVectorJobDescribeCmd, []string{"job-abc"}) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, "job-abc") +} + +func TestRunOchreVectorQuery_PrintsTable(t *testing.T) { + withVectorService(t, &fakeVectorService{query: []handlers_ochrevector.QueryResult{ + {Chunk: "hello world", SourceKey: "docs/a.txt", Score: 0.9}, + }}, nil) + + cmd := *ochreVectorQueryCmd + require.NoError(t, cmd.Flags().Set("index", testIndexID)) + require.NoError(t, cmd.Flags().Set("text", "hello")) + + var out string + code := withOchreExitCapture(t, func() { + out = captureStdout(t, func() { runOchreVectorQuery(&cmd, nil) }) + }) + require.Equal(t, -1, code) + require.Contains(t, out, "docs/a.txt") +} + +func TestRunOchreVectorQuery_ConnectFailureExits1(t *testing.T) { + withVectorService(t, nil, errors.New("dial nats: connection refused")) + + cmd := *ochreVectorQueryCmd + require.NoError(t, cmd.Flags().Set("index", testIndexID)) + require.NoError(t, cmd.Flags().Set("text", "hello")) + + code := withOchreExitCapture(t, func() { runOchreVectorQuery(&cmd, nil) }) + require.Equal(t, 1, code) +} diff --git a/docs/service-interfaces.yaml b/docs/service-interfaces.yaml index e99b4e2d1..d4f7e916a 100644 --- a/docs/service-interfaces.yaml +++ b/docs/service-interfaces.yaml @@ -57,6 +57,7 @@ services: - spinifex/handlers/acm - spinifex/handlers/iam - spinifex/handlers/rds + - spinifex/handlers/ochrevector subscribes: # EC2 instance lifecycle. - "ec2.RunInstances" @@ -125,6 +126,14 @@ services: # {account}.{db-instance}, so the subject is built at runtime and only the # prefix is declarable. - "rds.bus.*" + # Ochre vector store tenant surface (handlers/ochrevector/api.go), + # gated behind config.OchreVectorConfig.Enabled (default off). + - "ochre.vector.createIndex" + - "ochre.vector.deleteIndex" + - "ochre.vector.listIndexes" + - "ochre.vector.ingest" + - "ochre.vector.describeJob" + - "ochre.vector.query" publishes: # → vpcd - "vpc.create" diff --git a/go.mod b/go.mod index 94149bc6c..d8b6c7c02 100644 --- a/go.mod +++ b/go.mod @@ -126,6 +126,9 @@ require ( github.com/hashicorp/raft v1.7.3 // indirect github.com/hashicorp/raft-boltdb/v2 v2.3.1 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.10.0 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/josharian/native v1.1.0 // indirect github.com/klauspost/compress v1.19.0 // indirect diff --git a/go.sum b/go.sum index f347e2111..487f14078 100644 --- a/go.sum +++ b/go.sum @@ -1171,6 +1171,12 @@ github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLf github.com/influxdata/influxdb1-client v0.0.0-20191209144304-8bf82d3c094d/go.mod h1:qj24IKcXYK6Iy9ceXlo3Tc+vtHo9lIhSX5JddghvEPo= github.com/insomniacslk/dhcp v0.0.0-20260407060928-11b94ed970f2 h1:G3irkWwmpl0vH/nn83K2AHqLUZweC7XAONuwXy/w9Co= github.com/insomniacslk/dhcp v0.0.0-20260407060928-11b94ed970f2/go.mod h1:qfvBmyDNp+/liLEYWRvqny/PEz9hGe2Dz833eXILSmo= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jezek/xgb v1.0.0/go.mod h1:nrhwO0FX/enq75I7Y7G8iN1ubpSGZEiA3v9e9GyRFlk= github.com/jezek/xgb v1.1.1/go.mod h1:nrhwO0FX/enq75I7Y7G8iN1ubpSGZEiA3v9e9GyRFlk= github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k= diff --git a/scripts/images/rds-postgres/manifest.conf b/scripts/images/rds-postgres/manifest.conf index d3f5400a4..761a0a446 100644 --- a/scripts/images/rds-postgres/manifest.conf +++ b/scripts/images/rds-postgres/manifest.conf @@ -10,8 +10,10 @@ IMAGE_SIZE="2G" # icu-data-full is deliberate: the cluster uses an ICU locale provider (musl has # no libc locales) and stock english-only data would reject other collations. # blkid is a hard requirement — without it rds-datadir refuses to run rather -# than risk formatting a customer's datadir. -APK_PACKAGES="postgresql18 postgresql18-client postgresql18-contrib postgresql18-openrc postgresql-common-openrc icu-data-full ca-certificates cloud-utils-growpart e2fsprogs-extra xfsprogs blkid openrc" +# than risk formatting a customer's datadir. postgresql-pgvector ships the vector +# extension for the Ochre appliance; it stays inert on tenant DBs until a +# CREATE EXTENSION runs, so baking it into the shared image is harmless. +APK_PACKAGES="postgresql18 postgresql18-client postgresql18-contrib postgresql-pgvector postgresql18-openrc postgresql-common-openrc icu-data-full ca-certificates cloud-utils-growpart e2fsprogs-extra xfsprogs blkid openrc" # rds-agent is the only built binary (static, FIPS boot like the host build). BUILD_COMMANDS="CGO_ENABLED=0 GOFIPS140=v1.0.0 go build -ldflags \"-s -w\" -o bin/rds-agent ./cmd/rds-agent" diff --git a/scripts/images/rds-postgres/rds-init b/scripts/images/rds-postgres/rds-init index db5b5b97c..a906b15b3 100644 --- a/scripts/images/rds-postgres/rds-init +++ b/scripts/images/rds-postgres/rds-init @@ -85,6 +85,11 @@ MOUNTS_FILE="${RDS_MOUNTS_FILE:-/proc/mounts}" # AWS, so a customer's `GRANT rds_superuser TO ...` works here too. MASTER_GROUP_ROLE="rds_superuser" +# The Ochre vector store appliance is the one platform (non-tenant) consumer of +# this image, provisioned with this fixed master name. Tenant instances take a +# different master, so the pgvector install below is scoped to the appliance. +OCHRE_APPLIANCE_MASTER="ochre_vector_admin" + # NAMEDATALEN-1: the server truncates a longer identifier, so the database # created would not be the one the customer asked for. MAX_DB_NAME_LEN=63 @@ -152,6 +157,11 @@ load_handoff() { ;; esac + # Scopes the platform pgvector install to the appliance; a tenant master + # takes a different name and skips it. + RDS_IS_OCHRE_APPLIANCE=0 + [ "${_master_lower}" = "${OCHRE_APPLIANCE_MASTER}" ] && RDS_IS_OCHRE_APPLIANCE=1 + # The same rule the control plane applies, for the same reason the reserved # name is re-checked: an unusable name reaches initdb, fails at CREATE # DATABASE and spends the one-shot password on the way. psql quotes the @@ -496,6 +506,9 @@ bootstrap_master() { if [ "${_rc}" -eq 0 ] && [ -n "${RDS_DB_NAME}" ]; then create_database || _rc=$? fi + if [ "${_rc}" -eq 0 ]; then + create_platform_extensions || _rc=$? + fi # `set -e` is suspended for a function called from an AND-OR list, so the # stop status is captured rather than propagated. A failed stop must also not @@ -631,6 +644,23 @@ CREATE DATABASE :"dbname" OWNER :"master"; SQL } +# pgvector is baked in but left uninstalled so tenant databases stay free of it. +# The appliance master is NOSUPERUSER, so the bootstrap superuser installs the +# extension into a read-only extensions schema and grants the master the database +# and role privileges its daemon needs to create a per-account schema and role. +# Scoped to the appliance master alone; tenant masters skip this entirely. +create_platform_extensions() { + [ "${RDS_IS_OCHRE_APPLIANCE:-0}" = "1" ] || return 0 + log "configuring pgvector and platform privileges for the Ochre vector appliance" + psql_bootstrap < ochreStartupMaxBackoff { + backoff = ochreStartupMaxBackoff + } + } + return nil, lastErr +} + +// ensureAndConnectOchreApplianceOnce runs a single Ensure-then-Connect +// attempt, bounding Ensure by ochreApplianceLaunchTimeout the same way the +// original one-shot call did; Connect is bounded by d.ctx alone, since it +// does its own network dial rather than a long poll loop. +func (d *Daemon) ensureAndConnectOchreApplianceOnce(appliance *handlers_ochrevector.Appliance) (handlers_ochrevector.VectorBackend, error) { + ensureCtx, cancel := context.WithTimeout(d.ctx, ochreApplianceLaunchTimeout) + defer cancel() + if _, err := appliance.Ensure(ensureCtx); err != nil { + return nil, fmt.Errorf("ensure platform appliance: %w", err) + } + backend, err := appliance.Connect(d.ctx) + if err != nil { + return nil, fmt.Errorf("connect to platform appliance: %w", err) + } + return backend, nil +} diff --git a/spinifex/daemon/ochre_deps_test.go b/spinifex/daemon/ochre_deps_test.go new file mode 100644 index 000000000..8a7c566b2 --- /dev/null +++ b/spinifex/daemon/ochre_deps_test.go @@ -0,0 +1,48 @@ +// Exercises unexported ochre daemon wiring with no exported surface to +// drive it through. +// +//test:in-package +package daemon + +import ( + "context" + "testing" + + "github.com/mulgadc/spinifex/spinifex/config" + handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestStartOchreVector_DisabledSkipsConstruction pins the config gate's +// default-off behavior (D-series Stage 5b): with OchreVector.Enabled false +// (the zero value, i.e. every daemon that has not opted in), startOchreVector +// must construct and subscribe nothing, leaving d.ochreVectorService nil and +// d.natsSubscriptions untouched. No JetStream/NATS connection, ctx or +// natsSubscriptions map is set on d here, so any attempt to construct or +// register past the Enabled check would panic (nil natsConn, nil ctx, or a +// nil-map write) -- the test passing at all is itself part of the pin. +func TestStartOchreVector_DisabledSkipsConstruction(t *testing.T) { + d := &Daemon{ + config: &config.Config{}, + } + + d.startOchreVector() + + assert.Nil(t, d.ochreVectorService) + assert.Nil(t, d.natsSubscriptions, "disabled path must not touch the subscriptions map") +} + +// TestHandleOchreApplianceTeardown_NilApplianceRefuses proves the handler +// refuses with a clear error rather than a nil-pointer panic when the +// appliance never came up (disabled, still starting, or already torn down by +// an earlier call) -- it must never touch d.ochreVectorService in that case. +func TestHandleOchreApplianceTeardown_NilApplianceRefuses(t *testing.T) { + d := &Daemon{} + + out, err := d.handleOchreApplianceTeardown(context.Background(), &handlers_ochrevector.TeardownApplianceRequest{}, "") + + require.Error(t, err) + assert.Nil(t, out) + assert.Contains(t, err.Error(), "not enabled or not up") +} diff --git a/spinifex/daemon/vm_adapters.go b/spinifex/daemon/vm_adapters.go index e9f41cb42..11298de6f 100644 --- a/spinifex/daemon/vm_adapters.go +++ b/spinifex/daemon/vm_adapters.go @@ -620,6 +620,7 @@ func (d *Daemon) buildVMManagerDeps() vm.Deps { DetachDelay: d.detachDelay, DeviceDeletedTimeout: d.deviceDeletedTimeout, ConsumeCleanShutdownMarker: d.consumeCleanShutdownMarker(), + BackingStoreReady: func() bool { return d.checkPredastoreReady() && d.checkViperblockReady() }, } } diff --git a/spinifex/gateway/bedrock/embeddings.go b/spinifex/gateway/bedrock/embeddings.go new file mode 100644 index 000000000..830159090 --- /dev/null +++ b/spinifex/gateway/bedrock/embeddings.go @@ -0,0 +1,156 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +const embeddingsPath = "/v1/embeddings" + +// DefaultEmbeddingModel is Ochre's CPU-served embedding model. Callers pin a +// dimension per vector index; the adapter itself never hardcodes one and +// returns whatever the endpoint yields. +const DefaultEmbeddingModel = "nomic-embed-text-v1.5" + +// Embedder batch-embeds text against a self-hosted OpenAI/TEI-compatible +// endpoint, returning vectors ordered to match the caller's input slice. +type Embedder interface { + Embed(ctx context.Context, modelID string, inputs []string) ([][]float32, error) +} + +type embeddingsRequest struct { + Model string `json:"model"` + Input []string `json:"input"` +} + +type embeddingsDatum struct { + Index int `json:"index"` + Embedding []float32 `json:"embedding"` +} + +type embeddingsUsage struct { + PromptTokens int64 `json:"prompt_tokens"` + TotalTokens int64 `json:"total_tokens"` +} + +type embeddingsResponse struct { + Data []embeddingsDatum `json:"data"` + Model string `json:"model"` + Usage embeddingsUsage `json:"usage"` +} + +// embeddingsProvider serves batched text embeddings over the OpenAI/TEI +// /v1/embeddings wire. httpClient and endpointResolver are injectable for +// tests, mirroring vllmProvider. +type embeddingsProvider struct { + endpointResolver EndpointResolver + httpClient *http.Client +} + +var _ Embedder = (*embeddingsProvider)(nil) + +// newEmbeddingsProvider constructs an embeddingsProvider resolving endpoints +// via the shared EndpointResolver — the same static/registry-backed resolver +// vllmProvider uses, so the CPU embeddings endpoint never touches the +// tenant GPU serving lifecycle. +func newEmbeddingsProvider(endpointResolver EndpointResolver) *embeddingsProvider { + return &embeddingsProvider{ + endpointResolver: endpointResolver, + httpClient: &http.Client{Timeout: providerHTTPTimeout}, + } +} + +// NewEmbedder is newEmbeddingsProvider's exported constructor, for callers +// outside this package (e.g. the daemon's Ochre vector store wiring) that +// need an Embedder without reaching into gateway_bedrock's unexported types. +func NewEmbedder(endpointResolver EndpointResolver) Embedder { + return newEmbeddingsProvider(endpointResolver) +} + +// Embed resolves modelID's endpoint and batch-embeds inputs in a single +// request, returning vectors ordered to match inputs regardless of the order +// the endpoint's response lists them in. An empty inputs slice returns an +// empty result without calling the endpoint. +func (p *embeddingsProvider) Embed(ctx context.Context, modelID string, inputs []string) ([][]float32, error) { + if len(inputs) == 0 { + return nil, nil + } + + baseURL, ok, err := p.endpointResolver.Endpoint(ctx, modelID) + if err != nil { + slog.Error("embeddings: endpoint resolution failed", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + if !ok { + return nil, errors.New(awserrors.ErrorModelNotReadyException) + } + + reqBody, err := json.Marshal(embeddingsRequest{Model: modelID, Input: inputs}) + if err != nil { + slog.Error("embeddings: failed to marshal request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+embeddingsPath, bytes.NewReader(reqBody)) + if err != nil { + slog.Error("embeddings: failed to build request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := p.httpClient.Do(httpReq) + if err != nil { + slog.Error("embeddings: request failed", "model", modelID, "endpoint", baseURL, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + slog.Error("embeddings: failed to read response body", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + slog.Error("embeddings: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody)) + return nil, errors.New(mapUpstreamStatus(resp.StatusCode)) + } + + var er embeddingsResponse + if err := json.Unmarshal(respBody, &er); err != nil { + slog.Error("embeddings: failed to parse response", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorModelErrorException) + } + + vectors, err := orderEmbeddings(len(inputs), er.Data) + if err != nil { + slog.Error("embeddings: malformed response shape", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorModelErrorException) + } + return vectors, nil +} + +// orderEmbeddings sorts a batch response's vectors back to the caller's +// input order by the endpoint-reported index, rather than assuming the +// endpoint preserves request order. A count mismatch, an out-of-range index, +// or a duplicate index is treated as a malformed response. +func orderEmbeddings(n int, data []embeddingsDatum) ([][]float32, error) { + if len(data) != n { + return nil, errors.New("embeddings count mismatch") + } + out := make([][]float32, n) + for _, d := range data { + if d.Index < 0 || d.Index >= n || out[d.Index] != nil { + return nil, errors.New("embeddings index out of range or duplicated") + } + out[d.Index] = d.Embedding + } + return out, nil +} diff --git a/spinifex/gateway/bedrock/embeddings_test.go b/spinifex/gateway/bedrock/embeddings_test.go new file mode 100644 index 000000000..9c841ba7e --- /dev/null +++ b/spinifex/gateway/bedrock/embeddings_test.go @@ -0,0 +1,173 @@ +// Exercises the unexported embeddings adapter internals with no exported +// surface to drive them through. +// +//test:in-package +package gateway_bedrock + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestEmbeddingsProvider_Embed_RequestShape(t *testing.T) { + var captured embeddingsRequest + var capturedPath, capturedMethod string + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedMethod = r.Method + capturedPath = r.URL.Path + if !assert.NoError(t, json.NewDecoder(r.Body).Decode(&captured)) { + http.Error(w, "decode request body", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}],"model":"nomic-embed-text-v1.5","usage":{"prompt_tokens":1,"total_tokens":1}}`)) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + _, err := p.Embed(context.Background(), modelID, []string{"hello world"}) + require.NoError(t, err) + + assert.Equal(t, http.MethodPost, capturedMethod) + assert.Equal(t, embeddingsPath, capturedPath) + assert.Equal(t, modelID, captured.Model) + assert.Equal(t, []string{"hello world"}, captured.Input) +} + +func TestEmbeddingsProvider_Embed_HappyPathReordersToInputOrder(t *testing.T) { + var captured embeddingsRequest + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !assert.NoError(t, json.NewDecoder(r.Body).Decode(&captured)) { + http.Error(w, "decode request body", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + // Deliberately out of order relative to the request's input batch. + _, _ = w.Write([]byte(`{ + "data": [ + {"index": 2, "embedding": [2.0, 2.1, 2.2]}, + {"index": 0, "embedding": [0.0, 0.1, 0.2]}, + {"index": 1, "embedding": [1.0, 1.1, 1.2]} + ], + "model": "nomic-embed-text-v1.5", + "usage": {"prompt_tokens": 3, "total_tokens": 3} + }`)) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + inputs := []string{"first", "second", "third"} + vectors, err := p.Embed(context.Background(), modelID, inputs) + require.NoError(t, err) + + assert.Equal(t, inputs, captured.Input) + require.Len(t, vectors, 3) + assert.Equal(t, []float32{0.0, 0.1, 0.2}, vectors[0]) + assert.Equal(t, []float32{1.0, 1.1, 1.2}, vectors[1]) + assert.Equal(t, []float32{2.0, 2.1, 2.2}, vectors[2]) + for _, v := range vectors { + assert.Len(t, v, 3) + } +} + +func TestEmbeddingsProvider_Embed_EmptyInputSkipsHTTPCall(t *testing.T) { + called := false + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + vectors, err := p.Embed(context.Background(), modelID, nil) + require.NoError(t, err) + assert.Empty(t, vectors) + assert.False(t, called, "Embed must not call the endpoint for an empty input batch") +} + +func TestEmbeddingsProvider_Embed_ResolverMissReturnsErrorWithoutHTTPCall(t *testing.T) { + called := false + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + p := newEmbeddingsProvider(NewStaticEndpointResolver(nil)) + p.httpClient = ts.Client() + + _, err := p.Embed(context.Background(), DefaultEmbeddingModel, []string{"hello"}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) + assert.False(t, called, "Embed must not call the endpoint when the resolver misses") +} + +func TestEmbeddingsProvider_Embed_NonOKStatusReturnsError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"rate limited"}`)) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + _, err := p.Embed(context.Background(), modelID, []string{"hello"}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) +} + +func TestEmbeddingsProvider_Embed_MalformedJSONReturnsError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{not-json`)) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + _, err := p.Embed(context.Background(), modelID, []string{"hello"}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelErrorException, err.Error()) +} + +func TestEmbeddingsProvider_Embed_CountMismatchReturnsError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1]}],"model":"nomic-embed-text-v1.5"}`)) + })) + defer ts.Close() + + modelID := DefaultEmbeddingModel + p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + _, err := p.Embed(context.Background(), modelID, []string{"hello", "world"}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelErrorException, err.Error()) +} diff --git a/spinifex/handlers/ochrevector/api.go b/spinifex/handlers/ochrevector/api.go new file mode 100644 index 000000000..02441e6ab --- /dev/null +++ b/spinifex/handlers/ochrevector/api.go @@ -0,0 +1,337 @@ +package handlers_ochrevector + +import ( + "context" + "fmt" + "time" + + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +// NATS subjects the daemon subscribes and NATSVectorService forwards to, +// following the "eks.MethodName" naming EKS uses for its own internal +// (non-AWS-SDK-shaped) subjects. +const ( + SubjectCreateIndex = "ochre.vector.createIndex" + SubjectDeleteIndex = "ochre.vector.deleteIndex" + SubjectListIndexes = "ochre.vector.listIndexes" + SubjectIngest = "ochre.vector.ingest" + SubjectDescribeJob = "ochre.vector.describeJob" + SubjectQuery = "ochre.vector.query" +) + +// CreateIndexRequest names a new index and its fixed, per-index properties +// (D6): Dimension and EmbeddingModel are stamped once at create and never +// change for the index's lifetime. +type CreateIndexRequest struct { + IndexID string `json:"indexId"` + Name string `json:"name"` + Dimension int `json:"dimension"` + EmbeddingModel string `json:"embeddingModel"` +} + +type CreateIndexResponse struct { + Index Record `json:"index"` +} + +// DeleteIndexRequest identifies the index to drop. Idempotent: deleting an +// already-absent index is a no-op success (Service.DeleteIndex). +type DeleteIndexRequest struct { + IndexID string `json:"indexId"` +} + +type DeleteIndexResponse struct{} + +// ListIndexesRequest has no fields: accountID comes from the caller context +// (D10), never the payload. +type ListIndexesRequest struct{} + +type ListIndexesResponse struct { + Indexes []Record `json:"indexes"` +} + +// IngestRequest starts an ingestion job against IndexID from Source. Source's +// EmbeddingModel/Dimension are ignored and overwritten server-side from the +// index's own registered values (see vectorService.Ingest) -- an index's +// embedding model is a property of the index (D6/D8), not of any one caller. +type IngestRequest struct { + IndexID string `json:"indexId"` + Source SourceSpec `json:"source"` +} + +type IngestResponse struct { + Job JobRecord `json:"job"` +} + +// DescribeJobRequest looks up one ingestion job's current record. +type DescribeJobRequest struct { + JobID string `json:"jobId"` +} + +type DescribeJobResponse struct { + Job JobRecord `json:"job"` +} + +// QueryRequest asks for the K nearest chunks to Text within IndexID, after +// Filter narrows the candidate set (D9). Text is embedded server-side against +// IndexID's own pinned model (D8) -- the request never carries a vector. +// K<=0 takes the backend's default; any K above the hard cap is clamped, not +// rejected (D8/D10). +type QueryRequest struct { + IndexID string `json:"indexId"` + Text string `json:"text"` + K int `json:"k,omitempty"` + Filter *Filter `json:"filter,omitempty"` +} + +type QueryResponse struct { + Results []QueryResult `json:"results"` +} + +// VectorService is the tenant vector-store surface: index CRUD, ingestion, +// job status, and similarity query. Every method's accountID argument is +// supplied by the transport alone (a NATS message header for the daemon +// subscription, D10) -- no request type above carries an account field, so a +// spoofed payload account can never widen a caller's own scope. +type VectorService interface { + CreateIndex(ctx context.Context, req *CreateIndexRequest, accountID string) (*CreateIndexResponse, error) + DeleteIndex(ctx context.Context, req *DeleteIndexRequest, accountID string) (*DeleteIndexResponse, error) + ListIndexes(ctx context.Context, req *ListIndexesRequest, accountID string) (*ListIndexesResponse, error) + Ingest(ctx context.Context, req *IngestRequest, accountID string) (*IngestResponse, error) + DescribeJob(ctx context.Context, req *DescribeJobRequest, accountID string) (*DescribeJobResponse, error) + Query(ctx context.Context, req *QueryRequest, accountID string) (*QueryResponse, error) +} + +// indexService is the CreateIndex/DeleteIndex/ListIndexes surface vectorService +// delegates to; *Service satisfies it structurally. +type indexService interface { + CreateIndex(ctx context.Context, accountID, indexID string, params CreateIndexParams) (*Record, error) + DeleteIndex(ctx context.Context, accountID, indexID string) error + ListIndexes(ctx context.Context, accountID string) ([]Record, error) +} + +var _ indexService = (*Service)(nil) + +// ingestStarter is the StartIngest surface vectorService delegates to; +// *IngestService satisfies it structurally. +type ingestStarter interface { + StartIngest(ctx context.Context, accountID, indexID string, source SourceSpec) (*JobRecord, error) +} + +var _ ingestStarter = (*IngestService)(nil) + +// jobGetter is the job-lookup surface DescribeJob delegates to; *JobStore +// satisfies it structurally. +type jobGetter interface { + Get(ctx context.Context, accountID, jobID string) (*JobRecord, error) +} + +var _ jobGetter = (*JobStore)(nil) + +// indexGetter is the single-record registry lookup Ingest and Query use to +// resolve an index's pinned embedding model/dimension (D6/D8); *Registry +// satisfies it structurally. +type indexGetter interface { + Get(ctx context.Context, accountID, indexID string) (*Record, error) +} + +var _ indexGetter = (*Registry)(nil) + +// maxResponseChunkChars bounds a single QueryResult's Chunk text in the wire +// response (D10): NATS's 1MB default payload cap divided generously across +// maxQueryK results, with headroom for JSON and metadata overhead. The full +// chunk is always recoverable from the source object via SourceKey/ +// SourceOffset, so truncating here loses no data. +const maxResponseChunkChars = 4000 + +// truncateChunkForResponse bounds chunk to maxResponseChunkChars runes, +// appending a marker so a caller can tell a truncated chunk from a complete +// one without re-fetching the source object to check. +func truncateChunkForResponse(chunk string) string { + runes := []rune(chunk) + if len(runes) <= maxResponseChunkChars { + return chunk + } + return string(runes[:maxResponseChunkChars]) + " …[truncated; see SourceKey/SourceOffset]" +} + +// vectorService is the daemon-side VectorService implementation: index CRUD +// delegates to indexService (the orchestrator owning the registry+backend +// create/delete transitions), ingest/describe to the ingestion job stores, +// and Query resolves the index's pinned embedding model via registry before +// embedding and searching directly against backend. +type vectorService struct { + index indexService + ingest ingestStarter + jobs jobGetter + registry indexGetter + backend VectorBackend + embedder Embedder +} + +var _ VectorService = (*vectorService)(nil) + +// NewVectorService constructs the daemon-side VectorService over its +// dependencies. index/ingest/jobs/registry/backend/embedder are accepted as +// the minimal interfaces above, so a *Service/*IngestService/*JobStore/ +// *Registry/VectorBackend/Embedder wired for the rest of the daemon can be +// passed directly. +func NewVectorService(index indexService, ingest ingestStarter, jobs jobGetter, registry indexGetter, backend VectorBackend, embedder Embedder) VectorService { + return &vectorService{index: index, ingest: ingest, jobs: jobs, registry: registry, backend: backend, embedder: embedder} +} + +func (s *vectorService) CreateIndex(ctx context.Context, req *CreateIndexRequest, accountID string) (*CreateIndexResponse, error) { + if req == nil || req.IndexID == "" { + return nil, fmt.Errorf("ochrevector: create index requires an indexId") + } + rec, err := s.index.CreateIndex(ctx, accountID, req.IndexID, CreateIndexParams{ + Name: req.Name, + Dimension: req.Dimension, + EmbeddingModel: req.EmbeddingModel, + }) + if err != nil { + return nil, fmt.Errorf("ochrevector: create index %s: %w", req.IndexID, err) + } + return &CreateIndexResponse{Index: *rec}, nil +} + +func (s *vectorService) DeleteIndex(ctx context.Context, req *DeleteIndexRequest, accountID string) (*DeleteIndexResponse, error) { + if req == nil || req.IndexID == "" { + return nil, fmt.Errorf("ochrevector: delete index requires an indexId") + } + if err := s.index.DeleteIndex(ctx, accountID, req.IndexID); err != nil { + return nil, fmt.Errorf("ochrevector: delete index %s: %w", req.IndexID, err) + } + return &DeleteIndexResponse{}, nil +} + +func (s *vectorService) ListIndexes(ctx context.Context, _ *ListIndexesRequest, accountID string) (*ListIndexesResponse, error) { + recs, err := s.index.ListIndexes(ctx, accountID) + if err != nil { + return nil, fmt.Errorf("ochrevector: list indexes: %w", err) + } + return &ListIndexesResponse{Indexes: recs}, nil +} + +// Ingest stamps Source's EmbeddingModel/Dimension from the index's own +// registered values before starting the job, so the stored source-spec (and +// every vector it produces) always matches the index it belongs to, +// regardless of what a caller sent. +func (s *vectorService) Ingest(ctx context.Context, req *IngestRequest, accountID string) (*IngestResponse, error) { + if req == nil || req.IndexID == "" { + return nil, fmt.Errorf("ochrevector: ingest requires an indexId") + } + idx, err := s.registry.Get(ctx, accountID, req.IndexID) + if err != nil { + return nil, fmt.Errorf("ochrevector: ingest: lookup index %s: %w", req.IndexID, err) + } + if idx == nil { + return nil, fmt.Errorf("ochrevector: ingest: %w", ErrIndexNotFound) + } + + source := req.Source + source.EmbeddingModel = idx.EmbeddingModel + source.Dimension = idx.Dimension + + job, err := s.ingest.StartIngest(ctx, accountID, req.IndexID, source) + if err != nil { + return nil, fmt.Errorf("ochrevector: ingest: %w", err) + } + return &IngestResponse{Job: *job}, nil +} + +func (s *vectorService) DescribeJob(ctx context.Context, req *DescribeJobRequest, accountID string) (*DescribeJobResponse, error) { + if req == nil || req.JobID == "" { + return nil, fmt.Errorf("ochrevector: describe job requires a jobId") + } + job, err := s.jobs.Get(ctx, accountID, req.JobID) + if err != nil { + return nil, fmt.Errorf("ochrevector: describe job %s: %w", req.JobID, err) + } + if job == nil { + return nil, ErrJobNotFound + } + return &DescribeJobResponse{Job: *job}, nil +} + +// Query resolves IndexID's pinned embedding model from the registry (D8), +// embeds Text against it, then runs the similarity search directly against +// backend. A non-READY index (CREATING/DELETING/STALE) returns empty results +// rather than an error (D4): it is not yet, or no longer, queryable. +func (s *vectorService) Query(ctx context.Context, req *QueryRequest, accountID string) (*QueryResponse, error) { + if req == nil || req.IndexID == "" || req.Text == "" { + return nil, fmt.Errorf("ochrevector: query requires an indexId and text") + } + idx, err := s.registry.Get(ctx, accountID, req.IndexID) + if err != nil { + return nil, fmt.Errorf("ochrevector: query: lookup index %s: %w", req.IndexID, err) + } + if idx == nil { + return nil, fmt.Errorf("ochrevector: query: %w", ErrIndexNotFound) + } + if idx.State != StateReady { + return &QueryResponse{}, nil + } + + vectors, err := s.embedder.Embed(ctx, idx.EmbeddingModel, []string{req.Text}) + if err != nil { + return nil, fmt.Errorf("ochrevector: query: embed query text: %w", err) + } + if len(vectors) != 1 { + return nil, fmt.Errorf("ochrevector: query: embed returned %d vectors for 1 input", len(vectors)) + } + + results, err := s.backend.Query(ctx, accountID, req.IndexID, vectors[0], req.K, req.Filter) + if err != nil { + return nil, fmt.Errorf("ochrevector: query: %w", err) + } + for i := range results { + results[i].Chunk = truncateChunkForResponse(results[i].Chunk) + } + return &QueryResponse{Results: results}, nil +} + +// defaultVectorNATSTimeout bounds a request-reply round trip. Ingest itself +// only claims a job before replying (the run happens in the background), so +// this only needs to cover the claim/query latency, not a full ingestion run. +const defaultVectorNATSTimeout = 30 * time.Second + +// NATSVectorService is the NATS-forwarding VectorService client: any caller +// off the daemon's own process (the CLI, or a future KB-API wrapper) uses +// this instead of importing this package's registry/backend dependencies. +type NATSVectorService struct { + nc *nats.Conn + timeout time.Duration +} + +var _ VectorService = (*NATSVectorService)(nil) + +// NewNATSVectorService constructs a client bound to nc. +func NewNATSVectorService(nc *nats.Conn) *NATSVectorService { + return &NATSVectorService{nc: nc, timeout: defaultVectorNATSTimeout} +} + +func (c *NATSVectorService) CreateIndex(ctx context.Context, req *CreateIndexRequest, accountID string) (*CreateIndexResponse, error) { + return utils.NATSRequest[CreateIndexResponse](ctx, c.nc, SubjectCreateIndex, req, c.timeout, accountID) +} + +func (c *NATSVectorService) DeleteIndex(ctx context.Context, req *DeleteIndexRequest, accountID string) (*DeleteIndexResponse, error) { + return utils.NATSRequest[DeleteIndexResponse](ctx, c.nc, SubjectDeleteIndex, req, c.timeout, accountID) +} + +func (c *NATSVectorService) ListIndexes(ctx context.Context, req *ListIndexesRequest, accountID string) (*ListIndexesResponse, error) { + return utils.NATSRequest[ListIndexesResponse](ctx, c.nc, SubjectListIndexes, req, c.timeout, accountID) +} + +func (c *NATSVectorService) Ingest(ctx context.Context, req *IngestRequest, accountID string) (*IngestResponse, error) { + return utils.NATSRequest[IngestResponse](ctx, c.nc, SubjectIngest, req, c.timeout, accountID) +} + +func (c *NATSVectorService) DescribeJob(ctx context.Context, req *DescribeJobRequest, accountID string) (*DescribeJobResponse, error) { + return utils.NATSRequest[DescribeJobResponse](ctx, c.nc, SubjectDescribeJob, req, c.timeout, accountID) +} + +func (c *NATSVectorService) Query(ctx context.Context, req *QueryRequest, accountID string) (*QueryResponse, error) { + return utils.NATSRequest[QueryResponse](ctx, c.nc, SubjectQuery, req, c.timeout, accountID) +} diff --git a/spinifex/handlers/ochrevector/api_test.go b/spinifex/handlers/ochrevector/api_test.go new file mode 100644 index 000000000..6c9b2a0e3 --- /dev/null +++ b/spinifex/handlers/ochrevector/api_test.go @@ -0,0 +1,386 @@ +// Exercises unexported ochrevector API handler internals with no +// exported surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + apiAccountA = "111111111111" + apiAccountB = "222222222222" +) + +// apiTestSetup wires a real Service/IngestService/JobStore/Registry over an +// embedded JetStream server, plus a fakeBackend and stubEmbedder, so +// vectorService is exercised against the same collaborators it runs with in +// production -- only the backend and embedder are doubles. +type apiTestSetup struct { + svc VectorService + registry *Registry + backend *fakeBackend + embedder *stubEmbedder +} + +func newAPITestSetup(t *testing.T) *apiTestSetup { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + + registry := NewRegistry(js) + backend := newFakeBackend() + indexSvc := NewService(registry, backend) + + jobs := NewJobStore(js) + embedder := &stubEmbedder{} + ingestSvc := NewIngestService(jobs, registry, backend, nil, embedder) + + return &apiTestSetup{ + svc: NewVectorService(indexSvc, ingestSvc, jobs, registry, backend, embedder), + registry: registry, + backend: backend, + embedder: embedder, + } +} + +func TestVectorService_CreateIndex_HappyPath(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + + out, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{ + IndexID: "idx-one", Name: "kb1", Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5", + }, apiAccountA) + require.NoError(t, err) + assert.Equal(t, StateReady, out.Index.State) + assert.Equal(t, apiAccountA, out.Index.AccountID) + assert.True(t, s.backend.hasIndex(apiAccountA, "idx-one")) +} + +func TestVectorService_CreateIndex_MissingIndexIDErrors(t *testing.T) { + s := newAPITestSetup(t) + _, err := s.svc.CreateIndex(context.Background(), &CreateIndexRequest{Name: "kb1", Dimension: 768}, apiAccountA) + require.Error(t, err) +} + +func TestVectorService_CreateIndex_ExistsErrorPropagates(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + req := &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 768} + + _, err := s.svc.CreateIndex(ctx, req, apiAccountA) + require.NoError(t, err) + + _, err = s.svc.CreateIndex(ctx, req, apiAccountA) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexExists) +} + +func TestVectorService_DeleteIndex_HappyPath(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 768}, apiAccountA) + require.NoError(t, err) + + _, err = s.svc.DeleteIndex(ctx, &DeleteIndexRequest{IndexID: "idx-one"}, apiAccountA) + require.NoError(t, err) + + rec, err := s.registry.Get(ctx, apiAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, rec) +} + +func TestVectorService_DeleteIndex_AbsentIsNoop(t *testing.T) { + s := newAPITestSetup(t) + _, err := s.svc.DeleteIndex(context.Background(), &DeleteIndexRequest{IndexID: "no-such-index"}, apiAccountA) + require.NoError(t, err) +} + +func TestVectorService_ListIndexes_ScopedPerAccount(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-a", Name: "a", Dimension: 8}, apiAccountA) + require.NoError(t, err) + _, err = s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-b", Name: "b", Dimension: 8}, apiAccountB) + require.NoError(t, err) + + out, err := s.svc.ListIndexes(ctx, &ListIndexesRequest{}, apiAccountA) + require.NoError(t, err) + require.Len(t, out.Indexes, 1) + assert.Equal(t, "idx-a", out.Indexes[0].ID) +} + +// TestVectorService_Ingest_StampsIndexPinnedModel proves Ingest overwrites +// whatever EmbeddingModel/Dimension a caller sends with the index's own +// registered values (D6/D8): the stored source-spec must always match the +// index it belongs to, not whatever a stale or spoofed caller supplied. +func TestVectorService_Ingest_StampsIndexPinnedModel(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{ + IndexID: "idx-one", Name: "kb1", Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5", + }, apiAccountA) + require.NoError(t, err) + + out, err := s.svc.Ingest(ctx, &IngestRequest{ + IndexID: "idx-one", + Source: SourceSpec{Bucket: "docs", Prefix: "kb/", EmbeddingModel: "some-other-model", Dimension: 3}, + }, apiAccountA) + require.NoError(t, err) + assert.Equal(t, "nomic-embed-text-v1.5", out.Job.Source.EmbeddingModel) + assert.Equal(t, 768, out.Job.Source.Dimension) + assert.Equal(t, JobStatePending, out.Job.State) +} + +func TestVectorService_Ingest_MissingIndexErrors(t *testing.T) { + s := newAPITestSetup(t) + _, err := s.svc.Ingest(context.Background(), &IngestRequest{IndexID: "no-such-index", Source: SourceSpec{Bucket: "docs"}}, apiAccountA) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexNotFound) +} + +func TestVectorService_DescribeJob_HappyPath(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 8}, apiAccountA) + require.NoError(t, err) + ingestOut, err := s.svc.Ingest(ctx, &IngestRequest{IndexID: "idx-one", Source: SourceSpec{Bucket: "docs"}}, apiAccountA) + require.NoError(t, err) + + out, err := s.svc.DescribeJob(ctx, &DescribeJobRequest{JobID: ingestOut.Job.ID}, apiAccountA) + require.NoError(t, err) + assert.Equal(t, ingestOut.Job.ID, out.Job.ID) +} + +func TestVectorService_DescribeJob_AbsentReturnsErrJobNotFound(t *testing.T) { + s := newAPITestSetup(t) + _, err := s.svc.DescribeJob(context.Background(), &DescribeJobRequest{JobID: "no-such-job"}, apiAccountA) + require.Error(t, err) + assert.ErrorIs(t, err, ErrJobNotFound) +} + +// TestVectorService_Query_UsesIndexPinnedModel proves Query embeds Text +// against the index's own registered model, not any model the caller might +// have sent -- QueryRequest has no such field at all, but this also confirms +// the value actually reaching Embedder is the registry's. +func TestVectorService_Query_UsesIndexPinnedModel(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{ + IndexID: "idx-one", Name: "kb1", Dimension: 2, EmbeddingModel: "nomic-embed-text-v1.5", + }, apiAccountA) + require.NoError(t, err) + s.backend.queryResults = []QueryResult{{Chunk: "hello", SourceKey: "docs/a.txt", Score: 0.9}} + + out, err := s.svc.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "what is this about?"}, apiAccountA) + require.NoError(t, err) + require.Len(t, out.Results, 1) + assert.Equal(t, "hello", out.Results[0].Chunk) + assert.Equal(t, 1, s.embedder.callCount()) +} + +// TestVectorService_Query_NonReadyIndexReturnsEmptyNotError proves a +// CREATING/DELETING/STALE index reports empty results rather than an error +// (D4): it is not yet, or no longer, queryable, but a query against it is not +// a caller mistake. +func TestVectorService_Query_NonReadyIndexReturnsEmptyNotError(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + now := time.Now().UTC() + require.NoError(t, s.registry.Reserve(ctx, apiAccountA, Record{ + ID: "idx-one", Dimension: 2, EmbeddingModel: "m", State: StateCreating, CreatedAt: now, UpdatedAt: now, + })) + + out, err := s.svc.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) + require.NoError(t, err) + assert.Empty(t, out.Results) + assert.Equal(t, 0, s.embedder.callCount(), "a non-READY index must not even reach the embedder") +} + +func TestVectorService_Query_MissingIndexErrors(t *testing.T) { + s := newAPITestSetup(t) + _, err := s.svc.Query(context.Background(), &QueryRequest{IndexID: "no-such-index", Text: "hello"}, apiAccountA) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexNotFound) +} + +func TestVectorService_Query_EmbedderErrorPropagates(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 2}, apiAccountA) + require.NoError(t, err) + s.embedder.failAll = true + + _, err = s.svc.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) + require.Error(t, err) +} + +// TestVectorService_Query_TruncatesLongChunkText proves the NATS payload +// guard (D10): a chunk larger than maxResponseChunkChars is bounded in the +// response, with a marker pointing back at the source object rather than +// silently dropping the excess. +func TestVectorService_Query_TruncatesLongChunkText(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 2}, apiAccountA) + require.NoError(t, err) + + longChunk := strings.Repeat("a", maxResponseChunkChars+500) + s.backend.queryResults = []QueryResult{{Chunk: longChunk, SourceKey: "docs/a.txt", SourceOffset: 42, Score: 0.5}} + + out, err := s.svc.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) + require.NoError(t, err) + require.Len(t, out.Results, 1) + assert.Less(t, len(out.Results[0].Chunk), len(longChunk)) + assert.Contains(t, out.Results[0].Chunk, "truncated") + assert.Equal(t, "docs/a.txt", out.Results[0].SourceKey) + assert.Equal(t, 42, out.Results[0].SourceOffset) +} + +func TestVectorService_Query_ShortChunkUntouched(t *testing.T) { + s := newAPITestSetup(t) + ctx := context.Background() + _, err := s.svc.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 2}, apiAccountA) + require.NoError(t, err) + s.backend.queryResults = []QueryResult{{Chunk: "short chunk", SourceKey: "docs/a.txt", Score: 0.5}} + + out, err := s.svc.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) + require.NoError(t, err) + require.Len(t, out.Results, 1) + assert.Equal(t, "short chunk", out.Results[0].Chunk) +} + +// stubVectorNATSHandler mirrors daemon.handleNATSRequest's unmarshal -> +// service -> marshal/error -> respond wire contract, without importing the +// daemon package. It exists purely so this test exercises the exact bytes +// NATSVectorService and the real daemon subscription (Stage 5b) put on the +// wire. +func stubVectorNATSHandler[I any, O any](serviceFn func(context.Context, *I, string) (*O, error)) nats.MsgHandler { + return func(msg *nats.Msg) { + accountID := utils.AccountIDFromMsg(msg) + input := new(I) + if errResp := utils.UnmarshalJsonPayload(input, msg.Data); errResp != nil { + _ = msg.Respond(errResp) + return + } + output, err := serviceFn(context.Background(), input, accountID) + if err != nil { + payload := utils.GenerateErrorPayloadWithMessage(awserrors.ValidErrorCodeFromError(err), err.Error()) + _ = msg.Respond(payload) + return + } + data, err := json.Marshal(output) + if err != nil { + _ = msg.Respond(utils.GenerateErrorPayloadWithMessage(awserrors.ErrorServerInternal, err.Error())) + return + } + _ = msg.Respond(data) + } +} + +// subscribeVectorService registers every VectorService subject on the +// "spinifex-workers" queue group, exactly as daemon.subscribeAll will for the +// real service in Stage 5b, so the round trip below exercises the production +// subject names and queue group. +func subscribeVectorService(t *testing.T, nc *nats.Conn, svc VectorService) { + t.Helper() + subs := []struct { + subject string + handler nats.MsgHandler + }{ + {SubjectCreateIndex, stubVectorNATSHandler(svc.CreateIndex)}, + {SubjectDeleteIndex, stubVectorNATSHandler(svc.DeleteIndex)}, + {SubjectListIndexes, stubVectorNATSHandler(svc.ListIndexes)}, + {SubjectIngest, stubVectorNATSHandler(svc.Ingest)}, + {SubjectDescribeJob, stubVectorNATSHandler(svc.DescribeJob)}, + {SubjectQuery, stubVectorNATSHandler(svc.Query)}, + } + for _, s := range subs { + sub, err := nc.QueueSubscribe(s.subject, "spinifex-workers", s.handler) + require.NoError(t, err) + t.Cleanup(func() { _ = sub.Unsubscribe() }) + } +} + +func TestNATSVectorService_RoundTrip(t *testing.T) { + _, nc, js := testutil.StartTestJetStream(t) + + registry := NewRegistry(js) + backend := newFakeBackend() + indexSvc := NewService(registry, backend) + jobs := NewJobStore(js) + embedder := &stubEmbedder{} + ingestSvc := NewIngestService(jobs, registry, backend, nil, embedder) + svc := NewVectorService(indexSvc, ingestSvc, jobs, registry, backend, embedder) + + subscribeVectorService(t, nc, svc) + client := NewNATSVectorService(nc) + ctx := context.Background() + + createOut, err := client.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 2, EmbeddingModel: "m"}, apiAccountA) + require.NoError(t, err) + assert.Equal(t, StateReady, createOut.Index.State) + + listOut, err := client.ListIndexes(ctx, &ListIndexesRequest{}, apiAccountA) + require.NoError(t, err) + require.Len(t, listOut.Indexes, 1) + + ingestOut, err := client.Ingest(ctx, &IngestRequest{IndexID: "idx-one", Source: SourceSpec{Bucket: "docs"}}, apiAccountA) + require.NoError(t, err) + assert.NotEmpty(t, ingestOut.Job.ID) + + descOut, err := client.DescribeJob(ctx, &DescribeJobRequest{JobID: ingestOut.Job.ID}, apiAccountA) + require.NoError(t, err) + assert.Equal(t, ingestOut.Job.ID, descOut.Job.ID) + + backend.queryResults = []QueryResult{{Chunk: "hi", SourceKey: "docs/a.txt", Score: 0.9}} + queryOut, err := client.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) + require.NoError(t, err) + require.Len(t, queryOut.Results, 1) + assert.Equal(t, "hi", queryOut.Results[0].Chunk) + + _, err = client.DeleteIndex(ctx, &DeleteIndexRequest{IndexID: "idx-one"}, apiAccountA) + require.NoError(t, err) + + listOut, err = client.ListIndexes(ctx, &ListIndexesRequest{}, apiAccountA) + require.NoError(t, err) + assert.Empty(t, listOut.Indexes) +} + +// TestNATSVectorService_AccountIsolation proves accountID travels only via +// the NATS header the transport sets, never the JSON payload (D10): account +// B's client never sees account A's index even though both use the same +// wire and subject names. +func TestNATSVectorService_AccountIsolation(t *testing.T) { + _, nc, js := testutil.StartTestJetStream(t) + + registry := NewRegistry(js) + backend := newFakeBackend() + indexSvc := NewService(registry, backend) + jobs := NewJobStore(js) + embedder := &stubEmbedder{} + ingestSvc := NewIngestService(jobs, registry, backend, nil, embedder) + svc := NewVectorService(indexSvc, ingestSvc, jobs, registry, backend, embedder) + + subscribeVectorService(t, nc, svc) + client := NewNATSVectorService(nc) + ctx := context.Background() + + _, err := client.CreateIndex(ctx, &CreateIndexRequest{IndexID: "idx-one", Name: "kb1", Dimension: 2}, apiAccountA) + require.NoError(t, err) + + listOut, err := client.ListIndexes(ctx, &ListIndexesRequest{}, apiAccountB) + require.NoError(t, err) + assert.Empty(t, listOut.Indexes) +} diff --git a/spinifex/handlers/ochrevector/appliance.go b/spinifex/handlers/ochrevector/appliance.go new file mode 100644 index 000000000..5eb19f151 --- /dev/null +++ b/spinifex/handlers/ochrevector/appliance.go @@ -0,0 +1,541 @@ +package handlers_ochrevector + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math/big" + "net/url" + "sync" + "time" + + handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +// applianceBucket is the cluster-replicated KV bucket holding the platform +// Postgres appliance's singleton record (D2): endpoint, credential and +// lifecycle state for the one control-plane-owned pgvector instance. +const applianceBucket = "ochre-vector-platform" + +// applianceBucketHistory keeps one revision per key: the appliance record is +// one JSON document mutated in place, not a series. +const applianceBucketHistory = 1 + +// appliancePostgresKey is the fixed key every daemon races to Create, +// making the create itself the cluster-wide single-writer mutex. +const appliancePostgresKey = "postgres" + +// ApplianceIdentifier is the appliance's reserved, deployment-wide identifier +// (D2): a system resource under utils.GlobalAccountID, never a +// caller-supplied or generated id. +const ApplianceIdentifier = "ochre-vector-pg" + +// applianceMasterUsername is the fixed master role name the control plane +// connects as; it never has login privileges granted to any tenant. +const applianceMasterUsername = "ochre_vector_admin" + +// AppliancePostgresDatabase is the single database the appliance serves; +// every account's schema (D3) lives inside it. +const AppliancePostgresDatabase = "postgres" + +// Appliance lifecycle states (D2), analogous to the index registry's +// CREATING/READY but named for a long-lived singleton rather than a +// per-account resource. +const ( + ApplianceStateProvisioning = "PROVISIONING" + ApplianceStateAvailable = "AVAILABLE" +) + +// SubjectTeardownAppliance is the operator-only NATS subject for destroying +// the platform appliance singleton (D2): reached solely via `spx admin ochre +// appliance teardown`, never through the tenant-facing awsgw surface. +const SubjectTeardownAppliance = "ochre.appliance.teardown" + +// applianceStaleAfter bounds how long a PROVISIONING record may sit before a +// losing caller treats it as abandoned by a crashed winner and resumes the +// launch itself, rather than waiting on it forever. +const applianceStaleAfter = 5 * time.Minute + +// appliancePollInterval paces a loser's wait loop between staleness/state +// re-checks while a winner is still provisioning. +const appliancePollInterval = 2 * time.Second + +// appliancePasswordLength and appliancePasswordAlphabet define the generated +// master password: 32 characters from an alphanumeric-only alphabet, so the +// password never needs percent-encoding in a DSN or connection URL. +const appliancePasswordLength = 32 + +const appliancePasswordAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + +// ApplianceRecord is the persisted singleton record for the platform Postgres +// appliance. EncryptedPassword is AES-256-GCM ciphertext (base64, via +// handlers_iam.EncryptSecret/DecryptSecret) — the plaintext master password +// is never stored, logged, or held anywhere outside a launch/Connect call. +// +// No exported Appliance method ever returns an ApplianceRecord: Ensure +// returns the password-free ApplianceConnInfo, and Connect resolves and +// discards the plaintext internally. The exported fields exist only so this +// package's own JSON encode/decode round-trips them. +type ApplianceRecord struct { + Identifier string `json:"identifier"` + Endpoint string `json:"endpoint,omitempty"` + Port int `json:"port,omitempty"` + MasterUsername string `json:"masterUsername"` + EncryptedPassword string `json:"encryptedPassword"` + State string `json:"state"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// TeardownApplianceRequest has no fields: the appliance is a fixed, +// deployment-wide singleton (D2), never identified by a caller-supplied ID. +type TeardownApplianceRequest struct{} + +type TeardownApplianceResponse struct{} + +// ApplianceConnInfo is the appliance's connection endpoint: safe to log, +// return over an API, or hand to a caller freely, since it never carries the +// master password. +type ApplianceConnInfo struct { + Identifier string + Endpoint string + Port int + MasterUsername string +} + +// ApplianceLauncher is the seam to the real VM launch machinery: this stage +// fakes it in tests, a later stage wires it to handlers/rds's +// CreateDBInstance/LaunchDBInstanceVM. Launch must be identifier-idempotent +// downstream, since Ensure's re-entrant recovery path may call it again for +// an identifier it already launched. +type ApplianceLauncher interface { + Launch(ctx context.Context, identifier, masterUsername, masterPassword string) (endpoint string, port int, err error) + + // Delete removes identifier's backing instance. Idempotent: an identifier + // already gone (or never launched) is a no-op success, so Teardown never + // fails on a resource it already reclaimed. + Delete(ctx context.Context, identifier string) error +} + +// Appliance owns the singleton platform Postgres appliance: its encrypted +// credential store and the re-entrant Ensure state machine over it (D2). +type Appliance struct { + js jetstream.JetStream + masterKey []byte + launcher ApplianceLauncher + + mu sync.Mutex + kv jetstream.KeyValue + + // hostPortDeps and eniID back this daemon's routed presence in the + // appliance's VPC. A nil hostPortDeps.HostPort (WithHostPort never called) + // makes Connect skip all host-port work, so existing callers are unaffected. + hostPortDeps HostPortDeps + eniID string + + // registry and jobs are optional Teardown collaborators (WithStores): a + // nil pair (the zero value, and every test constructing an Appliance + // directly) skips the metadata purge, unchanged from before it existed. + registry *Registry + jobs *JobStore +} + +// NewAppliance constructs an Appliance over js, encrypting/decrypting the +// master password with masterKey and launching via launcher. masterKey must +// be non-empty: an appliance with no key to encrypt the credential it is +// about to generate is a configuration error, not a runtime one to defer. +func NewAppliance(js jetstream.JetStream, masterKey []byte, launcher ApplianceLauncher) (*Appliance, error) { + if len(masterKey) == 0 { + return nil, errors.New("ochrevector: appliance requires a master key to encrypt the platform postgres credential; none provided") + } + if launcher == nil { + return nil, errors.New("ochrevector: appliance requires an ApplianceLauncher") + } + return &Appliance{js: js, masterKey: masterKey, launcher: launcher}, nil +} + +// WithHostPort attaches this daemon's VPC host-port collaborators, so every +// future Connect first ensures a routed presence in the appliance's VPC +// before dialing it. Optional: an Appliance that never calls this performs +// no host-port work, which is what every caller not wired for the VPC fix +// (and every test that constructs an Appliance directly) gets by default. +func (a *Appliance) WithHostPort(deps HostPortDeps) *Appliance { + a.mu.Lock() + a.hostPortDeps = deps + a.mu.Unlock() + return a +} + +// WithStores attaches the index registry and job store Teardown purges +// alongside the appliance's own record, so a rebuilt appliance starts with +// no stale metadata pointing at data that no longer exists. Optional: an +// Appliance that never calls this skips the purge (back-compat default). +func (a *Appliance) WithStores(registry *Registry, jobs *JobStore) *Appliance { + a.mu.Lock() + a.registry = registry + a.jobs = jobs + a.mu.Unlock() + return a +} + +// TeardownHostPort removes this daemon's own VPC host port, if Connect ever +// installed one. Not per-endpoint -- the port belongs to this node, not to +// any one Connect call -- so this is meant for daemon shutdown, not for +// unwinding every failed or short-lived Connect. +func (a *Appliance) TeardownHostPort() error { + a.mu.Lock() + deps := a.hostPortDeps + eniID := a.eniID + a.mu.Unlock() + return removeApplianceHostPort(deps, eniID) +} + +// Teardown removes the singleton platform appliance -- RDS instance, host +// port, then (if WithStores was called) index/job metadata, then the KV +// record -- so a rebuilt appliance starts coherent. Every step runs +// regardless of an earlier one's failure, and every failure is joined into +// the returned error. Idempotent overall: tearing down an already-absent +// appliance is a no-op success. Does not re-provision -- the daemon's own +// Ensure does that on next startup, not this call. +func (a *Appliance) Teardown(ctx context.Context) error { + var errs []error + + if err := a.launcher.Delete(ctx, ApplianceIdentifier); err != nil { + errs = append(errs, fmt.Errorf("ochrevector: delete appliance db instance: %w", err)) + } + + if err := a.TeardownHostPort(); err != nil { + slog.Warn("ochrevector: appliance teardown: host port removal failed", "err", err) + errs = append(errs, fmt.Errorf("ochrevector: remove appliance host port: %w", err)) + } + + a.mu.Lock() + registry, jobs := a.registry, a.jobs + a.mu.Unlock() + // Purged after the data-bearing RDS instance is gone but before the + // appliance record: a crash here still leaves a record an operator can + // retry teardown against, re-attempting the (idempotent) purge too. + if registry != nil { + if err := registry.PurgeAll(ctx); err != nil { + errs = append(errs, fmt.Errorf("ochrevector: purge index registry: %w", err)) + } + } + if jobs != nil { + if err := jobs.PurgeAll(ctx); err != nil { + errs = append(errs, fmt.Errorf("ochrevector: purge ingestion jobs: %w", err)) + } + } + + kv, err := a.bucket(ctx) + if err != nil { + errs = append(errs, err) + return errors.Join(errs...) + } + if err := kv.Delete(ctx, appliancePostgresKey); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + errs = append(errs, fmt.Errorf("ochrevector: delete appliance record: %w", err)) + } + + return errors.Join(errs...) +} + +// ensureHostPort gives this daemon a routed presence in the appliance's VPC, +// called from Connect strictly before the DSN is built and dialed: the port +// binds an iface-id only after the ENI's logical switch port is programmed, +// which ensureDaemonENI does first. Returns the appliance's real dial IP, or +// "" as a no-op when WithHostPort was never called, so Connect's pre-VPC-fix +// dial path (rec.Endpoint) is unchanged for those callers. +func (a *Appliance) ensureHostPort(ctx context.Context, rec *ApplianceRecord) (string, error) { + a.mu.Lock() + deps := a.hostPortDeps + a.mu.Unlock() + if deps.HostPort == nil { + return "", nil + } + dialIP, eniID, err := ensureApplianceHostPort(ctx, deps, rec.Identifier) + if err != nil { + return "", fmt.Errorf("ochrevector: ensure daemon host port: %w", err) + } + a.mu.Lock() + a.eniID = eniID + a.mu.Unlock() + return dialIP, nil +} + +// bucket lazily opens (or creates) the appliance KV bucket, caching the +// handle, mirroring Registry.bucket. +func (a *Appliance) bucket(ctx context.Context) (jetstream.KeyValue, error) { + if a.js == nil { + return nil, errors.New("ochrevector: appliance has no JetStream client configured") + } + a.mu.Lock() + defer a.mu.Unlock() + if a.kv != nil { + return a.kv, nil + } + kv, err := kvutil.GetOrCreateBucket(ctx, a.js, applianceBucket, applianceBucketHistory) + if err != nil { + return nil, err + } + a.kv = kv + return kv, nil +} + +// Ensure is the re-entrant singleton-appliance orchestrator (D2): exactly one +// caller across a multi-node control plane wins the atomic KV claim and +// drives the launch; every other caller, concurrent or later, either waits +// for that launch or resumes a crashed one, and no caller ever generates a +// second password or identifier for the appliance. +func (a *Appliance) Ensure(ctx context.Context) (ApplianceConnInfo, error) { + kv, err := a.bucket(ctx) + if err != nil { + return ApplianceConnInfo{}, err + } + + rec, plaintext, err := newApplianceRecord(a.masterKey) + if err != nil { + return ApplianceConnInfo{}, err + } + data, err := json.Marshal(rec) + if err != nil { + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: encode appliance record: %w", err) + } + + rev, err := kv.Create(ctx, appliancePostgresKey, data) + switch { + case err == nil: + // Winner: the single-writer claim succeeded, so this caller alone + // drives the launch and the promotion to AVAILABLE. + return a.launchAndPromote(ctx, kv, rec, rev, plaintext) + case errors.Is(err, jetstream.ErrKeyExists): + return a.waitOrResume(ctx, kv) + default: + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: claim appliance: %w", err) + } +} + +// launchAndPromote runs the winner's path: launch the real instance, then +// flip the record to AVAILABLE. A launch failure leaves the record +// PROVISIONING rather than rolling it back — retrying immediately with a +// fresh claim would mint a second password for the same singleton appliance, +// which Ensure must never do; the staleness-triggered resume path is the +// only recovery. +func (a *Appliance) launchAndPromote(ctx context.Context, kv jetstream.KeyValue, rec ApplianceRecord, rev uint64, plaintext string) (ApplianceConnInfo, error) { + endpoint, port, err := a.launcher.Launch(ctx, rec.Identifier, rec.MasterUsername, plaintext) + if err != nil { + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: launch appliance: %w", err) + } + return a.promote(ctx, kv, rec, rev, endpoint, port) +} + +// promote flips rec to AVAILABLE via a revision-guarded (CAS) update. If the +// update loses a race to a concurrent resume of the same stale record, a +// record that has already reached AVAILABLE is this caller's success too, +// not a conflict to surface. +func (a *Appliance) promote(ctx context.Context, kv jetstream.KeyValue, rec ApplianceRecord, rev uint64, endpoint string, port int) (ApplianceConnInfo, error) { + rec.State = ApplianceStateAvailable + rec.Endpoint = endpoint + rec.Port = port + rec.UpdatedAt = time.Now().UTC() + data, err := json.Marshal(rec) + if err != nil { + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: encode appliance record: %w", err) + } + if _, err := kv.Update(ctx, appliancePostgresKey, data, rev); err != nil { + if current, _, getErr := a.getRecord(ctx, kv); getErr == nil && current != nil && current.State == ApplianceStateAvailable { + return connInfo(*current), nil + } + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: promote appliance to available: %w", err) + } + return connInfo(rec), nil +} + +// waitOrResume is the losing caller's path: read the winner's record and +// either return its endpoint once AVAILABLE, resume a stale crashed +// provision, or wait and re-check, bounded by ctx. +func (a *Appliance) waitOrResume(ctx context.Context, kv jetstream.KeyValue) (ApplianceConnInfo, error) { + for { + rec, entry, err := a.getRecord(ctx, kv) + if err != nil { + return ApplianceConnInfo{}, err + } + if rec == nil { + // The record vanished between our lost claim and this read (e.g. + // an operator delete); retry the claim from scratch. + return a.Ensure(ctx) + } + switch rec.State { + case ApplianceStateAvailable: + return connInfo(*rec), nil + case ApplianceStateProvisioning: + if time.Since(rec.UpdatedAt) > applianceStaleAfter { + return a.resume(ctx, kv, *rec, entry.Revision()) + } + default: + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: appliance record in unexpected state %q", rec.State) + } + select { + case <-ctx.Done(): + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: wait for appliance provisioning: %w", ctx.Err()) + case <-time.After(appliancePollInterval): + } + } +} + +// resume re-drives the launch for a stale PROVISIONING record left by a +// crashed winner, reusing rec's existing identifier and stored password +// rather than generating either anew (D2), then promotes it. +func (a *Appliance) resume(ctx context.Context, kv jetstream.KeyValue, rec ApplianceRecord, rev uint64) (ApplianceConnInfo, error) { + password, err := a.decryptPassword(&rec) + if err != nil { + return ApplianceConnInfo{}, err + } + endpoint, port, err := a.launcher.Launch(ctx, rec.Identifier, rec.MasterUsername, password) + if err != nil { + return ApplianceConnInfo{}, fmt.Errorf("ochrevector: resume stale appliance provisioning: %w", err) + } + return a.promote(ctx, kv, rec, rev, endpoint, port) +} + +// Connect builds a ready *pgxBackend for the current AVAILABLE appliance +// record: the sole accessor that decrypts the master password, and only for +// the moment it takes to compose the DSN and open the pool. The password +// never leaves this call. +func (a *Appliance) Connect(ctx context.Context) (*pgxBackend, error) { + kv, err := a.bucket(ctx) + if err != nil { + return nil, err + } + rec, _, err := a.getRecord(ctx, kv) + if err != nil { + return nil, err + } + if rec == nil { + return nil, errors.New("ochrevector: appliance not yet provisioned") + } + if rec.State != ApplianceStateAvailable { + return nil, fmt.Errorf("ochrevector: appliance not available (state %q)", rec.State) + } + + dialIP, err := a.ensureHostPort(ctx, rec) + if err != nil { + return nil, err + } + + password, err := a.decryptPassword(rec) + if err != nil { + return nil, err + } + // dialIP is only ever non-empty when the host-port path actually ran + // (WithHostPort configured); otherwise this is the unchanged pre-fix + // dial target. + dialHost := rec.Endpoint + if dialIP != "" { + dialHost = dialIP + } + dsn := buildDSN(dialHost, rec.Port, rec.MasterUsername, password) + + backend, err := NewPgxBackend(ctx, dsn) + if err != nil { + return nil, err + } + if err := backend.Init(ctx); err != nil { + backend.Close() + return nil, err + } + return backend, nil +} + +// getRecord reads and decodes the appliance record, returning (nil, nil, nil) +// when absent. +func (a *Appliance) getRecord(ctx context.Context, kv jetstream.KeyValue) (*ApplianceRecord, jetstream.KeyValueEntry, error) { + entry, err := kv.Get(ctx, appliancePostgresKey) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return nil, nil, nil + } + return nil, nil, fmt.Errorf("ochrevector: get appliance record: %w", err) + } + var rec ApplianceRecord + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return nil, nil, fmt.Errorf("ochrevector: decode appliance record: %w", err) + } + return &rec, entry, nil +} + +// decryptPassword resolves rec's encrypted master password to plaintext. The +// error path never includes the ciphertext or any decrypted material. +func (a *Appliance) decryptPassword(rec *ApplianceRecord) (string, error) { + plaintext, err := handlers_iam.DecryptSecret(rec.EncryptedPassword, a.masterKey) + if err != nil { + return "", fmt.Errorf("ochrevector: decrypt appliance master password: %w", err) + } + return plaintext, nil +} + +// newApplianceRecord builds a fresh PROVISIONING record with a newly +// generated master password, returning both the record (ciphertext only) and +// the plaintext for the caller's immediate use in the winning launch call. +func newApplianceRecord(masterKey []byte) (ApplianceRecord, string, error) { + password, err := generateMasterPassword() + if err != nil { + return ApplianceRecord{}, "", err + } + encrypted, err := handlers_iam.EncryptSecret(password, masterKey) + if err != nil { + return ApplianceRecord{}, "", fmt.Errorf("ochrevector: encrypt appliance master password: %w", err) + } + now := time.Now().UTC() + return ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + State: ApplianceStateProvisioning, + CreatedAt: now, + UpdatedAt: now, + }, password, nil +} + +// generateMasterPassword returns a strong, DSN/URL-safe random password: +// appliancePasswordLength characters drawn uniformly from +// appliancePasswordAlphabet via crypto/rand, so it never needs +// percent-encoding when composed into a connection URL. +func generateMasterPassword() (string, error) { + alphabetLen := big.NewInt(int64(len(appliancePasswordAlphabet))) + result := make([]byte, appliancePasswordLength) + for i := range result { + n, err := rand.Int(rand.Reader, alphabetLen) + if err != nil { + return "", fmt.Errorf("ochrevector: generate appliance master password: %w", err) + } + result[i] = appliancePasswordAlphabet[n.Int64()] + } + return string(result), nil +} + +// connInfo projects rec to its password-free connection info. +func connInfo(rec ApplianceRecord) ApplianceConnInfo { + return ApplianceConnInfo{ + Identifier: rec.Identifier, + Endpoint: rec.Endpoint, + Port: rec.Port, + MasterUsername: rec.MasterUsername, + } +} + +// buildDSN composes a postgres connection URL, URL-encoding user/password via +// net/url so even an unexpected special character round-trips safely rather +// than corrupting the DSN. +func buildDSN(endpoint string, port int, username, password string) string { + u := url.URL{ + Scheme: "postgres", + User: url.UserPassword(username, password), + Host: fmt.Sprintf("%s:%d", endpoint, port), + Path: "/" + AppliancePostgresDatabase, + } + return u.String() +} diff --git a/spinifex/handlers/ochrevector/appliance_test.go b/spinifex/handlers/ochrevector/appliance_test.go new file mode 100644 index 000000000..acb6e57d5 --- /dev/null +++ b/spinifex/handlers/ochrevector/appliance_test.go @@ -0,0 +1,791 @@ +// Exercises unexported ochrevector appliance internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "net/url" + "sync" + "testing" + "time" + + handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeLauncher is the test double for ApplianceLauncher: it counts +// invocations and records every identifier/password it was called with, so +// tests can assert exactly-once launch and password reuse without any real +// VM. +type fakeLauncher struct { + mu sync.Mutex + calls int + identifiers []string + passwords []string + endpoint string + port int + err error + + deleteCalls []string + deleteErr error +} + +var _ ApplianceLauncher = (*fakeLauncher)(nil) + +func (f *fakeLauncher) Launch(_ context.Context, identifier, _ string, masterPassword string) (string, int, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.calls++ + f.identifiers = append(f.identifiers, identifier) + f.passwords = append(f.passwords, masterPassword) + if f.err != nil { + return "", 0, f.err + } + return f.endpoint, f.port, nil +} + +func (f *fakeLauncher) Delete(_ context.Context, identifier string) error { + f.mu.Lock() + defer f.mu.Unlock() + f.deleteCalls = append(f.deleteCalls, identifier) + return f.deleteErr +} + +func (f *fakeLauncher) deleteCallCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.deleteCalls) +} + +func (f *fakeLauncher) callCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.calls +} + +func (f *fakeLauncher) lastPassword() string { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.passwords) == 0 { + return "" + } + return f.passwords[len(f.passwords)-1] +} + +func testMasterKey(t *testing.T) []byte { + t.Helper() + key := make([]byte, 32) + _, err := rand.Read(key) + require.NoError(t, err) + return key +} + +func TestNewAppliance_RequiresMasterKey(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + _, err := NewAppliance(js, nil, &fakeLauncher{}) + assert.Error(t, err) +} + +func TestNewAppliance_RequiresLauncher(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + _, err := NewAppliance(js, testMasterKey(t), nil) + assert.Error(t, err) +} + +// TestCredentialRoundTrip proves the generated master password survives +// encrypt-then-decrypt unchanged, and that the persisted ciphertext is never +// the plaintext itself. +func TestCredentialRoundTrip(t *testing.T) { + masterKey := testMasterKey(t) + + password, err := generateMasterPassword() + require.NoError(t, err) + assert.Len(t, password, appliancePasswordLength) + + encrypted, err := handlers_iam.EncryptSecret(password, masterKey) + require.NoError(t, err) + assert.NotEqual(t, password, encrypted) + assert.NotContains(t, encrypted, password) + + rec := ApplianceRecord{EncryptedPassword: encrypted} + appliance := &Appliance{masterKey: masterKey} + decrypted, err := appliance.decryptPassword(&rec) + require.NoError(t, err) + assert.Equal(t, password, decrypted) + + data, err := json.Marshal(rec) + require.NoError(t, err) + assert.NotContains(t, string(data), password) +} + +// TestEnsure_SingletonRace is the singleton-claim proof: N concurrent Ensure +// calls against one Appliance must produce exactly one launcher invocation +// and hand every caller back the same connection info. +func TestEnsure_SingletonRace(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + launcher := &fakeLauncher{endpoint: "10.0.0.5", port: 5432} + appliance, err := NewAppliance(js, testMasterKey(t), launcher) + require.NoError(t, err) + + const n = 5 + type result struct { + info ApplianceConnInfo + err error + } + results := make(chan result, n) + var wg sync.WaitGroup + for range n { + wg.Go(func() { + info, err := appliance.Ensure(context.Background()) + results <- result{info: info, err: err} + }) + } + wg.Wait() + close(results) + + var infos []ApplianceConnInfo + for r := range results { + require.NoError(t, r.err) + infos = append(infos, r.info) + } + require.Len(t, infos, n) + for _, info := range infos { + assert.Equal(t, infos[0], info) + } + assert.Equal(t, "10.0.0.5", infos[0].Endpoint) + assert.Equal(t, 1, launcher.callCount()) +} + +// TestEnsure_LoserReturnsAvailableWithoutRelaunch proves a second, later +// Ensure call against an already-AVAILABLE appliance returns immediately +// with the same connection info and never calls the launcher again. +func TestEnsure_LoserReturnsAvailableWithoutRelaunch(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + launcher := &fakeLauncher{endpoint: "10.0.0.6", port: 5432} + appliance, err := NewAppliance(js, testMasterKey(t), launcher) + require.NoError(t, err) + + first, err := appliance.Ensure(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, launcher.callCount()) + + second, err := appliance.Ensure(context.Background()) + require.NoError(t, err) + assert.Equal(t, first, second) + assert.Equal(t, 1, launcher.callCount()) +} + +// TestEnsure_StaleProvisioningIsResumed proves a PROVISIONING record left +// behind by a crashed winner (UpdatedAt older than applianceStaleAfter) is +// re-driven through the SAME identifier and stored password, reaching +// AVAILABLE, rather than a second claim being minted. +func TestEnsure_StaleProvisioningIsResumed(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.7", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + + seedPassword, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(seedPassword, masterKey) + require.NoError(t, err) + + stale := time.Now().UTC().Add(-2 * applianceStaleAfter) + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + State: ApplianceStateProvisioning, + CreatedAt: stale, + UpdatedAt: stale, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + info, err := appliance.Ensure(ctx) + require.NoError(t, err) + assert.Equal(t, "10.0.0.7", info.Endpoint) + assert.Equal(t, 1, launcher.callCount()) + assert.Equal(t, ApplianceIdentifier, launcher.identifiers[0]) + assert.Equal(t, seedPassword, launcher.lastPassword()) +} + +// TestEnsure_FreshProvisioningIsWaitedNotResumed proves a PROVISIONING +// record within the staleness grace period is waited on, bounded by the +// caller's context, and never resumed/relaunched. +func TestEnsure_FreshProvisioningIsWaitedNotResumed(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.8", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + + seedPassword, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(seedPassword, masterKey) + require.NoError(t, err) + + fresh := time.Now().UTC() + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + State: ApplianceStateProvisioning, + CreatedAt: fresh, + UpdatedAt: fresh, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + waitCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + _, err = appliance.Ensure(waitCtx) + assert.Error(t, err) + assert.Equal(t, 0, launcher.callCount()) +} + +// TestEnsure_PlaintextPasswordNeverPersisted proves the plaintext master +// password the launcher receives is never present in the raw bytes stored in +// the KV bucket. +func TestEnsure_PlaintextPasswordNeverPersisted(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + launcher := &fakeLauncher{endpoint: "10.0.0.9", port: 5432} + appliance, err := NewAppliance(js, testMasterKey(t), launcher) + require.NoError(t, err) + + ctx := context.Background() + _, err = appliance.Ensure(ctx) + require.NoError(t, err) + require.Equal(t, 1, launcher.callCount()) + + password := launcher.lastPassword() + require.NotEmpty(t, password) + + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + entry, err := kv.Get(ctx, appliancePostgresKey) + require.NoError(t, err) + assert.NotContains(t, string(entry.Value()), password) +} + +// TestEnsure_LaunchFailureLeavesRecordProvisioning proves a winner whose +// launch fails leaves the record PROVISIONING rather than rolling it back, +// and that a second, still-fresh Ensure call neither relaunches nor mints a +// new password for it. +func TestEnsure_LaunchFailureLeavesRecordProvisioning(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{err: errors.New("boom")} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + + ctx := context.Background() + _, err = appliance.Ensure(ctx) + assert.Error(t, err) + require.Equal(t, 1, launcher.callCount()) + firstPassword := launcher.lastPassword() + + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + rec, _, err := appliance.getRecord(ctx, kv) + require.NoError(t, err) + require.NotNil(t, rec) + assert.Equal(t, ApplianceStateProvisioning, rec.State) + + waitCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + _, err = appliance.Ensure(waitCtx) + assert.Error(t, err) + assert.Equal(t, 1, launcher.callCount()) + assert.Equal(t, firstPassword, launcher.lastPassword()) +} + +// TestResume_LaunchFailureLeavesRecordProvisioning proves a re-driven launch +// of a stale record that also fails leaves the record PROVISIONING with the +// same stored password, not a rollback or a fresh claim. +func TestResume_LaunchFailureLeavesRecordProvisioning(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{err: errors.New("boom")} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + + seedPassword, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(seedPassword, masterKey) + require.NoError(t, err) + + stale := time.Now().UTC().Add(-2 * applianceStaleAfter) + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + State: ApplianceStateProvisioning, + CreatedAt: stale, + UpdatedAt: stale, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + _, err = appliance.Ensure(ctx) + assert.Error(t, err) + assert.Equal(t, 1, launcher.callCount()) + assert.Equal(t, seedPassword, launcher.lastPassword()) + + rec, _, err := appliance.getRecord(ctx, kv) + require.NoError(t, err) + require.NotNil(t, rec) + assert.Equal(t, ApplianceStateProvisioning, rec.State) +} + +// TestConnect_NotProvisioned proves Connect refuses to build a backend when +// the appliance has no record yet, without needing any postgres to run. +func TestConnect_NotProvisioned(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + appliance, err := NewAppliance(js, testMasterKey(t), &fakeLauncher{}) + require.NoError(t, err) + + _, err = appliance.Connect(context.Background()) + assert.Error(t, err) +} + +// TestConnect_NotAvailable proves Connect refuses a PROVISIONING record. +func TestConnect_NotAvailable(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + appliance, err := NewAppliance(js, masterKey, &fakeLauncher{}) + require.NoError(t, err) + + password, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(password, masterKey) + require.NoError(t, err) + + now := time.Now().UTC() + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + State: ApplianceStateProvisioning, + CreatedAt: now, + UpdatedAt: now, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + _, err = appliance.Connect(ctx) + assert.Error(t, err) +} + +// seedAvailableAppliance seeds an AVAILABLE record so Connect gets past its +// state checks and reaches the host-port/dial path. +func seedAvailableAppliance(t *testing.T, appliance *Appliance, masterKey []byte, endpoint string, port int) { + t.Helper() + password, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(password, masterKey) + require.NoError(t, err) + + now := time.Now().UTC() + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + EncryptedPassword: encrypted, + Endpoint: endpoint, + Port: port, + State: ApplianceStateAvailable, + CreatedAt: now, + UpdatedAt: now, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) +} + +// TestConnect_NoHostPortDepsSkipsEnsure proves an Appliance that never calls +// WithHostPort performs no host-port work: Connect fails on the (unreachable +// in this test) pgx dial, not on anything host-port related, and the ENI +// side is left untouched. +func TestConnect_NoHostPortDepsSkipsEnsure(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + appliance, err := NewAppliance(js, masterKey, &fakeLauncher{}) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "127.0.0.1", 1) + + _, err = appliance.Connect(context.Background()) + require.Error(t, err) + assert.NotContains(t, err.Error(), "ensure daemon host port") +} + +// TestConnect_EnsuresHostPortBeforeDialing_PropagatesFailure proves Connect +// drives the daemon's VPC host port before it ever builds a DSN or dials +// Postgres: a host-port failure surfaces as that failure, not as a pgx +// connection error, and the ENI is still minted along the way. +func TestConnect_EnsuresHostPortBeforeDialing_PropagatesFailure(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + appliance, err := NewAppliance(js, masterKey, &fakeLauncher{}) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, testApplianceEndpoint, 5432) + + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(ApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + h.hostPort.failEnsure = errors.New("ovs-vsctl exploded") + appliance.WithHostPort(h.deps()) + + _, err = appliance.Connect(context.Background()) + require.Error(t, err) + assert.Contains(t, err.Error(), "ensure daemon host port") + assert.Contains(t, err.Error(), "ovs-vsctl exploded") + assert.Len(t, h.vpc.created, 1, "the ENI is minted before the host port install that then fails") +} + +// TestAppliance_TeardownHostPort proves TeardownHostPort removes the port +// Connect last installed, and is a no-op both before any Connect and for an +// Appliance that never opted into host-port management at all. +func TestAppliance_TeardownHostPort(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + appliance, err := NewAppliance(js, testMasterKey(t), &fakeLauncher{}) + require.NoError(t, err) + + assert.NoError(t, appliance.TeardownHostPort(), "never having connected must be a harmless no-op") + + h := newHostPortHarness() + appliance.WithHostPort(h.deps()) + assert.NoError(t, appliance.TeardownHostPort(), "no ENI ensured yet: still a no-op") + assert.Empty(t, h.hostPort.removals()) + + appliance.mu.Lock() + appliance.eniID = "eni-installed" + appliance.mu.Unlock() + + require.NoError(t, appliance.TeardownHostPort()) + assert.Equal(t, []string{"eni-installed"}, h.hostPort.removals()) +} + +// TestTeardown_DeletesRecordLauncherAndHostPort proves Teardown drives all +// three steps: the launcher's Delete for the fixed appliance identifier, the +// daemon's own host port removal, and the KV record. +func TestTeardown_DeletesRecordLauncherAndHostPort(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.30", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "10.0.0.30", 5432) + + h := newHostPortHarness() + appliance.WithHostPort(h.deps()) + appliance.mu.Lock() + appliance.eniID = "eni-installed" + appliance.mu.Unlock() + + require.NoError(t, appliance.Teardown(context.Background())) + + assert.Equal(t, []string{ApplianceIdentifier}, launcher.deleteCalls) + assert.Equal(t, []string{"eni-installed"}, h.hostPort.removals()) + + kv, err := appliance.bucket(context.Background()) + require.NoError(t, err) + rec, _, err := appliance.getRecord(context.Background(), kv) + require.NoError(t, err) + assert.Nil(t, rec, "the KV record must be gone after teardown") +} + +// TestTeardown_NoExistingApplianceIsANoOp proves tearing down an appliance +// that was never provisioned (or already torn down) succeeds rather than +// erroring on an absent KV record or an already-gone RDS instance. +func TestTeardown_NoExistingApplianceIsANoOp(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + launcher := &fakeLauncher{} + appliance, err := NewAppliance(js, testMasterKey(t), launcher) + require.NoError(t, err) + + assert.NoError(t, appliance.Teardown(context.Background())) + assert.Equal(t, 1, launcher.deleteCallCount()) + + // Tearing down again must still be a no-op success. + assert.NoError(t, appliance.Teardown(context.Background())) +} + +// TestTeardown_StillDeletesRecordWhenHostPortRemovalFails proves a host-port +// removal failure is best-effort: it does not stop the KV record from being +// deleted, but it is still surfaced in the joined error rather than swallowed. +func TestTeardown_StillDeletesRecordWhenHostPortRemovalFails(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.31", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "10.0.0.31", 5432) + + h := newHostPortHarness() + h.hostPort.failRemove = errors.New("ovs-vsctl exploded") + appliance.WithHostPort(h.deps()) + appliance.mu.Lock() + appliance.eniID = "eni-installed" + appliance.mu.Unlock() + + err = appliance.Teardown(context.Background()) + require.Error(t, err) + assert.Contains(t, err.Error(), "ovs-vsctl exploded") + + kv, err := appliance.bucket(context.Background()) + require.NoError(t, err) + rec, _, err := appliance.getRecord(context.Background(), kv) + require.NoError(t, err) + assert.Nil(t, rec, "the KV record must still be deleted despite the host-port failure") +} + +// TestTeardown_JoinsLauncherAndRecordErrors proves a launcher.Delete failure +// does not stop the KV record from being deleted, and both failures are +// joined into the returned error. +func TestTeardown_JoinsLauncherAndRecordErrors(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{deleteErr: errors.New("rds unreachable")} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "10.0.0.32", 5432) + + err = appliance.Teardown(context.Background()) + require.Error(t, err) + assert.Contains(t, err.Error(), "rds unreachable") + + kv, err := appliance.bucket(context.Background()) + require.NoError(t, err) + rec, _, err := appliance.getRecord(context.Background(), kv) + require.NoError(t, err) + assert.Nil(t, rec, "the KV record must still be deleted despite the launcher failure") +} + +// TestTeardown_WithStoresPurgesRegistryAndJobs proves an Appliance wired via +// WithStores purges every index and job record on Teardown, so a +// reprovisioned appliance never advertises metadata pointing at data that no +// longer exists. +func TestTeardown_WithStoresPurgesRegistryAndJobs(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.33", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "10.0.0.33", 5432) + + registry := NewRegistry(js) + jobs := NewJobStore(js) + ctx := context.Background() + now := time.Now().UTC() + require.NoError(t, registry.Reserve(ctx, "111111111111", Record{ID: "idx-one", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, jobs.Reserve(ctx, "111111111111", JobRecord{ID: "job-one", CreatedAt: now, UpdatedAt: now})) + + appliance.WithStores(registry, jobs) + require.NoError(t, appliance.Teardown(ctx)) + + remainingIndexes, err := registry.ListAll(ctx) + require.NoError(t, err) + assert.Empty(t, remainingIndexes, "Teardown must purge every index record") + + remainingJobs, err := jobs.ListAll(ctx) + require.NoError(t, err) + assert.Empty(t, remainingJobs, "Teardown must purge every job record") +} + +// TestTeardown_WithoutStoresStillWorks proves an Appliance that never calls +// WithStores (every pre-existing caller and test) still tears down cleanly: +// the registry/jobs purge is nil-safe, not a new required dependency. +func TestTeardown_WithoutStoresStillWorks(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + launcher := &fakeLauncher{endpoint: "10.0.0.34", port: 5432} + appliance, err := NewAppliance(js, masterKey, launcher) + require.NoError(t, err) + seedAvailableAppliance(t, appliance, masterKey, "10.0.0.34", 5432) + + require.NoError(t, appliance.Teardown(context.Background())) +} + +// TestBucket_RequiresJetStream proves bucket fails fast on a zero-value +// Appliance rather than panicking on a nil JetStream client. +func TestBucket_RequiresJetStream(t *testing.T) { + appliance := &Appliance{} + _, err := appliance.bucket(context.Background()) + assert.Error(t, err) +} + +// TestBuildDSN proves the DSN round-trips endpoint/port/username/password, +// percent-encoding special characters via net/url rather than concatenation. +func TestBuildDSN(t *testing.T) { + dsn := buildDSN("10.0.0.1", 5432, "ochre_vector_admin", "p@ss/word?") + parsed, err := url.Parse(dsn) + require.NoError(t, err) + assert.Equal(t, "postgres", parsed.Scheme) + assert.Equal(t, "10.0.0.1:5432", parsed.Host) + assert.Equal(t, "/"+AppliancePostgresDatabase, parsed.Path) + assert.Equal(t, "ochre_vector_admin", parsed.User.Username()) + pw, ok := parsed.User.Password() + assert.True(t, ok) + assert.Equal(t, "p@ss/word?", pw) +} + +// TestDecryptPassword_WrongKeyErrors proves decryptPassword fails rather +// than returning garbage when the appliance's masterKey does not match the +// key the record was encrypted with. +func TestDecryptPassword_WrongKeyErrors(t *testing.T) { + rightKey := testMasterKey(t) + wrongKey := testMasterKey(t) + + password, err := generateMasterPassword() + require.NoError(t, err) + encrypted, err := handlers_iam.EncryptSecret(password, rightKey) + require.NoError(t, err) + + appliance := &Appliance{masterKey: wrongKey} + _, err = appliance.decryptPassword(&ApplianceRecord{EncryptedPassword: encrypted}) + assert.Error(t, err) +} + +// TestNewApplianceRecord_EncryptFailureOnBadKeyLength proves newApplianceRecord +// surfaces an encrypt failure rather than silently persisting a plaintext or +// unencrypted password. +func TestNewApplianceRecord_EncryptFailureOnBadKeyLength(t *testing.T) { + _, _, err := newApplianceRecord([]byte("too-short")) + assert.Error(t, err) +} + +// TestGetRecord_MalformedJSON proves a corrupted record is a hard decode +// error, not a silently empty/zero record. +func TestGetRecord_MalformedJSON(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + appliance, err := NewAppliance(js, testMasterKey(t), &fakeLauncher{}) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Put(ctx, appliancePostgresKey, []byte("not json")) + require.NoError(t, err) + + _, _, err = appliance.getRecord(ctx, kv) + assert.Error(t, err) +} + +// TestWaitOrResume_UnexpectedStateErrors proves an appliance record in a +// state neither PROVISIONING nor AVAILABLE is a hard error, not a silent +// wait forever. +func TestWaitOrResume_UnexpectedStateErrors(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + masterKey := testMasterKey(t) + appliance, err := NewAppliance(js, masterKey, &fakeLauncher{}) + require.NoError(t, err) + + now := time.Now().UTC() + seed := ApplianceRecord{ + Identifier: ApplianceIdentifier, + MasterUsername: applianceMasterUsername, + State: "BOGUS", + CreatedAt: now, + UpdatedAt: now, + } + data, err := json.Marshal(seed) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + _, err = kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + _, err = appliance.waitOrResume(ctx, kv) + assert.Error(t, err) +} + +// TestWaitOrResume_VanishedRecordRetriesClaim proves a losing caller whose +// record disappears before it can be read (e.g. an operator delete) retries +// the claim from scratch rather than erroring. +func TestWaitOrResume_VanishedRecordRetriesClaim(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + launcher := &fakeLauncher{endpoint: "10.0.0.10", port: 5432} + appliance, err := NewAppliance(js, testMasterKey(t), launcher) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + + info, err := appliance.waitOrResume(ctx, kv) + require.NoError(t, err) + assert.Equal(t, "10.0.0.10", info.Endpoint) + assert.Equal(t, 1, launcher.callCount()) +} + +// TestPromote_ConflictFallsBackToCurrentAvailable proves promote treats a +// revision conflict against an already-AVAILABLE record (a concurrent +// resume won first) as success, not an error to surface. +func TestPromote_ConflictFallsBackToCurrentAvailable(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + appliance, err := NewAppliance(js, testMasterKey(t), &fakeLauncher{}) + require.NoError(t, err) + + rec, _, err := newApplianceRecord(appliance.masterKey) + require.NoError(t, err) + data, err := json.Marshal(rec) + require.NoError(t, err) + + ctx := context.Background() + kv, err := appliance.bucket(ctx) + require.NoError(t, err) + rev, err := kv.Create(ctx, appliancePostgresKey, data) + require.NoError(t, err) + + already := rec + already.State = ApplianceStateAvailable + already.Endpoint = "10.0.0.20" + already.Port = 5432 + alreadyData, err := json.Marshal(already) + require.NoError(t, err) + _, err = kv.Update(ctx, appliancePostgresKey, alreadyData, rev) + require.NoError(t, err) + + info, err := appliance.promote(ctx, kv, rec, rev, "10.0.0.21", 5432) + require.NoError(t, err) + assert.Equal(t, "10.0.0.20", info.Endpoint) +} diff --git a/spinifex/handlers/ochrevector/backend.go b/spinifex/handlers/ochrevector/backend.go new file mode 100644 index 000000000..072884f7a --- /dev/null +++ b/spinifex/handlers/ochrevector/backend.go @@ -0,0 +1,92 @@ +// Package handlers_ochrevector holds Ochre's tenant vector store: the +// pgvector backend, the KV index registry, and the create/delete/list/ +// reconcile orchestration over both. +package handlers_ochrevector + +import "context" + +// IndexSpec describes a single vector index at creation time. Dimension is +// fixed for the index's lifetime, stamped from the embedding model that +// produced it. +type IndexSpec struct { + ID string + Dimension int +} + +// VectorRow is one chunk's embedding plus its provenance, ready for +// ReplaceDocument to persist. Metadata is marshalled to jsonb as-is; a nil +// map persists as an empty jsonb object rather than SQL NULL. +type VectorRow struct { + Embedding []float32 + Chunk string + Metadata map[string]any + SourceOffset int +} + +// QueryResult is one nearest-neighbour hit from Query: its chunk text, +// stored metadata, source provenance, and cosine similarity Score in [0,1] +// (D8; 1 - cosine_distance). +type QueryResult struct { + Chunk string + Metadata map[string]any + SourceKey string + SourceOffset int + Score float32 +} + +// defaultQueryK and maxQueryK bound Query's k (D8/D10): k<=0 defaults to +// defaultQueryK, and any k above maxQueryK is silently clamped rather than +// rejected. +const ( + defaultQueryK = 10 + maxQueryK = 100 +) + +// clampQueryK applies Query's k default/clamp rule (D8). +func clampQueryK(k int) int { + if k <= 0 { + return defaultQueryK + } + if k > maxQueryK { + return maxQueryK + } + return k +} + +// VectorBackend is the swap-out seam behind every Postgres/pgvector +// operation: a future HA or replicated backend implements this interface +// without any caller change. Only the operations this stage needs are +// declared here; ingestion and query are later stages' extensions. +type VectorBackend interface { + // Init runs once-per-database setup (CREATE EXTENSION vector). Idempotent: + // safe to call on every daemon boot. + Init(ctx context.Context) error + + // EnsureAccount provisions accountID's schema, role and grants if they do + // not already exist. Idempotent: safe to call before every account-scoped + // operation. + EnsureAccount(ctx context.Context, accountID string) error + + // CreateIndex creates the backing table and HNSW index for spec under + // accountID's schema. Idempotent: safe to retry after a crash mid-create. + CreateIndex(ctx context.Context, accountID string, spec IndexSpec) error + + // DropIndex drops indexID's backing table under accountID's schema. + // Idempotent: dropping an already-absent index is a no-op success. + DropIndex(ctx context.Context, accountID, indexID string) error + + // ReplaceDocument atomically replaces every row previously stored for + // sourceKey under indexID with rows, in one transaction: delete-by-key + // then reinsert. A query never observes a half-replaced document, and a + // re-ingest of the same source key replaces rather than duplicates its + // chunks (D7). An empty rows deletes the key's existing rows and inserts + // nothing. + ReplaceDocument(ctx context.Context, accountID, indexID, sourceKey string, rows []VectorRow) error + + // Query runs a k-nearest-neighbour cosine similarity search against + // indexID's embeddings under accountID's schema, ranked nearest first + // (D8). k<=0 defaults to defaultQueryK; k is always clamped to + // maxQueryK, never errored on. A nil filter matches every row; a + // non-nil filter narrows results per its compiled WHERE clause (D9). + Query(ctx context.Context, accountID, indexID string, queryVector []float32, k int, filter *Filter) ([]QueryResult, error) +} diff --git a/spinifex/handlers/ochrevector/chunk.go b/spinifex/handlers/ochrevector/chunk.go new file mode 100644 index 000000000..368cc3d9d --- /dev/null +++ b/spinifex/handlers/ochrevector/chunk.go @@ -0,0 +1,186 @@ +package handlers_ochrevector + +import ( + "strings" + "unicode/utf8" +) + +// approxCharsPerToken is the documented, deliberately simple chars/token +// ratio behind DefaultChunkSize/DefaultChunkOverlap. There is no true +// tokenizer here -- sizing is in RUNES, not model tokens -- so both defaults +// are an approximation; a per-model tokenizer is a later refinement. +const approxCharsPerToken = 4 + +// DefaultChunkSize and DefaultChunkOverlap approximate ~512 tokens and ~64 +// overlap tokens (D10) in runes, via approxCharsPerToken. +const ( + DefaultChunkSize = 512 * approxCharsPerToken + DefaultChunkOverlap = 64 * approxCharsPerToken +) + +// chunkSeparators is the recursive-split hierarchy (D10): paragraph, line, +// sentence (three terminators), then word. A span still over size after +// every separator level falls back to a hard rune-count cut. +var chunkSeparators = []string{"\n\n", "\n", ". ", "! ", "? ", " "} + +// Chunk is one piece of split document text: Text is the exact substring +// (including any overlap carried over from the previous chunk), Offset is +// its starting position in the original text, measured in runes. +type Chunk struct { + Text string + Offset int +} + +// piece is one leaf-level split of the original text, before packing into +// size-bounded, overlapping Chunks. +type piece struct { + text string + offset int +} + +// ChunkText splits text into overlapping pieces of at most size runes each, +// recursing over chunkSeparators (paragraph -> line -> sentence -> word) so +// splits prefer natural boundaries over a mid-word or mid-sentence cut; any +// span with no such boundary within size (e.g. one long unbroken token) is +// hard-split by rune count. overlap runes from the tail of one chunk are +// carried into the start of the next, so adjacent chunks share context. +// +// size<=0 defaults to DefaultChunkSize; overlap<0 clamps to 0; overlap>=size +// clamps to size-1, since an overlap consuming (or exceeding) a whole chunk +// would stall the packer's forward progress. Whitespace-only or empty text +// returns nil. +func ChunkText(text string, size, overlap int) []Chunk { + if size <= 0 { + size = DefaultChunkSize + } + if overlap < 0 { + overlap = 0 + } + if overlap >= size { + overlap = size - 1 + } + if strings.TrimSpace(text) == "" { + return nil + } + + pieces := splitPieces(text, 0, size, overlap, 0) + return packChunks(pieces, size, overlap) +} + +// splitPieces recursively splits text on chunkSeparators[sepIdx:], returning +// the smallest number of ordered, non-overlapping pieces such that every +// piece is at most size runes -- preferring the coarsest separator level +// that achieves that. base is text's rune offset within the original +// document, so every returned piece's offset is correct with no further +// bookkeeping by the caller. overlap is threaded through only to size the +// hardSplit fallback's granularity (see hardSplit). +func splitPieces(text string, base, size, overlap, sepIdx int) []piece { + if utf8.RuneCountInString(text) <= size { + return []piece{{text: text, offset: base}} + } + if sepIdx >= len(chunkSeparators) { + return hardSplit(text, base, size, overlap) + } + + sep := chunkSeparators[sepIdx] + parts := strings.Split(text, sep) + if len(parts) == 1 { + // sep does not occur in text at all; skip straight to the next level + // rather than looping on a split that changes nothing. + return splitPieces(text, base, size, overlap, sepIdx+1) + } + + var out []piece + cursor := base + for i, part := range parts { + seg := part + if i < len(parts)-1 { + seg += sep + } + segLen := utf8.RuneCountInString(seg) + if segLen == 0 { + continue + } + if segLen > size { + out = append(out, splitPieces(seg, cursor, size, overlap, sepIdx+1)...) + } else { + out = append(out, piece{text: seg, offset: cursor}) + } + cursor += segLen + } + return out +} + +// hardSplit cuts text into contiguous rune spans, the terminal fallback once +// every separator level has been exhausted (e.g. one long unbroken token). +// Granularity is overlap when overlap>0, not size: pieces exactly size runes +// wide would leave packChunks nothing finer than a whole piece to back up +// by, so a requested overlap much smaller than size (the common case) would +// overshoot to a full extra piece. Sizing pieces to overlap instead lets +// packChunks land on (close to) the requested overlap even with no natural +// separator to split on. +func hardSplit(text string, base, size, overlap int) []piece { + granularity := size + if overlap > 0 { + granularity = overlap + } + runes := []rune(text) + out := make([]piece, 0, (len(runes)+granularity-1)/granularity) + for start := 0; start < len(runes); start += granularity { + end := min(start+granularity, len(runes)) + out = append(out, piece{text: string(runes[start:end]), offset: base + start}) + } + return out +} + +// packChunks greedily packs ordered, size-bounded pieces into Chunks near +// size runes each, backing up by ~overlap runes between adjacent chunks so +// they share trailing/leading context. Pieces are assumed contiguous and in +// original-text order (splitPieces' guarantee), so Chunk offsets come +// straight from the first piece in each pack and are strictly increasing. +func packChunks(pieces []piece, size, overlap int) []Chunk { + if len(pieces) == 0 { + return nil + } + + var chunks []Chunk + i := 0 + for i < len(pieces) { + start := i + length := 0 + j := i + for j < len(pieces) { + pl := utf8.RuneCountInString(pieces[j].text) + if length > 0 && length+pl > size { + break + } + length += pl + j++ + } + if j == start { + j = start + 1 // always include at least one piece, even if it alone exceeds size + } + + var sb strings.Builder + for _, p := range pieces[start:j] { + sb.WriteString(p.text) + } + chunks = append(chunks, Chunk{Text: sb.String(), Offset: pieces[start].offset}) + + if j >= len(pieces) { + break + } + + // Back up from j toward start to carry ~overlap runes of trailing + // pieces into the next chunk. k is floored at start+1 so the packer + // always advances -- even when overlap >= this chunk's whole span. + k := j + carried := 0 + for k > start+1 && carried < overlap { + carried += utf8.RuneCountInString(pieces[k-1].text) + k-- + } + i = k + } + return chunks +} diff --git a/spinifex/handlers/ochrevector/chunk_test.go b/spinifex/handlers/ochrevector/chunk_test.go new file mode 100644 index 000000000..0d148af66 --- /dev/null +++ b/spinifex/handlers/ochrevector/chunk_test.go @@ -0,0 +1,164 @@ +// Exercises the unexported chunker internals with no exported surface to +// drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestChunkText_Empty(t *testing.T) { + assert.Nil(t, ChunkText("", 50, 10)) +} + +func TestChunkText_WhitespaceOnly(t *testing.T) { + assert.Nil(t, ChunkText(" \n\t \n", 50, 10)) +} + +func TestChunkText_UnderSize(t *testing.T) { + text := "hello world, this is short" + chunks := ChunkText(text, 1000, 100) + require.Len(t, chunks, 1) + assert.Equal(t, text, chunks[0].Text) + assert.Equal(t, 0, chunks[0].Offset) +} + +func TestChunkText_ExactSize(t *testing.T) { + text := strings.Repeat("a", 50) + chunks := ChunkText(text, 50, 0) + require.Len(t, chunks, 1) + assert.Equal(t, text, chunks[0].Text) + assert.Equal(t, 0, chunks[0].Offset) +} + +// TestChunkText_OverSizeHardSplitNoOverlap proves the terminal hard-split +// fallback (no separator anywhere in the text) cuts into contiguous, +// size-bounded spans with exactly abutting offsets when overlap is 0. +func TestChunkText_OverSizeHardSplitNoOverlap(t *testing.T) { + text := strings.Repeat("a", 205) + chunks := ChunkText(text, 50, 0) + require.Len(t, chunks, 5) + + wantLens := []int{50, 50, 50, 50, 5} + wantOffsets := []int{0, 50, 100, 150, 200} + for i, c := range chunks { + assert.Equal(t, wantLens[i], utf8.RuneCountInString(c.Text), "chunk %d length", i) + assert.Equal(t, wantOffsets[i], c.Offset, "chunk %d offset", i) + } + // No separators to split on, but the pieces must still reconstruct the + // original text with no gaps or corruption. + assert.Equal(t, text, chunks[0].Text+chunks[1].Text+chunks[2].Text+chunks[3].Text+chunks[4].Text) +} + +// TestChunkText_OverlapExactOnHardSplit proves the hard-split fallback lands +// on exactly the requested overlap (not a whole extra piece) by sizing its +// granularity to overlap rather than size. +func TestChunkText_OverlapExactOnHardSplit(t *testing.T) { + // Each rune is distinguishable by position so an exact-overlap claim is + // actually verifiable, not just plausible. + var sb strings.Builder + for i := range 205 { + sb.WriteByte(byte('a' + i%26)) + } + text := sb.String() + + chunks := ChunkText(text, 50, 10) + require.GreaterOrEqual(t, len(chunks), 2) + + for i := 0; i < len(chunks)-1; i++ { + cur := []rune(chunks[i].Text) + next := []rune(chunks[i+1].Text) + require.GreaterOrEqual(t, len(cur), 10) + tail := string(cur[len(cur)-10:]) + head := string(next[:10]) + assert.Equal(t, tail, head, "chunk %d tail must equal chunk %d head (10-rune overlap)", i, i+1) + } +} + +// TestChunkText_OffsetMonotonicity proves every chunk's Offset strictly +// increases across a variety of separator levels (paragraphs, lines, +// sentences, words, and unbroken text). +func TestChunkText_OffsetMonotonicity(t *testing.T) { + text := strings.Repeat("Paragraph one.\n\nParagraph two has more words in it. And another sentence! And one more?\n\nline one\nline two\n\n", 5) + strings.Repeat("x", 300) + + chunks := ChunkText(text, 80, 20) + require.NotEmpty(t, chunks) + for i := 1; i < len(chunks); i++ { + assert.Greater(t, chunks[i].Offset, chunks[i-1].Offset, "offsets must strictly increase") + } +} + +// TestChunkText_ParagraphBoundaryPreferred proves a text with clean +// paragraph breaks under size splits along them rather than mid-word. +func TestChunkText_ParagraphBoundaryPreferred(t *testing.T) { + p1 := strings.Repeat("alpha ", 5) // 30 runes + p2 := strings.Repeat("beta ", 5) // 25 runes + text := p1 + "\n\n" + p2 + + chunks := ChunkText(text, 32, 0) + require.Len(t, chunks, 2) + assert.Equal(t, p1+"\n\n", chunks[0].Text) + assert.Equal(t, p2, chunks[1].Text) + assert.Equal(t, 0, chunks[0].Offset) + assert.Equal(t, utf8.RuneCountInString(p1+"\n\n"), chunks[1].Offset) +} + +// TestChunkText_SentenceAndWordFallback proves a single long line with no +// paragraph/line breaks falls through to sentence, then word, boundaries. +func TestChunkText_SentenceAndWordFallback(t *testing.T) { + text := "This is sentence one. This is sentence two. This is sentence three. This is sentence four." + chunks := ChunkText(text, 30, 5) + require.NotEmpty(t, chunks) + + // Every chunk's text must be an exact substring of the original at its + // reported offset -- i.e. no corruption from the recursive splitting. + runes := []rune(text) + for _, c := range chunks { + want := string(runes[c.Offset : c.Offset+utf8.RuneCountInString(c.Text)]) + assert.Equal(t, want, c.Text) + } +} + +// TestChunkText_NoSpacesLongWord proves a single token far longer than size +// (no separator matches at any level) still terminates and hard-splits. +func TestChunkText_NoSpacesLongWord(t *testing.T) { + text := strings.Repeat("x", 1000) + chunks := ChunkText(text, 100, 20) + require.NotEmpty(t, chunks) + for _, c := range chunks { + assert.LessOrEqual(t, utf8.RuneCountInString(c.Text), 100) + } +} + +// TestChunkText_DefaultsAreApproximationsInRunes documents that the exported +// defaults are rune counts, not real tokenizer output. +func TestChunkText_DefaultsAreApproximationsInRunes(t *testing.T) { + assert.Equal(t, 512*approxCharsPerToken, DefaultChunkSize) + assert.Equal(t, 64*approxCharsPerToken, DefaultChunkOverlap) + assert.Less(t, DefaultChunkOverlap, DefaultChunkSize) +} + +// TestChunkText_NegativeAndZeroInputsClamp proves size<=0 and overlap<0 (or +// overlap>=size) are normalized rather than left to misbehave or loop. +func TestChunkText_NegativeAndZeroInputsClamp(t *testing.T) { + text := strings.Repeat("word ", 50) + + // size<=0 defaults to DefaultChunkSize -- this text is well under it, so + // it must come back as a single chunk. + chunks := ChunkText(text, 0, 0) + require.Len(t, chunks, 1) + + // overlap<0 clamps to 0 rather than erroring. + chunks = ChunkText(text, 20, -5) + require.NotEmpty(t, chunks) + + // overlap>=size clamps to size-1 rather than stalling the packer forever. + chunks = ChunkText(text, 20, 100) + require.NotEmpty(t, chunks) +} diff --git a/spinifex/handlers/ochrevector/fake_backend_test.go b/spinifex/handlers/ochrevector/fake_backend_test.go new file mode 100644 index 000000000..4a6ae13e4 --- /dev/null +++ b/spinifex/handlers/ochrevector/fake_backend_test.go @@ -0,0 +1,153 @@ +// Provides an unexported fakeBackend used only by in-package ochrevector +// orchestration tests. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "errors" + "sync" +) + +// fakeBackend is an in-memory VectorBackend for orchestration tests: no +// Postgres, just enough bookkeeping to assert what Service asked for and to +// inject failures on demand. +type fakeBackend struct { + mu sync.Mutex + accounts map[string]bool + indexes map[string]IndexSpec // key: accountID+"/"+indexID + + docs map[string]map[string][]VectorRow // key: accountID+"/"+indexID -> sourceKey -> rows + replaceDocumentCalls map[string]int // key: accountID+"/"+indexID+"/"+sourceKey -> call count + + failEnsureAccount error + failCreateIndex error + failReplaceDocument error + failReplaceDocumentFor map[string]error // key: accountID+"/"+indexID+"/"+sourceKey -> error, for per-key failure injection + + failQuery error + queryResults []QueryResult // canned results Query returns, truncated to the clamped k + lastQueryK int // the clamped k from the most recent Query call + lastQueryFilter *Filter // the filter from the most recent Query call, for assertions +} + +var _ VectorBackend = (*fakeBackend)(nil) + +func newFakeBackend() *fakeBackend { + return &fakeBackend{ + accounts: map[string]bool{}, + indexes: map[string]IndexSpec{}, + docs: map[string]map[string][]VectorRow{}, + replaceDocumentCalls: map[string]int{}, + } +} + +func (f *fakeBackend) Init(_ context.Context) error { return nil } + +func (f *fakeBackend) EnsureAccount(_ context.Context, accountID string) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.failEnsureAccount != nil { + return f.failEnsureAccount + } + f.accounts[accountID] = true + return nil +} + +func (f *fakeBackend) CreateIndex(_ context.Context, accountID string, spec IndexSpec) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.failCreateIndex != nil { + return f.failCreateIndex + } + f.indexes[accountID+"/"+spec.ID] = spec + return nil +} + +func (f *fakeBackend) DropIndex(_ context.Context, accountID, indexID string) error { + f.mu.Lock() + defer f.mu.Unlock() + delete(f.indexes, accountID+"/"+indexID) + return nil +} + +func (f *fakeBackend) hasIndex(accountID, indexID string) bool { + f.mu.Lock() + defer f.mu.Unlock() + _, ok := f.indexes[accountID+"/"+indexID] + return ok +} + +func (f *fakeBackend) hasAccount(accountID string) bool { + f.mu.Lock() + defer f.mu.Unlock() + return f.accounts[accountID] +} + +// ReplaceDocument records rows under accountID+"/"+indexID+"/"+sourceKey, +// overwriting whatever was there before -- mirroring pgxBackend's +// delete-then-reinsert semantics without a real transaction, since a Go map +// assignment is already all-or-nothing from the caller's perspective. +func (f *fakeBackend) ReplaceDocument(_ context.Context, accountID, indexID, sourceKey string, rows []VectorRow) error { + f.mu.Lock() + defer f.mu.Unlock() + + docKey := accountID + "/" + indexID + callKey := docKey + "/" + sourceKey + f.replaceDocumentCalls[callKey]++ + + if f.failReplaceDocument != nil { + return f.failReplaceDocument + } + if err, ok := f.failReplaceDocumentFor[callKey]; ok { + return err + } + + if f.docs[docKey] == nil { + f.docs[docKey] = map[string][]VectorRow{} + } + f.docs[docKey][sourceKey] = append([]VectorRow(nil), rows...) + return nil +} + +// documentRows returns the current row set stored for sourceKey, or nil if +// ReplaceDocument was never called for it. +func (f *fakeBackend) documentRows(accountID, indexID, sourceKey string) []VectorRow { + f.mu.Lock() + defer f.mu.Unlock() + return f.docs[accountID+"/"+indexID][sourceKey] +} + +// replaceDocumentCallCount reports how many times ReplaceDocument was called +// for sourceKey, so a re-run test can assert "called again" without the row +// set having doubled. +func (f *fakeBackend) replaceDocumentCallCount(accountID, indexID, sourceKey string) int { + f.mu.Lock() + defer f.mu.Unlock() + return f.replaceDocumentCalls[accountID+"/"+indexID+"/"+sourceKey] +} + +// Query returns f.queryResults truncated to k's clamped value, recording k +// and filter for assertions -- it does not evaluate filter against +// queryResults, since exercising the real filter SQL semantics is the +// pgvector-backed integration test's job, not this in-memory double's. +func (f *fakeBackend) Query(_ context.Context, _, _ string, _ []float32, k int, filter *Filter) ([]QueryResult, error) { + f.mu.Lock() + defer f.mu.Unlock() + + f.lastQueryK = clampQueryK(k) + f.lastQueryFilter = filter + if f.failQuery != nil { + return nil, f.failQuery + } + + results := f.queryResults + if len(results) > f.lastQueryK { + results = results[:f.lastQueryK] + } + return append([]QueryResult(nil), results...), nil +} + +// errFakeBackend is a stand-in backend failure for rollback/error-path tests. +var errFakeBackend = errors.New("fake backend failure") diff --git a/spinifex/handlers/ochrevector/filter.go b/spinifex/handlers/ochrevector/filter.go new file mode 100644 index 000000000..4cfe1e0d0 --- /dev/null +++ b/spinifex/handlers/ochrevector/filter.go @@ -0,0 +1,342 @@ +package handlers_ochrevector + +import ( + "fmt" + "regexp" + "strconv" + "strings" +) + +// FilterOp is one Bedrock-KB-grade metadata filter operator (D9). Scope is +// exactly this set -- comparison, set, string and combinator operators -- +// so a Filter can only ever compile to one of the fixed SQL shapes below, +// never arbitrary SQL. +type FilterOp string + +// The Bedrock Knowledge Bases operator set (D9), verbatim: comparison, +// set, string, and combinator operators. .10 translates the KB API's +// operator names onto these directly. +const ( + FilterEquals FilterOp = "equals" + FilterNotEquals FilterOp = "notEquals" + FilterGreaterThan FilterOp = "greaterThan" + FilterGreaterThanOrEqual FilterOp = "greaterThanOrEquals" + FilterLessThan FilterOp = "lessThan" + FilterLessThanOrEqual FilterOp = "lessThanOrEquals" + FilterIn FilterOp = "in" + FilterNotIn FilterOp = "notIn" + FilterStartsWith FilterOp = "startsWith" + FilterStringContains FilterOp = "stringContains" + FilterListContains FilterOp = "listContains" + FilterAndAll FilterOp = "andAll" + FilterOrAll FilterOp = "orAll" +) + +// filterKeyPattern is the strict identifier allowlist a metadata key must +// pass before it is ever written into SQL text. A jsonb key has no +// parameterized form -- `metadata->>$1` is not valid Postgres, the key +// must appear in the query string itself -- so this regex is the only +// thing standing between a filter key and injection. It is deliberately +// conservative: letters, digits, underscore, must start with a letter or +// underscore, max 64 chars. +var filterKeyPattern = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]{0,63}$`) + +// Filter is one node of the Bedrock-KB metadata filter AST (D9): either a +// leaf comparison of Key against Value under Op, or a combinator (andAll/ +// orAll) over Children. Construct via the Equals/GreaterThan/.../AndAll +// helpers rather than the struct literal directly. +type Filter struct { + Op FilterOp + Key string + Value any + Children []*Filter +} + +// Equals builds an `equals` leaf filter: metadata[key] as text equals value. +func Equals(key string, value any) *Filter { return &Filter{Op: FilterEquals, Key: key, Value: value} } + +// NotEquals builds a `notEquals` leaf filter. +func NotEquals(key string, value any) *Filter { + return &Filter{Op: FilterNotEquals, Key: key, Value: value} +} + +// GreaterThan builds a `greaterThan` leaf filter: metadata[key] cast to +// numeric compared against value. +func GreaterThan(key string, value any) *Filter { + return &Filter{Op: FilterGreaterThan, Key: key, Value: value} +} + +// GreaterThanOrEquals builds a `greaterThanOrEquals` leaf filter. +func GreaterThanOrEquals(key string, value any) *Filter { + return &Filter{Op: FilterGreaterThanOrEqual, Key: key, Value: value} +} + +// LessThan builds a `lessThan` leaf filter. +func LessThan(key string, value any) *Filter { + return &Filter{Op: FilterLessThan, Key: key, Value: value} +} + +// LessThanOrEquals builds a `lessThanOrEquals` leaf filter. +func LessThanOrEquals(key string, value any) *Filter { + return &Filter{Op: FilterLessThanOrEqual, Key: key, Value: value} +} + +// In builds an `in` leaf filter: metadata[key] as text must equal one of values. +func In(key string, values []string) *Filter { + return &Filter{Op: FilterIn, Key: key, Value: values} +} + +// NotIn builds a `notIn` leaf filter: metadata[key] as text must equal none of values. +func NotIn(key string, values []string) *Filter { + return &Filter{Op: FilterNotIn, Key: key, Value: values} +} + +// StartsWith builds a `startsWith` leaf filter: metadata[key] as text +// starts with prefix. +func StartsWith(key, prefix string) *Filter { + return &Filter{Op: FilterStartsWith, Key: key, Value: prefix} +} + +// StringContains builds a `stringContains` leaf filter: metadata[key] as +// text contains substr. +func StringContains(key, substr string) *Filter { + return &Filter{Op: FilterStringContains, Key: key, Value: substr} +} + +// ListContains builds a `listContains` leaf filter: metadata[key], a jsonb +// array, contains value as a top-level element. +func ListContains(key, value string) *Filter { + return &Filter{Op: FilterListContains, Key: key, Value: value} +} + +// AndAll builds an `andAll` combinator over children, all of which must match. +func AndAll(children ...*Filter) *Filter { + return &Filter{Op: FilterAndAll, Children: children} +} + +// OrAll builds an `orAll` combinator over children, any of which may match. +func OrAll(children ...*Filter) *Filter { + return &Filter{Op: FilterOrAll, Children: children} +} + +// compile renders f into a parameterized SQL boolean expression starting at +// bind parameter $startParam, returning the expression, the bound args in +// order, and the next free parameter index for the caller to continue +// numbering from. Every metadata key is validated against filterKeyPattern +// before it reaches the returned SQL string; every value is returned only +// as a bound arg, never interpolated. +func (f *Filter) compile(startParam int) (whereSQL string, args []any, nextParam int, err error) { + if f == nil { + return "", nil, startParam, fmt.Errorf("ochrevector: nil filter") + } + switch f.Op { + case FilterEquals, FilterNotEquals, FilterStartsWith, FilterStringContains: + return f.compileText(startParam) + case FilterGreaterThan, FilterGreaterThanOrEqual, FilterLessThan, FilterLessThanOrEqual: + return f.compileNumeric(startParam) + case FilterIn, FilterNotIn: + return f.compileSet(startParam) + case FilterListContains: + return f.compileListContains(startParam) + case FilterAndAll, FilterOrAll: + return f.compileCombinator(startParam) + default: + return "", nil, startParam, fmt.Errorf("ochrevector: unknown filter operator %q", f.Op) + } +} + +// sanitizeFilterKey validates key against filterKeyPattern, the sole gate +// between a caller-supplied metadata key and raw SQL text. +func sanitizeFilterKey(key string) (string, error) { + if !filterKeyPattern.MatchString(key) { + return "", fmt.Errorf("ochrevector: invalid filter key %q", key) + } + return key, nil +} + +// compileText handles equals/notEquals/startsWith/stringContains: every one +// compares metadata[key] extracted as text (`->>'`) against a single bound +// text value. +func (f *Filter) compileText(startParam int) (string, []any, int, error) { + key, err := sanitizeFilterKey(f.Key) + if err != nil { + return "", nil, startParam, err + } + value, err := filterTextValue(f.Value) + if err != nil { + return "", nil, startParam, fmt.Errorf("ochrevector: filter %q on %q: %w", f.Op, f.Key, err) + } + + var sql string + switch f.Op { + case FilterEquals: + sql = fmt.Sprintf(`metadata->>'%s' = $%d`, key, startParam) + case FilterNotEquals: + sql = fmt.Sprintf(`metadata->>'%s' != $%d`, key, startParam) + case FilterStartsWith: + sql = fmt.Sprintf(`starts_with(metadata->>'%s', $%d)`, key, startParam) + case FilterStringContains: + sql = fmt.Sprintf(`position($%d IN metadata->>'%s') > 0`, startParam, key) + } + return sql, []any{value}, startParam + 1, nil +} + +// compileNumeric handles greaterThan(OrEquals)/lessThan(OrEquals): each +// casts metadata[key] to numeric and compares against a bound numeric value. +func (f *Filter) compileNumeric(startParam int) (string, []any, int, error) { + key, err := sanitizeFilterKey(f.Key) + if err != nil { + return "", nil, startParam, err + } + value, err := filterNumericValue(f.Value) + if err != nil { + return "", nil, startParam, fmt.Errorf("ochrevector: filter %q on %q: %w", f.Op, f.Key, err) + } + + var op string + switch f.Op { + case FilterGreaterThan: + op = ">" + case FilterGreaterThanOrEqual: + op = ">=" + case FilterLessThan: + op = "<" + case FilterLessThanOrEqual: + op = "<=" + } + sql := fmt.Sprintf(`(metadata->>'%s')::numeric %s $%d`, key, op, startParam) + return sql, []any{value}, startParam + 1, nil +} + +// compileSet handles in/notIn: metadata[key] as text compared against a +// single bound array parameter via ANY/ALL. +func (f *Filter) compileSet(startParam int) (string, []any, int, error) { + key, err := sanitizeFilterKey(f.Key) + if err != nil { + return "", nil, startParam, err + } + values, err := filterStringSlice(f.Value) + if err != nil { + return "", nil, startParam, fmt.Errorf("ochrevector: filter %q on %q: %w", f.Op, f.Key, err) + } + + var sql string + switch f.Op { + case FilterIn: + sql = fmt.Sprintf(`metadata->>'%s' = ANY($%d)`, key, startParam) + case FilterNotIn: + sql = fmt.Sprintf(`metadata->>'%s' != ALL($%d)`, key, startParam) + } + return sql, []any{values}, startParam + 1, nil +} + +// compileListContains handles listContains: metadata[key] is treated as a +// jsonb array and tested for value as a top-level element via the `?` +// jsonb-exists operator. +func (f *Filter) compileListContains(startParam int) (string, []any, int, error) { + key, err := sanitizeFilterKey(f.Key) + if err != nil { + return "", nil, startParam, err + } + value, ok := f.Value.(string) + if !ok { + return "", nil, startParam, fmt.Errorf("ochrevector: filter %q on %q requires a string value, got %T", f.Op, f.Key, f.Value) + } + sql := fmt.Sprintf(`metadata->'%s' ? $%d`, key, startParam) + return sql, []any{value}, startParam + 1, nil +} + +// compileCombinator handles andAll/orAll: every child is compiled in order, +// each consuming the next contiguous block of bind parameters, and the +// results are joined and parenthesized as one boolean expression. +func (f *Filter) compileCombinator(startParam int) (string, []any, int, error) { + if len(f.Children) == 0 { + return "", nil, startParam, fmt.Errorf("ochrevector: %q filter requires at least one child", f.Op) + } + joiner := " AND " + if f.Op == FilterOrAll { + joiner = " OR " + } + + parts := make([]string, 0, len(f.Children)) + var args []any + next := startParam + for _, child := range f.Children { + sql, childArgs, n, err := child.compile(next) + if err != nil { + return "", nil, startParam, err + } + parts = append(parts, sql) + args = append(args, childArgs...) + next = n + } + return "(" + strings.Join(parts, joiner) + ")", args, next, nil +} + +// filterTextValue renders v as the single text value a text-comparison +// filter binds. Scalars are formatted, not interpolated -- the returned +// string is always used as a bound arg, never written into SQL. +func filterTextValue(v any) (string, error) { + switch t := v.(type) { + case string: + return t, nil + case bool: + return strconv.FormatBool(t), nil + case int: + return strconv.Itoa(t), nil + case int32: + return strconv.FormatInt(int64(t), 10), nil + case int64: + return strconv.FormatInt(t, 10), nil + case float32: + return strconv.FormatFloat(float64(t), 'f', -1, 32), nil + case float64: + return strconv.FormatFloat(t, 'f', -1, 64), nil + default: + return "", fmt.Errorf("unsupported filter value type %T", v) + } +} + +// filterNumericValue coerces v to float64 for a numeric-comparison filter's +// bound parameter, rejecting non-numeric types. +func filterNumericValue(v any) (float64, error) { + switch t := v.(type) { + case float64: + return t, nil + case float32: + return float64(t), nil + case int: + return float64(t), nil + case int32: + return float64(t), nil + case int64: + return float64(t), nil + default: + return 0, fmt.Errorf("value must be numeric, got %T", v) + } +} + +// filterStringSlice coerces v to a []string for an in/notIn filter's bound +// array parameter, rejecting an empty or non-list value. +func filterStringSlice(v any) ([]string, error) { + var out []string + switch t := v.(type) { + case []string: + out = t + case []any: + out = make([]string, len(t)) + for i, item := range t { + s, err := filterTextValue(item) + if err != nil { + return nil, err + } + out[i] = s + } + default: + return nil, fmt.Errorf("value must be a list, got %T", v) + } + if len(out) == 0 { + return nil, fmt.Errorf("value must be a non-empty list") + } + return out, nil +} diff --git a/spinifex/handlers/ochrevector/filter_test.go b/spinifex/handlers/ochrevector/filter_test.go new file mode 100644 index 000000000..a76719627 --- /dev/null +++ b/spinifex/handlers/ochrevector/filter_test.go @@ -0,0 +1,281 @@ +// Exercises the unexported filter compiler internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestFilter_CompileEachOperator proves every Bedrock-KB operator (D9) +// compiles to its documented parameterized SQL shape, with args bound in +// order and nextParam advancing by exactly the params consumed. +func TestFilter_CompileEachOperator(t *testing.T) { + tests := []struct { + name string + filter *Filter + wantSQL string + wantArgs []any + wantNext int + startParam int + }{ + { + name: "equals", + filter: Equals("status", "active"), + wantSQL: `metadata->>'status' = $1`, + wantArgs: []any{"active"}, + wantNext: 2, + startParam: 1, + }, + { + name: "notEquals", + filter: NotEquals("status", "archived"), + wantSQL: `metadata->>'status' != $1`, + wantArgs: []any{"archived"}, + wantNext: 2, + startParam: 1, + }, + { + name: "greaterThan", + filter: GreaterThan("price", 10.5), + wantSQL: `(metadata->>'price')::numeric > $1`, + wantArgs: []any{10.5}, + wantNext: 2, + startParam: 1, + }, + { + name: "greaterThanOrEquals", + filter: GreaterThanOrEquals("price", 10), + wantSQL: `(metadata->>'price')::numeric >= $1`, + wantArgs: []any{float64(10)}, + wantNext: 2, + startParam: 1, + }, + { + name: "lessThan", + filter: LessThan("price", 100), + wantSQL: `(metadata->>'price')::numeric < $1`, + wantArgs: []any{float64(100)}, + wantNext: 2, + startParam: 1, + }, + { + name: "lessThanOrEquals", + filter: LessThanOrEquals("price", 100), + wantSQL: `(metadata->>'price')::numeric <= $1`, + wantArgs: []any{float64(100)}, + wantNext: 2, + startParam: 1, + }, + { + name: "in", + filter: In("tag", []string{"a", "b", "c"}), + wantSQL: `metadata->>'tag' = ANY($1)`, + wantArgs: []any{[]string{"a", "b", "c"}}, + wantNext: 2, + startParam: 1, + }, + { + name: "notIn", + filter: NotIn("tag", []string{"x", "y"}), + wantSQL: `metadata->>'tag' != ALL($1)`, + wantArgs: []any{[]string{"x", "y"}}, + wantNext: 2, + startParam: 1, + }, + { + name: "startsWith", + filter: StartsWith("path", "/docs/"), + wantSQL: `starts_with(metadata->>'path', $1)`, + wantArgs: []any{"/docs/"}, + wantNext: 2, + startParam: 1, + }, + { + name: "stringContains", + filter: StringContains("body", "quarterly"), + wantSQL: `position($1 IN metadata->>'body') > 0`, + wantArgs: []any{"quarterly"}, + wantNext: 2, + startParam: 1, + }, + { + name: "listContains", + filter: ListContains("owners", "alice"), + wantSQL: `metadata->'owners' ? $1`, + wantArgs: []any{"alice"}, + wantNext: 2, + startParam: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sql, args, next, err := tt.filter.compile(tt.startParam) + require.NoError(t, err) + assert.Equal(t, tt.wantSQL, sql) + assert.Equal(t, tt.wantArgs, args) + assert.Equal(t, tt.wantNext, next) + }) + } +} + +// TestFilter_CompileStartsAtArbitraryParam proves compile numbers its bound +// params from whatever startParam the caller passes -- the property Query +// depends on to continue numbering after the query vector's $1. +func TestFilter_CompileStartsAtArbitraryParam(t *testing.T) { + sql, args, next, err := Equals("status", "active").compile(5) + require.NoError(t, err) + assert.Equal(t, `metadata->>'status' = $5`, sql) + assert.Equal(t, []any{"active"}, args) + assert.Equal(t, 6, next) +} + +// TestFilter_CompileAndAll proves andAll joins every child with AND, +// parenthesized, and that param numbering stays contiguous across children. +func TestFilter_CompileAndAll(t *testing.T) { + f := AndAll( + Equals("category", "handbook"), + GreaterThan("version", 2), + StartsWith("path", "/kb/"), + ) + sql, args, next, err := f.compile(2) + require.NoError(t, err) + assert.Equal(t, + `(metadata->>'category' = $2 AND (metadata->>'version')::numeric > $3 AND starts_with(metadata->>'path', $4))`, + sql) + assert.Equal(t, []any{"handbook", float64(2), "/kb/"}, args) + assert.Equal(t, 5, next) +} + +// TestFilter_CompileOrAll proves orAll joins every child with OR. +func TestFilter_CompileOrAll(t *testing.T) { + f := OrAll(Equals("category", "handbook"), Equals("category", "faq")) + sql, args, next, err := f.compile(1) + require.NoError(t, err) + assert.Equal(t, `(metadata->>'category' = $1 OR metadata->>'category' = $2)`, sql) + assert.Equal(t, []any{"handbook", "faq"}, args) + assert.Equal(t, 3, next) +} + +// TestFilter_CompileNestedCombinators proves a combinator nested inside +// another compiles correctly and keeps param numbering contiguous through +// the nesting. +func TestFilter_CompileNestedCombinators(t *testing.T) { + f := AndAll( + OrAll(Equals("category", "handbook"), Equals("category", "faq")), + LessThanOrEquals("version", 5), + ) + sql, args, next, err := f.compile(1) + require.NoError(t, err) + assert.Equal(t, + `((metadata->>'category' = $1 OR metadata->>'category' = $2) AND (metadata->>'version')::numeric <= $3)`, + sql) + assert.Equal(t, []any{"handbook", "faq", float64(5)}, args) + assert.Equal(t, 4, next) +} + +// TestFilter_CombinatorRequiresChildren proves andAll/orAll with no +// children is a compile error, not an always-true/always-false SQL +// fragment (which could silently drop a filter caller intended). +func TestFilter_CombinatorRequiresChildren(t *testing.T) { + _, _, _, err := AndAll().compile(1) + require.Error(t, err) + + _, _, _, err = OrAll().compile(1) + require.Error(t, err) +} + +// TestFilter_NumericOperatorsRejectNonNumericValue proves a numeric +// comparison with a non-numeric value is a compile error, not a query that +// reaches Postgres and fails a ::numeric cast at execution time. +func TestFilter_NumericOperatorsRejectNonNumericValue(t *testing.T) { + _, _, _, err := GreaterThan("price", "not-a-number").compile(1) + require.Error(t, err) +} + +// TestFilter_InRejectsEmptyList proves in/notIn reject an empty value list. +func TestFilter_InRejectsEmptyList(t *testing.T) { + _, _, _, err := In("tag", nil).compile(1) + require.Error(t, err) + + _, _, _, err = NotIn("tag", []string{}).compile(1) + require.Error(t, err) +} + +// TestFilter_ListContainsRejectsNonStringValue proves listContains rejects +// a non-string value rather than silently stringifying it. +func TestFilter_ListContainsRejectsNonStringValue(t *testing.T) { + _, _, _, err := (&Filter{Op: FilterListContains, Key: "owners", Value: 42}).compile(1) + require.Error(t, err) +} + +// TestFilter_UnknownOperatorErrors proves an operator outside the Bedrock +// KB set is a compile error -- the filter grammar is closed, not +// extensible to arbitrary SQL via a crafted Op string. +func TestFilter_UnknownOperatorErrors(t *testing.T) { + _, _, _, err := (&Filter{Op: "dropTable", Key: "k", Value: "v"}).compile(1) + require.Error(t, err) +} + +// TestFilter_InjectionAttemptInKeyIsRejected is the core security property +// of D9: a jsonb key has no parameterized form, so it can only ever reach +// SQL text after passing filterKeyPattern. An injection payload as a key +// must be rejected outright, for every operator that takes a key. +func TestFilter_InjectionAttemptInKeyIsRejected(t *testing.T) { + maliciousKeys := []string{ + "status'; DROP TABLE idx_one; --", + "status' OR '1'='1", + "a.b", + "a b", + "a\"b", + "", + strings.Repeat("a", 65), // over the 64-char allowlist cap + } + + for _, key := range maliciousKeys { + t.Run(key, func(t *testing.T) { + _, _, _, err := Equals(key, "x").compile(1) + require.Error(t, err, "key %q must be rejected", key) + + _, _, _, err = GreaterThan(key, 1).compile(1) + require.Error(t, err, "key %q must be rejected", key) + + _, _, _, err = ListContains(key, "x").compile(1) + require.Error(t, err, "key %q must be rejected", key) + }) + } +} + +// TestFilter_MaliciousValueStaysBoundArg is D9's other invariant: a value +// is never interpolated into the SQL string, no matter its content -- it +// only ever appears as an element of the returned args slice. +func TestFilter_MaliciousValueStaysBoundArg(t *testing.T) { + const payload = `'; DROP TABLE idx_one; --` + + sql, args, _, err := Equals("status", payload).compile(1) + require.NoError(t, err) + assert.Equal(t, `metadata->>'status' = $1`, sql, "the SQL text must be unaffected by the value") + assert.NotContains(t, sql, payload, "the payload must never appear in the SQL string") + require.Len(t, args, 1) + assert.Equal(t, payload, args[0], "the payload must be preserved intact as a bound arg") + + sql, args, _, err = StringContains("body", payload).compile(1) + require.NoError(t, err) + assert.Equal(t, `position($1 IN metadata->>'body') > 0`, sql) + assert.NotContains(t, sql, payload) + require.Len(t, args, 1) + assert.Equal(t, payload, args[0]) + + sql, args, _, err = In("tag", []string{payload, "safe"}).compile(1) + require.NoError(t, err) + assert.Equal(t, `metadata->>'tag' = ANY($1)`, sql) + assert.NotContains(t, sql, payload) + require.Len(t, args, 1) + assert.Equal(t, []string{payload, "safe"}, args[0]) +} diff --git a/spinifex/handlers/ochrevector/hostport.go b/spinifex/handlers/ochrevector/hostport.go new file mode 100644 index 000000000..268545a8e --- /dev/null +++ b/spinifex/handlers/ochrevector/hostport.go @@ -0,0 +1,242 @@ +package handlers_ochrevector + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/netip" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// The appliance's reachable address is the customer-ENI IP RDS places in the +// system account's default VPC, an OVN logical switch on br-int. The daemon +// host has no route to it, so a healthy appliance is otherwise never +// reachable -- the same host-to-guest gap the Bedrock fix closed for +// InvokeModel. +const ( + ochreDaemonPortRoleTagKey = "spinifex:ochrevector-role" + ochreDaemonPortRole = "daemon-port" + ochreDaemonNodeTagKey = "spinifex:ochrevector-node" +) + +// hostPortPlumber installs and removes the host-side port backing the +// daemon's own ENI. Narrowed to what this package uses so it does not import +// the vm or network/host packages. +type hostPortPlumber interface { + EnsureVPCHostPort(eniID, mac, addr string) error + RemoveVPCHostPort(eniID string) error +} + +// vpcProvisioner is the narrow EC2 surface ensureDaemonENI and subnet +// resolution need: mint-or-describe the daemon's own ENI, and list the +// subnets in the appliance's VPC. +type vpcProvisioner interface { + CreateNetworkInterface(ctx context.Context, input *ec2.CreateNetworkInterfaceInput, accountID string) (*ec2.CreateNetworkInterfaceOutput, error) + DescribeNetworkInterfaces(ctx context.Context, input *ec2.DescribeNetworkInterfacesInput, accountID string) (*ec2.DescribeNetworkInterfacesOutput, error) + DescribeSubnets(ctx context.Context, input *ec2.DescribeSubnetsInput, accountID string) (*ec2.DescribeSubnetsOutput, error) +} + +// HostPortDeps bundles the collaborators the daemon's appliance-VPC host +// port needs. A zero HostPortDeps (HostPort nil) is a valid "do no +// host-port work" configuration, so Appliance instances that never opt in +// (every unit test, or a build with the fix not wired up) are unaffected. +type HostPortDeps struct { + VPC vpcProvisioner + HostPort hostPortPlumber + // NodeID names the daemon replica, so each node's appliance-VPC port is + // a distinct ENI rather than one they contend for. + NodeID string +} + +// daemonENI is the daemon's own network interface in the appliance's VPC, +// unattached to any VM: it exists so ovn-controller has a logical switch +// port to bind the host's internal port to. +type daemonENI struct { + id string + ip string + mac string +} + +// daemonPortDescription is what idempotence keys off. ENIs carry no tag +// filter in DescribeNetworkInterfaces, but description is filterable, so the +// node's identity goes there and the tags are for operator visibility only. +func daemonPortDescription(nodeID string) string { + return "Ochre vector daemon port for node " + nodeID +} + +// ensureDaemonENI describe-or-creates this node's ENI in subnetID. Mirrors +// handlers/bedrock's helper of the same shape: description-filtered +// describe first, create only on a miss, so a restart reuses the same +// address rather than leaking one per launch. +func ensureDaemonENI(ctx context.Context, vpcSvc vpcProvisioner, subnetID, nodeID string) (*daemonENI, error) { + desc := daemonPortDescription(nodeID) + out, err := vpcSvc.DescribeNetworkInterfaces(ctx, &ec2.DescribeNetworkInterfacesInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("subnet-id"), Values: aws.StringSlice([]string{subnetID})}, + {Name: aws.String("description"), Values: aws.StringSlice([]string{desc})}, + }, + }, utils.GlobalAccountID) + if err != nil { + return nil, fmt.Errorf("ochrevector: describe daemon ENI in subnet %s: %w", subnetID, err) + } + if out != nil { + for _, ni := range out.NetworkInterfaces { + eni := &daemonENI{ + id: aws.StringValue(ni.NetworkInterfaceId), + ip: aws.StringValue(ni.PrivateIpAddress), + mac: aws.StringValue(ni.MacAddress), + } + // A record missing either field cannot back a host port, so fall + // through and mint a usable one rather than fail the connect. + if eni.id != "" && eni.ip != "" && eni.mac != "" { + return eni, nil + } + slog.WarnContext(ctx, "ochrevector: ignoring incomplete daemon ENI record", + "node", nodeID, "eni", eni.id) + } + } + + created, err := vpcSvc.CreateNetworkInterface(ctx, &ec2.CreateNetworkInterfaceInput{ + SubnetId: aws.String(subnetID), + Description: aws.String(desc), + TagSpecifications: []*ec2.TagSpecification{{ + ResourceType: aws.String("network-interface"), + Tags: []*ec2.Tag{ + {Key: aws.String(ochreDaemonPortRoleTagKey), Value: aws.String(ochreDaemonPortRole)}, + {Key: aws.String(ochreDaemonNodeTagKey), Value: aws.String(nodeID)}, + }, + }}, + }, utils.GlobalAccountID) + if err != nil { + return nil, fmt.Errorf("ochrevector: create daemon ENI in subnet %s: %w", subnetID, err) + } + if created == nil || created.NetworkInterface == nil { + return nil, fmt.Errorf("ochrevector: create daemon ENI in subnet %s: no interface returned", subnetID) + } + ni := created.NetworkInterface + eni := &daemonENI{ + id: aws.StringValue(ni.NetworkInterfaceId), + ip: aws.StringValue(ni.PrivateIpAddress), + mac: aws.StringValue(ni.MacAddress), + } + if eni.id == "" || eni.ip == "" || eni.mac == "" { + return nil, fmt.Errorf("ochrevector: create daemon ENI in subnet %s: incomplete interface returned", subnetID) + } + return eni, nil +} + +// applianceInstanceTagKey mirrors the unexported rdsInstanceTagKey in +// handlers/rds/launch.go, which stamps it on the appliance's customer ENI at +// launch. Keep this literal in sync with that const; the e2e harness +// (tests/e2e/harness/rds.go) duplicates it the same way. +const applianceInstanceTagKey = "spinifex:rds-db-instance" + +// resolveApplianceTarget finds the appliance's own customer ENI by the tag +// handlers/rds stamps on it at launch, returning the real private IP to dial +// and the subnet it lives in -- never the vanity endpoint hostname, which is +// unresolvable and unroutable from the daemon host. +func resolveApplianceTarget(ctx context.Context, deps HostPortDeps, identifier string) (dialIP, subnetID, subnetCIDR string, err error) { + if deps.VPC == nil { + return "", "", "", errors.New("ochrevector: no VPC provider configured; cannot resolve the appliance's ENI") + } + + out, err := deps.VPC.DescribeNetworkInterfaces(ctx, &ec2.DescribeNetworkInterfacesInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("tag:" + applianceInstanceTagKey), Values: aws.StringSlice([]string{identifier})}, + }, + }, utils.GlobalAccountID) + if err != nil { + return "", "", "", fmt.Errorf("ochrevector: describe appliance ENI for %s: %w", identifier, err) + } + var ni *ec2.NetworkInterface + if out != nil { + for _, cand := range out.NetworkInterfaces { + if aws.StringValue(cand.PrivateIpAddress) != "" && aws.StringValue(cand.SubnetId) != "" { + ni = cand + break + } + } + } + if ni == nil { + return "", "", "", fmt.Errorf("ochrevector: no customer ENI found for appliance %s", identifier) + } + dialIP = aws.StringValue(ni.PrivateIpAddress) + subnetID = aws.StringValue(ni.SubnetId) + + subOut, err := deps.VPC.DescribeSubnets(ctx, &ec2.DescribeSubnetsInput{ + Filters: []*ec2.Filter{{Name: aws.String("subnet-id"), Values: aws.StringSlice([]string{subnetID})}}, + }, utils.GlobalAccountID) + if err != nil { + return "", "", "", fmt.Errorf("ochrevector: describe subnet %s: %w", subnetID, err) + } + if subOut == nil || len(subOut.Subnets) == 0 || aws.StringValue(subOut.Subnets[0].CidrBlock) == "" { + return "", "", "", fmt.Errorf("ochrevector: subnet %s has no cidr block", subnetID) + } + subnetCIDR = aws.StringValue(subOut.Subnets[0].CidrBlock) + return dialIP, subnetID, subnetCIDR, nil +} + +// ensureApplianceHostPort gives this daemon a routed presence in the subnet +// the platform appliance's real ENI lives in, so that address (not the +// unroutable vanity hostname) becomes reachable. One ENI per node, not per +// connect: describe-or-create means a restart reuses the same address +// rather than leaking one per launch. Returns the appliance's dial IP and +// the daemon's own ENI id, so the caller can build the DSN and later remove +// the port. +func ensureApplianceHostPort(ctx context.Context, deps HostPortDeps, identifier string) (dialIP, eniID string, err error) { + if deps.HostPort == nil { + return "", "", errors.New("ochrevector: no host-port plumber configured; the daemon cannot reach the platform appliance") + } + if deps.NodeID == "" { + return "", "", errors.New("ochrevector: no node id; cannot identify this daemon's appliance-VPC port") + } + if deps.VPC == nil { + return "", "", errors.New("ochrevector: no VPC provider configured; cannot mint the daemon's appliance-VPC port") + } + + dialIP, subnetID, subnetCIDR, err := resolveApplianceTarget(ctx, deps, identifier) + if err != nil { + return "", "", err + } + subnet, err := netip.ParsePrefix(subnetCIDR) + if err != nil { + return "", "", fmt.Errorf("ochrevector: parse appliance subnet CIDR %q: %w", subnetCIDR, err) + } + + eni, err := ensureDaemonENI(ctx, deps.VPC, subnetID, deps.NodeID) + if err != nil { + return "", "", err + } + ip, err := netip.ParseAddr(eni.ip) + if err != nil { + return "", "", fmt.Errorf("ochrevector: daemon ENI %s has unparseable address %q: %w", eni.id, eni.ip, err) + } + + // The ENI's own address at the SUBNET's prefix length. A /32 would + // address the port and still leave the appliance unreachable; this + // prefix is what installs the connected route to it. + addr := netip.PrefixFrom(ip, subnet.Bits()) + if err := deps.HostPort.EnsureVPCHostPort(eni.id, eni.mac, addr.String()); err != nil { + return "", "", fmt.Errorf("ochrevector: install daemon host port for ENI %s: %w", eni.id, err) + } + slog.InfoContext(ctx, "ochrevector: daemon appliance-VPC port ready", + "node", deps.NodeID, "eni", eni.id, "addr", addr) + return dialIP, eni.id, nil +} + +// removeApplianceHostPort tears down the port ensureApplianceHostPort +// installed. A no-op when there is nothing to remove, so it is safe to call +// on a daemon that never reached a successful ensure. +func removeApplianceHostPort(deps HostPortDeps, eniID string) error { + if deps.HostPort == nil || eniID == "" { + return nil + } + if err := deps.HostPort.RemoveVPCHostPort(eniID); err != nil { + return fmt.Errorf("ochrevector: remove daemon host port for ENI %s: %w", eniID, err) + } + return nil +} diff --git a/spinifex/handlers/ochrevector/hostport_test.go b/spinifex/handlers/ochrevector/hostport_test.go new file mode 100644 index 000000000..54d17d6dd --- /dev/null +++ b/spinifex/handlers/ochrevector/hostport_test.go @@ -0,0 +1,485 @@ +// Exercises unexported host/port resolution internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + "sync" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + testApplianceIdentifier = "ochre-vector-pg" + testApplianceSubnet = "subnet-appliance-0001" + testApplianceEndpoint = "10.244.1.9" + testApplianceENIID = "eni-appliance-0001" +) + +// fakeVPC is an in-memory vpcProvisioner: every CreateNetworkInterface call +// returns a fresh ENI ID/IP, and records are retained so +// DescribeNetworkInterfaces can answer both the tag-filtered appliance-ENI +// lookup and the description-filtered daemon-ENI describe-or-create. +// subnets backs DescribeSubnets. +type fakeVPC struct { + mu sync.Mutex + nextID int + created []string + records map[string]*ec2.NetworkInterface + subnets []*ec2.Subnet +} + +var _ vpcProvisioner = (*fakeVPC)(nil) + +func (f *fakeVPC) CreateNetworkInterface(_ context.Context, in *ec2.CreateNetworkInterfaceInput, _ string) (*ec2.CreateNetworkInterfaceOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.nextID++ + id := f.nextID + eniID := fmt.Sprintf("eni-%d", id) + f.created = append(f.created, eniID) + ni := &ec2.NetworkInterface{ + NetworkInterfaceId: aws.String(eniID), + PrivateIpAddress: aws.String(fmt.Sprintf("10.244.1.%d", id%250+10)), + MacAddress: aws.String("02:00:00:00:00:01"), + SubnetId: in.SubnetId, + Description: in.Description, + } + if f.records == nil { + f.records = map[string]*ec2.NetworkInterface{} + } + f.records[eniID] = ni + return &ec2.CreateNetworkInterfaceOutput{NetworkInterface: ni}, nil +} + +// putRecord seeds an ENI the fake did not mint, so a test can present the +// lookup with a record it would not otherwise produce. +func (f *fakeVPC) putRecord(eniID, subnetID, description, ip, mac string) { + f.mu.Lock() + defer f.mu.Unlock() + if f.records == nil { + f.records = map[string]*ec2.NetworkInterface{} + } + f.records[eniID] = &ec2.NetworkInterface{ + NetworkInterfaceId: aws.String(eniID), + SubnetId: aws.String(subnetID), + Description: aws.String(description), + PrivateIpAddress: aws.String(ip), + MacAddress: aws.String(mac), + } +} + +// putApplianceENI seeds the appliance's own customer ENI, tagged the way +// handlers/rds tags it at launch, so resolveApplianceTarget's tag-filtered +// lookup can find it. +func (f *fakeVPC) putApplianceENI(eniID, identifier, subnetID, ip string) { + f.mu.Lock() + defer f.mu.Unlock() + if f.records == nil { + f.records = map[string]*ec2.NetworkInterface{} + } + f.records[eniID] = &ec2.NetworkInterface{ + NetworkInterfaceId: aws.String(eniID), + SubnetId: aws.String(subnetID), + PrivateIpAddress: aws.String(ip), + TagSet: []*ec2.Tag{ + {Key: aws.String(applianceInstanceTagKey), Value: aws.String(identifier)}, + }, + } +} + +// tagValue returns the value of key in tags, or "" if absent. +func tagValue(tags []*ec2.Tag, key string) string { + for _, t := range tags { + if aws.StringValue(t.Key) == key { + return aws.StringValue(t.Value) + } + } + return "" +} + +func (f *fakeVPC) DescribeNetworkInterfaces(_ context.Context, in *ec2.DescribeNetworkInterfacesInput, _ string) (*ec2.DescribeNetworkInterfacesOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + var out []*ec2.NetworkInterface + for _, ni := range f.records { + match := true + for _, filter := range in.Filters { + name := aws.StringValue(filter.Name) + var got string + switch { + case name == "subnet-id": + got = aws.StringValue(ni.SubnetId) + case name == "description": + got = aws.StringValue(ni.Description) + case strings.HasPrefix(name, "tag:"): + got = tagValue(ni.TagSet, strings.TrimPrefix(name, "tag:")) + default: + return nil, fmt.Errorf("fakeVPC: unsupported ENI filter %q", name) + } + if !slices.Contains(aws.StringValueSlice(filter.Values), got) { + match = false + break + } + } + if match { + out = append(out, ni) + } + } + return &ec2.DescribeNetworkInterfacesOutput{NetworkInterfaces: out}, nil +} + +func (f *fakeVPC) DescribeSubnets(_ context.Context, in *ec2.DescribeSubnetsInput, _ string) (*ec2.DescribeSubnetsOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + if len(in.Filters) == 0 { + return &ec2.DescribeSubnetsOutput{Subnets: f.subnets}, nil + } + var out []*ec2.Subnet + for _, sn := range f.subnets { + match := true + for _, filter := range in.Filters { + var got string + switch aws.StringValue(filter.Name) { + case "subnet-id": + got = aws.StringValue(sn.SubnetId) + case "vpc-id": + got = aws.StringValue(sn.VpcId) + default: + return nil, fmt.Errorf("fakeVPC: unsupported subnet filter %q", aws.StringValue(filter.Name)) + } + if !slices.Contains(aws.StringValueSlice(filter.Values), got) { + match = false + break + } + } + if match { + out = append(out, sn) + } + } + return &ec2.DescribeSubnetsOutput{Subnets: out}, nil +} + +// fakeHostPort records the host-side ports ensureApplianceHostPort/ +// removeApplianceHostPort asked for. +type fakeHostPort struct { + mu sync.Mutex + ensured []hostPortCall + removed []string + failEnsure error + failRemove error +} + +var _ hostPortPlumber = (*fakeHostPort)(nil) + +type hostPortCall struct { + eniID string + mac string + addr string +} + +func (f *fakeHostPort) EnsureVPCHostPort(eniID, mac, addr string) error { + f.mu.Lock() + f.ensured = append(f.ensured, hostPortCall{eniID: eniID, mac: mac, addr: addr}) + f.mu.Unlock() + return f.failEnsure +} + +func (f *fakeHostPort) RemoveVPCHostPort(eniID string) error { + f.mu.Lock() + f.removed = append(f.removed, eniID) + f.mu.Unlock() + return f.failRemove +} + +func (f *fakeHostPort) calls() []hostPortCall { + f.mu.Lock() + defer f.mu.Unlock() + return slices.Clone(f.ensured) +} + +func (f *fakeHostPort) removals() []string { + f.mu.Lock() + defer f.mu.Unlock() + return slices.Clone(f.removed) +} + +// hostPortHarness bundles every HostPortDeps fake so a test can reach into +// any of them. +type hostPortHarness struct { + vpc *fakeVPC + hostPort *fakeHostPort +} + +func newHostPortHarness() *hostPortHarness { + return &hostPortHarness{ + vpc: &fakeVPC{}, + hostPort: &fakeHostPort{}, + } +} + +func (h *hostPortHarness) deps() HostPortDeps { + return HostPortDeps{ + VPC: h.vpc, + HostPort: h.hostPort, + NodeID: testNodeID, + } +} + +// putSubnet seeds a subnet in the harness's VPC covering the given CIDR. +func (h *hostPortHarness) putSubnet(subnetID, cidr string) { + h.vpc.mu.Lock() + h.vpc.subnets = append(h.vpc.subnets, &ec2.Subnet{ + SubnetId: aws.String(subnetID), + CidrBlock: aws.String(cidr), + }) + h.vpc.mu.Unlock() +} + +// putApplianceENI seeds the appliance's own tagged customer ENI in subnetID, +// so resolveApplianceTarget's tag-filtered describe finds it. +func (h *hostPortHarness) putApplianceENI(identifier, subnetID, ip string) { + h.vpc.putApplianceENI(testApplianceENIID, identifier, subnetID, ip) +} + +const testNodeID = "node-a" + +func TestEnsureApplianceHostPort_CreatesENIAndHostPort(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + + dialIP, eniID, err := ensureApplianceHostPort(t.Context(), h.deps(), testApplianceIdentifier) + require.NoError(t, err) + + assert.Equal(t, testApplianceEndpoint, dialIP, "the dial IP should be the appliance's own ENI address") + require.Len(t, h.vpc.created, 1, "one daemon ENI should be minted") + require.NotEmpty(t, eniID) + calls := h.hostPort.calls() + require.Len(t, calls, 1) + assert.Equal(t, eniID, calls[0].eniID) + assert.Equal(t, "02:00:00:00:00:01", calls[0].mac) +} + +// The port must carry the ENI address at the SUBNET's prefix length: a /32 +// addresses the port and still leaves the appliance unreachable, which is +// the entire reason the port exists. +func TestEnsureApplianceHostPort_UsesSubnetPrefixLength(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + + _, _, err := ensureApplianceHostPort(t.Context(), h.deps(), testApplianceIdentifier) + require.NoError(t, err) + + calls := h.hostPort.calls() + require.Len(t, calls, 1) + assert.True(t, strings.HasSuffix(calls[0].addr, "/24"), + "host port address %q should carry the subnet's /24, not a host route", calls[0].addr) + assert.False(t, strings.HasPrefix(calls[0].addr, "10.244.1.0/"), + "host port address %q must be the ENI's own address, not the network address", calls[0].addr) +} + +// One ENI per node, not one per connect: without describe-or-create every +// Connect would leak an address out of the appliance's subnet. +func TestEnsureApplianceHostPort_ReusesTheNodesExistingENI(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + deps := h.deps() + + firstIP, firstENI, err := ensureApplianceHostPort(t.Context(), deps, testApplianceIdentifier) + require.NoError(t, err) + secondIP, secondENI, err := ensureApplianceHostPort(t.Context(), deps, testApplianceIdentifier) + require.NoError(t, err) + + assert.Equal(t, firstIP, secondIP, "the dial IP should resolve the same both times") + assert.Equal(t, firstENI, secondENI, "the second ensure should reuse the first ENI") + assert.Len(t, h.vpc.created, 1, "the second ensure should reuse the first ENI") + calls := h.hostPort.calls() + require.Len(t, calls, 2, "the host port is still re-driven, since a reboot clears it") + assert.Equal(t, calls[0], calls[1], "both ensures should install the same port") +} + +// Two daemons sharing one address would fight over the same OVN logical +// port, so the node identity has to reach the ENI lookup. +func TestEnsureApplianceHostPort_DistinctENIPerNode(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + a := h.deps() + a.NodeID = "node-a" + b := h.deps() + b.NodeID = "node-b" + + _, eniA, err := ensureApplianceHostPort(t.Context(), a, testApplianceIdentifier) + require.NoError(t, err) + _, eniB, err := ensureApplianceHostPort(t.Context(), b, testApplianceIdentifier) + require.NoError(t, err) + + assert.NotEqual(t, eniA, eniB) + assert.Len(t, h.vpc.created, 2, "each node should own its own ENI") +} + +func TestEnsureApplianceHostPort_RequiresPlumberNodeIDAndVPC(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + + noPlumber := h.deps() + noPlumber.HostPort = nil + _, _, err := ensureApplianceHostPort(t.Context(), noPlumber, testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "host-port plumber") + + noNode := h.deps() + noNode.NodeID = "" + _, _, err = ensureApplianceHostPort(t.Context(), noNode, testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "node id") + + noVPC := h.deps() + noVPC.VPC = nil + _, _, err = ensureApplianceHostPort(t.Context(), noVPC, testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "VPC provider") + + assert.Empty(t, h.vpc.created, "no refusal should have minted an ENI") +} + +func TestEnsureApplianceHostPort_PropagatesPlumberFailure(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + h.hostPort.failEnsure = errors.New("ovs-vsctl exploded") + + _, _, err := ensureApplianceHostPort(t.Context(), h.deps(), testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "ovs-vsctl exploded") +} + +// An ENI record with no address or MAC cannot back a host port, so the +// lookup must fall through and mint a usable one rather than install a +// broken port. +func TestEnsureApplianceHostPort_SkipsIncompleteENIRecords(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + h.vpc.putRecord("eni-broken", testApplianceSubnet, daemonPortDescription(testNodeID), "", "") + + _, eniID, err := ensureApplianceHostPort(t.Context(), h.deps(), testApplianceIdentifier) + require.NoError(t, err) + assert.NotEqual(t, "eni-broken", eniID) + calls := h.hostPort.calls() + require.Len(t, calls, 1) + assert.NotEmpty(t, calls[0].mac) +} + +func TestDaemonPortDescriptionCarriesNodeID(t *testing.T) { + assert.Contains(t, daemonPortDescription("node-a"), "node-a") + assert.NotEqual(t, daemonPortDescription("node-a"), daemonPortDescription("node-b")) +} + +// --- tag-filtered appliance ENI resolution --- + +func TestResolveApplianceTarget_FindsTheTaggedApplianceENI(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + + dialIP, subnetID, cidr, err := resolveApplianceTarget(t.Context(), h.deps(), testApplianceIdentifier) + require.NoError(t, err) + assert.Equal(t, testApplianceEndpoint, dialIP) + assert.Equal(t, testApplianceSubnet, subnetID) + assert.Equal(t, "10.244.1.0/24", cidr) +} + +// An ENI tagged for a different identifier must never match. +func TestResolveApplianceTarget_IgnoresOtherIdentifiersTags(t *testing.T) { + h := newHostPortHarness() + h.putSubnet(testApplianceSubnet, "10.244.1.0/24") + h.putApplianceENI("some-other-db", testApplianceSubnet, testApplianceEndpoint) + + _, _, _, err := resolveApplianceTarget(t.Context(), h.deps(), testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "no customer ENI found") +} + +func TestResolveApplianceTarget_ErrorsWhenSubnetHasNoCIDR(t *testing.T) { + h := newHostPortHarness() + h.putApplianceENI(testApplianceIdentifier, testApplianceSubnet, testApplianceEndpoint) + + _, _, _, err := resolveApplianceTarget(t.Context(), h.deps(), testApplianceIdentifier) + require.Error(t, err) +} + +func TestResolveApplianceTarget_PropagatesDescribeFailure(t *testing.T) { + h := newHostPortHarness() + failing := &failingVPC{fakeVPC: h.vpc, failDescribeENIs: errors.New("nats timeout")} + deps := h.deps() + deps.VPC = failing + + _, _, _, err := resolveApplianceTarget(t.Context(), deps, testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "nats timeout") +} + +func TestResolveApplianceTarget_RequiresVPC(t *testing.T) { + h := newHostPortHarness() + noVPC := h.deps() + noVPC.VPC = nil + _, _, _, err := resolveApplianceTarget(t.Context(), noVPC, testApplianceIdentifier) + require.Error(t, err) + assert.Contains(t, err.Error(), "VPC provider") +} + +// failingVPC wraps a fakeVPC to inject a DescribeNetworkInterfaces failure, +// so resolveApplianceTarget's error propagation can be tested without a real +// EC2 fault. +type failingVPC struct { + *fakeVPC + + failDescribeENIs error +} + +var _ vpcProvisioner = (*failingVPC)(nil) + +func (f *failingVPC) DescribeNetworkInterfaces(ctx context.Context, in *ec2.DescribeNetworkInterfacesInput, accountID string) (*ec2.DescribeNetworkInterfacesOutput, error) { + if f.failDescribeENIs != nil { + return nil, f.failDescribeENIs + } + return f.fakeVPC.DescribeNetworkInterfaces(ctx, in, accountID) +} + +// --- teardown --- + +func TestRemoveApplianceHostPort_NoopWithoutHostPortOrENI(t *testing.T) { + h := newHostPortHarness() + require.NoError(t, removeApplianceHostPort(HostPortDeps{}, "eni-1")) + require.NoError(t, removeApplianceHostPort(h.deps(), "")) + assert.Empty(t, h.hostPort.removals()) +} + +func TestRemoveApplianceHostPort_RemovesTheInstalledPort(t *testing.T) { + h := newHostPortHarness() + require.NoError(t, removeApplianceHostPort(h.deps(), "eni-42")) + assert.Equal(t, []string{"eni-42"}, h.hostPort.removals()) +} + +func TestRemoveApplianceHostPort_PropagatesPlumberFailure(t *testing.T) { + h := newHostPortHarness() + h.hostPort.failRemove = errors.New("br-int is gone") + err := removeApplianceHostPort(h.deps(), "eni-42") + require.Error(t, err) + assert.Contains(t, err.Error(), "br-int is gone") +} diff --git a/spinifex/handlers/ochrevector/identifiers.go b/spinifex/handlers/ochrevector/identifiers.go new file mode 100644 index 000000000..99eeb5cc6 --- /dev/null +++ b/spinifex/handlers/ochrevector/identifiers.go @@ -0,0 +1,65 @@ +package handlers_ochrevector + +import ( + "fmt" + "regexp" + + "github.com/jackc/pgx/v5" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// indexIDPattern is a conservative allowlist for anything that reaches a SQL +// identifier: utils.GenerateResourceID's own alphabet ("idx-<17 hex>") plus +// enough headroom that a caller-supplied id still round-trips safely rather +// than being silently rejected. +var indexIDPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]{0,62}$`) + +// validateAccountID rejects anything that is not a well-formed 12-digit +// AWS-style account id before it ever reaches a SQL identifier or +// search_path. +func validateAccountID(accountID string) error { + if !utils.IsAccountID(accountID) { + return fmt.Errorf("ochrevector: invalid account id %q", accountID) + } + return nil +} + +// validateIndexID rejects anything outside a conservative identifier +// allowlist before it ever reaches a SQL identifier. +func validateIndexID(indexID string) error { + if !indexIDPattern.MatchString(indexID) { + return fmt.Errorf("ochrevector: invalid index id %q", indexID) + } + return nil +} + +// schemaName returns accountID's schema name. Callers must validateAccountID +// first. +func schemaName(accountID string) string { + return "kb_" + accountID +} + +// roleName returns accountID's non-login role name. Callers must +// validateAccountID first. +func roleName(accountID string) string { + return "kb_" + accountID + "_role" +} + +// tableName returns indexID's backing table name. Callers must +// validateIndexID first. +func tableName(indexID string) string { + return "idx_" + indexID +} + +// hnswIndexName returns indexID's HNSW index name. Callers must +// validateIndexID first. +func hnswIndexName(indexID string) string { + return "idx_" + indexID + "_embedding_hnsw" +} + +// sanitizeIdent quotes name as a single SQL identifier via pgx's own +// sanitizer — the last line of defense against injection even though every +// caller has already passed name through an allowlist regex first. +func sanitizeIdent(name string) string { + return pgx.Identifier{name}.Sanitize() +} diff --git a/spinifex/handlers/ochrevector/identifiers_test.go b/spinifex/handlers/ochrevector/identifiers_test.go new file mode 100644 index 000000000..27f869c1b --- /dev/null +++ b/spinifex/handlers/ochrevector/identifiers_test.go @@ -0,0 +1,52 @@ +// Exercises unexported identifier validation internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "testing" + + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/stretchr/testify/assert" +) + +func TestValidateAccountID(t *testing.T) { + assert.NoError(t, validateAccountID(utils.GlobalAccountID)) + assert.NoError(t, validateAccountID("123456789012")) + + assert.Error(t, validateAccountID("")) + assert.Error(t, validateAccountID("not-an-account")) + assert.Error(t, validateAccountID("12345")) + assert.Error(t, validateAccountID("1234567890123")) +} + +func TestValidateIndexID(t *testing.T) { + assert.NoError(t, validateIndexID("idx-0123456789abcdef1")) + assert.NoError(t, validateIndexID("a")) + assert.NoError(t, validateIndexID("my_index-1")) + + assert.Error(t, validateIndexID("")) + assert.Error(t, validateIndexID("-leading-dash")) + assert.Error(t, validateIndexID("has a space")) + assert.Error(t, validateIndexID("has;semicolon")) + assert.Error(t, validateIndexID("has\"quote")) +} + +func TestSchemaAndRoleNames(t *testing.T) { + assert.Equal(t, "kb_"+utils.GlobalAccountID, schemaName(utils.GlobalAccountID)) + assert.Equal(t, "kb_"+utils.GlobalAccountID+"_role", roleName(utils.GlobalAccountID)) +} + +func TestTableAndIndexNames(t *testing.T) { + assert.Equal(t, "idx_myindex", tableName("myindex")) + assert.Equal(t, "idx_myindex_embedding_hnsw", hnswIndexName("myindex")) +} + +func TestSanitizeIdent(t *testing.T) { + assert.Equal(t, `"kb_000000000000"`, sanitizeIdent("kb_000000000000")) + // A name containing a double quote is escaped by doubling it, pgx's own + // defense against an identifier that reaches Sanitize despite the + // allowlist regex rejecting it earlier. + assert.Equal(t, `"weird""name"`, sanitizeIdent(`weird"name`)) +} diff --git a/spinifex/handlers/ochrevector/ingest.go b/spinifex/handlers/ochrevector/ingest.go new file mode 100644 index 000000000..83dd88103 --- /dev/null +++ b/spinifex/handlers/ochrevector/ingest.go @@ -0,0 +1,348 @@ +package handlers_ochrevector + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/s3" + "github.com/mulgadc/spinifex/spinifex/objectstore" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// Embedder is the local seam for batch text embedding: the same method set +// as gateway_bedrock.Embedder, restated here rather than imported so this +// package never depends on the gateway. The real embedder (an +// *gateway_bedrock adapter) satisfies this structurally; wiring it in is a +// later stage's job. +type Embedder interface { + Embed(ctx context.Context, modelID string, inputs []string) ([][]float32, error) +} + +const ( + // ingestListPageSize bounds each ListObjectsV2 page; StartIngest's caller + // may point a job at a bucket/prefix far larger than fits in memory at + // once, so RunJob pages through it via ContinuationToken. + ingestListPageSize = 1000 + + // maxIngestObjectBytes bounds a single source object read into memory + // during ingestion. An object over this size is a per-document failure + // (recorded, not fatal to the job), not an unbounded read. + maxIngestObjectBytes = 32 * 1024 * 1024 + + // ingestMaxEmbedRetries bounds the retry attempts for one document's + // embed call before it counts as a failure -- the "bounded retry" D7 + // requires before a fully-down embedder fails the job. + ingestMaxEmbedRetries = 3 + + // jobStaleAfter bounds how long a RUNNING job may sit before Reconcile + // treats it as abandoned by a crashed worker and re-drives it from the + // start -- safe because ReplaceDocument makes reprocessing an + // already-done document a no-op (D7). + jobStaleAfter = 5 * time.Minute +) + +// ingestRetryBackoff paces retries between embed attempts. A var, not a +// const, so tests can shrink it rather than paying the real delay. +var ingestRetryBackoff = 200 * time.Millisecond + +// IngestService orchestrates the ingestion job lifecycle: claiming a job, +// listing a source bucket/prefix, chunking and embedding each object, and +// replacing its rows in the vector backend. +type IngestService struct { + Jobs *JobStore + Registry *Registry + Backend VectorBackend + Store objectstore.ObjectStore + Embedder Embedder +} + +// NewIngestService constructs an IngestService over its dependencies. +func NewIngestService(jobs *JobStore, registry *Registry, backend VectorBackend, store objectstore.ObjectStore, embedder Embedder) *IngestService { + return &IngestService{Jobs: jobs, Registry: registry, Backend: backend, Store: store, Embedder: embedder} +} + +// StartIngest validates indexID exists for accountID, then reserves a new +// PENDING job for source. It does not run the job -- a scheduler/daemon +// calls RunJob later (Stage 4); tests call RunJob directly. +func (s *IngestService) StartIngest(ctx context.Context, accountID, indexID string, source SourceSpec) (*JobRecord, error) { + idxRec, err := s.Registry.Get(ctx, accountID, indexID) + if err != nil { + return nil, err + } + if idxRec == nil { + return nil, ErrIndexNotFound + } + + now := time.Now().UTC() + rec := JobRecord{ + ID: utils.GenerateResourceID("job"), + IndexID: indexID, + Source: source, + State: JobStatePending, + CreatedAt: now, + UpdatedAt: now, + } + if err := s.Jobs.Reserve(ctx, accountID, rec); err != nil { + return nil, err + } + rec.AccountID = accountID + return &rec, nil +} + +// RunJob is the synchronous ingestion worker: lists job.Source, and for each +// object, reads/chunks/embeds/replaces its rows. A per-document failure +// (unreadable object, oversize, exhausted embed retries) is skipped and +// recorded rather than failing the whole job; a fully-unavailable embedder +// (every attempted document's embed calls exhaust retries) fails the job +// instead of reaching READY with everything skipped (D7). On success, the +// job's source spec is recorded on the index registry record (D4). +func (s *IngestService) RunJob(ctx context.Context, job JobRecord) error { + if err := s.Jobs.SetState(ctx, job.AccountID, job.ID, JobStateRunning); err != nil { + return err + } + + keys, err := s.listObjectKeys(ctx, job.Source) + if err != nil { + return s.failJob(ctx, job, fmt.Errorf("list source objects: %w", err)) + } + + var failedDocs []FailedDoc + done := 0 + embedFailures := 0 + + for _, key := range keys { + docErr := s.ingestObject(ctx, job.AccountID, job.IndexID, key, job.Source) + if docErr != nil { + slog.WarnContext(ctx, "ochrevector: ingest document failed", "job", job.ID, "key", key, "err", docErr.err) + failedDocs = append(failedDocs, FailedDoc{SourceKey: key, Reason: docErr.Error()}) + if docErr.fromEmbedder { + embedFailures++ + } + continue + } + done++ + } + + // A job that ingests nothing must not report READY, regardless of why: + // an all-embedder-down job gets the specific message, but any other + // all-failing reason (every GetObject failing, etc.) fails the job too + // rather than reaching READY with done==0. + if len(keys) > 0 && done == 0 { + if embedFailures == len(keys) { + return s.failJob(ctx, job, fmt.Errorf("ochrevector: embedder unavailable for all %d documents", len(keys))) + } + return s.failJob(ctx, job, fmt.Errorf("ochrevector: no documents ingested out of %d", len(keys))) + } + + if err := s.Registry.AppendSourceSpec(ctx, job.AccountID, job.IndexID, job.Source); err != nil { + return s.failJob(ctx, job, fmt.Errorf("record source spec: %w", err)) + } + + return s.Jobs.Update(ctx, job.AccountID, job.ID, func(rec *JobRecord) { + rec.State = JobStateReady + rec.DocumentsTotal = len(keys) + rec.DocumentsDone = done + rec.FailedDocuments = failedDocs + rec.Error = "" + }) +} + +// failJob transitions job to FAILED with cause recorded, and returns cause +// (joined with any error updating the record itself) so the caller sees why. +func (s *IngestService) failJob(ctx context.Context, job JobRecord, cause error) error { + if updErr := s.Jobs.Update(ctx, job.AccountID, job.ID, func(rec *JobRecord) { + rec.State = JobStateFailed + rec.Error = cause.Error() + }); updErr != nil { + return errors.Join(cause, updErr) + } + return cause +} + +// Reconcile re-drives every RUNNING job stale past jobStaleAfter -- the +// crash-recovery sweep a scheduler calls on a timer (no scheduler here). +// Re-running RunJob from the start is safe: ReplaceDocument makes +// reprocessing any already-done document a no-op (D7). +func (s *IngestService) Reconcile(ctx context.Context) error { + jobs, err := s.Jobs.ListAll(ctx) + if err != nil { + return fmt.Errorf("ochrevector: reconcile ingest jobs: list all: %w", err) + } + now := time.Now().UTC() + var errs []error + for _, job := range jobs { + if job.State != JobStateRunning { + continue + } + if now.Sub(job.UpdatedAt) < jobStaleAfter { + continue + } + if err := s.RunJob(ctx, job); err != nil { + errs = append(errs, fmt.Errorf("ochrevector: reconcile job %s: %w", job.ID, err)) + } + } + return errors.Join(errs...) +} + +// Sweep is the scheduler tick the daemon runs on a timer: it claims and runs +// every PENDING job, and re-drives any RUNNING job a crashed worker abandoned +// (stale past jobStaleAfter). RunJob is idempotent, so a re-run never +// double-writes a document. Reconcile stays the narrower crash-recovery method. +func (s *IngestService) Sweep(ctx context.Context) error { + jobs, err := s.Jobs.ListAll(ctx) + if err != nil { + return fmt.Errorf("ochrevector: sweep ingest jobs: list all: %w", err) + } + now := time.Now().UTC() + var errs []error + for _, job := range jobs { + switch { + case job.State == JobStatePending: + case job.State == JobStateRunning && now.Sub(job.UpdatedAt) >= jobStaleAfter: + default: + continue + } + if err := s.RunJob(ctx, job); err != nil { + errs = append(errs, fmt.Errorf("ochrevector: sweep job %s: %w", job.ID, err)) + } + } + return errors.Join(errs...) +} + +// ingestDocError is one document's ingestion failure, tagged with whether it +// originated from the embedder (after exhausted retries) so RunJob can tell +// an isolated per-doc miss from every document failing the same way. +type ingestDocError struct { + err error + fromEmbedder bool +} + +func (e *ingestDocError) Error() string { return e.err.Error() } +func (e *ingestDocError) Unwrap() error { return e.err } + +// ingestObject reads, chunks, embeds and replaces one source object's rows. +func (s *IngestService) ingestObject(ctx context.Context, accountID, indexID, key string, source SourceSpec) *ingestDocError { + obj, err := s.Store.GetObject(ctx, &s3.GetObjectInput{Bucket: aws.String(source.Bucket), Key: aws.String(key)}) + if err != nil { + return &ingestDocError{err: fmt.Errorf("get object: %w", err)} + } + defer obj.Body.Close() + + data, err := io.ReadAll(io.LimitReader(obj.Body, maxIngestObjectBytes+1)) + if err != nil { + return &ingestDocError{err: fmt.Errorf("read object body: %w", err)} + } + if len(data) > maxIngestObjectBytes { + return &ingestDocError{err: fmt.Errorf("object exceeds max ingest size (%d bytes)", maxIngestObjectBytes)} + } + + chunkSize := source.ChunkSize + if chunkSize <= 0 { + chunkSize = DefaultChunkSize + } + chunkOverlap := source.ChunkOverlap + if chunkOverlap < 0 { + chunkOverlap = DefaultChunkOverlap + } + + chunks := ChunkText(string(data), chunkSize, chunkOverlap) + if len(chunks) == 0 { + // An empty/whitespace-only document is not an error, but any + // previously-ingested rows for this key must still be cleared. + if err := s.Backend.ReplaceDocument(ctx, accountID, indexID, key, nil); err != nil { + return &ingestDocError{err: fmt.Errorf("replace document: %w", err)} + } + return nil + } + + texts := make([]string, len(chunks)) + for i, c := range chunks { + texts[i] = c.Text + } + + vectors, err := s.embedWithRetry(ctx, source.EmbeddingModel, texts) + if err != nil { + return &ingestDocError{err: fmt.Errorf("embed: %w", err), fromEmbedder: true} + } + if len(vectors) != len(chunks) { + return &ingestDocError{err: fmt.Errorf("embed: vector count mismatch (%d vectors for %d chunks)", len(vectors), len(chunks)), fromEmbedder: true} + } + + rows := make([]VectorRow, len(chunks)) + for i, c := range chunks { + rows[i] = VectorRow{Embedding: vectors[i], Chunk: c.Text, Metadata: rowMetadata(source, key), SourceOffset: c.Offset} + } + + if err := s.Backend.ReplaceDocument(ctx, accountID, indexID, key, rows); err != nil { + return &ingestDocError{err: fmt.Errorf("replace document: %w", err)} + } + return nil +} + +// embedWithRetry retries a batch Embed call up to ingestMaxEmbedRetries +// times, the bounded retry D7 requires before a document's embed failure +// counts against the job. +func (s *IngestService) embedWithRetry(ctx context.Context, modelID string, inputs []string) ([][]float32, error) { + var lastErr error + for attempt := range ingestMaxEmbedRetries { + vectors, err := s.Embedder.Embed(ctx, modelID, inputs) + if err == nil { + return vectors, nil + } + lastErr = err + if attempt == ingestMaxEmbedRetries-1 { + break + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(ingestRetryBackoff): + } + } + return nil, lastErr +} + +// rowMetadata stamps source's static per-source tags onto one ingested +// row's metadata, plus the originating source key -- without this, D9 +// filters have nothing to match against, since a row's metadata would +// otherwise always be empty. Returns a fresh map per call so rows never +// alias the same underlying map. +func rowMetadata(source SourceSpec, key string) map[string]any { + meta := make(map[string]any, len(source.Metadata)+1) + for k, v := range source.Metadata { + meta[k] = v + } + meta["sourceKey"] = key + return meta +} + +// listObjectKeys pages through job.Source's bucket/prefix via +// ContinuationToken, returning every object key found. +func (s *IngestService) listObjectKeys(ctx context.Context, source SourceSpec) ([]string, error) { + var keys []string + var token *string + for { + out, err := s.Store.ListObjectsV2(ctx, &s3.ListObjectsV2Input{ + Bucket: aws.String(source.Bucket), + Prefix: aws.String(source.Prefix), + MaxKeys: aws.Int64(ingestListPageSize), + ContinuationToken: token, + }) + if err != nil { + return nil, err + } + for _, obj := range out.Contents { + keys = append(keys, aws.StringValue(obj.Key)) + } + if out.IsTruncated == nil || !*out.IsTruncated || out.NextContinuationToken == nil { + break + } + token = out.NextContinuationToken + } + return keys, nil +} diff --git a/spinifex/handlers/ochrevector/ingest_test.go b/spinifex/handlers/ochrevector/ingest_test.go new file mode 100644 index 000000000..a26f84b9d --- /dev/null +++ b/spinifex/handlers/ochrevector/ingest_test.go @@ -0,0 +1,552 @@ +// Exercises unexported ingest pipeline internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "encoding/json" + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/s3" + "github.com/mulgadc/spinifex/spinifex/objectstore" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + ingestAccountA = "111111111111" + ingestBucket = "docs" + ingestPrefix = "kb/" +) + +// stubEmbedder is a controllable Embedder test double: it can fail every +// call, fail only calls whose input contains a marker substring, and always +// returns a deterministic (non-random) vector per input so ingest_test can +// assert exact row contents. +type stubEmbedder struct { + mu sync.Mutex + calls int + failAll bool + failIfContains string +} + +var _ Embedder = (*stubEmbedder)(nil) + +func (e *stubEmbedder) Embed(_ context.Context, _ string, inputs []string) ([][]float32, error) { + e.mu.Lock() + e.calls++ + e.mu.Unlock() + + if e.failAll { + return nil, errors.New("stub embedder: unavailable") + } + if e.failIfContains != "" { + for _, in := range inputs { + if strings.Contains(in, e.failIfContains) { + return nil, errors.New("stub embedder: rejected input") + } + } + } + vectors := make([][]float32, len(inputs)) + for i, in := range inputs { + vectors[i] = []float32{float32(len(in)), float32(i)} + } + return vectors, nil +} + +func (e *stubEmbedder) callCount() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls +} + +// newIngestTestSetup wires an IngestService over fresh in-memory stores, with +// accountA already owning a READY index "idx-one". +func newIngestTestSetup(t *testing.T) (*IngestService, *Registry, *fakeBackend, *objectstore.MemoryObjectStore, *stubEmbedder) { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + + registry := NewRegistry(js) + now := time.Now().UTC() + require.NoError(t, registry.Reserve(context.Background(), ingestAccountA, Record{ + ID: "idx-one", Dimension: 2, EmbeddingModel: "stub-embed", State: StateCreating, CreatedAt: now, UpdatedAt: now, + })) + require.NoError(t, registry.SetState(context.Background(), ingestAccountA, "idx-one", StateReady)) + + backend := newFakeBackend() + store := objectstore.NewMemoryObjectStore() + embedder := &stubEmbedder{} + jobs := NewJobStore(js) + + svc := NewIngestService(jobs, registry, backend, store, embedder) + return svc, registry, backend, store, embedder +} + +func testSource() SourceSpec { + return SourceSpec{Bucket: ingestBucket, Prefix: ingestPrefix, ChunkSize: 100, ChunkOverlap: 10, EmbeddingModel: "stub-embed", Dimension: 2} +} + +func putObject(t *testing.T, store *objectstore.MemoryObjectStore, key, body string) { + t.Helper() + _, err := store.PutObject(context.Background(), &s3.PutObjectInput{ + Bucket: aws.String(ingestBucket), + Key: aws.String(key), + Body: strings.NewReader(body), + }) + require.NoError(t, err) +} + +func TestStartIngest_MissingIndexErrors(t *testing.T) { + svc, _, _, _, _ := newIngestTestSetup(t) + ctx := context.Background() + + _, err := svc.StartIngest(ctx, ingestAccountA, "idx-does-not-exist", testSource()) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexNotFound) +} + +func TestStartIngest_ReturnsPendingJobWithoutRunning(t *testing.T) { + svc, _, backend, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + putObject(t, store, ingestPrefix+"doc1.txt", "hello world") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NotNil(t, job) + assert.Equal(t, JobStatePending, job.State) + assert.Equal(t, ingestAccountA, job.AccountID) + assert.Equal(t, "idx-one", job.IndexID) + + // StartIngest must not itself run the job. + assert.Equal(t, 0, embedder.callCount()) + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", ingestPrefix+"doc1.txt")) +} + +func TestRunJob_HappyPath(t *testing.T) { + svc, registry, backend, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + + putObject(t, store, ingestPrefix+"doc1.txt", "the quick brown fox jumps over the lazy dog") + putObject(t, store, ingestPrefix+"doc2.txt", "a second, different document body") + putObject(t, store, ingestPrefix+"doc3.txt", "and a third one, for good measure") + + source := testSource() + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", source) + require.NoError(t, err) + + require.NoError(t, svc.RunJob(ctx, *job)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 3, got.DocumentsTotal) + assert.Equal(t, 3, got.DocumentsDone) + assert.Empty(t, got.FailedDocuments) + assert.Empty(t, got.Error) + + for _, key := range []string{ingestPrefix + "doc1.txt", ingestPrefix + "doc2.txt", ingestPrefix + "doc3.txt"} { + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key), "key %s", key) + rows := backend.documentRows(ingestAccountA, "idx-one", key) + require.NotEmpty(t, rows, "key %s must have rows", key) + for _, row := range rows { + assert.NotEmpty(t, row.Embedding) + assert.NotEmpty(t, row.Chunk) + } + } + assert.Positive(t, embedder.callCount()) + + idxRec, err := registry.Get(ctx, ingestAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, idxRec) + require.Len(t, idxRec.SourceSpecs, 1, "source spec must be appended exactly once") + assert.Equal(t, source, idxRec.SourceSpecs[0]) +} + +// TestRunJob_IdempotentRerun proves running the same job twice keeps the +// backend's per-key row set replaced, not doubled, and appends the source +// spec only once despite two successful runs. +func TestRunJob_IdempotentRerun(t *testing.T) { + svc, registry, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "the quick brown fox jumps over the lazy dog") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + + require.NoError(t, svc.RunJob(ctx, *job)) + firstRows := backend.documentRows(ingestAccountA, "idx-one", key) + require.NotEmpty(t, firstRows) + + // Re-fetch the job (state now READY) and run it again, as Reconcile + // would for a stale RUNNING job. + rerun, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NoError(t, svc.RunJob(ctx, *rerun)) + + assert.Equal(t, 2, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key), "ReplaceDocument is called again on rerun") + secondRows := backend.documentRows(ingestAccountA, "idx-one", key) + assert.Len(t, secondRows, len(firstRows), "the row set must be replaced, not doubled") + + idxRec, err := registry.Get(ctx, ingestAccountA, "idx-one") + require.NoError(t, err) + require.Len(t, idxRec.SourceSpecs, 1, "an identical source spec must not be appended twice") +} + +// TestRunJob_PerDocumentFailureIsSkippedAndRecorded proves one document +// whose embed call fails is recorded in FailedDocuments while the rest of +// the job proceeds to READY. +func TestRunJob_PerDocumentFailureIsSkippedAndRecorded(t *testing.T) { + svc, _, backend, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + origBackoff := ingestRetryBackoff + ingestRetryBackoff = time.Millisecond + t.Cleanup(func() { ingestRetryBackoff = origBackoff }) + + goodKey := ingestPrefix + "good.txt" + badKey := ingestPrefix + "bad.txt" + putObject(t, store, goodKey, "a perfectly fine document body") + putObject(t, store, badKey, "this document contains POISON and will fail to embed") + embedder.failIfContains = "POISON" + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.RunJob(ctx, *job)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 2, got.DocumentsTotal) + assert.Equal(t, 1, got.DocumentsDone) + require.Len(t, got.FailedDocuments, 1) + assert.Equal(t, badKey, got.FailedDocuments[0].SourceKey) + assert.NotEmpty(t, got.FailedDocuments[0].Reason) + + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", goodKey)) + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", badKey)) +} + +// TestRunJob_EmbedderFullyDownFailsJob proves that when every document's +// embed call fails, the job reaches FAILED (after the bounded retry) rather +// than READY with everything skipped. +func TestRunJob_EmbedderFullyDownFailsJob(t *testing.T) { + svc, _, backend, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + origBackoff := ingestRetryBackoff + ingestRetryBackoff = time.Millisecond + t.Cleanup(func() { ingestRetryBackoff = origBackoff }) + + putObject(t, store, ingestPrefix+"doc1.txt", "document one") + putObject(t, store, ingestPrefix+"doc2.txt", "document two") + embedder.failAll = true + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + + err = svc.RunJob(ctx, *job) + require.Error(t, err) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateFailed, got.State) + assert.NotEmpty(t, got.Error) + + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", ingestPrefix+"doc1.txt")) + // Every doc must have been retried up to the bound before giving up. + assert.Equal(t, 2*ingestMaxEmbedRetries, embedder.callCount()) +} + +func TestRunJob_EmptyDocumentClearsExistingRows(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "empty.txt" + putObject(t, store, key, " \n\t ") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.RunJob(ctx, *job)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 1, got.DocumentsDone) + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key)) + assert.Empty(t, backend.documentRows(ingestAccountA, "idx-one", key)) +} + +// forceJobUpdatedAt seeds jobID's UpdatedAt directly via the underlying KV +// bucket, bypassing JobStore.Update (which always refreshes UpdatedAt to +// now), so a test can simulate a record a crashed worker left stale -- +// mirroring appliance_test.go's raw-KV seeding for the same reason. +func forceJobUpdatedAt(t *testing.T, store *JobStore, accountID, jobID string, when time.Time) { + t.Helper() + ctx := context.Background() + kv, err := store.bucket(ctx) + require.NoError(t, err) + entry, err := kv.Get(ctx, jobKey(accountID, jobID)) + require.NoError(t, err) + var rec JobRecord + require.NoError(t, json.Unmarshal(entry.Value(), &rec)) + rec.UpdatedAt = when + data, err := json.Marshal(rec) + require.NoError(t, err) + _, err = kv.Update(ctx, jobKey(accountID, jobID), data, entry.Revision()) + require.NoError(t, err) +} + +func TestIngestService_Reconcile_RedrivesStaleRunningJob(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "a document that a crashed worker never finished") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, job.ID, JobStateRunning)) + + // Force the record's UpdatedAt stale, as a crashed worker would leave it. + forceJobUpdatedAt(t, svc.Jobs, ingestAccountA, job.ID, time.Now().UTC().Add(-2*jobStaleAfter)) + + require.NoError(t, svc.Reconcile(ctx)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key)) +} + +func TestIngestService_Reconcile_LeavesFreshRunningJobAlone(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "still in flight") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, job.ID, JobStateRunning)) + + require.NoError(t, svc.Reconcile(ctx)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateRunning, got.State, "a RUNNING job within the grace period must be left alone") + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key)) +} + +// failingGetObjectStore wraps a MemoryObjectStore so ListObjectsV2 lists +// real objects but every GetObject call fails, simulating a +// non-embedder-related reason every document fails for (e.g. a source +// bucket gone unreachable mid-job). +type failingGetObjectStore struct { + *objectstore.MemoryObjectStore +} + +func (f *failingGetObjectStore) GetObject(_ context.Context, _ *s3.GetObjectInput) (*s3.GetObjectOutput, error) { + return nil, errors.New("simulated GetObject failure") +} + +// TestRunJob_AllDocumentsFailNonEmbedReason_FailsJob proves a job whose +// every document fails for a reason other than the embedder (every +// GetObject failing here) reaches FAILED, not READY with done==0 -- the +// old check only caught the all-embed-failures case, so a job that ingests +// literally nothing for any other reason used to sail through to READY. +func TestRunJob_AllDocumentsFailNonEmbedReason_FailsJob(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + ctx := context.Background() + + registry := NewRegistry(js) + now := time.Now().UTC() + require.NoError(t, registry.Reserve(ctx, ingestAccountA, Record{ + ID: "idx-one", Dimension: 2, EmbeddingModel: "stub-embed", State: StateCreating, CreatedAt: now, UpdatedAt: now, + })) + require.NoError(t, registry.SetState(ctx, ingestAccountA, "idx-one", StateReady)) + + mem := objectstore.NewMemoryObjectStore() + putObject(t, mem, ingestPrefix+"doc1.txt", "one") + putObject(t, mem, ingestPrefix+"doc2.txt", "two") + store := &failingGetObjectStore{MemoryObjectStore: mem} + + backend := newFakeBackend() + embedder := &stubEmbedder{} + jobs := NewJobStore(js) + svc := NewIngestService(jobs, registry, backend, store, embedder) + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + + err = svc.RunJob(ctx, *job) + require.Error(t, err) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateFailed, got.State) + assert.NotEmpty(t, got.Error) + assert.Equal(t, 0, got.DocumentsDone) + + idxRec, err := registry.Get(ctx, ingestAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, idxRec) + assert.Empty(t, idxRec.SourceSpecs, "a job that ingested nothing must not append its source spec") + + // The embedder was never even reached -- these are GetObject failures. + assert.Equal(t, 0, embedder.callCount()) +} + +// TestRunJob_StampsSourceMetadataOntoRows proves ingestObject stamps +// SourceSpec.Metadata's static per-source tags, plus the source key, onto +// every ingested row -- without this, D9 filters have nothing to match. +func TestRunJob_StampsSourceMetadataOntoRows(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "a document with static tags to filter on") + + source := testSource() + source.Metadata = map[string]string{"category": "handbook", "team": "platform"} + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", source) + require.NoError(t, err) + require.NoError(t, svc.RunJob(ctx, *job)) + + rows := backend.documentRows(ingestAccountA, "idx-one", key) + require.NotEmpty(t, rows) + for _, row := range rows { + assert.Equal(t, "handbook", row.Metadata["category"]) + assert.Equal(t, "platform", row.Metadata["team"]) + assert.Equal(t, key, row.Metadata["sourceKey"], "the source key must also be stamped, for a filter to narrow by originating document") + } + + // A second document's rows carry the same source tags but its own key -- + // a filter on category=handbook selects both; a filter on sourceKey + // selects the one document. + otherKey := ingestPrefix + "doc2.txt" + putObject(t, store, otherKey, "a second document, same source tags") + job2, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", source) + require.NoError(t, err) + require.NoError(t, svc.RunJob(ctx, *job2)) + + otherRows := backend.documentRows(ingestAccountA, "idx-one", otherKey) + require.NotEmpty(t, otherRows) + for _, row := range otherRows { + assert.Equal(t, "handbook", row.Metadata["category"], "both documents share the source's static tags") + assert.Equal(t, otherKey, row.Metadata["sourceKey"], "each document's rows carry its own source key") + } +} + +func TestIngestService_Reconcile_IgnoresPendingAndTerminalJobs(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + putObject(t, store, ingestPrefix+"doc1.txt", "pending, never run") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + forceJobUpdatedAt(t, svc.Jobs, ingestAccountA, job.ID, time.Now().UTC().Add(-2*jobStaleAfter)) + + require.NoError(t, svc.Reconcile(ctx)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + assert.Equal(t, JobStatePending, got.State, "Reconcile only re-drives RUNNING jobs, not PENDING ones") + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", ingestPrefix+"doc1.txt")) +} + +// TestIngestService_Sweep_RunsPendingJobToReady proves Sweep is what actually +// drives a fresh PENDING job to completion -- the gap Reconcile deliberately +// leaves open, since nothing else ever calls RunJob for a PENDING job. +func TestIngestService_Sweep_RunsPendingJobToReady(t *testing.T) { + svc, _, backend, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "a pending job the scheduler must claim and run") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.Equal(t, JobStatePending, job.State) + + require.NoError(t, svc.Sweep(ctx)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 1, got.DocumentsDone) + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key)) + assert.Positive(t, embedder.callCount()) +} + +// TestIngestService_Sweep_RedrivesStaleRunningJob proves Sweep also covers +// Reconcile's crash-recovery case, so a single scheduler tick handles both. +func TestIngestService_Sweep_RedrivesStaleRunningJob(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + key := ingestPrefix + "doc1.txt" + putObject(t, store, key, "a stale running job the scheduler must re-drive") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, job.ID, JobStateRunning)) + forceJobUpdatedAt(t, svc.Jobs, ingestAccountA, job.ID, time.Now().UTC().Add(-2*jobStaleAfter)) + + require.NoError(t, svc.Sweep(ctx)) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 1, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", key)) +} + +// TestIngestService_Sweep_LeavesFreshRunningAndTerminalJobsAlone proves Sweep +// does not touch a RUNNING job still within its grace period, nor a job +// already in a terminal state (READY or FAILED). +func TestIngestService_Sweep_LeavesFreshRunningAndTerminalJobsAlone(t *testing.T) { + svc, _, backend, store, _ := newIngestTestSetup(t) + ctx := context.Background() + + freshKey := ingestPrefix + "fresh.txt" + putObject(t, store, freshKey, "still in flight, must not be touched") + freshJob, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, freshJob.ID, JobStateRunning)) + + readyJob, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, readyJob.ID, JobStateReady)) + + failedJob, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource()) + require.NoError(t, err) + require.NoError(t, svc.Jobs.SetState(ctx, ingestAccountA, failedJob.ID, JobStateFailed)) + + require.NoError(t, svc.Sweep(ctx)) + + gotFresh, err := svc.Jobs.Get(ctx, ingestAccountA, freshJob.ID) + require.NoError(t, err) + require.NotNil(t, gotFresh) + assert.Equal(t, JobStateRunning, gotFresh.State, "a RUNNING job within the grace period must be left alone") + + gotReady, err := svc.Jobs.Get(ctx, ingestAccountA, readyJob.ID) + require.NoError(t, err) + require.NotNil(t, gotReady) + assert.Equal(t, JobStateReady, gotReady.State) + + gotFailed, err := svc.Jobs.Get(ctx, ingestAccountA, failedJob.ID) + require.NoError(t, err) + require.NotNil(t, gotFailed) + assert.Equal(t, JobStateFailed, gotFailed.State) + + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", freshKey)) +} diff --git a/spinifex/handlers/ochrevector/jobs.go b/spinifex/handlers/ochrevector/jobs.go new file mode 100644 index 000000000..ca5a3488c --- /dev/null +++ b/spinifex/handlers/ochrevector/jobs.go @@ -0,0 +1,274 @@ +package handlers_ochrevector + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +// jobsBucket is the cluster-replicated KV bucket holding ingestion job +// records, per-account, mirroring registryBucket's pattern. +const jobsBucket = "ochre-vector-jobs" + +// jobsBucketHistory keeps one revision per key: a job record is one JSON +// document mutated in place, not a series. +const jobsBucketHistory = 1 + +// Ingestion job lifecycle states. +const ( + JobStatePending = "PENDING" + JobStateRunning = "RUNNING" + JobStateReady = "READY" + JobStateFailed = "FAILED" +) + +// ErrJobExists reports that Reserve lost the single-writer claim on a job id +// already reserved for that account. +var ErrJobExists = errors.New("ochrevector: job already exists") + +// ErrJobNotFound reports that an operation targeted an account+job pair with +// no job record. +var ErrJobNotFound = errors.New("ochrevector: job not found") + +// FailedDoc records one document an ingestion job skipped rather than +// letting fail the whole job (D7): its source key and why it was skipped. +type FailedDoc struct { + SourceKey string `json:"sourceKey"` + Reason string `json:"reason"` +} + +// JobRecord is one ingestion job: PENDING (reserved, not yet run) -> +// RUNNING -> READY (completed, possibly with per-document failures) or +// FAILED (the job itself could not complete, e.g. embedder fully down). +type JobRecord struct { + ID string `json:"id"` + AccountID string `json:"accountId"` + IndexID string `json:"indexId"` + Source SourceSpec `json:"source"` + State string `json:"state"` + DocumentsTotal int `json:"documentsTotal"` + DocumentsDone int `json:"documentsDone"` + FailedDocuments []FailedDoc `json:"failedDocuments,omitempty"` + Error string `json:"error,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// JobStore persists JobRecords in the ochre-vector-jobs JetStream KV bucket, +// mirroring Registry's lazy get-or-create bucket, key "accountID/id", +// per-account prefix-scan list. +type JobStore struct { + js jetstream.JetStream + + mu sync.Mutex + kv jetstream.KeyValue +} + +// NewJobStore constructs a JobStore over js. +func NewJobStore(js jetstream.JetStream) *JobStore { + return &JobStore{js: js} +} + +// bucket lazily opens (or creates) the jobs KV bucket, caching the handle. +func (s *JobStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { + if s.js == nil { + return nil, errors.New("ochrevector: job store has no JetStream client configured") + } + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := kvutil.GetOrCreateBucket(ctx, s.js, jobsBucket, jobsBucketHistory) + if err != nil { + return nil, err + } + s.kv = kv + return kv, nil +} + +// jobKey scopes every record to its owning account, so a foreign account's +// raw-id guess can never collide with another tenant's key. +func jobKey(accountID, jobID string) string { + return accountID + "/" + jobID +} + +// Reserve atomically claims rec.ID for accountID: the create-only KV write is +// the single-writer mutex (mirrors Registry.Reserve), so two concurrent +// starts of the same job id race safely and exactly one wins. rec.AccountID +// is stamped from accountID, overriding whatever the caller set. +func (s *JobStore) Reserve(ctx context.Context, accountID string, rec JobRecord) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + rec.AccountID = accountID + data, err := json.Marshal(rec) + if err != nil { + return fmt.Errorf("ochrevector: encode job record %s: %w", rec.ID, err) + } + if _, err := kv.Create(ctx, jobKey(accountID, rec.ID), data); err != nil { + if errors.Is(err, jetstream.ErrKeyExists) { + return ErrJobExists + } + return fmt.Errorf("ochrevector: reserve job %s: %w", rec.ID, err) + } + return nil +} + +// Get reads accountID's record for jobID, returning (nil, nil) when absent. +func (s *JobStore) Get(ctx context.Context, accountID, jobID string) (*JobRecord, error) { + kv, err := s.bucket(ctx) + if err != nil { + return nil, err + } + return getJobRecord(ctx, kv, jobKey(accountID, jobID)) +} + +// getJobRecord reads and decodes one record, returning (nil, nil) when absent. +func getJobRecord(ctx context.Context, kv jetstream.KeyValue, key string) (*JobRecord, error) { + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return nil, nil + } + return nil, fmt.Errorf("ochrevector: kv get %s: %w", key, err) + } + var rec JobRecord + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return nil, fmt.Errorf("ochrevector: decode %s: %w", key, err) + } + return &rec, nil +} + +// List returns every job record owned by accountID, so one tenant's listing +// never surfaces another's. +func (s *JobStore) List(ctx context.Context, accountID string) ([]JobRecord, error) { + kv, err := s.bucket(ctx) + if err != nil { + return nil, err + } + return listJobRecords(ctx, kv, accountID+"/") +} + +// ListAll returns every job record across every account, for the ingestion +// reconciler's crash-recovery sweep. +func (s *JobStore) ListAll(ctx context.Context) ([]JobRecord, error) { + kv, err := s.bucket(ctx) + if err != nil { + return nil, err + } + return listJobRecords(ctx, kv, "") +} + +// listJobRecords walks every key with the given prefix ("" matches +// everything) and decodes each into a JobRecord, skipping any key that +// disappears between the key listing and the read. +func listJobRecords(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]JobRecord, error) { + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return nil, nil + } + return nil, fmt.Errorf("ochrevector: list keys: %w", err) + } + var out []JobRecord + for _, key := range keys { + if !strings.HasPrefix(key, prefix) { + continue + } + rec, err := getJobRecord(ctx, kv, key) + if err != nil { + return nil, err + } + if rec != nil { + out = append(out, *rec) + } + } + return out, nil +} + +// Update reads accountID's jobID record, applies fn to mutate it in place, +// refreshes UpdatedAt, and writes it back with a revision-guarded (CAS) +// update, mirroring Registry.SetState's pattern for the general case of an +// arbitrary field mutation (progress counters, failed-document list, error). +func (s *JobStore) Update(ctx context.Context, accountID, jobID string, fn func(rec *JobRecord)) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + key := jobKey(accountID, jobID) + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return ErrJobNotFound + } + return fmt.Errorf("ochrevector: kv get %s: %w", key, err) + } + var rec JobRecord + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return fmt.Errorf("ochrevector: decode %s: %w", key, err) + } + fn(&rec) + rec.UpdatedAt = time.Now().UTC() + data, err := json.Marshal(rec) + if err != nil { + return fmt.Errorf("ochrevector: encode job record %s: %w", rec.ID, err) + } + if _, err := kv.Update(ctx, key, data, entry.Revision()); err != nil { + return fmt.Errorf("ochrevector: update job %s: %w", jobID, err) + } + return nil +} + +// SetState is Update specialized to the common single-field state +// transition, mirroring Registry.SetState. +func (s *JobStore) SetState(ctx context.Context, accountID, jobID, state string) error { + return s.Update(ctx, accountID, jobID, func(rec *JobRecord) { + rec.State = state + }) +} + +// PurgeAll deletes every job record across every account, so a rebuilt +// appliance's job history never references data that no longer exists. +// Idempotent: an empty bucket is a no-op success. +func (s *JobStore) PurgeAll(ctx context.Context) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return nil + } + return fmt.Errorf("ochrevector: list keys: %w", err) + } + for _, key := range keys { + if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("ochrevector: purge job %s: %w", key, err) + } + } + return nil +} + +// Delete removes accountID's jobID record. Idempotent: deleting an +// already-absent record is a no-op success. +func (s *JobStore) Delete(ctx context.Context, accountID, jobID string) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + key := jobKey(accountID, jobID) + if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("ochrevector: delete job %s: %w", jobID, err) + } + return nil +} diff --git a/spinifex/handlers/ochrevector/jobs_test.go b/spinifex/handlers/ochrevector/jobs_test.go new file mode 100644 index 000000000..494f2a219 --- /dev/null +++ b/spinifex/handlers/ochrevector/jobs_test.go @@ -0,0 +1,202 @@ +// Exercises the unexported job store internals with no exported surface +// to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + jobAccountA = "111111111111" + jobAccountB = "222222222222" +) + +func newJobStoreTestStore(t *testing.T) *JobStore { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + return NewJobStore(js) +} + +func TestJobStore_ReserveClaimAndCollision(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + rec := JobRecord{ID: "job-one", IndexID: "idx-one", State: JobStatePending, CreatedAt: now, UpdatedAt: now} + require.NoError(t, store.Reserve(ctx, jobAccountA, rec)) + + // A second reservation of the same id under the same account loses the + // single-writer claim. + err := store.Reserve(ctx, jobAccountA, rec) + require.Error(t, err) + assert.ErrorIs(t, err, ErrJobExists) + + // The same id under a different account is a distinct key and succeeds. + require.NoError(t, store.Reserve(ctx, jobAccountB, rec)) + + got, err := store.Get(ctx, jobAccountA, "job-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, jobAccountA, got.AccountID, "Reserve must stamp AccountID from the accountID argument") + assert.Equal(t, "idx-one", got.IndexID) +} + +func TestJobStore_GetNotFound(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + got, err := store.Get(ctx, jobAccountA, "does-not-exist") + require.NoError(t, err) + assert.Nil(t, got) +} + +func TestJobStore_SetStateTransitions(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + rec := JobRecord{ID: "job-one", State: JobStatePending, CreatedAt: now, UpdatedAt: now} + require.NoError(t, store.Reserve(ctx, jobAccountA, rec)) + + require.NoError(t, store.SetState(ctx, jobAccountA, "job-one", JobStateRunning)) + got, err := store.Get(ctx, jobAccountA, "job-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateRunning, got.State) + assert.True(t, got.UpdatedAt.After(now) || got.UpdatedAt.Equal(now), "SetState must refresh UpdatedAt") + + require.NoError(t, store.SetState(ctx, jobAccountA, "job-one", JobStateReady)) + got, err = store.Get(ctx, jobAccountA, "job-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) +} + +func TestJobStore_SetStateNotFound(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + err := store.SetState(ctx, jobAccountA, "does-not-exist", JobStateRunning) + require.Error(t, err) + assert.ErrorIs(t, err, ErrJobNotFound) +} + +func TestJobStore_Update(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-one", CreatedAt: now, UpdatedAt: now})) + + err := store.Update(ctx, jobAccountA, "job-one", func(rec *JobRecord) { + rec.State = JobStateReady + rec.DocumentsTotal = 3 + rec.DocumentsDone = 2 + rec.FailedDocuments = []FailedDoc{{SourceKey: "bad.txt", Reason: "boom"}} + }) + require.NoError(t, err) + + got, err := store.Get(ctx, jobAccountA, "job-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateReady, got.State) + assert.Equal(t, 3, got.DocumentsTotal) + assert.Equal(t, 2, got.DocumentsDone) + require.Len(t, got.FailedDocuments, 1) + assert.Equal(t, "bad.txt", got.FailedDocuments[0].SourceKey) +} + +func TestJobStore_UpdateNotFound(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + err := store.Update(ctx, jobAccountA, "does-not-exist", func(rec *JobRecord) {}) + require.Error(t, err) + assert.ErrorIs(t, err, ErrJobNotFound) +} + +func TestJobStore_Delete(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-one", CreatedAt: now, UpdatedAt: now})) + + require.NoError(t, store.Delete(ctx, jobAccountA, "job-one")) + got, err := store.Get(ctx, jobAccountA, "job-one") + require.NoError(t, err) + assert.Nil(t, got) + + // Deleting an already-absent record is a no-op success. + require.NoError(t, store.Delete(ctx, jobAccountA, "job-one")) +} + +// TestJobStore_PurgeAll proves PurgeAll clears every job record across every +// account, not just one -- the appliance-teardown use case it exists for. +func TestJobStore_PurgeAll(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-one", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, store.Reserve(ctx, jobAccountB, JobRecord{ID: "job-two", CreatedAt: now, UpdatedAt: now})) + + require.NoError(t, store.PurgeAll(ctx)) + + all, err := store.ListAll(ctx) + require.NoError(t, err) + assert.Empty(t, all) +} + +// TestJobStore_PurgeAllOnEmptyBucketIsANoOp proves PurgeAll succeeds against +// a job store with no records yet -- a bucket kv.Keys reports as ErrNoKeysFound. +func TestJobStore_PurgeAllOnEmptyBucketIsANoOp(t *testing.T) { + store := newJobStoreTestStore(t) + require.NoError(t, store.PurgeAll(context.Background())) +} + +func TestJobStore_ListAccountScoped(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-a1", IndexID: "idx-a", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-a2", IndexID: "idx-a", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, store.Reserve(ctx, jobAccountB, JobRecord{ID: "job-b1", IndexID: "idx-b", CreatedAt: now, UpdatedAt: now})) + + listA, err := store.List(ctx, jobAccountA) + require.NoError(t, err) + assert.Len(t, listA, 2) + + listB, err := store.List(ctx, jobAccountB) + require.NoError(t, err) + require.Len(t, listB, 1, "account B must not see account A's jobs") + assert.Equal(t, "job-b1", listB[0].ID) +} + +func TestJobStore_ListAll(t *testing.T) { + store := newJobStoreTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, store.Reserve(ctx, jobAccountA, JobRecord{ID: "job-a1", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, store.Reserve(ctx, jobAccountB, JobRecord{ID: "job-b1", CreatedAt: now, UpdatedAt: now})) + + all, err := store.ListAll(ctx) + require.NoError(t, err) + assert.Len(t, all, 2) +} + +func TestJobStore_BucketRequiresJetStream(t *testing.T) { + store := &JobStore{} + _, err := store.bucket(context.Background()) + assert.Error(t, err) +} diff --git a/spinifex/handlers/ochrevector/pgx_backend.go b/spinifex/handlers/ochrevector/pgx_backend.go new file mode 100644 index 000000000..49fbe1c30 --- /dev/null +++ b/spinifex/handlers/ochrevector/pgx_backend.go @@ -0,0 +1,350 @@ +package handlers_ochrevector + +import ( + "context" + "encoding/json" + "fmt" + "strconv" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// hnswM and hnswEfConstruction are the fixed HNSW build parameters (D10); +// ef_search stays runtime-tunable per query in a later stage. +const ( + hnswM = 16 + hnswEfConstruction = 64 +) + +// extensionsSchema holds pgvector, out of public and out of every account +// schema. Account roles get USAGE (read-only) on it and carry it on their +// search_path so the vector type and its opclasses resolve unqualified. +const extensionsSchema = "extensions" + +// pgxBackend is the pgx/v5 VectorBackend implementation. One pool is shared +// by every account (D3); it never carries a pool-wide search_path, since a +// shared pool serves every account and a stale connection-level setting +// would leak across them. Every account-scoped operation instead opens a +// transaction and sets ROLE and search_path LOCAL to that transaction alone. +type pgxBackend struct { + pool *pgxpool.Pool +} + +var _ VectorBackend = (*pgxBackend)(nil) + +// NewPgxBackend connects a pooled pgx/v5 client to dsn. The pool is shared +// across every account; per-account isolation is enforced per-operation, not +// at the pool level (see pgxBackend). +func NewPgxBackend(ctx context.Context, dsn string) (*pgxBackend, error) { + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + return nil, fmt.Errorf("ochrevector: connect postgres: %w", err) + } + return &pgxBackend{pool: pool}, nil +} + +// Close releases the pool's connections. Safe to call once, at shutdown. +func (b *pgxBackend) Close() { + b.pool.Close() +} + +// Init ensures the vector extension exists in the extensions schema. Both are +// database-scoped, so this runs once at daemon boot. In production the appliance +// bootstrap has already installed them as the superuser, so these IF-NOT-EXISTS +// statements are a no-op; against a superuser test DSN they create them. +func (b *pgxBackend) Init(ctx context.Context) error { + if _, err := b.pool.Exec(ctx, fmt.Sprintf(`CREATE SCHEMA IF NOT EXISTS %s`, extensionsSchema)); err != nil { + return fmt.Errorf("ochrevector: create extensions schema: %w", err) + } + if _, err := b.pool.Exec(ctx, fmt.Sprintf(`CREATE EXTENSION IF NOT EXISTS vector SCHEMA %s`, extensionsSchema)); err != nil { + return fmt.Errorf("ochrevector: create vector extension: %w", err) + } + return nil +} + +// EnsureAccount creates accountID's schema and a non-login role granted +// USAGE and CREATE on it alone, with no access to any other schema. Safe to +// call before every account-scoped operation: every statement here is +// idempotent (CREATE ... IF NOT EXISTS, or a pre-check before CREATE ROLE, +// which has no IF NOT EXISTS form). +func (b *pgxBackend) EnsureAccount(ctx context.Context, accountID string) error { + if err := validateAccountID(accountID); err != nil { + return err + } + schema := sanitizeIdent(schemaName(accountID)) + role := sanitizeIdent(roleName(accountID)) + + tx, err := b.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("ochrevector: begin ensure-account tx for %s: %w", accountID, err) + } + defer func() { _ = tx.Rollback(ctx) }() + + // #nosec G201 -- schema is an identifier sanitized by pgx.Identifier and + // validated against indexIDPattern-equivalent account-id validation above; + // no user value is ever concatenated here. + if _, err := tx.Exec(ctx, fmt.Sprintf(`CREATE SCHEMA IF NOT EXISTS %s`, schema)); err != nil { + return fmt.Errorf("ochrevector: create schema for account %s: %w", accountID, err) + } + + var exists bool + if err := tx.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = $1)`, roleName(accountID)).Scan(&exists); err != nil { + return fmt.Errorf("ochrevector: check role for account %s: %w", accountID, err) + } + if !exists { + // #nosec G201 -- role is a sanitized, validated identifier; CREATE ROLE + // has no parameterized form and no IF NOT EXISTS, hence the pre-check. + if _, err := tx.Exec(ctx, fmt.Sprintf(`CREATE ROLE %s NOLOGIN`, role)); err != nil { + return fmt.Errorf("ochrevector: create role for account %s: %w", accountID, err) + } + } + + // #nosec G201 -- schema/role are sanitized, validated identifiers. + if _, err := tx.Exec(ctx, fmt.Sprintf(`GRANT USAGE, CREATE ON SCHEMA %s TO %s`, schema, role)); err != nil { + return fmt.Errorf("ochrevector: grant schema to role for account %s: %w", accountID, err) + } + // The account role never gets USAGE on any other schema (including + // public), so cross-account access requires an explicit grant this code + // never makes — isolation holds even against a query bug elsewhere. + // #nosec G201 -- role is a sanitized, validated identifier. + if _, err := tx.Exec(ctx, fmt.Sprintf(`REVOKE ALL ON SCHEMA public FROM %s`, role)); err != nil { + return fmt.Errorf("ochrevector: revoke public schema from role for account %s: %w", accountID, err) + } + + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("ochrevector: commit ensure-account for account %s: %w", accountID, err) + } + return nil +} + +// withAccountTx runs fn inside a transaction scoped to accountID: SET LOCAL +// ROLE to the account's role and SET LOCAL search_path to its schema, so +// every statement fn issues is enforced by Postgres' own grants — even a +// query bug in fn cannot reach another account's schema, because the role +// itself has no grant there. +func (b *pgxBackend) withAccountTx(ctx context.Context, accountID string, fn func(ctx context.Context, tx pgx.Tx) error) error { + if err := validateAccountID(accountID); err != nil { + return err + } + tx, err := b.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("ochrevector: begin account tx for %s: %w", accountID, err) + } + defer func() { _ = tx.Rollback(ctx) }() + + role := sanitizeIdent(roleName(accountID)) + schema := sanitizeIdent(schemaName(accountID)) + // #nosec G201 -- role is a sanitized, validated identifier; SET does not + // accept bound parameters. + if _, err := tx.Exec(ctx, fmt.Sprintf(`SET LOCAL ROLE %s`, role)); err != nil { + return fmt.Errorf("ochrevector: set local role for account %s: %w", accountID, err) + } + // #nosec G201 -- schema is a sanitized, validated identifier and + // extensionsSchema is a fixed constant; SET does not accept bound parameters. + // extensions trails the account schema so the account's own objects win name + // resolution and only the shared vector type falls through to it. + if _, err := tx.Exec(ctx, fmt.Sprintf(`SET LOCAL search_path = %s, %s`, schema, extensionsSchema)); err != nil { + return fmt.Errorf("ochrevector: set local search_path for account %s: %w", accountID, err) + } + + if err := fn(ctx, tx); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("ochrevector: commit account tx for account %s: %w", accountID, err) + } + return nil +} + +// CreateIndex creates spec's backing table and HNSW index under accountID's +// schema. Both statements are IF NOT EXISTS, so a retry after a crash +// mid-create (the Reconcile path) is safe rather than erroring on a +// partially-created index. +func (b *pgxBackend) CreateIndex(ctx context.Context, accountID string, spec IndexSpec) error { + if err := validateIndexID(spec.ID); err != nil { + return err + } + if spec.Dimension <= 0 { + return fmt.Errorf("ochrevector: index %s: dimension must be positive, got %d", spec.ID, spec.Dimension) + } + table := sanitizeIdent(tableName(spec.ID)) + hnswIdx := sanitizeIdent(hnswIndexName(spec.ID)) + + return b.withAccountTx(ctx, accountID, func(ctx context.Context, tx pgx.Tx) error { + // #nosec G201 -- table is a sanitized, validated identifier; Dimension + // is a validated positive int, not a caller string — vector(N) has no + // parameterized form. + createTable := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( + id bigserial PRIMARY KEY, + embedding vector(%d) NOT NULL, + chunk text NOT NULL, + metadata jsonb NOT NULL DEFAULT '{}'::jsonb, + source_key text NOT NULL, + source_offset integer NOT NULL DEFAULT 0 + )`, table, spec.Dimension) + if _, err := tx.Exec(ctx, createTable); err != nil { + return fmt.Errorf("ochrevector: create table for index %s: %w", spec.ID, err) + } + + // #nosec G201 -- hnswIdx/table are sanitized, validated identifiers. + createIndex := fmt.Sprintf( + `CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw (embedding vector_cosine_ops) WITH (m = %d, ef_construction = %d)`, + hnswIdx, table, hnswM, hnswEfConstruction) + if _, err := tx.Exec(ctx, createIndex); err != nil { + return fmt.Errorf("ochrevector: create hnsw index for index %s: %w", spec.ID, err) + } + return nil + }) +} + +// DropIndex drops indexID's backing table under accountID's schema. Idempotent: +// dropping an already-absent index is a no-op success. +func (b *pgxBackend) DropIndex(ctx context.Context, accountID, indexID string) error { + if err := validateIndexID(indexID); err != nil { + return err + } + table := sanitizeIdent(tableName(indexID)) + + return b.withAccountTx(ctx, accountID, func(ctx context.Context, tx pgx.Tx) error { + // #nosec G201 -- table is a sanitized, validated identifier. + if _, err := tx.Exec(ctx, fmt.Sprintf(`DROP TABLE IF EXISTS %s`, table)); err != nil { + return fmt.Errorf("ochrevector: drop table for index %s: %w", indexID, err) + } + return nil + }) +} + +// ReplaceDocument deletes every row for sourceKey then reinserts rows, in one +// transaction (D7): a query mid-ingest never sees a half-replaced document, +// and a re-ingest of the same key replaces rather than accumulates. +func (b *pgxBackend) ReplaceDocument(ctx context.Context, accountID, indexID, sourceKey string, rows []VectorRow) error { + if err := validateIndexID(indexID); err != nil { + return err + } + table := sanitizeIdent(tableName(indexID)) + + return b.withAccountTx(ctx, accountID, func(ctx context.Context, tx pgx.Tx) error { + // #nosec G201 -- table is a sanitized, validated identifier; sourceKey + // is bound as a parameter. + if _, err := tx.Exec(ctx, fmt.Sprintf(`DELETE FROM %s WHERE source_key = $1`, table), sourceKey); err != nil { + return fmt.Errorf("ochrevector: delete existing rows for %s: %w", sourceKey, err) + } + + // #nosec G201 -- table is a sanitized, validated identifier; every + // value below is bound as a parameter. The embedding has no pgx + // native type here, so it is bound as pgvector's own "[v1,v2,...]" + // text form and cast server-side via ::vector. + insert := fmt.Sprintf(`INSERT INTO %s (embedding, chunk, metadata, source_key, source_offset) VALUES ($1::vector, $2, $3::jsonb, $4, $5)`, table) + for _, row := range rows { + metadata := row.Metadata + if metadata == nil { + metadata = map[string]any{} + } + metaJSON, err := json.Marshal(metadata) + if err != nil { + return fmt.Errorf("ochrevector: encode metadata for %s: %w", sourceKey, err) + } + if _, err := tx.Exec(ctx, insert, encodeVector(row.Embedding), row.Chunk, metaJSON, sourceKey, row.SourceOffset); err != nil { + return fmt.Errorf("ochrevector: insert row for %s: %w", sourceKey, err) + } + } + return nil + }) +} + +// Query runs a k-nearest-neighbour cosine similarity search against +// indexID's embedding column under accountID's schema (D8), optionally +// narrowed by filter (D9). The query vector is always bound param $1; a +// non-nil filter's own bound params continue contiguously from $2, and the +// LIMIT is the last param, so numbering matches the compiled WHERE clause +// exactly regardless of whether a filter is present. +func (b *pgxBackend) Query(ctx context.Context, accountID, indexID string, queryVector []float32, k int, filter *Filter) ([]QueryResult, error) { + if err := validateIndexID(indexID); err != nil { + return nil, err + } + if len(queryVector) == 0 { + return nil, fmt.Errorf("ochrevector: query %s: empty query vector", indexID) + } + table := sanitizeIdent(tableName(indexID)) + + args := []any{encodeVector(queryVector)} + whereSQL := "" + limitParam := 2 + if filter != nil { + sql, filterArgs, next, err := filter.compile(2) + if err != nil { + return nil, fmt.Errorf("ochrevector: query %s: compile filter: %w", indexID, err) + } + whereSQL = " WHERE " + sql + args = append(args, filterArgs...) + limitParam = next + } + args = append(args, clampQueryK(k)) + + // #nosec G201 -- table is a sanitized, validated identifier; whereSQL + // comes only from filter.compile, whose metadata keys pass an + // identifier allowlist before ever reaching SQL text and whose values + // are always appended to args as bound params, never interpolated here. + query := fmt.Sprintf( + `SELECT chunk, metadata, source_key, source_offset, 1 - (embedding <=> $1::vector) AS score FROM %s%s ORDER BY embedding <=> $1::vector LIMIT $%d`, + table, whereSQL, limitParam) + + var results []QueryResult + err := b.withAccountTx(ctx, accountID, func(ctx context.Context, tx pgx.Tx) error { + if filter != nil { + // A selective filter under HNSW needs iterative scan to still + // surface k results rather than stopping short on the index's + // default candidate list (D9); with no filter the default + // (non-iterative) scan is left in place. + if _, err := tx.Exec(ctx, `SET LOCAL hnsw.iterative_scan = 'relaxed_order'`); err != nil { + return fmt.Errorf("set iterative scan: %w", err) + } + } + + rows, err := tx.Query(ctx, query, args...) + if err != nil { + return fmt.Errorf("query: %w", err) + } + defer rows.Close() + + for rows.Next() { + var ( + res QueryResult + metaJSON []byte + ) + if err := rows.Scan(&res.Chunk, &metaJSON, &res.SourceKey, &res.SourceOffset, &res.Score); err != nil { + return fmt.Errorf("scan row: %w", err) + } + if len(metaJSON) > 0 { + if err := json.Unmarshal(metaJSON, &res.Metadata); err != nil { + return fmt.Errorf("decode metadata: %w", err) + } + } + results = append(results, res) + } + return rows.Err() + }) + if err != nil { + return nil, fmt.Errorf("ochrevector: query %s: %w", indexID, err) + } + return results, nil +} + +// encodeVector renders embedding in pgvector's own text input format +// ("[v1,v2,...]"). Bound as an ordinary string parameter and cast +// server-side via ::vector, so no pgvector client dependency is needed for +// this one value encoding. +func encodeVector(embedding []float32) string { + var sb strings.Builder + sb.WriteByte('[') + for i, v := range embedding { + if i > 0 { + sb.WriteByte(',') + } + sb.WriteString(strconv.FormatFloat(float64(v), 'f', -1, 32)) + } + sb.WriteByte(']') + return sb.String() +} diff --git a/spinifex/handlers/ochrevector/pgx_backend_test.go b/spinifex/handlers/ochrevector/pgx_backend_test.go new file mode 100644 index 000000000..cb5e72a2b --- /dev/null +++ b/spinifex/handlers/ochrevector/pgx_backend_test.go @@ -0,0 +1,203 @@ +// Exercises the unexported pgx backend internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "fmt" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// testDSNEnv names the environment variable carrying a real Postgres/pgvector +// DSN. Unset in CI and by default locally, so this test skips cleanly rather +// than requiring a database (no testcontainers dependency is introduced). +const testDSNEnv = "OCHREVECTOR_TEST_DSN" + +const ( + pgTestAccountA = "111111111111" + pgTestAccountB = "222222222222" +) + +// TestPgxBackend_Integration exercises pgxBackend against a real pgvector +// instance: Init/EnsureAccount/CreateIndex/DropIndex, plus the isolation +// invariant that one account's role cannot read another account's schema. +func TestPgxBackend_Integration(t *testing.T) { + dsn := os.Getenv(testDSNEnv) + if dsn == "" { + t.Skipf("%s not set; skipping pgvector integration test", testDSNEnv) + } + + ctx := context.Background() + backend, err := NewPgxBackend(ctx, dsn) + require.NoError(t, err) + t.Cleanup(backend.Close) + + require.NoError(t, backend.Init(ctx)) + require.NoError(t, backend.EnsureAccount(ctx, pgTestAccountA)) + require.NoError(t, backend.EnsureAccount(ctx, pgTestAccountB)) + + const indexID = "idx-integrationtest01" + t.Cleanup(func() { + _ = backend.DropIndex(context.Background(), pgTestAccountA, indexID) + }) + + require.NoError(t, backend.CreateIndex(ctx, pgTestAccountA, IndexSpec{ID: indexID, Dimension: 8})) + + // Re-running CreateIndex must not error (IF NOT EXISTS), the property + // Reconcile's forward-retry depends on. + require.NoError(t, backend.CreateIndex(ctx, pgTestAccountA, IndexSpec{ID: indexID, Dimension: 8})) + + // Account B's role has no grant on account A's schema, so a query + // against A's table under B's role must fail even though it is + // syntactically well-formed. + tx, err := backend.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + + _, err = tx.Exec(ctx, fmt.Sprintf("SET LOCAL ROLE %s", sanitizeIdent(roleName(pgTestAccountB)))) + require.NoError(t, err) + + _, err = tx.Exec(ctx, fmt.Sprintf("SELECT 1 FROM %s.%s LIMIT 1", + sanitizeIdent(schemaName(pgTestAccountA)), sanitizeIdent(tableName(indexID)))) + assert.Error(t, err, "account B's role must not be able to read account A's table") + + require.NoError(t, backend.DropIndex(ctx, pgTestAccountA, indexID)) + + // Dropping an already-absent index is a no-op success. + require.NoError(t, backend.DropIndex(ctx, pgTestAccountA, indexID)) +} + +// TestPgxBackend_ReplaceDocument_Integration proves ReplaceDocument's +// delete-then-reinsert replaces a source key's rows on re-run rather than +// accumulating duplicates (D7), against a real pgvector instance. +func TestPgxBackend_ReplaceDocument_Integration(t *testing.T) { + dsn := os.Getenv(testDSNEnv) + if dsn == "" { + t.Skipf("%s not set; skipping pgvector integration test", testDSNEnv) + } + + ctx := context.Background() + backend, err := NewPgxBackend(ctx, dsn) + require.NoError(t, err) + t.Cleanup(backend.Close) + + require.NoError(t, backend.Init(ctx)) + require.NoError(t, backend.EnsureAccount(ctx, pgTestAccountA)) + + const indexID = "idx-replacedoctest01" + t.Cleanup(func() { + _ = backend.DropIndex(context.Background(), pgTestAccountA, indexID) + }) + require.NoError(t, backend.CreateIndex(ctx, pgTestAccountA, IndexSpec{ID: indexID, Dimension: 4})) + + const sourceKey = "docs/one.txt" + firstRows := []VectorRow{ + {Embedding: []float32{1, 0, 0, 0}, Chunk: "first chunk", Metadata: map[string]any{"k": "v"}, SourceOffset: 0}, + {Embedding: []float32{0, 1, 0, 0}, Chunk: "second chunk", SourceOffset: 100}, + } + require.NoError(t, backend.ReplaceDocument(ctx, pgTestAccountA, indexID, sourceKey, firstRows)) + + count, err := countDocumentRows(ctx, t, backend, pgTestAccountA, indexID, sourceKey) + require.NoError(t, err) + assert.Equal(t, 2, count) + + // Re-ingesting the same key with a DIFFERENT chunk count must replace, + // not accumulate: a stale query would otherwise see both the old and + // new chunk sets at once. + secondRows := []VectorRow{ + {Embedding: []float32{0, 0, 1, 0}, Chunk: "only chunk now", SourceOffset: 0}, + } + require.NoError(t, backend.ReplaceDocument(ctx, pgTestAccountA, indexID, sourceKey, secondRows)) + + count, err = countDocumentRows(ctx, t, backend, pgTestAccountA, indexID, sourceKey) + require.NoError(t, err) + assert.Equal(t, 1, count, "ReplaceDocument must replace, not accumulate, rows for the same source key") + + // A different source key's rows are unaffected by another key's replace. + require.NoError(t, backend.ReplaceDocument(ctx, pgTestAccountA, indexID, "docs/two.txt", []VectorRow{ + {Embedding: []float32{0, 0, 0, 1}, Chunk: "unrelated document", SourceOffset: 0}, + })) + count, err = countDocumentRows(ctx, t, backend, pgTestAccountA, indexID, sourceKey) + require.NoError(t, err) + assert.Equal(t, 1, count) +} + +// TestPgxBackend_Query_Integration proves Query against a real pgvector +// instance: nearest-first ordering, a score in [0,1] (D8), and a metadata +// filter (D9) narrowing the result set correctly. +func TestPgxBackend_Query_Integration(t *testing.T) { + dsn := os.Getenv(testDSNEnv) + if dsn == "" { + t.Skipf("%s not set; skipping pgvector integration test", testDSNEnv) + } + + ctx := context.Background() + backend, err := NewPgxBackend(ctx, dsn) + require.NoError(t, err) + t.Cleanup(backend.Close) + + require.NoError(t, backend.Init(ctx)) + require.NoError(t, backend.EnsureAccount(ctx, pgTestAccountA)) + + const indexID = "idx-querytest01" + t.Cleanup(func() { + _ = backend.DropIndex(context.Background(), pgTestAccountA, indexID) + }) + require.NoError(t, backend.CreateIndex(ctx, pgTestAccountA, IndexSpec{ID: indexID, Dimension: 3})) + + rows := []struct { + key string + row VectorRow + }{ + {"docs/near.txt", VectorRow{Embedding: []float32{1, 0, 0}, Chunk: "closest to the query", Metadata: map[string]any{"category": "handbook"}}}, + {"docs/mid.txt", VectorRow{Embedding: []float32{0.7, 0.7, 0}, Chunk: "somewhat related", Metadata: map[string]any{"category": "faq"}}}, + {"docs/far.txt", VectorRow{Embedding: []float32{0, 0, 1}, Chunk: "unrelated", Metadata: map[string]any{"category": "faq"}}}, + } + for _, r := range rows { + require.NoError(t, backend.ReplaceDocument(ctx, pgTestAccountA, indexID, r.key, []VectorRow{r.row})) + } + + // Unfiltered: nearest first, every score in [0,1]. + results, err := backend.Query(ctx, pgTestAccountA, indexID, []float32{1, 0, 0}, 3, nil) + require.NoError(t, err) + require.Len(t, results, 3) + assert.Equal(t, "docs/near.txt", results[0].SourceKey, "the identical vector must rank first") + for _, res := range results { + assert.GreaterOrEqual(t, res.Score, float32(0)) + assert.LessOrEqual(t, res.Score, float32(1)) + } + assert.GreaterOrEqual(t, results[0].Score, results[1].Score, "results must be nearest-first") + assert.GreaterOrEqual(t, results[1].Score, results[2].Score, "results must be nearest-first") + + // Filtered: category=faq narrows to the two faq-tagged rows only. + filtered, err := backend.Query(ctx, pgTestAccountA, indexID, []float32{1, 0, 0}, 10, Equals("category", "faq")) + require.NoError(t, err) + require.Len(t, filtered, 2) + for _, res := range filtered { + assert.Equal(t, "faq", res.Metadata["category"]) + } +} + +// countDocumentRows counts sourceKey's rows in indexID's table under +// accountID's schema, using the account's own role (as query paths will). +func countDocumentRows(ctx context.Context, t *testing.T, backend *pgxBackend, accountID, indexID, sourceKey string) (int, error) { + t.Helper() + tx, err := backend.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + + _, err = tx.Exec(ctx, fmt.Sprintf("SET LOCAL ROLE %s", sanitizeIdent(roleName(accountID)))) + require.NoError(t, err) + _, err = tx.Exec(ctx, fmt.Sprintf("SET LOCAL search_path = %s", sanitizeIdent(schemaName(accountID)))) + require.NoError(t, err) + + var count int + err = tx.QueryRow(ctx, fmt.Sprintf("SELECT count(*) FROM %s WHERE source_key = $1", sanitizeIdent(tableName(indexID))), sourceKey).Scan(&count) + return count, err +} diff --git a/spinifex/handlers/ochrevector/query_test.go b/spinifex/handlers/ochrevector/query_test.go new file mode 100644 index 000000000..a9379081e --- /dev/null +++ b/spinifex/handlers/ochrevector/query_test.go @@ -0,0 +1,108 @@ +// Exercises unexported query clamping internals with no exported surface +// to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestClampQueryK proves Query's k default/clamp rule directly (D8/D10): +// k<=0 defaults to 10, k in range passes through, and any k above 100 is +// silently clamped rather than erroring. +func TestClampQueryK(t *testing.T) { + tests := []struct { + name string + k int + want int + }{ + {"zero defaults", 0, defaultQueryK}, + {"negative defaults", -5, defaultQueryK}, + {"in range passes through", 25, 25}, + {"at cap passes through", maxQueryK, maxQueryK}, + {"over cap clamps", 500, maxQueryK}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, clampQueryK(tt.k)) + }) + } +} + +// TestFakeBackend_Query_DefaultsAndClampsK proves the backend seam itself +// applies the same default/clamp rule Query documents, using the in-memory +// fake so this is a pure-Go test with no Postgres dependency. +func TestFakeBackend_Query_DefaultsAndClampsK(t *testing.T) { + backend := newFakeBackend() + ctx := context.Background() + + canned := make([]QueryResult, 150) + for i := range canned { + canned[i] = QueryResult{Chunk: "chunk", SourceKey: "docs/one.txt", Score: 1.0} + } + backend.queryResults = canned + + // k<=0 defaults to 10. + got, err := backend.Query(ctx, "111111111111", "idx-one", []float32{1, 0}, 0, nil) + require.NoError(t, err) + assert.Len(t, got, defaultQueryK) + assert.Equal(t, defaultQueryK, backend.lastQueryK) + + // A k above the hard cap is clamped, not rejected. + got, err = backend.Query(ctx, "111111111111", "idx-one", []float32{1, 0}, 1000, nil) + require.NoError(t, err) + assert.Len(t, got, maxQueryK) + assert.Equal(t, maxQueryK, backend.lastQueryK) + + // An in-range k passes straight through. + got, err = backend.Query(ctx, "111111111111", "idx-one", []float32{1, 0}, 30, nil) + require.NoError(t, err) + assert.Len(t, got, 30) +} + +// TestFakeBackend_Query_ResultsPlumbedThrough proves the exact QueryResult +// values set on the backend come back unchanged through the Query call -- +// the plumbing this stage adds, independent of any real ANN ranking. +func TestFakeBackend_Query_ResultsPlumbedThrough(t *testing.T) { + backend := newFakeBackend() + ctx := context.Background() + + want := []QueryResult{ + {Chunk: "first chunk", Metadata: map[string]any{"category": "handbook"}, SourceKey: "docs/a.txt", SourceOffset: 0, Score: 0.95}, + {Chunk: "second chunk", Metadata: map[string]any{"category": "faq"}, SourceKey: "docs/b.txt", SourceOffset: 128, Score: 0.42}, + } + backend.queryResults = want + + got, err := backend.Query(ctx, "111111111111", "idx-one", []float32{0.1, 0.2, 0.3}, 5, nil) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +// TestFakeBackend_Query_PropagatesBackendError proves a backend failure +// surfaces to the caller rather than being swallowed. +func TestFakeBackend_Query_PropagatesBackendError(t *testing.T) { + backend := newFakeBackend() + backend.failQuery = errFakeBackend + + _, err := backend.Query(context.Background(), "111111111111", "idx-one", []float32{1}, 5, nil) + require.Error(t, err) + assert.ErrorIs(t, err, errFakeBackend) +} + +// TestFakeBackend_Query_RecordsFilterForAssertions proves the filter a +// caller passes reaches the backend call, so an orchestration-level test +// can assert the right filter was requested even though the fake does not +// evaluate it. +func TestFakeBackend_Query_RecordsFilterForAssertions(t *testing.T) { + backend := newFakeBackend() + f := Equals("category", "handbook") + + _, err := backend.Query(context.Background(), "111111111111", "idx-one", []float32{1}, 5, f) + require.NoError(t, err) + assert.Same(t, f, backend.lastQueryFilter) +} diff --git a/spinifex/handlers/ochrevector/rds_launcher.go b/spinifex/handlers/ochrevector/rds_launcher.go new file mode 100644 index 000000000..ac82b46fb --- /dev/null +++ b/spinifex/handlers/ochrevector/rds_launcher.go @@ -0,0 +1,156 @@ +package handlers_ochrevector + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/rds" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +// rdsCreateSubject, rdsDescribeSubject and rdsDeleteSubject mirror +// handlers/rds's SubjectCreateDBInstance/SubjectDescribeDBInstances/ +// SubjectDeleteDBInstance. Restated as literals rather than imported, so +// this launcher stays decoupled from the rds handler package (it talks to +// RDS only over NATS, like any other client). +const ( + rdsCreateSubject = "rds.CreateDBInstance" + rdsDescribeSubject = "rds.DescribeDBInstances" + rdsDeleteSubject = "rds.DeleteDBInstance" +) + +// rdsAvailableStatus mirrors handlers/rds's StatusAvailable ("available"), +// restated for the same reason as the subjects above. +const rdsAvailableStatus = "available" + +// The platform appliance's own fixed sizing: the smallest supported RDS +// instance class and the minimum allocated storage, since it serves +// control-plane vector data rather than a tenant workload. +const ( + applianceEngine = "postgres" + applianceInstanceClass = "db.t3.micro" + applianceAllocatedStorage = 20 +) + +// rdsCreateTimeout and rdsDescribeTimeout bound one NATS round trip each, +// mirroring handlers/rds's own NATSService client timeouts for the same +// two calls. +const ( + rdsCreateTimeout = 5 * time.Minute + rdsDescribeTimeout = 30 * time.Second + rdsDeleteTimeout = 5 * time.Minute +) + +// launchPollInterval paces Launch's wait between DescribeDBInstances polls +// while the appliance instance boots toward "available". A var, not a const, +// so tests can shrink it rather than paying the real delay (mirrors +// ingest.go's ingestRetryBackoff). +var launchPollInterval = 5 * time.Second + +// rdsLauncher is the real ApplianceLauncher: it provisions the platform +// Postgres appliance by NATS request to RDS's CreateDBInstance, then polls +// DescribeDBInstances until the instance reports available, bounded by the +// caller's context and timeout. +type rdsLauncher struct { + nc *nats.Conn + timeout time.Duration +} + +var _ ApplianceLauncher = (*rdsLauncher)(nil) + +// NewRDSApplianceLauncher constructs an ApplianceLauncher over nc, bounding +// the whole launch (create-and-poll-to-available) by timeout. +func NewRDSApplianceLauncher(nc *nats.Conn, timeout time.Duration) *rdsLauncher { + return &rdsLauncher{nc: nc, timeout: timeout} +} + +// Launch provisions identifier as a new RDS DB instance under masterUsername/ +// masterPassword, then waits for it to become available. It is +// identifier-idempotent: a create that fails because identifier already +// exists (Ensure's re-entrant recovery path calling Launch again for an +// identifier it already launched) is treated as "already launching", not a +// hard failure, and Launch proceeds straight to polling. The password is +// never logged. +func (l *rdsLauncher) Launch(ctx context.Context, identifier, masterUsername, masterPassword string) (string, int, error) { + ctx, cancel := context.WithTimeout(ctx, l.timeout) + defer cancel() + + if err := l.create(ctx, identifier, masterUsername, masterPassword); err != nil { + return "", 0, err + } + return l.pollUntilAvailable(ctx, identifier) +} + +// create sends the CreateDBInstance request. An AWS DBInstanceAlreadyExists +// error is swallowed: the caller only reaches here because it holds (or is +// resuming) the appliance's single-writer claim, so an identifier collision +// means a previous attempt already started this same launch. +func (l *rdsLauncher) create(ctx context.Context, identifier, masterUsername, masterPassword string) error { + input := &rds.CreateDBInstanceInput{ + DBInstanceIdentifier: aws.String(identifier), + Engine: aws.String(applianceEngine), + DBInstanceClass: aws.String(applianceInstanceClass), + AllocatedStorage: aws.Int64(applianceAllocatedStorage), + MasterUsername: aws.String(masterUsername), + MasterUserPassword: aws.String(masterPassword), + } + _, err := utils.NATSRequest[rds.CreateDBInstanceOutput](ctx, l.nc, rdsCreateSubject, input, rdsCreateTimeout, utils.GlobalAccountID) + if err != nil { + if strings.Contains(err.Error(), awserrors.ErrorDBInstanceAlreadyExists) { + return nil + } + return fmt.Errorf("ochrevector: create appliance db instance %s: %w", identifier, err) + } + return nil +} + +// Delete removes identifier's backing RDS DB instance, skipping any final +// snapshot: the appliance is a system resource, not customer data worth +// retaining past its own teardown. Idempotent: an identifier already gone +// (or never created) is a no-op success. +func (l *rdsLauncher) Delete(ctx context.Context, identifier string) error { + ctx, cancel := context.WithTimeout(ctx, l.timeout) + defer cancel() + + input := &rds.DeleteDBInstanceInput{ + DBInstanceIdentifier: aws.String(identifier), + SkipFinalSnapshot: aws.Bool(true), + } + _, err := utils.NATSRequest[rds.DeleteDBInstanceOutput](ctx, l.nc, rdsDeleteSubject, input, rdsDeleteTimeout, utils.GlobalAccountID) + if err != nil { + if strings.Contains(err.Error(), awserrors.ErrorDBInstanceNotFound) { + return nil + } + return fmt.Errorf("ochrevector: delete appliance db instance %s: %w", identifier, err) + } + return nil +} + +// pollUntilAvailable polls DescribeDBInstances for identifier until it +// reports StatusAvailable with a resolved endpoint, or ctx is done. +func (l *rdsLauncher) pollUntilAvailable(ctx context.Context, identifier string) (string, int, error) { + for { + out, err := utils.NATSRequest[rds.DescribeDBInstancesOutput](ctx, l.nc, rdsDescribeSubject, + &rds.DescribeDBInstancesInput{DBInstanceIdentifier: aws.String(identifier)}, rdsDescribeTimeout, utils.GlobalAccountID) + if err != nil { + return "", 0, fmt.Errorf("ochrevector: describe appliance db instance %s: %w", identifier, err) + } + if len(out.DBInstances) == 1 { + inst := out.DBInstances[0] + if aws.StringValue(inst.DBInstanceStatus) == rdsAvailableStatus && + inst.Endpoint != nil && aws.StringValue(inst.Endpoint.Address) != "" { + return aws.StringValue(inst.Endpoint.Address), int(aws.Int64Value(inst.Endpoint.Port)), nil + } + } + select { + case <-ctx.Done(): + return "", 0, fmt.Errorf("ochrevector: wait for appliance db instance %s to become available: %w", identifier, ctx.Err()) + case <-time.After(launchPollInterval): + } + } +} diff --git a/spinifex/handlers/ochrevector/rds_launcher_test.go b/spinifex/handlers/ochrevector/rds_launcher_test.go new file mode 100644 index 000000000..b396d6d76 --- /dev/null +++ b/spinifex/handlers/ochrevector/rds_launcher_test.go @@ -0,0 +1,149 @@ +// Exercises the unexported RDS launcher internals with no exported +// surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "encoding/json" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/rds" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestRDSApplianceLauncher_LaunchCreatesAndPollsToAvailable pins Launch's +// happy path: a well-formed CreateDBInstance request, then a poll loop over +// DescribeDBInstances that keeps going while the instance is still +// "creating" and returns the endpoint once it reports "available". +func TestRDSApplianceLauncher_LaunchCreatesAndPollsToAvailable(t *testing.T) { + orig := launchPollInterval + launchPollInterval = 5 * time.Millisecond + defer func() { launchPollInterval = orig }() + + _, nc := testutil.StartTestNATS(t) + defer nc.Close() + + var mu sync.Mutex + createCalls := 0 + var lastCreate rds.CreateDBInstanceInput + + createSub, err := nc.Subscribe(rdsCreateSubject, func(msg *nats.Msg) { + mu.Lock() + createCalls++ + _ = json.Unmarshal(msg.Data, &lastCreate) + mu.Unlock() + out, _ := json.Marshal(rds.CreateDBInstanceOutput{}) + _ = msg.Respond(out) + }) + require.NoError(t, err) + defer func() { _ = createSub.Unsubscribe() }() + + describeCalls := 0 + describeSub, err := nc.Subscribe(rdsDescribeSubject, func(msg *nats.Msg) { + mu.Lock() + describeCalls++ + call := describeCalls + mu.Unlock() + + status := "creating" + var endpoint *rds.Endpoint + if call >= 2 { + status = rdsAvailableStatus + endpoint = &rds.Endpoint{Address: aws.String("appliance.internal"), Port: aws.Int64(5432)} + } + out, _ := json.Marshal(rds.DescribeDBInstancesOutput{ + DBInstances: []*rds.DBInstance{{ + DBInstanceIdentifier: aws.String("ochre-vector-pg"), + DBInstanceStatus: aws.String(status), + Endpoint: endpoint, + }}, + }) + _ = msg.Respond(out) + }) + require.NoError(t, err) + defer func() { _ = describeSub.Unsubscribe() }() + + launcher := NewRDSApplianceLauncher(nc, 5*time.Second) + endpoint, port, err := launcher.Launch(t.Context(), "ochre-vector-pg", "ochre_vector_admin", "s3cr3t-pw") + require.NoError(t, err) + assert.Equal(t, "appliance.internal", endpoint) + assert.Equal(t, 5432, port) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, 1, createCalls, "create must be requested exactly once") + assert.GreaterOrEqual(t, describeCalls, 2, "must poll past the first non-available observation") + assert.Equal(t, "ochre-vector-pg", aws.StringValue(lastCreate.DBInstanceIdentifier)) + assert.Equal(t, "postgres", aws.StringValue(lastCreate.Engine)) + assert.Equal(t, "ochre_vector_admin", aws.StringValue(lastCreate.MasterUsername)) + assert.Equal(t, "s3cr3t-pw", aws.StringValue(lastCreate.MasterUserPassword)) +} + +// TestRDSApplianceLauncher_LaunchIsIdempotentOnAlreadyExists pins Launch's +// re-entrant recovery path (Appliance.resume calling Launch a second time for +// an identifier it already launched): a DBInstanceAlreadyExists response from +// CreateDBInstance must not fail the launch, only the poll result matters. +func TestRDSApplianceLauncher_LaunchIsIdempotentOnAlreadyExists(t *testing.T) { + _, nc := testutil.StartTestNATS(t) + defer nc.Close() + + createSub, err := nc.Subscribe(rdsCreateSubject, func(msg *nats.Msg) { + // Mirrors handlers/rds's actual CreateDBInstance collision error + // (awserrors.Errorf), whose Error() text carries the code as a + // ": " suffix -- the idempotency check below depends on that + // suffix being present, so a hand-written message would not exercise + // the real production shape. + rdsErr := awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists, "DB instance %s already exists", "ochre-vector-pg") + payload := utils.GenerateErrorPayloadWithMessage(awserrors.ValidErrorCodeFromError(rdsErr), rdsErr.Error()) + _ = msg.Respond(payload) + }) + require.NoError(t, err) + defer func() { _ = createSub.Unsubscribe() }() + + describeSub, err := nc.Subscribe(rdsDescribeSubject, func(msg *nats.Msg) { + out, _ := json.Marshal(rds.DescribeDBInstancesOutput{ + DBInstances: []*rds.DBInstance{{ + DBInstanceIdentifier: aws.String("ochre-vector-pg"), + DBInstanceStatus: aws.String(rdsAvailableStatus), + Endpoint: &rds.Endpoint{Address: aws.String("appliance.internal"), Port: aws.Int64(5432)}, + }}, + }) + _ = msg.Respond(out) + }) + require.NoError(t, err) + defer func() { _ = describeSub.Unsubscribe() }() + + launcher := NewRDSApplianceLauncher(nc, 5*time.Second) + endpoint, port, err := launcher.Launch(t.Context(), "ochre-vector-pg", "ochre_vector_admin", "s3cr3t-pw") + require.NoError(t, err) + assert.Equal(t, "appliance.internal", endpoint) + assert.Equal(t, 5432, port) +} + +// TestRDSApplianceLauncher_LaunchPropagatesOtherCreateErrors pins that a +// create failure unrelated to identifier collision (e.g. an invalid +// parameter) is NOT swallowed the way DBInstanceAlreadyExists is. +func TestRDSApplianceLauncher_LaunchPropagatesOtherCreateErrors(t *testing.T) { + _, nc := testutil.StartTestNATS(t) + defer nc.Close() + + createSub, err := nc.Subscribe(rdsCreateSubject, func(msg *nats.Msg) { + payload := utils.GenerateErrorPayloadWithMessage(awserrors.ErrorInvalidParameterValue, "DBInstanceClass is not supported") + _ = msg.Respond(payload) + }) + require.NoError(t, err) + defer func() { _ = createSub.Unsubscribe() }() + + launcher := NewRDSApplianceLauncher(nc, 5*time.Second) + _, _, err = launcher.Launch(t.Context(), "ochre-vector-pg", "ochre_vector_admin", "s3cr3t-pw") + assert.Error(t, err) +} diff --git a/spinifex/handlers/ochrevector/registry.go b/spinifex/handlers/ochrevector/registry.go new file mode 100644 index 000000000..cb1c0ff84 --- /dev/null +++ b/spinifex/handlers/ochrevector/registry.go @@ -0,0 +1,325 @@ +package handlers_ochrevector + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "maps" + "slices" + "strings" + "sync" + "time" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +// registryBucket is the cluster-replicated KV bucket holding index registry +// records (D4): the authoritative source for an index's existence, +// configuration and state, surviving a Postgres rebuild even though the +// vectors themselves do not. +const registryBucket = "ochre-vector-indexes" + +// registryBucketHistory keeps one revision per key: a registry record is one +// JSON document mutated in place, not a series. +const registryBucketHistory = 1 + +// Index lifecycle states (D4). +const ( + StateCreating = "CREATING" + StateReady = "READY" + StateDeleting = "DELETING" + StateStale = "STALE" +) + +// ErrIndexExists reports that Reserve lost the single-writer claim on an +// index id already reserved for that account. +var ErrIndexExists = errors.New("ochrevector: index already exists") + +// ErrIndexNotFound reports that an operation targeted an account+index pair +// with no registry record. +var ErrIndexNotFound = errors.New("ochrevector: index not found") + +// SourceSpec is one ingestion source bound to an index: where its documents +// come from and how they are chunked/embedded. Populated by the ingestion +// stage; persisted here so a rebuilt index can re-ingest without operator +// input (D4). +type SourceSpec struct { + Bucket string `json:"bucket"` + Prefix string `json:"prefix"` + ChunkSize int `json:"chunkSize"` + ChunkOverlap int `json:"chunkOverlap"` + EmbeddingModel string `json:"embeddingModel"` + Dimension int `json:"dimension"` + // Metadata is a static set of per-source tags stamped onto every row + // ingested from this source (see ingestObject), giving D9 filters + // something to match against beyond the source key. + Metadata map[string]string `json:"metadata,omitempty"` +} + +// sourceSpecEqual reports whether a and b are the same source spec. +// SourceSpec carries a map (Metadata), so it is not comparable with ==; +// AppendSourceSpec's dedup check uses this instead of slices.Contains. +func sourceSpecEqual(a, b SourceSpec) bool { + return a.Bucket == b.Bucket && + a.Prefix == b.Prefix && + a.ChunkSize == b.ChunkSize && + a.ChunkOverlap == b.ChunkOverlap && + a.EmbeddingModel == b.EmbeddingModel && + a.Dimension == b.Dimension && + maps.Equal(a.Metadata, b.Metadata) +} + +// Record is the authoritative registry entry for one vector index (D4). +type Record struct { + ID string `json:"id"` + AccountID string `json:"accountId"` + Name string `json:"name"` + Dimension int `json:"dimension"` + EmbeddingModel string `json:"embeddingModel"` + State string `json:"state"` + SourceSpecs []SourceSpec `json:"sourceSpecs,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// Registry persists index Records in the ochre-vector-indexes JetStream KV +// bucket, mirroring GuardrailStore's gateway-direct-KV pattern (lazy +// get-or-create bucket, key "accountID/id", per-account prefix-scan list). +type Registry struct { + js jetstream.JetStream + + mu sync.Mutex + kv jetstream.KeyValue +} + +// NewRegistry constructs a Registry over js. +func NewRegistry(js jetstream.JetStream) *Registry { + return &Registry{js: js} +} + +// bucket lazily opens (or creates) the registry KV bucket, caching the +// handle. +func (reg *Registry) bucket(ctx context.Context) (jetstream.KeyValue, error) { + if reg.js == nil { + return nil, errors.New("ochrevector: registry has no JetStream client configured") + } + reg.mu.Lock() + defer reg.mu.Unlock() + if reg.kv != nil { + return reg.kv, nil + } + kv, err := kvutil.GetOrCreateBucket(ctx, reg.js, registryBucket, registryBucketHistory) + if err != nil { + return nil, err + } + reg.kv = kv + return kv, nil +} + +// registryKey scopes every record to its owning account, so a foreign +// account's raw-id guess can never collide with another tenant's key. +func registryKey(accountID, indexID string) string { + return accountID + "/" + indexID +} + +// Reserve atomically claims rec.ID for accountID: the create-only KV write is +// the single-writer mutex (mirrors handlers/rds's identifier claim via +// jetstream.ErrKeyExists), so two concurrent creates of the same id race +// safely and exactly one wins. rec.AccountID is stamped from accountID, +// overriding whatever the caller set. +func (reg *Registry) Reserve(ctx context.Context, accountID string, rec Record) error { + kv, err := reg.bucket(ctx) + if err != nil { + return err + } + rec.AccountID = accountID + data, err := json.Marshal(rec) + if err != nil { + return fmt.Errorf("ochrevector: encode index record %s: %w", rec.ID, err) + } + if _, err := kv.Create(ctx, registryKey(accountID, rec.ID), data); err != nil { + if errors.Is(err, jetstream.ErrKeyExists) { + return ErrIndexExists + } + return fmt.Errorf("ochrevector: reserve index %s: %w", rec.ID, err) + } + return nil +} + +// Get reads accountID's record for indexID, returning (nil, nil) when absent. +func (reg *Registry) Get(ctx context.Context, accountID, indexID string) (*Record, error) { + kv, err := reg.bucket(ctx) + if err != nil { + return nil, err + } + return getRecord(ctx, kv, registryKey(accountID, indexID)) +} + +// getRecord reads and decodes one record, returning (nil, nil) when absent. +func getRecord(ctx context.Context, kv jetstream.KeyValue, key string) (*Record, error) { + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return nil, nil + } + return nil, fmt.Errorf("ochrevector: kv get %s: %w", key, err) + } + var rec Record + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return nil, fmt.Errorf("ochrevector: decode %s: %w", key, err) + } + return &rec, nil +} + +// List returns every index record owned by accountID, so one tenant's +// listing never surfaces another's. +func (reg *Registry) List(ctx context.Context, accountID string) ([]Record, error) { + kv, err := reg.bucket(ctx) + if err != nil { + return nil, err + } + return listRecords(ctx, kv, accountID+"/") +} + +// ListAll returns every index record across every account, for the +// reconciler's crash-recovery sweep. +func (reg *Registry) ListAll(ctx context.Context) ([]Record, error) { + kv, err := reg.bucket(ctx) + if err != nil { + return nil, err + } + return listRecords(ctx, kv, "") +} + +// listRecords walks every key with the given prefix ("" matches everything) +// and decodes each into a Record, skipping any key that disappears between +// the key listing and the read. +func listRecords(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]Record, error) { + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return nil, nil + } + return nil, fmt.Errorf("ochrevector: list keys: %w", err) + } + var out []Record + for _, key := range keys { + if !strings.HasPrefix(key, prefix) { + continue + } + rec, err := getRecord(ctx, kv, key) + if err != nil { + return nil, err + } + if rec != nil { + out = append(out, *rec) + } + } + return out, nil +} + +// SetState updates accountID's indexID record's State in place, returning +// ErrIndexNotFound if no such record exists. +func (reg *Registry) SetState(ctx context.Context, accountID, indexID, state string) error { + kv, err := reg.bucket(ctx) + if err != nil { + return err + } + key := registryKey(accountID, indexID) + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return ErrIndexNotFound + } + return fmt.Errorf("ochrevector: kv get %s: %w", key, err) + } + var rec Record + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return fmt.Errorf("ochrevector: decode %s: %w", key, err) + } + rec.State = state + rec.UpdatedAt = time.Now().UTC() + data, err := json.Marshal(rec) + if err != nil { + return fmt.Errorf("ochrevector: encode index record %s: %w", rec.ID, err) + } + if _, err := kv.Update(ctx, key, data, entry.Revision()); err != nil { + return fmt.Errorf("ochrevector: update state for index %s: %w", indexID, err) + } + return nil +} + +// AppendSourceSpec adds spec to accountID's indexID record's SourceSpecs +// unless an identical spec is already present, so a repeated ingest against +// an unchanged source never accumulates duplicate entries (D4: the stored +// source-spec set drives auto-repopulation after a Postgres rebuild). +func (reg *Registry) AppendSourceSpec(ctx context.Context, accountID, indexID string, spec SourceSpec) error { + kv, err := reg.bucket(ctx) + if err != nil { + return err + } + key := registryKey(accountID, indexID) + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return ErrIndexNotFound + } + return fmt.Errorf("ochrevector: kv get %s: %w", key, err) + } + var rec Record + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + return fmt.Errorf("ochrevector: decode %s: %w", key, err) + } + if slices.ContainsFunc(rec.SourceSpecs, func(s SourceSpec) bool { return sourceSpecEqual(s, spec) }) { + return nil + } + rec.SourceSpecs = append(rec.SourceSpecs, spec) + rec.UpdatedAt = time.Now().UTC() + data, err := json.Marshal(rec) + if err != nil { + return fmt.Errorf("ochrevector: encode index record %s: %w", rec.ID, err) + } + if _, err := kv.Update(ctx, key, data, entry.Revision()); err != nil { + return fmt.Errorf("ochrevector: append source spec for index %s: %w", indexID, err) + } + return nil +} + +// PurgeAll deletes every index record across every account, so a rebuilt +// appliance's registry never advertises a table that no longer exists. +// Idempotent: an empty bucket is a no-op success. +func (reg *Registry) PurgeAll(ctx context.Context) error { + kv, err := reg.bucket(ctx) + if err != nil { + return err + } + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return nil + } + return fmt.Errorf("ochrevector: list keys: %w", err) + } + for _, key := range keys { + if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("ochrevector: purge index %s: %w", key, err) + } + } + return nil +} + +// Delete removes accountID's indexID record. Idempotent: deleting an +// already-absent record is a no-op success. +func (reg *Registry) Delete(ctx context.Context, accountID, indexID string) error { + kv, err := reg.bucket(ctx) + if err != nil { + return err + } + key := registryKey(accountID, indexID) + if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("ochrevector: delete index %s: %w", indexID, err) + } + return nil +} diff --git a/spinifex/handlers/ochrevector/registry_test.go b/spinifex/handlers/ochrevector/registry_test.go new file mode 100644 index 000000000..79d9c7454 --- /dev/null +++ b/spinifex/handlers/ochrevector/registry_test.go @@ -0,0 +1,202 @@ +// Exercises the unexported registry internals with no exported surface +// to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + regAccountA = "111111111111" + regAccountB = "222222222222" +) + +func newRegistryTestStore(t *testing.T) *Registry { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + return NewRegistry(js) +} + +func TestRegistry_ReserveClaimAndCollision(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + rec := Record{ID: "idx-one", Name: "first", Dimension: 768, State: StateCreating, CreatedAt: now, UpdatedAt: now} + require.NoError(t, reg.Reserve(ctx, regAccountA, rec)) + + // A second reservation of the same id under the same account loses the + // single-writer claim. + err := reg.Reserve(ctx, regAccountA, rec) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexExists) + + // The same id under a different account is a distinct key and succeeds. + require.NoError(t, reg.Reserve(ctx, regAccountB, rec)) + + got, err := reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, regAccountA, got.AccountID, "Reserve must stamp AccountID from the accountID argument") + assert.Equal(t, "first", got.Name) +} + +func TestRegistry_GetNotFound(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + got, err := reg.Get(ctx, regAccountA, "does-not-exist") + require.NoError(t, err) + assert.Nil(t, got) +} + +func TestRegistry_SetStateTransitions(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + rec := Record{ID: "idx-one", Dimension: 768, State: StateCreating, CreatedAt: now, UpdatedAt: now} + require.NoError(t, reg.Reserve(ctx, regAccountA, rec)) + + require.NoError(t, reg.SetState(ctx, regAccountA, "idx-one", StateReady)) + got, err := reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, StateReady, got.State) + assert.True(t, got.UpdatedAt.After(now) || got.UpdatedAt.Equal(now), "SetState must refresh UpdatedAt") + + require.NoError(t, reg.SetState(ctx, regAccountA, "idx-one", StateDeleting)) + got, err = reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, StateDeleting, got.State) +} + +func TestRegistry_SetStateNotFound(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + err := reg.SetState(ctx, regAccountA, "does-not-exist", StateReady) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexNotFound) +} + +func TestRegistry_Delete(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-one", CreatedAt: now, UpdatedAt: now})) + + require.NoError(t, reg.Delete(ctx, regAccountA, "idx-one")) + got, err := reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, got) + + // Deleting an already-absent record is a no-op success. + require.NoError(t, reg.Delete(ctx, regAccountA, "idx-one")) +} + +// TestRegistry_PurgeAll proves PurgeAll clears every record across every +// account, not just one -- the appliance-teardown use case it exists for. +func TestRegistry_PurgeAll(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-one", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, reg.Reserve(ctx, regAccountB, Record{ID: "idx-two", CreatedAt: now, UpdatedAt: now})) + + require.NoError(t, reg.PurgeAll(ctx)) + + all, err := reg.ListAll(ctx) + require.NoError(t, err) + assert.Empty(t, all) +} + +// TestRegistry_PurgeAllOnEmptyBucketIsANoOp proves PurgeAll succeeds against +// a registry with no records yet -- a bucket kv.Keys reports as ErrNoKeysFound. +func TestRegistry_PurgeAllOnEmptyBucketIsANoOp(t *testing.T) { + reg := newRegistryTestStore(t) + require.NoError(t, reg.PurgeAll(context.Background())) +} + +func TestRegistry_ListAccountScoped(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-a1", Name: "a1", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-a2", Name: "a2", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, reg.Reserve(ctx, regAccountB, Record{ID: "idx-b1", Name: "b1", CreatedAt: now, UpdatedAt: now})) + + listA, err := reg.List(ctx, regAccountA) + require.NoError(t, err) + assert.Len(t, listA, 2) + + listB, err := reg.List(ctx, regAccountB) + require.NoError(t, err) + require.Len(t, listB, 1, "account B must not see account A's indexes") + assert.Equal(t, "b1", listB[0].Name) +} + +func TestRegistry_AppendSourceSpecIsIdempotent(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-one", CreatedAt: now, UpdatedAt: now})) + + spec := SourceSpec{Bucket: "docs", Prefix: "kb/", ChunkSize: 512, ChunkOverlap: 64, EmbeddingModel: "nomic-embed-text-v1.5", Dimension: 768} + require.NoError(t, reg.AppendSourceSpec(ctx, regAccountA, "idx-one", spec)) + + got, err := reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, got) + require.Len(t, got.SourceSpecs, 1) + assert.Equal(t, spec, got.SourceSpecs[0]) + + // Appending the identical spec again must not duplicate it. + require.NoError(t, reg.AppendSourceSpec(ctx, regAccountA, "idx-one", spec)) + got, err = reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.Len(t, got.SourceSpecs, 1) + + // A distinct spec (different prefix) is a second entry. + other := spec + other.Prefix = "kb2/" + require.NoError(t, reg.AppendSourceSpec(ctx, regAccountA, "idx-one", other)) + got, err = reg.Get(ctx, regAccountA, "idx-one") + require.NoError(t, err) + require.Len(t, got.SourceSpecs, 2) +} + +func TestRegistry_AppendSourceSpecNotFound(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + err := reg.AppendSourceSpec(ctx, regAccountA, "does-not-exist", SourceSpec{Bucket: "docs"}) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexNotFound) +} + +func TestRegistry_ListAll(t *testing.T) { + reg := newRegistryTestStore(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, reg.Reserve(ctx, regAccountA, Record{ID: "idx-a1", CreatedAt: now, UpdatedAt: now})) + require.NoError(t, reg.Reserve(ctx, regAccountB, Record{ID: "idx-b1", CreatedAt: now, UpdatedAt: now})) + + all, err := reg.ListAll(ctx) + require.NoError(t, err) + assert.Len(t, all, 2) +} diff --git a/spinifex/handlers/ochrevector/service.go b/spinifex/handlers/ochrevector/service.go new file mode 100644 index 000000000..2be5ffa29 --- /dev/null +++ b/spinifex/handlers/ochrevector/service.go @@ -0,0 +1,198 @@ +package handlers_ochrevector + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" +) + +// reconcileStuckAfter bounds how long a CREATING/DELETING record may sit +// before Reconcile treats it as abandoned by a crashed writer and resumes it. +const reconcileStuckAfter = 5 * time.Minute + +// rollbackTimeout bounds a rollback's own cleanup, run with the request +// context already cancelled or expired. +const rollbackTimeout = 30 * time.Second + +// CreateIndexParams is the caller-supplied half of a new index: everything +// the registry does not derive itself. +type CreateIndexParams struct { + Name string + Dimension int + EmbeddingModel string +} + +// Service orchestrates the Registry and a VectorBackend into the +// create/delete/list/reconcile index lifecycle (D4): the registry is always +// mutated first on create and last on delete, so a crash never leaves an +// orphan Postgres table with no registry entry, nor a registry entry that +// silently outlives its table's success. +type Service struct { + Registry *Registry + Backend VectorBackend +} + +// NewService constructs a Service over registry and backend. +func NewService(registry *Registry, backend VectorBackend) *Service { + return &Service{Registry: registry, Backend: backend} +} + +// CreateIndex reserves indexID in the registry (state CREATING), provisions +// the account's schema/role and the index's table in the backend, then flips +// the record to READY. A backend failure rolls back both the backend state +// and the reservation, so no half-created index survives the call. +func (s *Service) CreateIndex(ctx context.Context, accountID, indexID string, params CreateIndexParams) (*Record, error) { + if err := validateIndexID(indexID); err != nil { + return nil, err + } + if params.Dimension <= 0 { + return nil, fmt.Errorf("ochrevector: index %s: dimension must be positive, got %d", indexID, params.Dimension) + } + + now := time.Now().UTC() + rec := Record{ + ID: indexID, + Name: params.Name, + Dimension: params.Dimension, + EmbeddingModel: params.EmbeddingModel, + State: StateCreating, + CreatedAt: now, + UpdatedAt: now, + } + if err := s.Registry.Reserve(ctx, accountID, rec); err != nil { + return nil, err + } + rec.AccountID = accountID + + if err := s.provisionIndex(ctx, accountID, indexID, params.Dimension); err != nil { + s.rollbackCreate(ctx, accountID, indexID) + return nil, err + } + if err := s.Registry.SetState(ctx, accountID, indexID, StateReady); err != nil { + s.rollbackCreate(ctx, accountID, indexID) + return nil, err + } + + rec.State = StateReady + rec.UpdatedAt = time.Now().UTC() + return &rec, nil +} + +// provisionIndex runs the two backend calls a create needs: the account's +// schema/role (idempotent, safe even if it already exists) and the index's +// table. Both are safe to retry, so Reconcile calls this same path to resume +// a create a crash interrupted rather than rolling it back. +func (s *Service) provisionIndex(ctx context.Context, accountID, indexID string, dimension int) error { + if err := s.Backend.EnsureAccount(ctx, accountID); err != nil { + return fmt.Errorf("ochrevector: ensure account %s: %w", accountID, err) + } + if err := s.Backend.CreateIndex(ctx, accountID, IndexSpec{ID: indexID, Dimension: dimension}); err != nil { + return fmt.Errorf("ochrevector: create index %s: %w", indexID, err) + } + return nil +} + +// rollbackCreate unwinds a failed create: drop whatever table might have +// been created, then withdraw the reservation, so no half-state — a table +// without a registry entry, or a registry entry without a real table — +// survives. Best-effort: failures are logged, not propagated, since the +// caller has already failed and there is nothing further to roll back to. +func (s *Service) rollbackCreate(ctx context.Context, accountID, indexID string) { + rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout) + defer cancel() + if err := s.Backend.DropIndex(rbCtx, accountID, indexID); err != nil { + slog.WarnContext(rbCtx, "ochrevector: rollback drop index failed", "index", indexID, "account", accountID, "err", err) + } + if err := s.Registry.Delete(rbCtx, accountID, indexID); err != nil { + slog.WarnContext(rbCtx, "ochrevector: rollback delete registry record failed", "index", indexID, "account", accountID, "err", err) + } +} + +// DeleteIndex drops indexID's table in Postgres first, then removes its +// registry record (D4) — the ordering that guarantees an orphan table with no +// registry entry never happens, at the cost of a possible orphan *record* +// (state DELETING) if the process crashes between the two steps, which +// Reconcile finishes. Idempotent: deleting an already-absent index is a +// no-op success. +func (s *Service) DeleteIndex(ctx context.Context, accountID, indexID string) error { + rec, err := s.Registry.Get(ctx, accountID, indexID) + if err != nil { + return err + } + if rec == nil { + return nil + } + if err := s.Registry.SetState(ctx, accountID, indexID, StateDeleting); err != nil { + return err + } + if err := s.Backend.DropIndex(ctx, accountID, indexID); err != nil { + return fmt.Errorf("ochrevector: drop index %s: %w", indexID, err) + } + if err := s.Registry.Delete(ctx, accountID, indexID); err != nil { + return err + } + return nil +} + +// ListIndexes returns every index record accountID owns. +func (s *Service) ListIndexes(ctx context.Context, accountID string) ([]Record, error) { + return s.Registry.List(ctx, accountID) +} + +// Reconcile completes every record stuck in CREATING or DELETING for longer +// than reconcileStuckAfter — the re-entrant crash-recovery primitive a +// reaper calls on a schedule (no scheduler here). A record still within the +// grace period is left alone, so a create/delete in normal-latency flight is +// never disturbed by a concurrent Reconcile pass. +func (s *Service) Reconcile(ctx context.Context) error { + recs, err := s.Registry.ListAll(ctx) + if err != nil { + return fmt.Errorf("ochrevector: reconcile: list all: %w", err) + } + now := time.Now().UTC() + var errs []error + for _, rec := range recs { + if now.Sub(rec.UpdatedAt) < reconcileStuckAfter { + continue + } + switch rec.State { + case StateCreating: + if err := s.reconcileCreating(ctx, rec); err != nil { + errs = append(errs, err) + } + case StateDeleting: + if err := s.reconcileDeleting(ctx, rec); err != nil { + errs = append(errs, err) + } + } + } + return errors.Join(errs...) +} + +// reconcileCreating resumes a stuck CREATING record: EnsureAccount and +// CreateIndex are both idempotent, so this retries the same forward path +// rather than rolling back a create that may have mostly succeeded before +// the crash. +func (s *Service) reconcileCreating(ctx context.Context, rec Record) error { + if err := s.provisionIndex(ctx, rec.AccountID, rec.ID, rec.Dimension); err != nil { + return fmt.Errorf("ochrevector: reconcile creating index %s: %w", rec.ID, err) + } + if err := s.Registry.SetState(ctx, rec.AccountID, rec.ID, StateReady); err != nil { + return fmt.Errorf("ochrevector: reconcile creating index %s: %w", rec.ID, err) + } + return nil +} + +// reconcileDeleting resumes a stuck DELETING record: drop (idempotent), then +// remove the registry entry. +func (s *Service) reconcileDeleting(ctx context.Context, rec Record) error { + if err := s.Backend.DropIndex(ctx, rec.AccountID, rec.ID); err != nil { + return fmt.Errorf("ochrevector: reconcile deleting index %s: %w", rec.ID, err) + } + if err := s.Registry.Delete(ctx, rec.AccountID, rec.ID); err != nil { + return fmt.Errorf("ochrevector: reconcile deleting index %s: %w", rec.ID, err) + } + return nil +} diff --git a/spinifex/handlers/ochrevector/service_test.go b/spinifex/handlers/ochrevector/service_test.go new file mode 100644 index 000000000..87e67dbb4 --- /dev/null +++ b/spinifex/handlers/ochrevector/service_test.go @@ -0,0 +1,218 @@ +// Exercises the unexported service orchestration internals with no +// exported surface to drive them through. +// +//test:in-package +package handlers_ochrevector + +import ( + "context" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + svcAccountA = "111111111111" + svcAccountB = "222222222222" +) + +func newServiceTestSetup(t *testing.T) (*Service, *fakeBackend) { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + backend := newFakeBackend() + return NewService(NewRegistry(js), backend), backend +} + +func TestService_CreateIndex_HappyPath(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + + rec, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", CreateIndexParams{Name: "kb1", Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"}) + require.NoError(t, err) + require.NotNil(t, rec) + assert.Equal(t, StateReady, rec.State) + assert.Equal(t, svcAccountA, rec.AccountID) + + assert.True(t, backend.hasAccount(svcAccountA)) + assert.True(t, backend.hasIndex(svcAccountA, "idx-one")) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, StateReady, stored.State) +} + +func TestService_CreateIndex_ReserveCollision(t *testing.T) { + svc, _ := newServiceTestSetup(t) + ctx := context.Background() + params := CreateIndexParams{Name: "kb1", Dimension: 768} + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", params) + require.NoError(t, err) + + _, err = svc.CreateIndex(ctx, svcAccountA, "idx-one", params) + require.Error(t, err) + assert.ErrorIs(t, err, ErrIndexExists) + + // The first index survives the second (failed) create attempt untouched. + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, StateReady, stored.State) +} + +func TestService_CreateIndex_RollsBackOnBackendFailure(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + backend.failCreateIndex = errFakeBackend + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", CreateIndexParams{Name: "kb1", Dimension: 768}) + require.Error(t, err) + + // No half-state: no registry record and no backend index survive the + // failed create. + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, stored, "a failed create must not leave a registry record behind") + assert.False(t, backend.hasIndex(svcAccountA, "idx-one")) +} + +func TestService_CreateIndex_RollsBackOnEnsureAccountFailure(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + backend.failEnsureAccount = errFakeBackend + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", CreateIndexParams{Name: "kb1", Dimension: 768}) + require.Error(t, err) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, stored) +} + +func TestService_CreateIndex_RejectsNonPositiveDimension(t *testing.T) { + svc, _ := newServiceTestSetup(t) + ctx := context.Background() + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", CreateIndexParams{Name: "kb1", Dimension: 0}) + require.Error(t, err) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, stored, "an invalid create must never reach Reserve") +} + +func TestService_DeleteIndex_HappyPath(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-one", CreateIndexParams{Name: "kb1", Dimension: 768}) + require.NoError(t, err) + + require.NoError(t, svc.DeleteIndex(ctx, svcAccountA, "idx-one")) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-one") + require.NoError(t, err) + assert.Nil(t, stored) + assert.False(t, backend.hasIndex(svcAccountA, "idx-one")) +} + +func TestService_DeleteIndex_AbsentIsNoop(t *testing.T) { + svc, _ := newServiceTestSetup(t) + ctx := context.Background() + + require.NoError(t, svc.DeleteIndex(ctx, svcAccountA, "does-not-exist")) +} + +func TestService_ListIndexes_AccountScoped(t *testing.T) { + svc, _ := newServiceTestSetup(t) + ctx := context.Background() + + _, err := svc.CreateIndex(ctx, svcAccountA, "idx-a1", CreateIndexParams{Name: "a1", Dimension: 768}) + require.NoError(t, err) + _, err = svc.CreateIndex(ctx, svcAccountB, "idx-b1", CreateIndexParams{Name: "b1", Dimension: 768}) + require.NoError(t, err) + + listA, err := svc.ListIndexes(ctx, svcAccountA) + require.NoError(t, err) + require.Len(t, listA, 1) + assert.Equal(t, "idx-a1", listA[0].ID) + + listB, err := svc.ListIndexes(ctx, svcAccountB) + require.NoError(t, err) + require.Len(t, listB, 1, "account B must not see account A's indexes") + assert.Equal(t, "idx-b1", listB[0].ID) +} + +// seedStuckRecord reserves a record directly (bypassing CreateIndex/ +// DeleteIndex) with an UpdatedAt old enough for Reconcile to treat it as +// abandoned by a crashed writer. +func seedStuckRecord(t *testing.T, svc *Service, accountID, indexID, state string) { + t.Helper() + stuckAt := time.Now().UTC().Add(-2 * reconcileStuckAfter) + rec := Record{ + ID: indexID, + Name: "stuck", + Dimension: 768, + State: state, + CreatedAt: stuckAt, + UpdatedAt: stuckAt, + } + require.NoError(t, svc.Registry.Reserve(context.Background(), accountID, rec)) +} + +func TestService_Reconcile_FinishesStuckCreating(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + + // Simulates a crash after Reserve but before the backend ever ran. + seedStuckRecord(t, svc, svcAccountA, "idx-stuck-create", StateCreating) + assert.False(t, backend.hasIndex(svcAccountA, "idx-stuck-create")) + + require.NoError(t, svc.Reconcile(ctx)) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-stuck-create") + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, StateReady, stored.State) + assert.True(t, backend.hasIndex(svcAccountA, "idx-stuck-create")) +} + +func TestService_Reconcile_FinishesStuckDeleting(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + + // Simulates a crash after SetState(DELETING) but before DropIndex ran: + // the backend still has the table. + seedStuckRecord(t, svc, svcAccountA, "idx-stuck-delete", StateDeleting) + require.NoError(t, backend.CreateIndex(ctx, svcAccountA, IndexSpec{ID: "idx-stuck-delete", Dimension: 768})) + assert.True(t, backend.hasIndex(svcAccountA, "idx-stuck-delete")) + + require.NoError(t, svc.Reconcile(ctx)) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-stuck-delete") + require.NoError(t, err) + assert.Nil(t, stored) + assert.False(t, backend.hasIndex(svcAccountA, "idx-stuck-delete")) +} + +func TestService_Reconcile_LeavesFreshRecordsAlone(t *testing.T) { + svc, backend := newServiceTestSetup(t) + ctx := context.Background() + + now := time.Now().UTC() + require.NoError(t, svc.Registry.Reserve(ctx, svcAccountA, Record{ + ID: "idx-fresh", Dimension: 768, State: StateCreating, CreatedAt: now, UpdatedAt: now, + })) + + require.NoError(t, svc.Reconcile(ctx)) + + stored, err := svc.Registry.Get(ctx, svcAccountA, "idx-fresh") + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, StateCreating, stored.State, "a record still within the grace period must not be disturbed") + assert.False(t, backend.hasIndex(svcAccountA, "idx-fresh")) +} diff --git a/spinifex/network/host/vpc_hostport.go b/spinifex/network/host/vpc_hostport.go index a6be40432..2a258d05c 100644 --- a/spinifex/network/host/vpc_hostport.go +++ b/spinifex/network/host/vpc_hostport.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "net/netip" + "strings" "github.com/mulgadc/spinifex/spinifex/vm" ) @@ -75,6 +76,9 @@ func installVPCHostPort(ctx context.Context, r Runner, d VPCHostPort) error { if err := d.validate(); err != nil { return err } + if err := pruneCollidingVPCHostPorts(ctx, r, d.Name, d.Addr); err != nil { + return err + } if _, err := r.Run(ctx, "ovs-vsctl", "--may-exist", "add-port", "br-int", d.Name, "--", "set", "Interface", d.Name, "type=internal", "external_ids:iface-id="+d.IfaceID, "external_ids:attached-mac="+d.MAC); err != nil { @@ -95,6 +99,33 @@ func installVPCHostPort(ctx context.Context, r Runner, d VPCHostPort) error { return nil } +// pruneCollidingVPCHostPorts removes any daemon host port other than keep that +// already carries addr's host IP. A replaced appliance leaves its vhp- port and +// the connected route behind even after its ENI record is gone, so without this +// two ports would claim the same address and the kernel could route to the dead +// one. Only vhp- ports are pruned, never a guest tap sharing the address. +func pruneCollidingVPCHostPorts(ctx context.Context, r Runner, keep string, addr netip.Prefix) error { + out, err := r.Run(ctx, "ip", "-4", "-o", "addr", "show", "to", addr.Addr().String()+"/32") + if err != nil { + return fmt.Errorf("list interfaces holding %s: %w", addr.Addr(), err) + } + for line := range strings.SplitSeq(strings.TrimSpace(string(out)), "\n") { + fields := strings.Fields(line) + if len(fields) < 2 { + continue + } + dev := fields[1] + if dev == keep || !strings.HasPrefix(dev, vpcHostPortPrefix) { + continue + } + slog.Info("pruning stale VPC host port before install", "stale", dev, "keep", keep, "addr", addr) + if err := removeVPCHostPort(ctx, r, dev); err != nil { + return err + } + } + return nil +} + // removeVPCHostPort deletes the port, which takes its address with it. func removeVPCHostPort(ctx context.Context, r Runner, name string) error { if name == "" { diff --git a/spinifex/network/host/vpc_hostport_test.go b/spinifex/network/host/vpc_hostport_test.go index 5c27d9635..061668543 100644 --- a/spinifex/network/host/vpc_hostport_test.go +++ b/spinifex/network/host/vpc_hostport_test.go @@ -103,6 +103,48 @@ func TestInstallVPCHostPortDoesNotMaskHostAddress(t *testing.T) { } } +// A replaced appliance leaves its vhp- port behind at the reused subnet IP. +// Installing the new port must prune that stale namesake first, or two devs +// claim one address and the kernel can route to the dead one. +func TestInstallVPCHostPortPrunesCollidingPort(t *testing.T) { + s := newStubRunner() + s.expect("ovs-vsctl", nil, nil) + // Distinct prefixes so the stub matches deterministically: the readiness + // probe (`ip -4 -o addr show`) reports a stale vhp- namesake at our addr. + s.expect("ip -4", []byte("7: vhp-stalest inet 10.246.29.9/24 brd 10.246.29.255 scope global vhp-stalest\\ valid_lft forever preferred_lft forever\n"), nil) + s.expect("ip link", nil, nil) + s.expect("ip addr", nil, nil) + + d := testVPCHostPort() + if err := installVPCHostPort(context.Background(), s, d); err != nil { + t.Fatalf("installVPCHostPort: %v", err) + } + if !s.called("ovs-vsctl --if-exists del-port br-int vhp-stalest") { + t.Errorf("stale colliding vhp port was not pruned; calls: %v", s.calls) + } + if s.called("ovs-vsctl --if-exists del-port br-int " + d.Name) { + t.Errorf("must not prune the port being installed; calls: %v", s.calls) + } +} + +// Only our own vhp- ports are pruned: the port being installed (post-reboot it +// already holds the addr) and any foreign dev sharing it must be left alone. +func TestPruneCollidingVPCHostPortsLeavesKeepAndForeign(t *testing.T) { + s := newStubRunner() + s.expect("ovs-vsctl", nil, nil) + keep := VPCHostPortName("eni-0123456789abcdef0") + s.expect("ip -4", []byte( + "3: "+keep+" inet 10.246.29.9/24 scope global "+keep+"\n"+ + "9: eth0 inet 10.246.29.9/24 scope global eth0\n"), nil) + + if err := pruneCollidingVPCHostPorts(context.Background(), s, keep, netip.MustParsePrefix("10.246.29.9/24")); err != nil { + t.Fatalf("pruneCollidingVPCHostPorts: %v", err) + } + if s.called("ovs-vsctl --if-exists del-port") { + t.Errorf("nothing should be pruned (keep + foreign dev); calls: %v", s.calls) + } +} + func TestRemoveVPCHostPort(t *testing.T) { s := newStubRunner() s.expect("ovs-vsctl", nil, nil) diff --git a/spinifex/vm/manager.go b/spinifex/vm/manager.go index bf72c3603..f762e1bc8 100644 --- a/spinifex/vm/manager.go +++ b/spinifex/vm/manager.go @@ -71,6 +71,11 @@ type Deps struct { // ConsumeCleanShutdownMarker reports whether the previous run wrote a clean // shutdown marker, consuming it. Nil is treated as "no marker" (cautious recovery). ConsumeCleanShutdownMarker func() bool + + // BackingStoreReady reports whether this node's block-storage backend + // (predastore + viperblock) is confirmed up. Nil is treated as ready, + // preserving behaviour for tests and nodes with storage unwired. + BackingStoreReady func() bool } // Manager owns the in-memory map of running VMs on this node and every @@ -100,6 +105,13 @@ func (m *Manager) SetDeps(deps Deps) { m.deps = deps } // NodeID returns the node identifier the manager was constructed with. func (m *Manager) NodeID() string { return m.deps.NodeID } +// backingStoreReady reports whether the block-storage backend is confirmed up. +// A nil probe is treated as ready, so an unreachable NBD socket is only ever +// escalated to a kill when the store as a whole is known to be serving. +func (m *Manager) backingStoreReady() bool { + return m.deps.BackingStoreReady == nil || m.deps.BackingStoreReady() +} + // Get returns the VM for id (and true) or (nil, false). func (m *Manager) Get(id string) (*VM, bool) { m.mu.Lock() diff --git a/spinifex/vm/restore.go b/spinifex/vm/restore.go index 67546344a..55f5e0902 100644 --- a/spinifex/vm/restore.go +++ b/spinifex/vm/restore.go @@ -149,14 +149,26 @@ func (m *Manager) classifyRestoredInstances() []*VM { } if isInstanceProcessRunning(instance) { - if !AreVolumeSocketsValid(instance) { + socketsValid := AreVolumeSocketsValid(instance) + if !socketsValid && m.backingStoreReady() { + // Sockets are genuinely stale under a healthy store: a real + // orphan, so reap it and relaunch below. slog.Warn("QEMU alive but NBD sockets are stale, killing orphaned process for relaunch", "instance", instance.ID) if !killOrphanedQEMU(instance) { continue } } else { - slog.Info("Instance QEMU process still alive, reconnecting", "instance", instance.ID) + // Sockets valid, or the backing store is not yet ready. The + // latter makes an unreachable socket a transient dependency + // gap, never grounds to destroy a running VM: keep it and + // reconnect QMP, which needs neither predastore nor viperblock. + if socketsValid { + slog.Info("Instance QEMU process still alive, reconnecting", "instance", instance.ID) + } else { + slog.Warn("NBD sockets unreachable but backing store not ready; reconnecting instead of killing running QEMU", + "instance", instance.ID) + } if err := m.reconnectInstance(instance); err != nil { slog.Error("Failed to reconnect to running instance, marking recovery-failed to preserve user data", "instanceId", instance.ID, "err", err) diff --git a/spinifex/vm/restore_test.go b/spinifex/vm/restore_test.go index a2b7139bd..c4deba22c 100644 --- a/spinifex/vm/restore_test.go +++ b/spinifex/vm/restore_test.go @@ -101,6 +101,24 @@ func classifyTestManager(t *testing.T) (*Manager, *fakeStateStore, *fakeResource return m, store, rc } +// backingStoreReady gates the restore-path kill: an unreachable NBD socket is +// only reaped as an orphan when the store is confirmed up. A nil probe keeps +// today's behaviour so existing classifier tests and storage-unwired nodes are +// unaffected. +func TestManager_backingStoreReady(t *testing.T) { + m := NewManager() + assert.True(t, m.backingStoreReady(), + "a nil probe must default to ready so a missing wiring never blocks recovery") + + m.SetDeps(Deps{BackingStoreReady: func() bool { return false }}) + assert.False(t, m.backingStoreReady(), + "a store reporting not-ready must gate the kill so a transient gap cannot force-kill a live VM") + + m.SetDeps(Deps{BackingStoreReady: func() bool { return true }}) + assert.True(t, m.backingStoreReady(), + "a ready store must let genuine orphan reaping proceed") +} + func TestClassifyRestoredInstances_TerminatedMigrates(t *testing.T) { m, store, _ := classifyTestManager(t) v := &VM{ID: "i-term", Status: StateTerminated, InstanceType: "t3.micro"}