feat: Ochre vector store (tenant pgvector, embeddings, ingestion, operator teardown) (#788)

* feat: add /v1/embeddings gateway adapter

* feat: Ochre vector store pgvector backend, KV index registry, orchestration

Adds handlers_ochrevector: the VectorBackend interface + pgx/v5 impl
with per-account schema/role isolation enforced via SET LOCAL ROLE +
search_path inside every transaction, the KV index registry
(CREATING/READY/DELETING/STALE), and the create/delete/list/reconcile
orchestration tying them together with crash-safe rollback and
re-entrant recovery of stuck records.

* feat: singleton appliance-ensure state machine for ochre vector store

Adds the encrypted control-plane credential store and re-entrant
Ensure state machine for the platform pgvector appliance (D2): a
fixed-key KV claim mirroring Registry.Reserve elects exactly one
winner to launch the appliance via an injected ApplianceLauncher,
losers wait or resume a stale crashed provision, and the master
password is generated once, stored only AES-256-GCM encrypted, and
decrypted solely inside Connect to build a pgxBackend.

* feat: Ochre vector store async ingestion (chunker, job store, worker)

Adds Stage 3: a deterministic recursive chunk splitter, a per-account
ingestion job-record store mirroring the index registry, the
transactional per-document replace on the backend, and the
IngestService worker that lists a predastore source, chunks and
embeds each object, and replaces its rows idempotently.

* feat: Ochre vector store similarity query and metadata filter grammar

Adds VectorBackend.Query (cosine kNN, k default/cap, iterative HNSW
scan under a filter) and a Bedrock-KB-shaped metadata Filter AST that
compiles to parameterized SQL over jsonb, closing off arbitrary SQL by
allowlisting jsonb keys and always binding values as params. Also
fixes RunJob reporting READY on a job that ingested nothing regardless
of failure cause, and stamps SourceSpec.Metadata onto ingested rows so
the new filters have data to match against.

* feat: Ochre vector store request/response API, VectorService, and CLI

Adds the daemon-facing VectorService interface (index CRUD, ingest,
describe job, query), its server implementation over the existing
registry/backend/embedder, a NATS-forwarding client, and the
`spx admin ochre vector` operator CLI. Query resolves and embeds
against each index's own pinned model; ingest stamps the same
model/dimension onto every source spec server-side. accountID is
always taken from the caller context, never the request payload.

Daemon subscription wiring and manifest registration are a later
change; this is pure request/response plumbing with no construction
into daemon.go or awsgw.go.

* feat: wire Ochre vector store VectorService into the daemon

Constructs the platform Postgres appliance and VectorService behind a new
OchreVector config section that defaults off, so existing daemons are
unaffected. Adds a real RDS-backed ApplianceLauncher that provisions the
appliance over rds.CreateDBInstance/DescribeDBInstances, registers the six
ochre.vector.* NATS subjects mirroring the bedrock endpoint service pattern,
and degrades to a disabled (nil) service on any boot-time appliance failure
rather than blocking daemon startup.

* fix: fix Ochre vector store boot ordering and self-register subjects

startOchreVector ran before subscribeAll, so its appliance launch fired
rds.CreateDBInstance/DescribeDBInstances requests with no responder yet
registered, timing out the launch on a first-enable boot. It also ran
synchronously, blocking daemon startup up to the appliance launch timeout.

Move it to a background goroutine started after subscribeAll returns, and
have it register its own ochre.vector.* subjects once the VectorService is
actually constructed, via a new registerNatsSubs helper shared with
subscribeAll's boot-time table. Guard d.natsSubscriptions with d.mu, since
this background registration can now race the shutdown path's
unsubscribe-all sweep.

* feat: bake pgvector into the shared rds-postgres image

* fix: reach the ochre appliance over a routed daemon port on its ENI IP

The daemon dialled the RDS appliance over rec.Endpoint, a vanity spx3.net
hostname the control-plane host can neither resolve nor route to, so a healthy
appliance was never reachable and the vector store disabled itself.

Resolve the appliance's customer-ENI private IP read-only via
ec2.DescribeNetworkInterfaces filtered on the spinifex:rds-db-instance tag, dial
that IP for the DSN instead of the hostname, and give the daemon its own routed
ENI in the appliance's subnet bound to an OVS host port so the IP is reachable.
Retry appliance startup with bounded backoff instead of hard-disabling on the
first transient.

* fix: install pgvector at bootstrap for the Ochre vector appliance

The daemon connects to the appliance as its NOSUPERUSER master, and the
image's pgvector is not a trusted extension, so CREATE EXTENSION vector
failed with permission denied. rds-init now installs it via the bootstrap
superuser, gated on the appliance's fixed master name so tenant instances
are unaffected.

* fix: grant the Ochre appliance master CREATE on its database

The daemon's EnsureAccount runs CREATE SCHEMA per account as the NOSUPERUSER
master, which does not own the postgres database, so bootstrap now grants it
CREATE on the database alongside the pgvector install. Both are gated to the
appliance master, so tenant instances are unaffected.

* fix: make the Ochre appliance's pgvector reachable under per-account isolation

The daemon connects as a NOSUPERUSER master and runs every account operation
under SET ROLE with a search_path pinned to the account schema. Two gaps blocked
that against a real non-superuser Postgres: the master could not SET ROLE into
the roles it creates, and the account role could not see the vector type.

Bootstrap now grants the appliance master createrole_self_grant and installs
pgvector into a read-only extensions schema granted to PUBLIC; the backend
carries that schema on each account search_path so the vector type resolves
without weakening per-account write isolation. All appliance-scoped.

* fix: prune stale VPC host port before installing the daemon appliance port

A replaced Ochre appliance is handed the same subnet IP, so a new daemon ENI
mints a fresh vhp- port at an address a superseded port still holds. The old
port survives at the OS/OVS layer even after its ENI record is wiped, so two
interfaces claim one address and the kernel can route to the dead one
(EHOSTUNREACH). installVPCHostPort now prunes any other vhp- port holding the
target address before installing.

* fix: gate VM restore's stale-socket kill on backing-store readiness

On daemon start, a live VM whose NBD sockets fail to dial was force-killed as an
orphan. While predastore/viperblock is still coming up after a restart those
sockets are transiently unreachable, so a healthy VM (the Ochre appliance was
the observed victim) was destroyed and its RDS record flipped to failed. The
kill now only fires when the backing store is confirmed up; otherwise the live
QEMU is kept and QMP reconnected. Genuine orphan reaping under a healthy store
is unchanged.

* fix: drive ochre ingest jobs from a daemon-side sweep scheduler

* feat: add operator teardown verb for the ochre platform appliance

Adds Appliance.Teardown to remove the singleton platform Postgres
appliance (RDS instance, host port, KV record), an admin-only NATS
subject and daemon handler to drive it, and `spx admin ochre appliance
teardown --confirm` to invoke it without exposing the action through
the tenant awsgw surface.

* fix: register ochre appliance teardown before Connect so it survives a broken appliance

* fix: purge ochre index and job metadata on appliance teardown

* test: mark ochre in-package test files for the external-test gate
This commit is contained in:
Josh Siviour
2026-08-13 15:54:48 +10:00
committed by GitHub
parent 560b437f3b
commit d52245001e
49 changed files with 9367 additions and 16 deletions
+87
View File
@@ -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)
}
@@ -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)
}
+450
View File
@@ -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 <index-id>",
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 <job-id>",
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)
}
+375
View File
@@ -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)
}
+9
View File
@@ -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"
+3
View File
@@ -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
+6
View File
@@ -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=
+4 -2
View File
@@ -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"
+30
View File
@@ -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 <<SQL
CREATE SCHEMA IF NOT EXISTS extensions;
CREATE EXTENSION IF NOT EXISTS vector SCHEMA extensions;
GRANT USAGE ON SCHEMA extensions TO PUBLIC;
GRANT CREATE ON DATABASE postgres TO "${OCHRE_APPLIANCE_MASTER}";
ALTER ROLE "${OCHRE_APPLIANCE_MASTER}" SET createrole_self_grant = 'set, inherit';
SQL
}
# Drops the master password from the handoff once applied, so it is not left
# readable for the life of the boot. The rest is tmpfs and goes at reboot.
consume_secrets() {
@@ -313,6 +313,15 @@ if run_ok "initialize"; then
&& pass "initialize: password passed through the environment" || fail "initialize: password not in psql env"
grep -q 'CREATE DATABASE' "${PSQL_CALLS}" \
&& pass "initialize: initial database created" || fail "initialize: no CREATE DATABASE"
grep -q 'CREATE EXTENSION' "${PSQL_CALLS}" \
&& fail "initialize: a tenant instance installed a platform extension" \
|| pass "initialize: tenant instance installs no platform extension"
grep -q 'GRANT CREATE ON DATABASE' "${PSQL_CALLS}" \
&& fail "initialize: a tenant instance was granted database CREATE" \
|| pass "initialize: tenant instance gets no platform database grant"
grep -q 'createrole_self_grant' "${PSQL_CALLS}" \
&& fail "initialize: a tenant master got createrole_self_grant" \
|| pass "initialize: tenant master gets no createrole_self_grant"
# The master role is administrative but not a PostgreSQL superuser: a
# superuser reaches outside the database (COPY FROM PROGRAM, pg_read_file,
@@ -369,6 +378,29 @@ if run_ok "attach"; then
&& pass "attach: TLS survives a handoff without a new cert" || fail "attach: TLS turned off"
fi
# --- Case 3b: the Ochre appliance master installs pgvector ---
# No DBName, as the appliance is launched: create_database is skipped and the
# gated extension install is the only platform-specific bootstrap step.
reset_state
MASTER_USER=ochre_vector_admin
write_handoff initialize 's3cr3t' ''
write_parameters
if run_ok "ochre-appliance"; then
grep -q 'CREATE EXTENSION IF NOT EXISTS vector SCHEMA extensions' "${PSQL_CALLS}" \
&& pass "ochre-appliance: pgvector installed into the extensions schema" \
|| fail "ochre-appliance: pgvector not installed into the extensions schema"
grep -q 'GRANT USAGE ON SCHEMA extensions TO PUBLIC' "${PSQL_CALLS}" \
&& pass "ochre-appliance: extensions schema readable by account roles" \
|| fail "ochre-appliance: extensions schema not granted to PUBLIC"
grep -q 'GRANT CREATE ON DATABASE postgres TO "ochre_vector_admin"' "${PSQL_CALLS}" \
&& pass "ochre-appliance: master granted CREATE on the database" \
|| fail "ochre-appliance: master not granted CREATE on the database"
grep -q 'createrole_self_grant' "${PSQL_CALLS}" \
&& pass "ochre-appliance: master can set into the roles it creates" \
|| fail "ochre-appliance: master missing createrole_self_grant"
fi
unset MASTER_USER
# --- Case 3a: a parameter group cannot log the master password ---
# The customer's parameters are installed before the master role is applied and
# the bootstrap postmaster does not override config_file, so log_statement = all
+32 -11
View File
@@ -99,17 +99,18 @@ type Config struct {
DataDir string `json:"DataDir" mapstructure:"data_dir"`
Services []string `json:"Services" mapstructure:"services"` // Which services this node runs locally
Daemon DaemonConfig `json:"Daemon" mapstructure:"daemon"`
NATS NATSConfig `json:"NATS" mapstructure:"nats"`
Predastore PredastoreConfig `json:"Predastore" mapstructure:"predastore"`
Viperblock ViperblockConfig `json:"Viperblock" mapstructure:"viperblock"`
EBS EBSConfig `json:"EBS" mapstructure:"ebs"`
AWSGW AWSGWConfig `json:"AWSGW" mapstructure:"awsgw"`
VPCD VPCDConfig `json:"VPCD" mapstructure:"vpcd"`
Northstar NorthstarConfig `json:"Northstar" mapstructure:"northstar"`
RDS RDSConfig `json:"RDS" mapstructure:"rds"`
ACM ACMConfig `json:"ACM" mapstructure:"acm"`
Bedrock BedrockConfig `json:"Bedrock" mapstructure:"bedrock"`
Daemon DaemonConfig `json:"Daemon" mapstructure:"daemon"`
NATS NATSConfig `json:"NATS" mapstructure:"nats"`
Predastore PredastoreConfig `json:"Predastore" mapstructure:"predastore"`
Viperblock ViperblockConfig `json:"Viperblock" mapstructure:"viperblock"`
EBS EBSConfig `json:"EBS" mapstructure:"ebs"`
AWSGW AWSGWConfig `json:"AWSGW" mapstructure:"awsgw"`
VPCD VPCDConfig `json:"VPCD" mapstructure:"vpcd"`
Northstar NorthstarConfig `json:"Northstar" mapstructure:"northstar"`
RDS RDSConfig `json:"RDS" mapstructure:"rds"`
ACM ACMConfig `json:"ACM" mapstructure:"acm"`
Bedrock BedrockConfig `json:"Bedrock" mapstructure:"bedrock"`
OchreVector OchreVectorConfig `json:"OchreVector" mapstructure:"ochre_vector"`
BaseDir string `json:"BaseDir" mapstructure:"base_dir"`
WalDir string `json:"WalDir" mapstructure:"wal_dir"`
@@ -253,6 +254,26 @@ type BedrockConfig struct {
// (10.248.0.0/14) and EKS control-plane (10.252.0.0/14) supernets.
const BedrockDefaultSystemVPCSupernet = "10.244.0.0/14"
// OchreVectorConfig gates the Ochre vector store's platform Postgres
// appliance and VectorService NATS surface. Disabled by default: an unset
// section constructs nothing and registers no ochre.vector.* subjects, so an
// existing deployment is unaffected until an operator opts in.
type OchreVectorConfig struct {
Enabled bool `json:"Enabled" mapstructure:"enabled"`
// EmbeddingsEndpoint is the base URL of the self-hosted OpenAI/TEI-
// compatible embeddings endpoint (gateway_bedrock.Embedder) serving
// EmbeddingModel.
EmbeddingsEndpoint string `json:"EmbeddingsEndpoint" mapstructure:"embeddings_endpoint"`
// EmbeddingModel is the model id sent to EmbeddingsEndpoint and stamped
// on every index record. Empty takes gateway_bedrock.DefaultEmbeddingModel
// ("nomic-embed-text-v1.5").
EmbeddingModel string `json:"EmbeddingModel" mapstructure:"embedding_model"`
// PostgresImage names the appliance's Postgres image. Currently
// informational only: RDS's CreateDBInstanceInput has no image-selection
// field, so this is not yet threaded into the appliance launch.
PostgresImage string `json:"PostgresImage" mapstructure:"postgres_image"`
}
// ACMConfig holds the operator-level ACM certificate-issuance configuration.
// Deliberately small: four keys, nothing deployment-specific and nothing
// derivable. In particular there is no allowed-domains list (public modes are
+69 -1
View File
@@ -58,6 +58,7 @@ import (
handlers_eks "github.com/mulgadc/spinifex/spinifex/handlers/eks"
handlers_elbv2 "github.com/mulgadc/spinifex/spinifex/handlers/elbv2"
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
handlers_rds "github.com/mulgadc/spinifex/spinifex/handlers/rds"
"github.com/mulgadc/spinifex/spinifex/instancetypes"
"github.com/mulgadc/spinifex/spinifex/kvutil"
@@ -160,6 +161,8 @@ type Daemon struct {
bedrockReaper *handlers_bedrock.Reaper
acmService *handlers_acm.ACMServiceImpl
acmRenewalWorker *handlers_acm.Worker
ochreVectorService handlers_ochrevector.VectorService
ochreAppliance *handlers_ochrevector.Appliance
ecrMetaService *handlers_ecr.MetaServiceImpl
routeTableService *handlers_ec2_routetable.RouteTableServiceImpl
natGatewayService *handlers_ec2_natgw.NatGatewayServiceImpl
@@ -1128,6 +1131,12 @@ func (d *Daemon) subscribeAll() error {
)
}
// Ochre vector store tenant surface is deliberately NOT registered here:
// its subjects only exist once the platform appliance has finished
// launching (startOchreVector, run in the background after this method
// returns), and registering them requires this same table-driven
// mechanism after the fact — see registerNatsSubs.
// ECR gateway → daemon subscriptions. The daemon owns the per-account
// JetStream KV metadata; blob/manifest bytes never traverse these subjects.
if d.ecrMetaService != nil {
@@ -1173,6 +1182,17 @@ func (d *Daemon) subscribeAll() error {
natsSub{dhcp.TopicOwnerCheck, d.handleDHCPOwnerCheck, "spinifex-workers"},
)
return d.registerNatsSubs(subs)
}
// registerNatsSubs turns each entry in subs into a live NATS subscription and
// records it in d.natsSubscriptions. Shared by subscribeAll's boot-time table
// and any subject registered later, off the main boot goroutine (e.g. Ochre's
// VectorService, whose subjects only exist once its platform appliance has
// finished launching) — the map write is mutex-guarded because that later
// registration can race the shutdown path's unsubscribe-all sweep, unlike
// subscribeAll's own single-goroutine, boot-time-only writes.
func (d *Daemon) registerNatsSubs(subs []natsSub) error {
for _, s := range subs {
var sub *nats.Subscription
var err error
@@ -1185,7 +1205,9 @@ func (d *Daemon) subscribeAll() error {
if err != nil {
return fmt.Errorf("failed to subscribe to %s: %w", s.topic, err)
}
d.mu.Lock()
d.natsSubscriptions[s.topic] = sub
d.mu.Unlock()
slog.Info("Subscribed to NATS topic", "topic", s.topic, "queue", s.queueGroup)
}
return nil
@@ -1865,13 +1887,51 @@ func (d *Daemon) startCluster() error {
return fmt.Errorf("failed to subscribe to NATS topics: %w", err)
}
// Ochre vector store platform appliance + VectorService, backgrounded:
// launching the appliance can block on a cold RDS VM boot for minutes, and
// its own launch request needs the rds.* responders subscribeAll just
// registered above, so it cannot run any earlier. Gated on
// config.OchreVector.Enabled (default off); any failure along the way is
// logged and leaves the feature's subjects unregistered rather than
// failing startCluster or blocking daemon boot. Tracked on shutdownWg so
// shutdown waits for it to observe d.ctx cancellation and unwind instead
// of leaking; d.ctx is what bounds every NATS call it makes.
d.shutdownWg.Go(func() {
defer func() {
if r := recover(); r != nil {
slog.Error("Ochre vector store startup goroutine panicked", "recover", r)
}
}()
d.startOchreVector()
})
// Remove this node's appliance-VPC host port on shutdown. Same config gate
// as startOchreVector, so a disabled daemon spawns nothing. d.ochreAppliance
// stays nil until the retry loop succeeds, so early shutdown tears down nothing.
if d.config.OchreVector.Enabled {
d.shutdownWg.Go(func() {
<-d.ctx.Done()
d.mu.Lock()
appliance := d.ochreAppliance
d.mu.Unlock()
if appliance == nil {
return
}
if err := appliance.TeardownHostPort(); err != nil {
slog.Warn("Ochre vector store: daemon host port teardown failed", "err", err)
}
})
}
// DNS record writer: a queue-group consumer of dns.recordset.change. Every
// node subscribes, so the writer locks each zone before its
// read-modify-write. No-op when northstar S3 is not configured.
if sub, err := d.dnsWriter.Subscribe(d.natsConn); err != nil {
return fmt.Errorf("failed to subscribe DNS record writer: %w", err)
} else if sub != nil {
d.mu.Lock()
d.natsSubscriptions[handlers_dns.SubjectRecordsetChange] = sub
d.mu.Unlock()
slog.Info("Subscribed DNS record writer", "subject", handlers_dns.SubjectRecordsetChange, "queue", handlers_dns.QueueGroup)
}
@@ -2557,7 +2617,15 @@ func (d *Daemon) setupShutdown() {
d.eksService.Shutdown()
}
for _, sub := range d.natsSubscriptions {
// Snapshotted under d.mu rather than ranged over live: a background
// registration (e.g. Ochre's VectorService, registered off the boot
// goroutine once its appliance finishes launching) can still be
// writing to this map concurrently with shutdown.
d.mu.Lock()
subsSnapshot := make(map[string]*nats.Subscription, len(d.natsSubscriptions))
maps.Copy(subsSnapshot, d.natsSubscriptions)
d.mu.Unlock()
for _, sub := range subsSnapshot {
slog.Info("Unsubscribing from NATS", "subject", sub.Subject)
if err := sub.Unsubscribe(); err != nil {
if errors.Is(err, nats.ErrBadSubscription) {
+269
View File
@@ -0,0 +1,269 @@
package daemon
import (
"context"
"errors"
"fmt"
"log/slog"
"path/filepath"
"time"
"github.com/mulgadc/spinifex/spinifex/admin"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
"github.com/mulgadc/spinifex/spinifex/network/host"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/nats-io/nats.go/jetstream"
)
// ochreApplianceLaunchTimeout bounds the platform Postgres appliance's whole
// create-and-poll-to-available launch: generous enough for a cold RDS VM
// boot plus initdb on the smallest instance class.
const ochreApplianceLaunchTimeout = 10 * time.Minute
// ochreStartupMaxAttempts bounds how many times startOchreVector retries the
// Ensure-then-Connect pair. A still-provisioning appliance or a host port
// not yet routed is transient, but the retry still needs a ceiling so a
// genuine misconfiguration eventually gives up instead of retrying forever.
const ochreStartupMaxAttempts = 5
// ochreStartupInitialBackoff and ochreStartupMaxBackoff bound the wait
// between attempts: it doubles from the initial value and is capped so a
// run of early failures cannot stall an eventually-successful attempt for
// an unbounded amount of time.
const (
ochreStartupInitialBackoff = 15 * time.Second
ochreStartupMaxBackoff = 3 * time.Minute
)
// startOchreVector wires the Ochre vector store's VectorService when
// config.OchreVectorConfig.Enabled is set. Disabled (the default) leaves
// d.ochreVectorService nil, so subscribeAll registers no ochre.vector.*
// subject and daemon behavior is byte-for-byte unchanged. Any failure below
// — JetStream, the master key, or the platform appliance itself — is logged
// and leaves d.ochreVectorService nil rather than failing startCluster: the
// vector store is a feature dependency, never a daemon-boot one.
func (d *Daemon) startOchreVector() {
cfg := d.config.OchreVector
if !cfg.Enabled {
return
}
js, err := jetstream.New(d.natsConn)
if err != nil {
slog.Warn("Ochre vector store disabled: JetStream unavailable", "err", err)
return
}
masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key"))
if err != nil {
slog.Warn("Ochre vector store disabled: master key unavailable", "err", err)
return
}
launcher := handlers_ochrevector.NewRDSApplianceLauncher(d.natsConn, ochreApplianceLaunchTimeout)
appliance, err := handlers_ochrevector.NewAppliance(js, masterKey, launcher)
if err != nil {
slog.Warn("Ochre vector store disabled: appliance construction failed", "err", err)
return
}
// Give the daemon a routed presence in the appliance's system-VPC subnet
// before Connect dials it: a tag-filtered ENI describe resolves the real
// dial IP and subnet, never the unroutable vanity endpoint hostname.
appliance.WithHostPort(handlers_ochrevector.HostPortDeps{
VPC: d.vpcService,
HostPort: host.NewOVSPlumber(),
NodeID: d.node,
})
// registry and jobs need only js, never a live backend, so they are built
// (and attached) here rather than after Connect: Teardown must be able to
// purge them even for an appliance that never comes up. Reused below,
// once connected, rather than reconstructed.
registry := handlers_ochrevector.NewRegistry(js)
jobs := handlers_ochrevector.NewJobStore(js)
appliance.WithStores(registry, jobs)
// Publish the appliance and register the operator teardown subject BEFORE
// Connect: recovering a broken singleton is teardown's whole purpose, so it
// must stay reachable even when the appliance is unreachable and the connect
// loop below fails. Teardown needs only js+launcher+KV, never a live backend.
d.mu.Lock()
d.ochreAppliance = appliance
d.mu.Unlock()
if err := d.registerNatsSubs([]natsSub{
{handlers_ochrevector.SubjectTeardownAppliance, handleNATSRequest(d.handleOchreApplianceTeardown), "spinifex-workers"},
}); err != nil {
slog.Error("Ochre vector store: failed to register appliance teardown subject", "err", err)
}
backend, err := d.connectOchreAppliance(appliance)
if err != nil {
slog.Warn("Ochre vector store disabled: platform appliance not reachable after retrying",
"attempts", ochreStartupMaxAttempts, "err", err)
return
}
embedModel := cfg.EmbeddingModel
if embedModel == "" {
embedModel = gateway_bedrock.DefaultEmbeddingModel
}
embedder := gateway_bedrock.NewEmbedder(gateway_bedrock.NewStaticEndpointResolver(map[string]string{
embedModel: cfg.EmbeddingsEndpoint,
}))
store := objectstore.NewS3ObjectStoreFromConfig(admin.DialTarget(d.config.Predastore.Host),
d.config.Predastore.Region, d.config.Predastore.AccessKey, d.config.Predastore.SecretKey)
// registry and jobs were already built above (and attached to appliance
// via WithStores) before Connect was attempted; reused here, not rebuilt.
service := handlers_ochrevector.NewService(registry, backend)
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, store, embedder)
vectorService := handlers_ochrevector.NewVectorService(service, ingest, jobs, registry, backend, embedder)
// A shutdown that lands in the gap between Connect succeeding above and
// this check does not leave anything to unwind: nothing has been
// registered yet, so there are no subjects or map entries to clean up --
// this is a best-effort skip of a subscribe attempt that would very
// likely fail on a closing NATS connection anyway, not a correctness
// requirement.
if d.ctx.Err() != nil {
slog.Warn("Ochre vector store: daemon shutting down before appliance came up; not registering subjects")
return
}
// Registered here rather than in subscribeAll: these six subjects only
// exist once the appliance above is actually connected, which can be
// minutes after subscribeAll already ran. registerNatsSubs is the same
// table-driven mechanism subscribeAll itself uses, so a queue-group
// registration here is indistinguishable from one made at boot.
subs := []natsSub{
{handlers_ochrevector.SubjectCreateIndex, handleNATSRequest(vectorService.CreateIndex), "spinifex-workers"},
{handlers_ochrevector.SubjectDeleteIndex, handleNATSRequest(vectorService.DeleteIndex), "spinifex-workers"},
{handlers_ochrevector.SubjectListIndexes, handleNATSRequest(vectorService.ListIndexes), "spinifex-workers"},
{handlers_ochrevector.SubjectIngest, handleNATSRequest(vectorService.Ingest), "spinifex-workers"},
{handlers_ochrevector.SubjectDescribeJob, handleNATSRequest(vectorService.DescribeJob), "spinifex-workers"},
{handlers_ochrevector.SubjectQuery, handleNATSRequest(vectorService.Query), "spinifex-workers"},
}
if err := d.registerNatsSubs(subs); err != nil {
slog.Error("Ochre vector store: failed to register NATS subjects", "err", err)
return
}
// Assigned last, and only after every subject above is live: a reader
// (there is none today beyond observability) must never see a non-nil
// service whose subjects are not yet actually serving.
d.ochreVectorService = vectorService
go d.runOchreIngestScheduler(ingest)
slog.Info("Ochre vector store enabled")
}
// ochreIngestSweepInterval paces the scheduler tick that drives PENDING ingest
// jobs and re-drives stale RUNNING ones.
const ochreIngestSweepInterval = 15 * time.Second
// runOchreIngestScheduler drives the ingestion sweep on a timer until the
// daemon context is cancelled. An initial sweep keeps submit latency low; the
// ticker then re-drives PENDING and crash-abandoned RUNNING jobs.
func (d *Daemon) runOchreIngestScheduler(ingest *handlers_ochrevector.IngestService) {
if err := ingest.Sweep(d.ctx); err != nil {
slog.Warn("Ochre vector store: initial ingest sweep", "err", err)
}
ticker := time.NewTicker(ochreIngestSweepInterval)
defer ticker.Stop()
for {
select {
case <-d.ctx.Done():
return
case <-ticker.C:
if err := ingest.Sweep(d.ctx); err != nil {
slog.Warn("Ochre vector store: ingest sweep", "err", err)
}
}
}
}
// handleOchreApplianceTeardown answers SubjectTeardownAppliance: destroys the
// platform appliance singleton, then clears this daemon's own references so
// nothing keeps issuing queries against a torn-down backend. A daemon whose
// appliance never came up (disabled, still starting, or already torn down by
// an earlier call) has nothing to act on, which is reported as an error
// rather than silently accepted -- an operator asking to tear down expects
// one to have existed.
func (d *Daemon) handleOchreApplianceTeardown(ctx context.Context, _ *handlers_ochrevector.TeardownApplianceRequest, _ string) (*handlers_ochrevector.TeardownApplianceResponse, error) {
d.mu.Lock()
appliance := d.ochreAppliance
d.mu.Unlock()
if appliance == nil {
return nil, errors.New("ochrevector: platform appliance is not enabled or not up on this node")
}
if err := appliance.Teardown(ctx); err != nil {
return nil, fmt.Errorf("ochrevector: teardown platform appliance: %w", err)
}
d.mu.Lock()
d.ochreAppliance = nil
d.ochreVectorService = nil
d.mu.Unlock()
return &handlers_ochrevector.TeardownApplianceResponse{}, nil
}
// connectOchreAppliance drives Ensure then Connect as one retryable unit, up
// to ochreStartupMaxAttempts times with exponential backoff between
// attempts. Previously a single provisioning timeout -- the appliance still
// booting, or its host port not yet routed -- hard-disabled the feature for
// the daemon's entire lifetime; this bounds that to a handful of attempts
// instead of exactly one, while still eventually giving up on a persistent
// failure rather than retrying forever.
func (d *Daemon) connectOchreAppliance(appliance *handlers_ochrevector.Appliance) (handlers_ochrevector.VectorBackend, error) {
backoff := ochreStartupInitialBackoff
var lastErr error
for attempt := 1; attempt <= ochreStartupMaxAttempts; attempt++ {
if d.ctx.Err() != nil {
return nil, d.ctx.Err()
}
backend, err := d.ensureAndConnectOchreApplianceOnce(appliance)
if err == nil {
return backend, nil
}
lastErr = err
if attempt == ochreStartupMaxAttempts {
break
}
slog.Warn("Ochre vector store: appliance not ready, retrying",
"attempt", attempt, "max_attempts", ochreStartupMaxAttempts, "backoff", backoff, "err", lastErr)
select {
case <-d.ctx.Done():
return nil, d.ctx.Err()
case <-time.After(backoff):
}
backoff *= 2
if backoff > 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
}
+48
View File
@@ -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")
}
+1
View File
@@ -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() },
}
}
+156
View File
@@ -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
}
+173
View File
@@ -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())
}
+337
View File
@@ -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)
}
+386
View File
@@ -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)
}
+541
View File
@@ -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()
}
@@ -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)
}
+92
View File
@@ -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)
}
+186
View File
@@ -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
}
+164
View File
@@ -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)
}
@@ -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")
+342
View File
@@ -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
}
@@ -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])
}
+242
View File
@@ -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
}
@@ -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")
}
@@ -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()
}
@@ -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`))
}
+348
View File
@@ -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
}
@@ -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))
}
+274
View File
@@ -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
}
+202
View File
@@ -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)
}
@@ -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()
}
@@ -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
}
+108
View File
@@ -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)
}
@@ -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):
}
}
}
@@ -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
// ": <code>" 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)
}
+325
View File
@@ -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
}
@@ -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)
}
+198
View File
@@ -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
}
@@ -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"))
}
+31
View File
@@ -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 == "" {
@@ -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)
+12
View File
@@ -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()
+14 -2
View File
@@ -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)
+18
View File
@@ -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"}