diff --git a/.docker/docker-compose-infra-oriole-pgvector-override.yml b/.docker/docker-compose-infra-oriole-pgvector-override.yml new file mode 100644 index 000000000..787adbca7 --- /dev/null +++ b/.docker/docker-compose-infra-oriole-pgvector-override.yml @@ -0,0 +1,6 @@ +services: + tenant_db: + build: + context: . + dockerfile: .docker/orioledb-pgvector.Dockerfile + image: supabase-storage-orioledb-pgvector:pg17-local diff --git a/.docker/orioledb-pgvector.Dockerfile b/.docker/orioledb-pgvector.Dockerfile new file mode 100644 index 000000000..99a75f9f2 --- /dev/null +++ b/.docker/orioledb-pgvector.Dockerfile @@ -0,0 +1,22 @@ +ARG ORIOLEDB_IMAGE=orioledb/orioledb:latest-pg17 +FROM ${ORIOLEDB_IMAGE} + +ARG PGVECTOR_VERSION=0.8.2 + +USER root + +RUN apk add --no-cache --virtual .pgvector-build-deps \ + build-base \ + clang \ + git \ + llvm \ + && git clone --depth 1 --branch "v${PGVECTOR_VERSION}" \ + https://github.com/pgvector/pgvector.git /tmp/pgvector \ + && cd /tmp/pgvector \ + && make clean \ + && make OPTFLAGS="" \ + && make install \ + && rm -rf /tmp/pgvector \ + && apk del .pgvector-build-deps + +USER postgres diff --git a/.env.acceptance.sample b/.env.acceptance.sample index 2b034b439..35951ae32 100644 --- a/.env.acceptance.sample +++ b/.env.acceptance.sample @@ -24,7 +24,8 @@ ACCEPTANCE_ENABLE_RENDER=true # ACCEPTANCE_RLS_READ_OBJECT points at an existing object readable by authenticated role. ACCEPTANCE_ENABLE_RLS_SETUP=true # ACCEPTANCE_RLS_READ_OBJECT=authenticated/casestudy.png -# Vector acceptance requires a configured S3 Vectors-compatible service. +# Vector acceptance requires either the local pgvector provider or a configured +# S3 Vectors-compatible service. ACCEPTANCE_ENABLE_VECTOR=false ACCEPTANCE_ENABLE_ICEBERG=true ACCEPTANCE_ENABLE_WIRE=true diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 2c51dca74..175b3c7e8 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -81,7 +81,8 @@ jobs: env: ACCEPTANCE_ADMIN_URL: ${{ matrix.tenancy == 'multitenant' && 'http://127.0.0.1:5001' || '' }} ACCEPTANCE_ENABLE_ADMIN: ${{ matrix.tenancy == 'multitenant' && 'true' || 'false' }} - ACCEPTANCE_INFRA_RESTART_SCRIPT: ${{ matrix.database == 'oriole' && 'infra:restart:ci:oriole' || 'infra:restart:ci' }} + ACCEPTANCE_ENABLE_VECTOR: "true" + ACCEPTANCE_INFRA_RESTART_SCRIPT: ${{ matrix.database == 'oriole' && 'infra:restart:ci:oriole:pgvector' || 'infra:restart:ci' }} ACCEPTANCE_PROFILE: ${{ inputs.profile || 'full' }} ACCEPTANCE_X_FORWARDED_HOST: ${{ matrix.tenancy == 'multitenant' && 'bjhaohmqunupljrqypxz.local.dev' || '' }} MULTI_TENANT: ${{ matrix.tenancy == 'multitenant' && 'true' || 'false' }} @@ -89,6 +90,9 @@ jobs: REQUEST_X_FORWARDED_HOST_REGEXP: ${{ matrix.tenancy == 'multitenant' && '^([a-z]{20})[.]local[.](?:com|dev)$' || '' }} STORAGE_PUBLIC_URL: ${{ matrix.tenancy == 'multitenant' && 'http://127.0.0.1:5000' || '' }} STORAGE_BACKEND: ${{ matrix.storage_backend }} + VECTOR_BUCKET_PROVIDER: pgvector + VECTOR_DATABASE_URL: ${{ matrix.tenancy == 'single' && 'postgresql://postgres:postgres@127.0.0.1:5432/postgres' || '' }} + VECTOR_STORE_MIGRATIONS_ENABLED: "true" run: | mkdir -p data coverage/acceptance chmod -R 777 data diff --git a/acceptance/API_COVERAGE.md b/acceptance/API_COVERAGE.md index 1e2a2aae4..4c665d93a 100644 --- a/acceptance/API_COVERAGE.md +++ b/acceptance/API_COVERAGE.md @@ -42,7 +42,7 @@ the configured target, and the selected profile includes the spec: | Render | `ACCEPTANCE_ENABLE_RENDER=true` | public, authenticated, and signed image transformation routes, `webp` output format, non-image input errors, invalid transformation validation | | RLS | `ACCEPTANCE_ENABLE_RLS_SETUP=true` plus anon/authenticated keys and RLS resource config | authenticated allow and anon deny for read/write on configured policies | | Path edges | Derived from `ACCEPTANCE_TARGET` and `STORAGE_BACKEND` | list-v2 preservation for object names with empty path segments; local S3/MinIO backends skip this case directly | -| Vector | `ACCEPTANCE_ENABLE_VECTOR=true` | vector bucket pagination, index pagination, put/get/list/query/delete lifecycle | +| Vector | `ACCEPTANCE_ENABLE_VECTOR=true` with local pgvector or S3 Vectors configured | vector bucket pagination, index pagination, vector list pagination, metadata filter keys, non-filterable metadata rejection, default distance omission, cosine and euclidean query behavior, put/get/list/query/delete lifecycle | | Iceberg | `ACCEPTANCE_ENABLE_ICEBERG=true` | analytics bucket, catalog config, namespace (create/list/load/head/drop, missing load/drop, upsert on re-create, drop blocked when non-empty), table create/list/page-size/load/head/drop, missing load/drop, commit success/conflict | ## Intentionally Gated diff --git a/acceptance/README.md b/acceptance/README.md index f0f8b2f4f..9f6853888 100644 --- a/acceptance/README.md +++ b/acceptance/README.md @@ -98,6 +98,15 @@ Local CI also enables admin acceptance for multitenant matrix entries. Path-edge from the local storage backend, so empty path segment object names are exercised only on backends that can store them. +Local CI enables vector acceptance on PostgreSQL and OrioleDB matrix rows using the pgvector-backed +local provider, covering both S3/file storage backends and single/multitenant modes. OrioleDB rows +use the locally built OrioleDB+pgvector image. Single-tenant rows create and migrate a dedicated +`storage_vectors` database from `VECTOR_DATABASE_URL`. Multitenant pgvector rows provision the local +tenant with the configured tenant database URL and pool URL. Multitenant pgvector index DDL reuses +the active tenant transaction connection; single-tenant pgvector and S3 Vectors index creation keep +physical side effects outside retried metadata transactions and clean up committed metadata on +post-commit failures. + ## GitHub Environments The workflow dispatch `acceptance_environment` input uses `local` for the managed local run. Any @@ -125,7 +134,7 @@ secrets as environment secrets. | `ACCEPTANCE_ENABLE_CDN` | Enables CDN purge tests. Managed local runs provide a purge stub by default. | | `ACCEPTANCE_ENABLE_RENDER` | Enables image transformation tests. | | `ACCEPTANCE_ENABLE_RLS_SETUP` | Enables RLS tests; requires service, anon, authenticated keys and bucket/prefix policy resources. | -| `ACCEPTANCE_ENABLE_VECTOR` | Enables vector bucket API tests. Requires a configured S3 Vectors-compatible service. | +| `ACCEPTANCE_ENABLE_VECTOR` | Enables vector bucket API tests. Requires local pgvector or a configured S3 Vectors target. | | `ACCEPTANCE_ENABLE_ICEBERG` | Enables Iceberg catalog API tests. | | `ACCEPTANCE_ENABLE_WIRE` | Enables wire-level tests outside the `wire` / `full` profiles. | | `ACCEPTANCE_RLS_BUCKET` | Bucket used by opt-in RLS tests. Defaults to local dummy `bucket2`. | diff --git a/acceptance/scripts/run-managed-local.ts b/acceptance/scripts/run-managed-local.ts index 81b2490a7..75c5de597 100644 --- a/acceptance/scripts/run-managed-local.ts +++ b/acceptance/scripts/run-managed-local.ts @@ -277,7 +277,11 @@ function closeHttpServer(server: HttpServer): Promise { function resolveInfraRestartScript() { const script = acceptanceEnv('ACCEPTANCE_INFRA_RESTART_SCRIPT') ?? 'infra:restart:ci' - const allowed = new Set(['infra:restart:ci', 'infra:restart:ci:oriole']) + const allowed = new Set([ + 'infra:restart:ci', + 'infra:restart:ci:oriole', + 'infra:restart:ci:oriole:pgvector', + ]) if (!allowed.has(script)) { throw new Error(`Unsupported ACCEPTANCE_INFRA_RESTART_SCRIPT: ${script}`) diff --git a/acceptance/specs/vector.test.ts b/acceptance/specs/vector.test.ts index e53267319..f4fdbf43d 100644 --- a/acceptance/specs/vector.test.ts +++ b/acceptance/specs/vector.test.ts @@ -13,8 +13,23 @@ interface VectorListIndexesResponse { nextToken?: string } -interface VectorListResponse { - vectors?: Array<{ key?: string }> +interface VectorResponse { + distanceMetric?: string + nextToken?: string + vectors?: Array<{ + data?: { float32?: number[] } + distance?: number + key?: string + metadata?: Record + }> +} + +function fixedLengthVectorKey(prefix: string, length: number): string { + if (prefix.length > length) { + throw new Error(`Vector key prefix is longer than ${length} characters`) + } + + return `${prefix}${'x'.repeat(length - prefix.length)}` } describeAcceptance( @@ -35,7 +50,23 @@ describeAcceptance( const secondaryVectorBucketName = `${vectorPrefix}-b` const indexName = `idx-${suffix}-a` const secondaryIndexName = `idx-${suffix}-b` + const defaultPageIndexName = `bulk-${suffix}` const vectorKeys = [`vec-a-${suffix}`, `vec-b-${suffix}`] + const maxLengthVectorKey = fixedLengthVectorKey(`vec-max-${suffix}-`, 1024) + const tooLongVectorKey = fixedLengthVectorKey(`vec-too-long-${suffix}-`, 1025) + const defaultPageVectorKeys = [ + maxLengthVectorKey, + ...Array.from({ length: 127 }, (_, i) => `bulk-${i.toString().padStart(3, '0')}-${suffix}`), + ] + const euclideanVectorKeys = [ + `vec-origin-${suffix}`, + `vec-near-${suffix}`, + `vec-far-${suffix}`, + ] + const euclideanDistractorVectorKeys = Array.from( + { length: 32 }, + (_, i) => `vec-far-${i.toString().padStart(2, '0')}-${suffix}` + ) try { await client.request('POST', '/vector/CreateVectorBucket', { @@ -104,6 +135,9 @@ describeAcceptance( dimension: 2, distanceMetric: 'cosine', indexName, + metadataConfiguration: { + nonFilterableMetadataKeys: ['private-note'], + }, vectorBucketName, }, expectedStatus: 200, @@ -114,7 +148,7 @@ describeAcceptance( body: { dataType: 'float32', dimension: 2, - distanceMetric: 'cosine', + distanceMetric: 'euclidean', indexName: secondaryIndexName, vectorBucketName, }, @@ -122,6 +156,18 @@ describeAcceptance( token, }) + await client.request('POST', '/vector/CreateIndex', { + body: { + dataType: 'float32', + dimension: 2, + distanceMetric: 'cosine', + indexName: defaultPageIndexName, + vectorBucketName, + }, + expectedStatus: 200, + token, + }) + const indexes = await client.request( 'POST', '/vector/ListIndexes', @@ -179,6 +225,10 @@ describeAcceptance( key: vectorKeys[0], metadata: { group: 'acceptance', + 'private-note': `private-${suffix}-a`, + role: 'primary', + score: 0.75, + 'user-id': `user-${suffix}-a`, }, }, { @@ -188,6 +238,9 @@ describeAcceptance( key: vectorKeys[1], metadata: { group: 'acceptance', + role: 'secondary', + score: 10, + 'user-id': `user-${suffix}-b`, }, }, ], @@ -196,9 +249,10 @@ describeAcceptance( token, }) - const vectors = await client.request('POST', '/vector/ListVectors', { + const vectors = await client.request('POST', '/vector/ListVectors', { body: { indexName, + maxResults: 1, returnData: true, returnMetadata: true, vectorBucketName, @@ -206,11 +260,137 @@ describeAcceptance( expectedStatus: 200, token, }) - expect(vectors.json?.vectors?.map((vector) => vector.key)).toEqual( + expect(vectors.json?.vectors).toHaveLength(1) + expect(vectors.json?.nextToken).toBeTruthy() + + const vectorsSecondPage = await client.request( + 'POST', + '/vector/ListVectors', + { + body: { + indexName, + maxResults: 1, + nextToken: vectors.json?.nextToken, + returnData: true, + returnMetadata: true, + vectorBucketName, + }, + expectedStatus: 200, + token, + } + ) + expect( + [...(vectors.json?.vectors ?? []), ...(vectorsSecondPage.json?.vectors ?? [])].map( + (vector) => vector.key + ) + ).toEqual(expect.arrayContaining(vectorKeys)) + expect(vectorsSecondPage.json?.nextToken).toBeUndefined() + + await client.request('POST', '/vector/PutVectors', { + body: { + indexName: defaultPageIndexName, + vectorBucketName, + vectors: defaultPageVectorKeys.map((key, i) => ({ + data: { + float32: [i, 0], + }, + key, + })), + }, + expectedStatus: 200, + token, + }) + + const defaultPageVectors = await client.request( + 'POST', + '/vector/ListVectors', + { + body: { + indexName: defaultPageIndexName, + vectorBucketName, + }, + expectedStatus: 200, + token, + } + ) + expect(defaultPageVectors.json?.vectors).toHaveLength(defaultPageVectorKeys.length) + expect(defaultPageVectors.json?.nextToken).toBeUndefined() + expect(defaultPageVectors.json?.vectors?.map((vector) => vector.key)).toEqual( + expect.arrayContaining(defaultPageVectorKeys) + ) + + await client.request('POST', '/vector/PutVectors', { + body: { + indexName: defaultPageIndexName, + vectorBucketName, + vectors: [ + { + data: { + float32: [0, 0], + }, + }, + ], + }, + expectedStatus: 400, + token, + }) + + await client.request('POST', '/vector/PutVectors', { + body: { + indexName: defaultPageIndexName, + vectorBucketName, + vectors: [ + { + data: { + float32: [0, 0], + }, + key: tooLongVectorKey, + }, + ], + }, + expectedStatus: 400, + token, + }) + + const largePageVectors = await client.request( + 'POST', + '/vector/ListVectors', + { + body: { + indexName, + maxResults: 1000, + vectorBucketName, + }, + expectedStatus: 200, + token, + } + ) + expect(largePageVectors.json?.vectors?.map((vector) => vector.key)).toEqual( expect.arrayContaining(vectorKeys) ) - const fetched = await client.request('POST', '/vector/GetVectors', { + const segmentedVectors = await Promise.all( + [0, 1].map((segmentIndex) => + client.request('POST', '/vector/ListVectors', { + body: { + indexName, + maxResults: 1000, + segmentCount: 2, + segmentIndex, + vectorBucketName, + }, + expectedStatus: 200, + token, + }) + ) + ) + const segmentedKeys = segmentedVectors.flatMap( + (segment) => segment.json?.vectors?.map((vector) => vector.key ?? '') ?? [] + ) + expect(new Set(segmentedKeys).size).toBe(segmentedKeys.length) + expect(segmentedKeys).toEqual(expect.arrayContaining(vectorKeys)) + + const fetched = await client.request('POST', '/vector/GetVectors', { body: { indexName, keys: vectorKeys, @@ -224,8 +404,27 @@ describeAcceptance( expect(fetched.json?.vectors?.map((vector) => vector.key)).toEqual( expect.arrayContaining(vectorKeys) ) + const fetchedPrimary = fetched.json?.vectors?.find((vector) => vector.key === vectorKeys[0]) + expect(fetchedPrimary?.metadata).toMatchObject({ + group: 'acceptance', + role: 'primary', + score: 0.75, + 'user-id': `user-${suffix}-a`, + }) + expect(fetchedPrimary?.data?.float32?.[0]).toBeCloseTo(1, 5) + expect(fetchedPrimary?.data?.float32?.[1]).toBeCloseTo(0, 5) + + await client.request('POST', '/vector/GetVectors', { + body: { + indexName, + keys: Array.from({ length: 101 }, (_, i) => `missing-${i}-${suffix}`), + vectorBucketName, + }, + expectedStatus: 400, + token, + }) - const query = await client.request('POST', '/vector/QueryVectors', { + const query = await client.request('POST', '/vector/QueryVectors', { body: { filter: { group: 'acceptance', @@ -244,6 +443,135 @@ describeAcceptance( }) expect(query.json?.vectors?.map((vector) => vector.key)).toContain(vectorKeys[0]) + const hyphenKeyFilter = await client.request( + 'POST', + '/vector/QueryVectors', + { + body: { + filter: { + 'user-id': `user-${suffix}-a`, + }, + indexName, + queryVector: { + float32: [1, 0], + }, + returnMetadata: true, + topK: 2, + vectorBucketName, + }, + expectedStatus: 200, + token, + } + ) + expect(hyphenKeyFilter.json?.vectors?.map((vector) => vector.key)).toEqual([vectorKeys[0]]) + expect(hyphenKeyFilter.json?.vectors?.[0]?.distance).toBeUndefined() + + await client.request('POST', '/vector/QueryVectors', { + body: { + filter: { + 'private-note': `private-${suffix}-a`, + }, + indexName, + queryVector: { + float32: [1, 0], + }, + topK: 2, + vectorBucketName, + }, + expectedStatus: 400, + token, + }) + + const numericFilter = await client.request('POST', '/vector/QueryVectors', { + body: { + filter: { + score: { + $gte: 5, + }, + }, + indexName, + queryVector: { + float32: [1, 0], + }, + returnMetadata: true, + topK: 2, + vectorBucketName, + }, + expectedStatus: 200, + token, + }) + expect(numericFilter.json?.vectors?.map((vector) => vector.key)).toEqual([vectorKeys[1]]) + + await client.request('POST', '/vector/PutVectors', { + body: { + indexName: secondaryIndexName, + vectorBucketName, + vectors: [ + { + data: { + float32: [0, 0], + }, + key: euclideanVectorKeys[0], + }, + { + data: { + float32: [3, 4], + }, + key: euclideanVectorKeys[1], + }, + { + data: { + float32: [8, 15], + }, + key: euclideanVectorKeys[2], + }, + ...euclideanDistractorVectorKeys.map((key, i) => ({ + data: { + float32: [100 + i, -100 - i], + }, + key, + })), + ], + }, + expectedStatus: 200, + token, + }) + + const euclideanQuery = await client.request( + 'POST', + '/vector/QueryVectors', + { + body: { + indexName: secondaryIndexName, + queryVector: { + float32: [0, 0], + }, + returnDistance: true, + topK: 3, + vectorBucketName, + }, + expectedStatus: 200, + token, + } + ) + expect(euclideanQuery.json?.distanceMetric).toBe('euclidean') + expect(euclideanQuery.json?.vectors?.map((vector) => vector.key)).toEqual( + euclideanVectorKeys + ) + expect(euclideanQuery.json?.vectors?.[0]?.distance).toBeCloseTo(0, 5) + expect(euclideanQuery.json?.vectors?.[1]?.distance).toBeCloseTo(5, 3) + expect(euclideanQuery.json?.vectors?.[2]?.distance).toBeCloseTo(17, 3) + + await client.request('POST', '/vector/DeleteVectors', { + body: { + indexName, + keys: Array.from({ length: 501 }, (_, i) => `missing-${i}-${suffix}`), + vectorBucketName, + }, + expectedStatus: 400, + token, + }) + await client.request('POST', '/vector/DeleteVectors', { body: { indexName, @@ -274,6 +602,16 @@ describeAcceptance( token, }) .catch(() => undefined) + await client + .request('POST', '/vector/DeleteIndex', { + body: { + indexName: defaultPageIndexName, + vectorBucketName, + }, + expectedStatus: [200, 400, 404], + token, + }) + .catch(() => undefined) await client .request('POST', '/vector/DeleteVectorBucket', { body: { diff --git a/package.json b/package.json index a70879947..e83f41142 100644 --- a/package.json +++ b/package.json @@ -29,6 +29,9 @@ "test:multigres": "npm run infra:restart:multigres && npm run test:dummy-data && npm run test:integration", "test:coverage": "npm run infra:restart && npm run test:dummy-data && npm run test:integration:coverage", "test:coverage:ci": "npm run infra:restart:ci && npm run test:dummy-data && npm run test:integration:coverage", + "test:oriole:ci": "npm run infra:restart:ci:oriole && npm run test:dummy-data && npm run test:integration", + "test:vector:oriole:pgvector": "VECTOR_BUCKET_PROVIDER=pgvector VECTOR_STORE_MIGRATIONS_ENABLED=true VECTOR_DATABASE_URL=postgresql://postgres:postgres@127.0.0.1:5432/postgres npm run infra:restart:oriole:pgvector && npm run test:dummy-data && VECTOR_BUCKET_PROVIDER=pgvector VECTOR_STORE_MIGRATIONS_ENABLED=true VECTOR_DATABASE_URL=postgresql://postgres:postgres@127.0.0.1:5432/postgres npm run test:integration -- src/test/pgvector-adapter.test.ts src/test/vectors-pgvector.test.ts", + "test:multigres:ci": "npm run infra:restart:ci:multigres && npm run test:dummy-data && npm run test:integration", "infra:stop": "docker compose --project-directory . --profile monitoring -f ./.docker/docker-compose-infra.yml down --remove-orphans", "infra:start": "docker compose --project-directory . --profile monitoring -f ./.docker/docker-compose-infra.yml up -d && sleep 5 && npm run migration:run", "infra:start:ci": "docker compose --project-directory . -f ./.docker/docker-compose-infra.yml up -d && sleep 5 && npm run migration:run", @@ -36,14 +39,16 @@ "infra:start:ci:oriole": "docker compose --project-directory . -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-oriole-override.yml up -d && sleep 5 && npm run migration:run", "infra:start:multigres": "docker compose --project-directory . --profile monitoring -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-multigres-override.yml up -d && docker compose --project-directory . --profile monitoring -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-multigres-override.yml up -d --wait --wait-timeout 180 tenant_db && npm run migration:run", "infra:start:ci:multigres": "docker compose --project-directory . -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-multigres-override.yml up -d && docker compose --project-directory . -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-multigres-override.yml up -d --wait --wait-timeout 180 tenant_db && npm run migration:run", + "infra:start:oriole:pgvector": "docker compose --project-directory . --profile monitoring -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-oriole-override.yml -f ./.docker/docker-compose-infra-oriole-pgvector-override.yml up -d --build && sleep 5 && npm run migration:run", + "infra:start:ci:oriole:pgvector": "docker compose --project-directory . -f ./.docker/docker-compose-infra.yml -f ./.docker/docker-compose-infra-oriole-override.yml -f ./.docker/docker-compose-infra-oriole-pgvector-override.yml up -d --build && sleep 5 && npm run migration:run", "infra:restart": "npm run infra:stop && npm run infra:start", "infra:restart:ci": "npm run infra:stop && npm run infra:start:ci", "infra:restart:oriole": "npm run infra:stop && npm run infra:start:oriole", "infra:restart:ci:oriole": "npm run infra:stop && npm run infra:start:ci:oriole", "infra:restart:multigres": "npm run infra:stop && npm run infra:start:multigres", "infra:restart:ci:multigres": "npm run infra:stop && npm run infra:start:ci:multigres", - "test:oriole:ci": "npm run infra:restart:ci:oriole && npm run test:dummy-data && npm run test:integration", - "test:multigres:ci": "npm run infra:restart:ci:multigres && npm run test:dummy-data && npm run test:integration" + "infra:restart:oriole:pgvector": "npm run infra:stop && npm run infra:start:oriole:pgvector", + "infra:restart:ci:oriole:pgvector": "npm run infra:stop && npm run infra:start:ci:oriole:pgvector" }, "dependencies": { "@aws-sdk/client-ecs": "^3.1041.0", diff --git a/src/http/error-handler.test.ts b/src/http/error-handler.test.ts new file mode 100644 index 000000000..d7fd188af --- /dev/null +++ b/src/http/error-handler.test.ts @@ -0,0 +1,70 @@ +import { ErrorCode } from '@internal/errors' +import Fastify from 'fastify' +import { setErrorHandler } from './error-handler' + +describe('setErrorHandler', () => { + it('maps Fastify schema validation failures to InvalidRequest', async () => { + const app = Fastify() + setErrorHandler(app) + + app.post( + '/validated', + { + schema: { + body: { + type: 'object', + properties: { + count: { type: 'number' }, + }, + required: ['count'], + additionalProperties: false, + }, + }, + }, + async () => ({ ok: true }) + ) + + try { + const response = await app.inject({ + method: 'POST', + url: '/validated', + payload: { count: 'not-a-number' }, + }) + + expect(response.statusCode).toBe(400) + expect(response.json()).toMatchObject({ + statusCode: '400', + code: ErrorCode.InvalidRequest, + }) + } finally { + await app.close() + } + }) + + it('uses the fallback status code in Fastify error payloads when statusCode is undefined', async () => { + const app = Fastify() + setErrorHandler(app) + + app.get('/undefined-status-code', async () => { + const error = new Error('boom') as Error & { statusCode?: number } + error.statusCode = undefined + throw error + }) + + try { + const response = await app.inject({ + method: 'GET', + url: '/undefined-status-code', + }) + + expect(response.statusCode).toBe(500) + expect(response.json()).toMatchObject({ + statusCode: '500', + code: ErrorCode.InternalError, + message: 'boom', + }) + } finally { + await app.close() + } + }) +}) diff --git a/src/http/error-handler.ts b/src/http/error-handler.ts index 9f181c013..2bea093d6 100644 --- a/src/http/error-handler.ts +++ b/src/http/error-handler.ts @@ -1,5 +1,11 @@ import { FastifyError } from '@fastify/error' -import { ErrorCode, isRenderableError, StorageBackendError, StorageError } from '@internal/errors' +import { + ErrorCode, + getErrorCode, + isRenderableError, + StorageBackendError, + StorageError, +} from '@internal/errors' import { isDatabaseSlowDownError } from '@internal/errors/database-error' import { FastifyInstance } from 'fastify' @@ -83,11 +89,17 @@ export const setErrorHandler = ( ) } - return reply.status(err.statusCode || 500).send( + const errorCode = getErrorCode(err) + const responseErrorCode = ( + errorCode === ErrorCode.UnknownError ? ErrorCode.InternalError : errorCode + ) as ErrorCode + const responseStatusCode = err.statusCode || 500 + + return reply.status(responseStatusCode).send( formatter({ - statusCode: `${err.statusCode}`, + statusCode: `${responseStatusCode}`, error: err.name, - code: ErrorCode.InternalError, + code: responseErrorCode, message: err.message, }) ) diff --git a/src/http/plugins/vector.ts b/src/http/plugins/vector.ts index adb80aa5b..f688fc872 100644 --- a/src/http/plugins/vector.ts +++ b/src/http/plugins/vector.ts @@ -10,6 +10,7 @@ import { } from '@internal/sharding' import { createS3VectorClient, + createVectorTransactionKnexResolver, KnexVectorMetadataDB, PgVectorStore, S3Vector, @@ -90,7 +91,9 @@ export const s3vector = fastifyPlugin(async function (fastify: FastifyInstance) // uses the singleton; s3 uses the singleton S3 client. let adapter: VectorStore if (vectorBucketProvider === 'pgvector') { - adapter = isMultitenant ? new PgVectorStore(db) : stPgVectorAdapter! + adapter = isMultitenant + ? new PgVectorStore(createVectorTransactionKnexResolver(db)) + : stPgVectorAdapter! } else { adapter = s3Adapter! } diff --git a/src/http/routes/vector/create-bucket.ts b/src/http/routes/vector/create-bucket.ts index 1ae0784ce..95a32b13f 100644 --- a/src/http/routes/vector/create-bucket.ts +++ b/src/http/routes/vector/create-bucket.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const createVectorBucket = { type: 'object', @@ -16,14 +17,17 @@ const createVectorBucket = { summary: 'Create a vector bucket', } as const -interface createVectorIndexRequest extends AuthenticatedRequest { +interface createVectorBucketRequest extends AuthenticatedRequest { Body: FromSchema<(typeof createVectorBucket)['body']> } export default async function routes(fastify: FastifyInstance) { - fastify.post( + const createVectorBucketValidator = compileNoCoercionValidator(createVectorBucket.body) + + fastify.post( '/CreateVectorBucket', { + validatorCompiler: createVectorBucketValidator, config: { operation: { type: ROUTE_OPERATIONS.CREATE_VECTOR_BUCKET }, }, diff --git a/src/http/routes/vector/create-index.ts b/src/http/routes/vector/create-index.ts index 1b9959382..a13b8e12b 100644 --- a/src/http/routes/vector/create-index.ts +++ b/src/http/routes/vector/create-index.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const createVectorIndex = { type: 'object', @@ -10,7 +11,13 @@ const createVectorIndex = { type: 'object', properties: { dataType: { type: 'string', enum: ['float32'] }, - dimension: { type: 'number', minimum: 1, maximum: 4096 }, + dimension: { + type: 'integer', + minimum: 1, + maximum: 4096, + description: + 'Vector dimensionality. S3 Vectors supports up to 4096. The local pgvector backend supports up to 4000 and will reject larger values.', + }, distanceMetric: { type: 'string', enum: ['cosine', 'euclidean'] }, indexName: { type: 'string', @@ -26,7 +33,10 @@ const createVectorIndex = { properties: { nonFilterableMetadataKeys: { type: 'array', - items: { type: 'string' }, + minItems: 1, + maxItems: 10, + uniqueItems: true, + items: { type: 'string', minLength: 1, maxLength: 63 }, }, }, }, @@ -42,9 +52,12 @@ interface createVectorIndexRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const createVectorIndexValidator = compileNoCoercionValidator(createVectorIndex.body) + fastify.post( '/CreateIndex', { + validatorCompiler: createVectorIndexValidator, config: { operation: { type: ROUTE_OPERATIONS.CREATE_VECTOR_INDEX }, }, diff --git a/src/http/routes/vector/delete-bucket.ts b/src/http/routes/vector/delete-bucket.ts index 31d88dfdf..9d74a2321 100644 --- a/src/http/routes/vector/delete-bucket.ts +++ b/src/http/routes/vector/delete-bucket.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const deleteVectorBucket = { type: 'object', @@ -21,9 +22,12 @@ interface deleteVectorIndexRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const deleteVectorBucketValidator = compileNoCoercionValidator(deleteVectorBucket.body) + fastify.post( '/DeleteVectorBucket', { + validatorCompiler: deleteVectorBucketValidator, config: { operation: { type: ROUTE_OPERATIONS.DELETE_VECTOR_BUCKET }, }, diff --git a/src/http/routes/vector/delete-index.ts b/src/http/routes/vector/delete-index.ts index 266e4ad50..a3aea84a7 100644 --- a/src/http/routes/vector/delete-index.ts +++ b/src/http/routes/vector/delete-index.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const deleteVectorIndex = { type: 'object', @@ -29,9 +30,12 @@ interface deleteVectorIndexRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const deleteVectorIndexValidator = compileNoCoercionValidator(deleteVectorIndex.body) + fastify.post( '/DeleteIndex', { + validatorCompiler: deleteVectorIndexValidator, config: { operation: { type: ROUTE_OPERATIONS.DELETE_VECTOR_INDEX }, }, diff --git a/src/http/routes/vector/delete-vectors.ts b/src/http/routes/vector/delete-vectors.ts index 6d02192dd..6e45779e4 100644 --- a/src/http/routes/vector/delete-vectors.ts +++ b/src/http/routes/vector/delete-vectors.ts @@ -1,8 +1,10 @@ import { ERRORS } from '@internal/errors' +import { MAX_DELETE_VECTOR_KEYS, MAX_VECTOR_KEY_LENGTH } from '@storage/protocols/vector/limits' import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const deleteVector = { body: { @@ -10,7 +12,12 @@ const deleteVector = { properties: { vectorBucketName: { type: 'string' }, indexName: { type: 'string' }, - keys: { type: 'array', items: { type: 'string' } }, + keys: { + type: 'array', + minItems: 1, + maxItems: MAX_DELETE_VECTOR_KEYS, + items: { type: 'string', minLength: 1, maxLength: MAX_VECTOR_KEY_LENGTH }, + }, }, required: ['vectorBucketName', 'indexName', 'keys'], }, @@ -22,9 +29,12 @@ interface deleteVectorRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const deleteVectorsValidator = compileNoCoercionValidator(deleteVector.body) + fastify.post( '/DeleteVectors', { + validatorCompiler: deleteVectorsValidator, config: { operation: { type: ROUTE_OPERATIONS.DELETE_VECTORS }, }, diff --git a/src/http/routes/vector/get-bucket.ts b/src/http/routes/vector/get-bucket.ts index cfc3611c7..804d036e5 100644 --- a/src/http/routes/vector/get-bucket.ts +++ b/src/http/routes/vector/get-bucket.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const getVectorBucket = { type: 'object', @@ -21,9 +22,12 @@ interface getVectorBucketRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const getVectorBucketValidator = compileNoCoercionValidator(getVectorBucket.body) + fastify.post( '/GetVectorBucket', { + validatorCompiler: getVectorBucketValidator, config: { operation: { type: ROUTE_OPERATIONS.GET_VECTOR_BUCKET }, }, diff --git a/src/http/routes/vector/get-index.ts b/src/http/routes/vector/get-index.ts index 89f5b6272..f5963c191 100644 --- a/src/http/routes/vector/get-index.ts +++ b/src/http/routes/vector/get-index.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const getVectorIndex = { type: 'object', @@ -29,9 +30,12 @@ interface getVectorIndexRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const getVectorIndexValidator = compileNoCoercionValidator(getVectorIndex.body) + fastify.post( '/GetIndex', { + validatorCompiler: getVectorIndexValidator, config: { operation: { type: ROUTE_OPERATIONS.GET_VECTOR_INDEX }, }, diff --git a/src/http/routes/vector/get-vectors.ts b/src/http/routes/vector/get-vectors.ts index 18bf97bd4..1116971a6 100644 --- a/src/http/routes/vector/get-vectors.ts +++ b/src/http/routes/vector/get-vectors.ts @@ -1,8 +1,10 @@ import { ERRORS } from '@internal/errors' +import { MAX_GET_VECTOR_KEYS, MAX_VECTOR_KEY_LENGTH } from '@storage/protocols/vector/limits' import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const getVectors = { type: 'object', @@ -10,7 +12,12 @@ const getVectors = { type: 'object', properties: { indexName: { type: 'string' }, - keys: { type: 'array', items: { type: 'string' } }, + keys: { + type: 'array', + minItems: 1, + maxItems: MAX_GET_VECTOR_KEYS, + items: { type: 'string', minLength: 1, maxLength: MAX_VECTOR_KEY_LENGTH }, + }, returnData: { type: 'boolean', default: false }, returnMetadata: { type: 'boolean', default: false }, vectorBucketName: { type: 'string' }, @@ -25,9 +32,12 @@ interface getVectorsRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const getVectorsValidator = compileNoCoercionValidator(getVectors.body) + fastify.post( '/GetVectors', { + validatorCompiler: getVectorsValidator, config: { operation: { type: ROUTE_OPERATIONS.GET_VECTORS }, }, diff --git a/src/http/routes/vector/list-buckets.ts b/src/http/routes/vector/list-buckets.ts index 31ee59736..719510173 100644 --- a/src/http/routes/vector/list-buckets.ts +++ b/src/http/routes/vector/list-buckets.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const listBucket = { type: 'object', @@ -22,9 +23,12 @@ interface listBucketRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const listBucketsValidator = compileNoCoercionValidator(listBucket.body) + fastify.post( '/ListVectorBuckets', { + validatorCompiler: listBucketsValidator, config: { operation: { type: ROUTE_OPERATIONS.LIST_VECTOR_BUCKETS }, }, diff --git a/src/http/routes/vector/list-indexes.ts b/src/http/routes/vector/list-indexes.ts index c0703f313..e051a1657 100644 --- a/src/http/routes/vector/list-indexes.ts +++ b/src/http/routes/vector/list-indexes.ts @@ -3,6 +3,7 @@ import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const listIndex = { type: 'object', @@ -24,9 +25,12 @@ interface listIndexRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const listIndexesValidator = compileNoCoercionValidator(listIndex.body) + fastify.post( '/ListIndexes', { + validatorCompiler: listIndexesValidator, config: { operation: { type: ROUTE_OPERATIONS.LIST_VECTOR_INDEXES }, }, diff --git a/src/http/routes/vector/list-vectors.ts b/src/http/routes/vector/list-vectors.ts index 1a8d0980d..2580a35e4 100644 --- a/src/http/routes/vector/list-vectors.ts +++ b/src/http/routes/vector/list-vectors.ts @@ -1,8 +1,10 @@ import { ERRORS } from '@internal/errors' +import { MAX_LIST_RESULTS, MAX_SEGMENT_COUNT } from '@storage/protocols/vector/limits' import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const listVectors = { type: 'object', @@ -19,12 +21,12 @@ const listVectors = { description: '3-63 chars, lowercase letters, numbers, hyphens, dots; must start/end with letter or number. Must be unique within the vector bucket.', }, - maxResults: { type: 'number', minimum: 1, maximum: 500 }, + maxResults: { type: 'integer', minimum: 1, maximum: MAX_LIST_RESULTS }, nextToken: { type: 'string' }, returnData: { type: 'boolean' }, returnMetadata: { type: 'boolean' }, - segmentCount: { type: 'number', minimum: 1, maximum: 16 }, - segmentIndex: { type: 'number', minimum: 0, maximum: 15 }, + segmentCount: { type: 'integer', minimum: 1, maximum: MAX_SEGMENT_COUNT }, + segmentIndex: { type: 'integer', minimum: 0, maximum: MAX_SEGMENT_COUNT - 1 }, }, required: ['vectorBucketName', 'indexName'], }, @@ -36,9 +38,12 @@ interface listVectorsRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const listVectorsValidator = compileNoCoercionValidator(listVectors.body) + fastify.post( '/ListVectors', { + validatorCompiler: listVectorsValidator, config: { operation: { type: ROUTE_OPERATIONS.LIST_VECTORS }, }, @@ -52,6 +57,26 @@ export default async function routes(fastify: FastifyInstance) { throw ERRORS.FeatureNotEnabled('vectorStore', 'Vector service not configured') } + const { segmentCount, segmentIndex } = request.body + if ( + (segmentCount === undefined && segmentIndex !== undefined) || + (segmentCount !== undefined && segmentIndex === undefined) + ) { + throw ERRORS.InvalidParameter('segmentCount/segmentIndex', { + message: 'segmentCount and segmentIndex must be provided together', + }) + } + + if ( + segmentCount !== undefined && + segmentIndex !== undefined && + segmentIndex >= segmentCount + ) { + throw ERRORS.InvalidParameter('segmentIndex', { + message: 'segmentIndex must be less than segmentCount', + }) + } + const indexResult = await request.s3Vector.listVectors(request.body) return response.send(indexResult) diff --git a/src/http/routes/vector/put-vectors.ts b/src/http/routes/vector/put-vectors.ts index caaba3958..b949ad39e 100644 --- a/src/http/routes/vector/put-vectors.ts +++ b/src/http/routes/vector/put-vectors.ts @@ -1,8 +1,18 @@ import { ERRORS } from '@internal/errors' +import { + MAX_PUT_VECTORS, + MAX_VECTOR_KEY_LENGTH, + MIN_VECTOR_DIMENSIONS, +} from '@storage/protocols/vector/limits' import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' + +const metadataPrimitive = { + anyOf: [{ type: 'string' }, { type: 'boolean' }, { type: 'number' }], +} as const const putVector = { body: { @@ -20,26 +30,36 @@ const putVector = { vectors: { type: 'array', minItems: 1, - maxItems: 500, + maxItems: MAX_PUT_VECTORS, items: { type: 'object', properties: { data: { type: 'object', properties: { - float32: { type: 'array', items: { type: 'number' } }, + float32: { + type: 'array', + minItems: MIN_VECTOR_DIMENSIONS, + items: { type: 'number' }, + }, }, required: ['float32'], }, metadata: { type: 'object', additionalProperties: { - oneOf: [{ type: 'string' }, { type: 'boolean' }, { type: 'number' }], + anyOf: [ + metadataPrimitive, + { + type: 'array', + items: metadataPrimitive, + }, + ], }, }, - key: { type: 'string' }, + key: { type: 'string', minLength: 1, maxLength: MAX_VECTOR_KEY_LENGTH }, }, - required: ['data'], + required: ['data', 'key'], }, }, }, @@ -53,10 +73,13 @@ interface putVectorRequest extends AuthenticatedRequest { } export default async function routes(fastify: FastifyInstance) { + const putVectorsValidator = compileNoCoercionValidator(putVector.body) + fastify.post( '/PutVectors', { bodyLimit: 20 * 1024 * 1024, // 20 MB + validatorCompiler: putVectorsValidator, config: { operation: { type: ROUTE_OPERATIONS.PUT_VECTORS }, }, @@ -76,7 +99,7 @@ export default async function routes(fastify: FastifyInstance) { vectors: request.body.vectors.map((v) => { return { ...v, - key: v.key || undefined, + key: v.key ?? undefined, } }), }) diff --git a/src/http/routes/vector/query-vectors.ts b/src/http/routes/vector/query-vectors.ts index c0083bd18..27d09da9d 100644 --- a/src/http/routes/vector/query-vectors.ts +++ b/src/http/routes/vector/query-vectors.ts @@ -1,9 +1,10 @@ import { ERRORS } from '@internal/errors' -import Ajv from 'ajv' -import { FastifyInstance, FastifySchema, FastifySchemaCompiler } from 'fastify' +import { MAX_QUERY_TOP_K, MIN_VECTOR_DIMENSIONS } from '@storage/protocols/vector/limits' +import { FastifyInstance } from 'fastify' import { FromSchema } from 'json-schema-to-ts' import { AuthenticatedRequest } from '../../types' import { ROUTE_OPERATIONS } from '../operations' +import { compileNoCoercionValidator } from './validation' const queryVectorPrimitiveSchema = { $id: 'queryVectorPrimitive', @@ -84,14 +85,23 @@ const queryVectorBodyBaseProperties = { queryVector: { type: 'object', properties: { - float32: { type: 'array', items: { type: 'number' } }, + float32: { + type: 'array', + minItems: MIN_VECTOR_DIMENSIONS, + items: { type: 'number' }, + }, }, required: ['float32'], additionalProperties: false, }, returnDistance: { type: 'boolean' }, returnMetadata: { type: 'boolean' }, - topK: { type: 'number' }, + topK: { + type: 'integer', + minimum: 1, + maximum: MAX_QUERY_TOP_K, + description: `Number of nearest-neighbor results to return, from 1 to ${MAX_QUERY_TOP_K}.`, + }, vectorBucketName: { type: 'string' }, } as const @@ -171,25 +181,14 @@ export default async function routes(fastify: FastifyInstance) { // schemas never appear in components.schemas of the emitted OpenAPI spec. fastify.addSchema(queryVectorBodyDocSchema) - // Strict validation runs through a separate Ajv instance (removeAdditional - // disabled so payloads aren't silently stripped before reaching the handler). - const ajvNoRemoval = new Ajv({ - allErrors: true, - removeAdditional: false, - coerceTypes: false, - }) - ajvNoRemoval.addSchema(queryVectorPrimitiveSchema) - ajvNoRemoval.addSchema(queryVectorFieldOperatorsSchema) - ajvNoRemoval.addSchema(queryVectorLogicalFilterSchema) - ajvNoRemoval.addSchema(queryVectorFilterSchema) - ajvNoRemoval.addSchema(queryVectorBodySchema) - - const validateBody = ajvNoRemoval.compile(queryVectorBodySchema) - - const queryVectorsValidator: FastifySchemaCompiler = () => (data: unknown) => { - if (validateBody(data)) return { value: data } - return { error: validateBody.errors ?? [] } - } + // Strict validation runs through a separate Ajv instance so vector filters + // keep their scalar types and payloads aren't silently stripped. + const queryVectorsValidator = compileNoCoercionValidator(queryVectorBodySchema, [ + queryVectorPrimitiveSchema, + queryVectorFieldOperatorsSchema, + queryVectorLogicalFilterSchema, + queryVectorFilterSchema, + ]) fastify.post( '/QueryVectors', diff --git a/src/http/routes/vector/validation.test.ts b/src/http/routes/vector/validation.test.ts new file mode 100644 index 000000000..f3ba25f43 --- /dev/null +++ b/src/http/routes/vector/validation.test.ts @@ -0,0 +1,48 @@ +import { compileNoCoercionValidator } from './validation' + +describe('compileNoCoercionValidator', () => { + it('validates referenced schemas without coercing scalar types', () => { + const primitiveSchema = { + $id: 'testVectorPrimitive', + oneOf: [{ type: 'string' }, { type: 'boolean' }, { type: 'number' }], + } as const + const bodySchema = { + type: 'object', + properties: { + flag: { type: 'boolean' }, + metadataValue: { $ref: 'testVectorPrimitive#' }, + score: { type: 'number' }, + }, + required: ['flag', 'metadataValue', 'score'], + additionalProperties: false, + } as const + + const validate = compileNoCoercionValidator(bodySchema, [primitiveSchema])({ + schema: bodySchema, + } as never) + const validBody = { + flag: true, + metadataValue: 0.75, + score: 0.75, + } + const invalidBody = { + flag: 'true', + metadataValue: '0.75', + score: '0.75', + } + + expect(validate(validBody)).toEqual({ value: validBody }) + expect(validate(invalidBody)).toEqual({ + error: expect.arrayContaining([ + expect.objectContaining({ + instancePath: '/flag', + keyword: 'type', + }), + expect.objectContaining({ + instancePath: '/score', + keyword: 'type', + }), + ]), + }) + }) +}) diff --git a/src/http/routes/vector/validation.ts b/src/http/routes/vector/validation.ts new file mode 100644 index 000000000..16e6c7a9d --- /dev/null +++ b/src/http/routes/vector/validation.ts @@ -0,0 +1,25 @@ +import Ajv, { type AnySchema } from 'ajv' +import { FastifySchema, FastifySchemaCompiler } from 'fastify' + +export function compileNoCoercionValidator( + schema: AnySchema, + refs: AnySchema[] = [] +): FastifySchemaCompiler { + const ajvNoCoercion = new Ajv({ + allErrors: true, + coerceTypes: false, + removeAdditional: false, + useDefaults: true, + }) + + for (const ref of refs) { + ajvNoCoercion.addSchema(ref) + } + + const validateBody = ajvNoCoercion.compile(schema) + + return () => (data: unknown) => { + if (validateBody(data)) return { value: data } + return { error: validateBody.errors ?? [] } + } +} diff --git a/src/internal/database/migrations/migrate.ts b/src/internal/database/migrations/migrate.ts index 310f118e0..f8427612f 100644 --- a/src/internal/database/migrations/migrate.ts +++ b/src/internal/database/migrations/migrate.ts @@ -412,11 +412,13 @@ async function ensureVectorDatabaseExists( maintenanceUrl: string, ssl: ClientConfig['ssl'] ): Promise { + let defaultAccessMethod = '' const client = await connect({ connectionString: maintenanceUrl, ssl, }) try { + defaultAccessMethod = await getDefaultAccessMethod(client) const exists = await client.query(`SELECT 1 FROM pg_database WHERE datname = $1`, [ VECTOR_DATABASE_NAME, ]) @@ -431,6 +433,42 @@ async function ensureVectorDatabaseExists( } finally { await client.end() } + + await configureVectorDatabaseAccessMethod({ + databaseUrl: deriveVectorDatabaseUrl(maintenanceUrl), + defaultAccessMethod, + ssl, + }) +} + +async function configureVectorDatabaseAccessMethod({ + databaseUrl, + defaultAccessMethod, + ssl, +}: { + databaseUrl: string + defaultAccessMethod: string + ssl: ClientConfig['ssl'] +}): Promise { + if (defaultAccessMethod !== 'orioledb') { + return + } + + const client = await connect({ + connectionString: databaseUrl, + ssl, + }) + try { + await client.query('CREATE EXTENSION IF NOT EXISTS orioledb') + await client.query( + `ALTER DATABASE "${VECTOR_DATABASE_NAME}" SET default_table_access_method = 'orioledb'` + ) + logSchema.info(logger, `[Migrations] Configured database ${VECTOR_DATABASE_NAME} for Oriole`, { + type: 'migrations', + }) + } finally { + await client.end() + } } /** diff --git a/src/internal/database/migrations/vector-store-migrations.test.ts b/src/internal/database/migrations/vector-store-migrations.test.ts new file mode 100644 index 000000000..b130503e6 --- /dev/null +++ b/src/internal/database/migrations/vector-store-migrations.test.ts @@ -0,0 +1,222 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { + mockClientConfigs, + mockPgClients, + mockLoadMigrationFilesCached, + mockRunMigrationStep, + mockValidateMigrationHashes, +} = vi.hoisted(() => ({ + mockClientConfigs: [] as Array>, + mockPgClients: [] as MockPgClient[], + mockLoadMigrationFilesCached: vi.fn(), + mockRunMigrationStep: vi.fn(), + mockValidateMigrationHashes: vi.fn(), +})) + +vi.mock('../../../config', () => ({ + MultitenantMigrationStrategy: { + ON_REQUEST: 'ON_REQUEST', + PROGRESSIVE: 'PROGRESSIVE', + FULL_FLEET: 'FULL_FLEET', + }, + getConfig: () => ({ + isMultitenant: false, + multitenantDatabaseUrl: '', + pgQueueEnable: false, + databaseSSLRootCert: '', + dbMigrationStrategy: 'ON_REQUEST', + dbAnonRole: 'anon', + dbAuthenticatedRole: 'authenticated', + dbSuperUser: 'postgres', + dbServiceRole: 'service_role', + dbInstallRoles: false, + dbRefreshMigrationHashesOnMismatch: false, + dbMigrationFreezeAt: undefined, + icebergShards: [], + multitenantDatabaseQueryTimeout: 1000, + vectorBucketProvider: 'pgvector', + }), +})) + +vi.mock('pg', () => ({ + Client: vi.fn(function MockClient(config: Record) { + const client = mockPgClients.shift() + if (!client) { + throw new Error(`Unexpected pg Client for ${String(config.connectionString)}`) + } + + mockClientConfigs.push(config) + return client + }), +})) + +vi.mock('postgres-migrations/dist/run-migration', () => ({ + runMigration: vi.fn(() => async (migration: unknown) => { + mockRunMigrationStep(migration) + return migration + }), +})) + +vi.mock('postgres-migrations/dist/validation', () => ({ + validateMigrationHashes: mockValidateMigrationHashes, +})) + +vi.mock('@storage/events', () => ({ + RunMigrationsOnTenants: class { + static batchSend = vi.fn() + }, + ResetMigrationsOnTenant: class { + static batchSend = vi.fn() + }, +})) + +vi.mock('../../monitoring', () => ({ + logger: {}, + logSchema: { + info: vi.fn(), + warning: vi.fn(), + error: vi.fn(), + }, +})) + +vi.mock('../multitenant-db', () => ({ + multitenantKnex: {}, +})) + +vi.mock('../tenant', () => ({ + getTenantConfig: vi.fn(), + TenantMigrationStatus: { + COMPLETED: 'COMPLETED', + FAILED: 'FAILED', + FAILED_STALE: 'FAILED_STALE', + }, +})) + +vi.mock('../pool', () => ({ + searchPath: ['storage', 'public'], +})) + +vi.mock('./files', () => ({ + lastLocalMigrationName: vi.fn(), + loadMigrationFilesCached: mockLoadMigrationFilesCached, + localMigrationFiles: vi.fn(), +})) + +vi.mock('./progressive', () => ({ + ProgressiveMigrations: class { + start() { + return undefined + } + }, +})) + +import { runVectorStoreMigrations } from './migrate' + +interface MockPgClient { + connect: ReturnType + end: ReturnType + on: ReturnType + queries: string[] + query: ReturnType +} + +function queryText(query: unknown): string { + if (query && typeof query === 'object') { + if ('text' in query && typeof query.text === 'string') { + return query.text + } + if ('sql' in query && typeof query.sql === 'string') { + return query.sql + } + } + + return String(query) +} + +function createMockPgClient(options: { + defaultAccessMethod?: string + databaseExists?: boolean + migrationTableExists?: boolean + schemaExists?: boolean +}): MockPgClient { + const queries: string[] = [] + + return { + connect: vi.fn().mockResolvedValue(undefined), + end: vi.fn().mockResolvedValue(undefined), + on: vi.fn(), + queries, + query: vi.fn(async (query: unknown) => { + const text = queryText(query) + queries.push(text) + + if (text === 'SHOW default_table_access_method') { + return { rows: [{ default_table_access_method: options.defaultAccessMethod ?? 'heap' }] } + } + + if (text.includes('SELECT 1 FROM pg_database')) { + return { rows: options.databaseExists ? [{ exists: 1 }] : [] } + } + + if (text.includes('pg_try_advisory_lock')) { + return { rows: [{ pg_try_advisory_lock: true }] } + } + + if (text.includes('pg_catalog.pg_class')) { + return { rows: [{ exists: options.migrationTableExists ?? false }] } + } + + if (text.includes('information_schema.schemata')) { + return { rows: [{ exists: options.schemaExists ?? false }] } + } + + return { rows: [] } + }), + } +} + +describe('runVectorStoreMigrations', () => { + beforeEach(() => { + vi.clearAllMocks() + mockClientConfigs.length = 0 + mockPgClients.length = 0 + mockLoadMigrationFilesCached.mockResolvedValue([ + { + id: 1, + name: 'create-vector-store', + hash: 'hash', + sql: 'CREATE TABLE vector_store_test(id int);', + contents: 'CREATE TABLE vector_store_test(id int);', + }, + ]) + }) + + it('configures a dedicated vector database for Oriole before running vector migrations', async () => { + const maintenanceClient = createMockPgClient({ + defaultAccessMethod: 'orioledb', + databaseExists: false, + }) + const vectorSetupClient = createMockPgClient({}) + const migrationClient = createMockPgClient({}) + mockPgClients.push(maintenanceClient, vectorSetupClient, migrationClient) + + await runVectorStoreMigrations({ + databaseUrl: 'postgresql://postgres:postgres@127.0.0.1:5432/postgres', + createDatabase: true, + waitForLock: false, + }) + + expect(mockClientConfigs.map((config) => config.connectionString)).toEqual([ + 'postgresql://postgres:postgres@127.0.0.1:5432/postgres', + 'postgresql://postgres:postgres@127.0.0.1:5432/storage_vectors', + 'postgresql://postgres:postgres@127.0.0.1:5432/storage_vectors', + ]) + expect(maintenanceClient.queries).toContain('CREATE DATABASE "storage_vectors"') + expect(vectorSetupClient.queries).toContain('CREATE EXTENSION IF NOT EXISTS orioledb') + expect(vectorSetupClient.queries).toContain( + 'ALTER DATABASE "storage_vectors" SET default_table_access_method = \'orioledb\'' + ) + expect(migrationClient.queries[0]).toBe("SET statement_timeout TO '12h'") + }) +}) diff --git a/src/internal/errors/codes.test.ts b/src/internal/errors/codes.test.ts index a96f14c22..56e597cf3 100644 --- a/src/internal/errors/codes.test.ts +++ b/src/internal/errors/codes.test.ts @@ -1,5 +1,5 @@ import { IcebergError, IcebergErrorType } from '@storage/protocols/iceberg/catalog/errors' -import { ErrorCode, normalizeRawError } from './codes' +import { ErrorCode, getErrorCode, normalizeRawError } from './codes' import { StorageBackendError } from './storage-error' describe('normalizeRawError', () => { @@ -109,3 +109,19 @@ describe('normalizeRawError', () => { expect(result.raw).toBe('Failed to stringify error') }) }) + +describe('getErrorCode', () => { + it.each([ + ['FST_ERR_VALIDATION', ErrorCode.InvalidRequest], + ['FST_ERR_CTP_EMPTY_JSON_BODY', ErrorCode.InvalidRequest], + ['FST_ERR_CTP_INVALID_JSON_BODY', ErrorCode.InvalidRequest], + ['FST_ERR_CTP_INVALID_CONTENT_LENGTH', ErrorCode.InvalidRequest], + ['FST_ERR_CTP_INVALID_MEDIA_TYPE', ErrorCode.InvalidMimeType], + ['FST_ERR_CTP_BODY_TOO_LARGE', ErrorCode.EntityTooLarge], + ])('maps Fastify error code %s to %s', (fastifyCode, expectedErrorCode) => { + const error = new Error('Fastify error') + Object.assign(error, { code: fastifyCode }) + + expect(getErrorCode(error)).toBe(expectedErrorCode) + }) +}) diff --git a/src/internal/errors/codes.ts b/src/internal/errors/codes.ts index bde07db7b..79dfd8ce6 100644 --- a/src/internal/errors/codes.ts +++ b/src/internal/errors/codes.ts @@ -579,7 +579,7 @@ function hasStringErrorCode(error: Error): error is Error & { code: string } { return 'code' in error && typeof error.code === 'string' } -function getErrorCode(error: Error): string { +export function getErrorCode(error: Error): string { if (error instanceof IcebergError && error.error) { return error.error } diff --git a/src/internal/sharding/strategy/catalog.ts b/src/internal/sharding/strategy/catalog.ts index d8d9a904d..778f787d9 100644 --- a/src/internal/sharding/strategy/catalog.ts +++ b/src/internal/sharding/strategy/catalog.ts @@ -260,6 +260,7 @@ export class ShardCatalog implements Sharder { await this.factory.withTransaction(async (store) => { const resv = await store.fetchReservationById(reservationId) if (!resv) return + if (resv.status !== 'pending') return await store.updateReservationStatus(reservationId, 'cancelled') }) } diff --git a/src/storage/protocols/vector/adapter/pgvector/errors.ts b/src/storage/protocols/vector/adapter/pgvector/errors.ts index a487124a9..47d0a9a7c 100644 --- a/src/storage/protocols/vector/adapter/pgvector/errors.ts +++ b/src/storage/protocols/vector/adapter/pgvector/errors.ts @@ -28,6 +28,7 @@ export async function handlePgVectorError( // 42P01 undefined_table → index/bucket missing // 42P07 duplicate_table → CREATE TABLE on already-existing index // 23505 unique_violation → key collision (we use upsert, but defensively map) + // 21000 cardinality_violation → duplicate source keys in an upsert batch // 22P02 invalid_text_representation → e.g. dimension mismatch on vector cast // 22023 invalid_parameter_value → pgvector dimension mismatch switch (e.code) { @@ -36,6 +37,10 @@ export async function handlePgVectorError( case '42P07': case '23505': throw ERRORS.S3VectorConflictException(resource.type, resource.name) + case '21000': + throw ERRORS.InvalidParameter(resource.name, { + message: e.message ?? 'duplicate keys in pgvector upsert batch', + }) case '22P02': case '22023': throw ERRORS.InvalidParameter(resource.name, { diff --git a/src/storage/protocols/vector/adapter/pgvector/filter.test.ts b/src/storage/protocols/vector/adapter/pgvector/filter.test.ts index 52fe13c9d..7401f4af9 100644 --- a/src/storage/protocols/vector/adapter/pgvector/filter.test.ts +++ b/src/storage/protocols/vector/adapter/pgvector/filter.test.ts @@ -1,29 +1,79 @@ -import { translateFilter } from './filter' +import { translateFilter, translateFilterForKnex } from './filter' + +function eqSql(column: string, field: number, scalar: number, array: number, value: number) { + return ( + `(${column}->>$${field} = $${scalar} OR ` + + `EXISTS (SELECT 1 FROM jsonb_array_elements(CASE WHEN jsonb_typeof(${column}->$${array}) = 'array' THEN ${column}->$${array} ELSE '[]'::jsonb END) AS elem(value) WHERE elem.value#>>'{}' = $${value}))` + ) +} + +function neSql( + column: string, + type: number, + array: number, + value: number, + field: number, + scalar: number +) { + return ( + `(CASE WHEN jsonb_typeof(${column}->$${type}) = 'array' THEN NOT ` + + `EXISTS (SELECT 1 FROM jsonb_array_elements(${column}->$${array}) AS elem(value) WHERE elem.value#>>'{}' = $${value}) ` + + `ELSE ${column}->>$${field} <> $${scalar} END)` + ) +} + +function inSql(column: string, field: number, scalar: number, array: number, value: number) { + return ( + `(${column}->>$${field} = ANY($${scalar}) OR ` + + `EXISTS (SELECT 1 FROM jsonb_array_elements(CASE WHEN jsonb_typeof(${column}->$${array}) = 'array' THEN ${column}->$${array} ELSE '[]'::jsonb END) AS elem(value) WHERE elem.value#>>'{}' = ANY($${value})))` + ) +} + +function ninSql( + column: string, + type: number, + array: number, + value: number, + field: number, + scalar: number +) { + return ( + `(CASE WHEN jsonb_typeof(${column}->$${type}) = 'array' THEN NOT ` + + `EXISTS (SELECT 1 FROM jsonb_array_elements(${column}->$${array}) AS elem(value) WHERE elem.value#>>'{}' = ANY($${value})) ` + + `ELSE ${column}->>$${field} <> ALL($${scalar}) END)` + ) +} describe('translateFilter (pgvector / JSONB)', () => { describe('implicit equality', () => { it('string', () => { expect(translateFilter({ category: 'foo' })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['category', 'foo'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['category', 'foo', 'category', 'foo'], }) }) it('number', () => { expect(translateFilter({ n: 5 })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['n', '5'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['n', '5', 'n', '5'], }) }) it('boolean', () => { expect(translateFilter({ active: true })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['active', 'true'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['active', 'true', 'active', 'true'], }) }) it('preserves embedded quotes via parameter (no manual escaping needed)', () => { expect(translateFilter({ title: "it's" })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['title', "it's"], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['title', "it's", 'title', "it's"], + }) + }) + it('matches primitive filters against list-valued metadata', () => { + expect(translateFilter({ category: 'documentary' })).toEqual({ + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['category', 'documentary', 'category', 'documentary'], }) }) }) @@ -31,26 +81,26 @@ describe('translateFilter (pgvector / JSONB)', () => { describe('accepts arbitrary metadata key strings', () => { it('keys with hyphens', () => { expect(translateFilter({ 'user-id': 'abc' })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['user-id', 'abc'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['user-id', 'abc', 'user-id', 'abc'], }) }) it('keys with dots', () => { expect(translateFilter({ 'my.key': 'v' })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['my.key', 'v'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['my.key', 'v', 'my.key', 'v'], }) }) it('keys with spaces', () => { expect(translateFilter({ 'a b c': 'v' })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['a b c', 'v'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['a b c', 'v', 'a b c', 'v'], }) }) it('keys that look hostile are still passed as parameters', () => { expect(translateFilter({ "'; DROP TABLE x;--": 'v' })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ["'; DROP TABLE x;--", 'v'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ["'; DROP TABLE x;--", 'v', "'; DROP TABLE x;--", 'v'], }) }) }) @@ -58,14 +108,14 @@ describe('translateFilter (pgvector / JSONB)', () => { describe('field operators', () => { it('$eq', () => { expect(translateFilter({ category: { $eq: 'foo' } })).toEqual({ - sql: 'metadata->>$1 = $2', - params: ['category', 'foo'], + sql: eqSql('metadata', 1, 2, 3, 4), + params: ['category', 'foo', 'category', 'foo'], }) }) it('$ne', () => { expect(translateFilter({ category: { $ne: 'foo' } })).toEqual({ - sql: 'metadata->>$1 <> $2', - params: ['category', 'foo'], + sql: neSql('metadata', 1, 2, 3, 4, 5), + params: ['category', 'category', 'foo', 'category', 'foo'], }) }) it('$gt guards with jsonb_typeof + numeric cast', () => { @@ -94,14 +144,14 @@ describe('translateFilter (pgvector / JSONB)', () => { }) it('$in uses ANY with array param', () => { expect(translateFilter({ tag: { $in: ['a', 'b', 'c'] } })).toEqual({ - sql: 'metadata->>$1 = ANY($2)', - params: ['tag', ['a', 'b', 'c']], + sql: inSql('metadata', 1, 2, 3, 4), + params: ['tag', ['a', 'b', 'c'], 'tag', ['a', 'b', 'c']], }) }) it('$nin uses <> ALL with array param', () => { expect(translateFilter({ tag: { $nin: ['a', 'b'] } })).toEqual({ - sql: 'metadata->>$1 <> ALL($2)', - params: ['tag', ['a', 'b']], + sql: ninSql('metadata', 1, 2, 3, 4, 5), + params: ['tag', 'tag', ['a', 'b'], 'tag', ['a', 'b']], }) }) it('$exists true uses jsonb_exists (function form avoids knex `?` collision)', () => { @@ -132,8 +182,8 @@ describe('translateFilter (pgvector / JSONB)', () => { describe('multi-field (implicit AND)', () => { it('joins clauses with AND', () => { expect(translateFilter({ a: 1, b: 'x' })).toEqual({ - sql: 'metadata->>$1 = $2 AND metadata->>$3 = $4', - params: ['a', '1', 'b', 'x'], + sql: `${eqSql('metadata', 1, 2, 3, 4)} AND ${eqSql('metadata', 5, 6, 7, 8)}`, + params: ['a', '1', 'a', '1', 'b', 'x', 'b', 'x'], }) }) }) @@ -141,20 +191,20 @@ describe('translateFilter (pgvector / JSONB)', () => { describe('logical operators', () => { it('$and', () => { expect(translateFilter({ $and: [{ a: 1 }, { b: 'x' }] })).toEqual({ - sql: '(metadata->>$1 = $2) AND (metadata->>$3 = $4)', - params: ['a', '1', 'b', 'x'], + sql: `(${eqSql('metadata', 1, 2, 3, 4)}) AND (${eqSql('metadata', 5, 6, 7, 8)})`, + params: ['a', '1', 'a', '1', 'b', 'x', 'b', 'x'], }) }) it('$or', () => { expect(translateFilter({ $or: [{ a: 1 }, { b: 'x' }] })).toEqual({ - sql: '(metadata->>$1 = $2) OR (metadata->>$3 = $4)', - params: ['a', '1', 'b', 'x'], + sql: `(${eqSql('metadata', 1, 2, 3, 4)}) OR (${eqSql('metadata', 5, 6, 7, 8)})`, + params: ['a', '1', 'a', '1', 'b', 'x', 'b', 'x'], }) }) it('nested $and within $or', () => { expect(translateFilter({ $or: [{ $and: [{ a: 1 }, { b: 2 }] }, { c: 3 }] })).toEqual({ - sql: '((metadata->>$1 = $2) AND (metadata->>$3 = $4)) OR (metadata->>$5 = $6)', - params: ['a', '1', 'b', '2', 'c', '3'], + sql: `((${eqSql('metadata', 1, 2, 3, 4)}) AND (${eqSql('metadata', 5, 6, 7, 8)})) OR (${eqSql('metadata', 9, 10, 11, 12)})`, + params: ['a', '1', 'a', '1', 'b', '2', 'b', '2', 'c', '3', 'c', '3'], }) }) it('deeply nested mix', () => { @@ -162,9 +212,9 @@ describe('translateFilter (pgvector / JSONB)', () => { translateFilter({ $and: [{ $or: [{ a: 1 }, { b: 2 }] }, { c: { $gte: 5 } }] }) ).toEqual({ sql: - '((metadata->>$1 = $2) OR (metadata->>$3 = $4)) AND ' + - "((jsonb_typeof(metadata->$5) = 'number' AND (metadata->>$6)::numeric >= $7))", - params: ['a', '1', 'b', '2', 'c', 'c', 5], + `((${eqSql('metadata', 1, 2, 3, 4)}) OR (${eqSql('metadata', 5, 6, 7, 8)})) AND ` + + "((jsonb_typeof(metadata->$9) = 'number' AND (metadata->>$10)::numeric >= $11))", + params: ['a', '1', 'a', '1', 'b', '2', 'b', '2', 'c', 'c', 5], }) }) }) @@ -172,8 +222,19 @@ describe('translateFilter (pgvector / JSONB)', () => { describe('column override', () => { it('uses provided column reference', () => { expect(translateFilter({ a: 1 }, 'v.metadata')).toEqual({ - sql: 'v.metadata->>$1 = $2', - params: ['a', '1'], + sql: eqSql('v.metadata', 1, 2, 3, 4), + params: ['a', '1', 'a', '1'], + }) + }) + }) + + describe('knex raw conversion', () => { + it('expands reused numbered placeholders into positional bindings', () => { + expect(translateFilterForKnex({ category: 'cats' })).toEqual({ + sql: + '(metadata->>? = ? OR ' + + "EXISTS (SELECT 1 FROM jsonb_array_elements(CASE WHEN jsonb_typeof(metadata->?) = 'array' THEN metadata->? ELSE '[]'::jsonb END) AS elem(value) WHERE elem.value#>>'{}' = ?))", + params: ['category', 'cats', 'category', 'category', 'cats'], }) }) }) diff --git a/src/storage/protocols/vector/adapter/pgvector/filter.ts b/src/storage/protocols/vector/adapter/pgvector/filter.ts index f12889a16..705aadc5a 100644 --- a/src/storage/protocols/vector/adapter/pgvector/filter.ts +++ b/src/storage/protocols/vector/adapter/pgvector/filter.ts @@ -91,17 +91,73 @@ function primitiveAsText(value: FilterPrimitive): string { return String(value) } +function jsonArrayContainsText(ctx: Ctx, fieldName: string, value: string): string { + const arrayNode = jsonValue(ctx, fieldName) + const valueParam = placeholder(ctx, value) + return `EXISTS (SELECT 1 FROM jsonb_array_elements(CASE WHEN jsonb_typeof(${arrayNode}) = 'array' THEN ${arrayNode} ELSE '[]'::jsonb END) AS elem(value) WHERE elem.value#>>'{}' = ${valueParam})` +} + +function jsonArrayContainsAnyText(ctx: Ctx, fieldName: string, values: string[]): string { + const arrayNode = jsonValue(ctx, fieldName) + const valuesParam = placeholder(ctx, values) + return `EXISTS (SELECT 1 FROM jsonb_array_elements(CASE WHEN jsonb_typeof(${arrayNode}) = 'array' THEN ${arrayNode} ELSE '[]'::jsonb END) AS elem(value) WHERE elem.value#>>'{}' = ANY(${valuesParam}))` +} + +function jsonArrayContainsTextAssumingArray(ctx: Ctx, fieldName: string, value: string): string { + const arrayNode = jsonValue(ctx, fieldName) + const valueParam = placeholder(ctx, value) + return `EXISTS (SELECT 1 FROM jsonb_array_elements(${arrayNode}) AS elem(value) WHERE elem.value#>>'{}' = ${valueParam})` +} + +function jsonArrayContainsAnyTextAssumingArray( + ctx: Ctx, + fieldName: string, + values: string[] +): string { + const arrayNode = jsonValue(ctx, fieldName) + const valuesParam = placeholder(ctx, values) + return `EXISTS (SELECT 1 FROM jsonb_array_elements(${arrayNode}) AS elem(value) WHERE elem.value#>>'{}' = ANY(${valuesParam}))` +} + +function jsonScalarOrArrayEquals(ctx: Ctx, fieldName: string, value: FilterPrimitive): string { + const textValue = primitiveAsText(value) + const lhs = jsonText(ctx, fieldName) + const scalarParam = placeholder(ctx, textValue) + return `(${lhs} = ${scalarParam} OR ${jsonArrayContainsText(ctx, fieldName, textValue)})` +} + +function jsonScalarAndArrayNotEquals(ctx: Ctx, fieldName: string, value: FilterPrimitive): string { + const textValue = primitiveAsText(value) + const typeNode = jsonValue(ctx, fieldName) + const arrayClause = jsonArrayContainsTextAssumingArray(ctx, fieldName, textValue) + const lhs = jsonText(ctx, fieldName) + const scalarParam = placeholder(ctx, textValue) + return `(CASE WHEN jsonb_typeof(${typeNode}) = 'array' THEN NOT ${arrayClause} ELSE ${lhs} <> ${scalarParam} END)` +} + +function jsonScalarOrArrayIn(ctx: Ctx, fieldName: string, values: string[]): string { + const lhs = jsonText(ctx, fieldName) + const scalarValuesParam = placeholder(ctx, values) + return `(${lhs} = ANY(${scalarValuesParam}) OR ${jsonArrayContainsAnyText(ctx, fieldName, values)})` +} + +function jsonScalarAndArrayNotIn(ctx: Ctx, fieldName: string, values: string[]): string { + const typeNode = jsonValue(ctx, fieldName) + const arrayClause = jsonArrayContainsAnyTextAssumingArray(ctx, fieldName, values) + const lhs = jsonText(ctx, fieldName) + const scalarValuesParam = placeholder(ctx, values) + return `(CASE WHEN jsonb_typeof(${typeNode}) = 'array' THEN NOT ${arrayClause} ELSE ${lhs} <> ALL(${scalarValuesParam}) END)` +} + function translateFieldOperator(ctx: Ctx, fieldName: string, op: string, raw: unknown): string { switch (op) { case '$eq': { const v = validatePrimitive(raw) - const lhs = jsonText(ctx, fieldName) - return `${lhs} = ${placeholder(ctx, primitiveAsText(v))}` + return jsonScalarOrArrayEquals(ctx, fieldName, v) } case '$ne': { const v = validatePrimitive(raw) - const lhs = jsonText(ctx, fieldName) - return `${lhs} <> ${placeholder(ctx, primitiveAsText(v))}` + return jsonScalarAndArrayNotEquals(ctx, fieldName, v) } case '$gt': case '$gte': @@ -128,8 +184,7 @@ function translateFieldOperator(ctx: Ctx, fieldName: string, op: string, raw: un }) } const values = (raw as FilterPrimitive[]).map((v) => primitiveAsText(validatePrimitive(v))) - const lhs = jsonText(ctx, fieldName) - return `${lhs} = ANY(${placeholder(ctx, values)})` + return jsonScalarOrArrayIn(ctx, fieldName, values) } case '$nin': { if (!Array.isArray(raw) || raw.length === 0) { @@ -138,8 +193,7 @@ function translateFieldOperator(ctx: Ctx, fieldName: string, op: string, raw: un }) } const values = (raw as FilterPrimitive[]).map((v) => primitiveAsText(validatePrimitive(v))) - const lhs = jsonText(ctx, fieldName) - return `${lhs} <> ALL(${placeholder(ctx, values)})` + return jsonScalarAndArrayNotIn(ctx, fieldName, values) } case '$exists': { if (typeof raw !== 'boolean') { @@ -185,8 +239,7 @@ function translateFieldClause( return parts.length === 1 ? parts[0] : `(${parts.join(' AND ')})` } const v = validatePrimitive(value) - const lhs = jsonText(ctx, fieldName) - return `${lhs} = ${placeholder(ctx, primitiveAsText(v))}` + return jsonScalarOrArrayEquals(ctx, fieldName, v) } function translateInternal(ctx: Ctx, filter: S3VectorFilter): string { @@ -253,3 +306,22 @@ export function translateFilter(filter: S3VectorFilter, column = 'metadata'): Tr const sql = translateInternal(ctx, filter) return { sql, params: ctx.params } } + +export function translateFilterForKnex( + filter: S3VectorFilter, + column = 'metadata' +): TranslatedFilter { + const translated = translateFilter(filter, column) + const params: unknown[] = [] + const sql = translated.sql.replace(/\$(\d+)/g, (placeholderRef, index) => { + const paramIndex = Number(index) - 1 + if (paramIndex < 0 || paramIndex >= translated.params.length) { + throw ERRORS.InvalidParameter('filter', { + message: `Invalid filter placeholder: ${placeholderRef}`, + }) + } + params.push(translated.params[paramIndex]) + return '?' + }) + return { sql, params } +} diff --git a/src/storage/protocols/vector/adapter/pgvector/index.ts b/src/storage/protocols/vector/adapter/pgvector/index.ts index 144893171..e538d2708 100644 --- a/src/storage/protocols/vector/adapter/pgvector/index.ts +++ b/src/storage/protocols/vector/adapter/pgvector/index.ts @@ -20,9 +20,19 @@ import { ERRORS } from '@internal/errors' import BaseTtlCache from '@isaacs/ttlcache' import type { DocumentType } from '@smithy/types' import type { Knex } from 'knex' +import { + MAX_DELETE_VECTOR_KEYS, + MAX_GET_VECTOR_KEYS, + MAX_LIST_RESULTS, + MAX_QUERY_TOP_K, + MAX_SEGMENT_COUNT, + validatePutVectors, + validateVectorKeys, +} from '../../limits' +import { paginateNPlusOne } from '../../pagination' import { VectorStore } from '../s3-vector' import { handlePgVectorError } from './errors' -import { S3VectorFilter, translateFilter } from './filter' +import { S3VectorFilter, translateFilterForKnex } from './filter' // Cache the resolved distance metric for ~5 min per (bucket, index). Avoids // hammering pg_index on every queryVectors call; stale entries (e.g., an @@ -40,19 +50,39 @@ const metricCache = new BaseTtlCache({ max: METRIC_CACHE_MAX, updateAgeOnGet: true, }) +type PgVectorTableCapabilityKind = 'bridged-hnsw' | 'standard' | 'unknown' +interface PgVectorTableCapability { + kind: PgVectorTableCapabilityKind + requiresManualUpsert: boolean + requiresExactQueryScan: boolean +} +const STANDARD_TABLE_CAPABILITY: PgVectorTableCapability = { + kind: 'standard', + requiresManualUpsert: false, + requiresExactQueryScan: false, +} +const BRIDGED_HNSW_TABLE_CAPABILITY: PgVectorTableCapability = { + kind: 'bridged-hnsw', + requiresManualUpsert: true, + requiresExactQueryScan: true, +} +const UNKNOWN_TABLE_CAPABILITY: PgVectorTableCapability = { + kind: 'unknown', + requiresManualUpsert: false, + requiresExactQueryScan: true, +} +const tableCapabilityCaches = new WeakMap>>() function metricCacheKey(bucket: string, index: string): string { return `${bucket}\x00${index}` } const SCHEMA = 'storage_vectors' +const MAX_DIMENSIONS = 4_000 +const DEFAULT_HNSW_EF_SEARCH = 40 -// Caps mirror the S3Vectors service limits so behaviour stays consistent -// across backends. We clamp/reject early to avoid forwarding negative or -// huge values into LIMIT — negative LIMITs are treated as unlimited by -// Postgres and large ones drive surprise CPU/memory cost. -const MAX_TOP_K = 30 -const MAX_LIST_RESULTS = 1_000 +// Manual OrioleDB upserts can still deadlock under concurrent writers. +const MANUAL_UPSERT_MAX_ATTEMPTS = 3 function validatePositiveInt(name: string, value: number, max: number): number { if (!Number.isInteger(value) || value < 1 || value > max) { @@ -63,18 +93,135 @@ function validatePositiveInt(name: string, value: number, max: number): number { return value } +function isRetryableWriteConflict(error: unknown): boolean { + const code = (error as { code?: unknown })?.code + return code === '40P01' || code === '40001' +} + +function isBridgedHnswUpsertError(error: unknown): boolean { + const parts = [ + (error as { message?: unknown })?.message, + (error as { detail?: unknown })?.detail, + ].filter((part): part is string => typeof part === 'string') + + return parts.some((part) => part.includes('unexpected self-updated tuple')) +} + +function validateListSegment(input: ListVectorsInput): + | { + segmentCount: number + segmentIndex: number + } + | undefined { + const hasSegmentCount = input.segmentCount !== undefined + const hasSegmentIndex = input.segmentIndex !== undefined + + if (hasSegmentCount !== hasSegmentIndex) { + throw ERRORS.InvalidParameter('segmentCount/segmentIndex', { + message: 'segmentCount and segmentIndex must be provided together', + }) + } + + if (!hasSegmentCount) { + return undefined + } + + const segmentCount = validatePositiveInt('segmentCount', input.segmentCount!, MAX_SEGMENT_COUNT) + const segmentIndex = input.segmentIndex! + if (!Number.isInteger(segmentIndex) || segmentIndex < 0 || segmentIndex >= segmentCount) { + throw ERRORS.InvalidParameter('segmentIndex', { + message: `segmentIndex must be an integer in [0, ${segmentCount - 1}], got: ${segmentIndex}`, + }) + } + + return { segmentCount, segmentIndex } +} + +function tableCapabilityCache(db: Knex): Map> { + let cache = tableCapabilityCaches.get(db) + if (!cache) { + cache = new Map() + tableCapabilityCaches.set(db, cache) + } + return cache +} + +function capabilityForAccessMethod(accessMethod: unknown): PgVectorTableCapability { + if (typeof accessMethod !== 'string') { + return UNKNOWN_TABLE_CAPABILITY + } + + return accessMethod === 'orioledb' ? BRIDGED_HNSW_TABLE_CAPABILITY : STANDARD_TABLE_CAPABILITY +} + +function forgetTableCapability(db: Knex, table: string): void { + tableCapabilityCaches.get(db)?.delete(table) +} + +async function resolveTableCapability(db: Knex, table: string): Promise { + const cache = tableCapabilityCache(db) + let capability = cache.get(table) + if (capability) { + return capability + } + + const capabilityProbe: Promise = db + .raw( + `SELECT am.amname + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + JOIN pg_am am ON am.oid = c.relam + WHERE n.nspname = ? AND c.relname = ? + LIMIT 1`, + [SCHEMA, table] + ) + .then((result: { rows?: Array> }) => { + if (result.rows?.[0]?.amname === undefined && cache.get(table) === capabilityProbe) { + cache.delete(table) + } + return capabilityForAccessMethod(result.rows?.[0]?.amname) + }) + .catch(() => { + // Do not cache probe failures. Query paths treat unknown tables as + // exact-scan for correctness, and the next request can retry. + if (cache.get(table) === capabilityProbe) { + cache.delete(table) + } + return UNKNOWN_TABLE_CAPABILITY + }) + cache.set(table, capabilityProbe) + + return capabilityProbe +} + /** * Knex handle abstraction: ST mode passes a singleton Knex; MT mode passes * a function that resolves to the tenant's pool on each call. We can't use a * bare `() => Knex` because `Knex` itself is callable, so the function form is * wrapped in a tagged object. */ -export type KnexResolver = Knex | { resolve: () => Knex } +export type KnexResolver = Knex | { resolve: () => Knex; root?: () => Knex } function resolveKnex(r: KnexResolver): Knex { return 'resolve' in r && typeof r.resolve === 'function' ? r.resolve() : (r as Knex) } +function resolveRootKnex(r: KnexResolver): Knex { + return 'resolve' in r && typeof r.resolve === 'function' + ? (r.root?.() ?? r.resolve()) + : (r as Knex) +} + +function hasRootKnexResolver(r: KnexResolver): boolean { + return 'resolve' in r && typeof r.resolve === 'function' && typeof r.root === 'function' +} + +async function withPgTransaction(db: Knex, fn: (trx: Knex) => Promise): Promise { + // Knex transaction handles create a savepoint when transaction() is called on + // them. Keep statement failures from poisoning the caller's outer transaction. + return db.transaction(async (trx) => fn(trx as unknown as Knex)) +} + function tableName(vectorBucketName: string, indexName: string): string { // Combined logical key may exceed Postgres' 63-char identifier limit, so we // hash. SHA-256 truncated to 24 hex chars keeps the table name well within @@ -87,7 +234,11 @@ function tableName(vectorBucketName: string, indexName: string): string { } function qualifiedTable(vectorBucketName: string, indexName: string): string { - return `${SCHEMA}.${tableName(vectorBucketName, indexName)}` + return qualifiedTableName(tableName(vectorBucketName, indexName)) +} + +function qualifiedTableName(table: string): string { + return `${SCHEMA}.${table}` } interface OpClassChoice { @@ -113,13 +264,66 @@ function toVectorLiteral(values: number[]): string { return `[${values.join(',')}]` } +type PgVectorMetadataPrimitive = string | boolean | number +type PgVectorMetadata = Record + +function validateMetadata(metadata: DocumentType | undefined, vectorKey: string): PgVectorMetadata { + if (metadata === undefined) { + return {} + } + + const invalidMetadata = (message: string) => + ERRORS.InvalidParameter('vectors.metadata', { + message: `Invalid metadata for vector "${vectorKey}": ${message}`, + }) + + if (metadata === null || Array.isArray(metadata) || typeof metadata !== 'object') { + throw invalidMetadata('metadata must be an object') + } + + const validatePrimitive = (value: unknown, fieldName: string): PgVectorMetadataPrimitive => { + if (typeof value === 'string' || typeof value === 'boolean') { + return value + } + + if (typeof value === 'number' && Number.isFinite(value)) { + return value + } + + throw invalidMetadata( + `metadata field "${fieldName}" must be a string, boolean, or finite number` + ) + } + + const validated: PgVectorMetadata = {} + for (const [key, value] of Object.entries(metadata)) { + if (Array.isArray(value)) { + validated[key] = value.map((item, index) => validatePrimitive(item, `${key}[${index}]`)) + continue + } + + validated[key] = validatePrimitive(value, key) + } + + return validated +} + export class PgVectorStore implements VectorStore { - constructor(private readonly knex: KnexResolver) {} + readonly maxDimensions = MAX_DIMENSIONS + readonly transactionalIndexOperations: boolean + + constructor(private readonly knex: KnexResolver) { + this.transactionalIndexOperations = hasRootKnexResolver(knex) + } private db(): Knex { return resolveKnex(this.knex) } + private rootDb(): Knex { + return resolveRootKnex(this.knex) + } + async createVectorIndex(command: CreateIndexCommandInput): Promise { if (!command.indexName || !command.vectorBucketName) { throw ERRORS.MissingParameter('indexName/vectorBucketName') @@ -136,9 +340,9 @@ export class PgVectorStore implements VectorStore { // at 3072) and S3Vectors' own 4096 cap after trivial truncation. We // validate at create-index time so we fail loudly with a clear error // rather than later at INDEX-build time with an opaque pgvector error. - if (!dimension || !Number.isInteger(dimension) || dimension < 1 || dimension > 4_000) { + if (!dimension || !Number.isInteger(dimension) || dimension < 1 || dimension > MAX_DIMENSIONS) { throw ERRORS.InvalidParameter('dimension', { - message: `Invalid dimension for pgvector HNSW: ${dimension} (must be 1..4000)`, + message: `Invalid dimension for pgvector HNSW: ${dimension} (must be 1..${MAX_DIMENSIONS})`, }) } @@ -147,10 +351,11 @@ export class PgVectorStore implements VectorStore { return handlePgVectorError( async () => { + const db = this.db() // Wrap DDL in a transaction so a failed CREATE INDEX (missing // opclass, permissions, transient error) rolls back the CREATE TABLE. // Otherwise the orphan table would block retries with "already exists". - await this.db().transaction(async (trx) => { + await withPgTransaction(db, async (trx) => { // Postgres doesn't allow parameter binding inside type modifiers // like `halfvec(N)` — N must be a literal at parse time. We've // validated `dimension` is an integer in [1, 4_000] above, so @@ -171,6 +376,8 @@ export class PgVectorStore implements VectorStore { [`${table}_hnsw`, table] ) }) + forgetTableCapability(db, table) + forgetTableCapability(this.rootDb(), table) // Prime the metric cache so subsequent queryVectors don't need a // round-trip lookup to pick the right distance operator. The two // names are validated non-empty at the top of this method. @@ -187,9 +394,13 @@ export class PgVectorStore implements VectorStore { async deleteVectorIndex(param: DeleteIndexCommandInput): Promise { const bucket = param.vectorBucketName! const index = param.indexName! + const table = tableName(bucket, index) return handlePgVectorError( async () => { - await this.db().raw(`DROP TABLE IF EXISTS ${SCHEMA}.??`, [tableName(bucket, index)]) + const db = this.db() + await db.raw(`DROP TABLE IF EXISTS ${SCHEMA}.??`, [table]) + forgetTableCapability(db, table) + forgetTableCapability(this.rootDb(), table) metricCache.delete(metricCacheKey(bucket, index)) return { $metadata: {} } as DeleteIndexCommandOutput }, @@ -200,11 +411,12 @@ export class PgVectorStore implements VectorStore { async putVectors(command: PutVectorsInput): Promise { const bucket = command.vectorBucketName! const index = command.indexName! - const vectors = command.vectors ?? [] - if (vectors.length === 0) return {} as PutVectorsOutput + const vectors = validatePutVectors(command.vectors) return handlePgVectorError( async () => { + const db = this.db() + const rows = vectors.map((v) => { if (!v.key) throw ERRORS.MissingParameter('vector.key') if (!v.data || !v.data.float32) throw ERRORS.MissingParameter('vector.data.float32') @@ -215,31 +427,141 @@ export class PgVectorStore implements VectorStore { // it as a JSON object so jsonb_to_recordset parses `metadata` as a // JSONB object (not a JSONB string). Otherwise `metadata->>'key'` // returns NULL because the column would hold a quoted string. - metadata: v.metadata ?? {}, + metadata: validateMetadata(v.metadata, v.key), } }) + const serializedRows = JSON.stringify([...rows].sort((a, b) => a.key.localeCompare(b.key))) + const table = tableName(bucket, index) + const qualified = qualifiedTableName(table) + const capability = await resolveTableCapability(db, table) + + if (capability.requiresManualUpsert) { + await this.putVectorsManually(db, qualified, serializedRows) + return {} as PutVectorsOutput + } + + try { + await db.raw( + `INSERT INTO ?? (key, embedding, metadata) + SELECT key, embedding::halfvec, metadata + FROM jsonb_to_recordset(?::jsonb) + AS x(key text, embedding text, metadata jsonb) + ON CONFLICT (key) DO UPDATE + SET embedding = EXCLUDED.embedding, + metadata = EXCLUDED.metadata`, + [qualified, serializedRows] + ) + } catch (e) { + if (!isBridgedHnswUpsertError(e)) { + throw e + } + + forgetTableCapability(db, table) + const refreshedCapability = await resolveTableCapability(db, table) + if (!refreshedCapability.requiresManualUpsert) { + throw e + } + + await this.putVectorsManually(db, qualified, serializedRows) + } - await this.db().raw( - `INSERT INTO ?? (key, embedding, metadata) - SELECT key, embedding::halfvec, metadata - FROM jsonb_to_recordset(?::jsonb) - AS x(key text, embedding text, metadata jsonb) - ON CONFLICT (key) DO UPDATE - SET embedding = EXCLUDED.embedding, - metadata = EXCLUDED.metadata`, - [qualifiedTable(bucket, index), JSON.stringify(rows)] - ) return {} as PutVectorsOutput }, { type: 'vectors', name: index } ) } + private async putVectorsManually(db: Knex, table: string, serializedRows: string): Promise { + // OrioleDB supports pgvector HNSW indexes through index bridging, but its + // bridged HNSW path can reject ON CONFLICT DO UPDATE with "unexpected + // self-updated tuple". Plain UPDATE, INSERT, and DO NOTHING work, so this + // fallback preserves upsert semantics without that conflict action. + for (let attempt = 1; attempt <= MANUAL_UPSERT_MAX_ATTEMPTS; attempt += 1) { + try { + await withPgTransaction(db, async (trx) => { + await trx.raw( + `WITH input AS ( + SELECT key, embedding::halfvec AS embedding, metadata + FROM jsonb_to_recordset(?::jsonb) + AS x(key text, embedding text, metadata jsonb) + ) + UPDATE ${table} AS target + SET embedding = input.embedding, + metadata = input.metadata + FROM input + WHERE target.key = input.key`, + [serializedRows] + ) + + await trx.raw( + `INSERT INTO ${table} (key, embedding, metadata) + SELECT key, embedding::halfvec, metadata + FROM jsonb_to_recordset(?::jsonb) + AS x(key text, embedding text, metadata jsonb) + ON CONFLICT (key) DO NOTHING`, + [serializedRows] + ) + + // Close the READ COMMITTED race where two writers both miss the first + // UPDATE for a new key, one INSERT wins, and the other INSERT does + // nothing. The final UPDATE applies the later writer's payload without + // using ON CONFLICT DO UPDATE, which OrioleDB bridged HNSW can reject. + await trx.raw( + `WITH input AS ( + SELECT key, embedding::halfvec AS embedding, metadata + FROM jsonb_to_recordset(?::jsonb) + AS x(key text, embedding text, metadata jsonb) + ) + UPDATE ${table} AS target + SET embedding = input.embedding, + metadata = input.metadata + FROM input + WHERE target.key = input.key`, + [serializedRows] + ) + }) + return + } catch (error) { + if (attempt === MANUAL_UPSERT_MAX_ATTEMPTS || !isRetryableWriteConflict(error)) { + throw error + } + } + } + } + + private async queryVectorsRaw( + db: Knex, + table: string, + sql: string, + params: unknown[], + topK: number + ): Promise<{ rows: unknown[] }> { + const capability = await resolveTableCapability(db, table) + if (!capability.requiresExactQueryScan) { + return withPgTransaction(db, async (trx): Promise<{ rows: unknown[] }> => { + await trx.raw(`SELECT set_config('hnsw.ef_search', ?, true)`, [ + String(Math.max(topK, DEFAULT_HNSW_EF_SEARCH)), + ]) + return trx.raw(sql, params) + }) + } + + // The same OrioleDB bridged HNSW path that rejects ON CONFLICT DO UPDATE + // can also miss rows inserted after the index was created. Use exact scan + // semantics for pools where we have observed that path. + return withPgTransaction(db, async (trx): Promise<{ rows: unknown[] }> => { + await trx.raw( + `SELECT set_config('enable_indexscan', 'off', true), + set_config('enable_bitmapscan', 'off', true)` + ) + return trx.raw(sql, params) + }) + } + async getVectors(input: GetVectorsCommandInput): Promise { const bucket = input.vectorBucketName! const index = input.indexName! - const keys = input.keys ?? [] - if (keys.length === 0) return { vectors: [] } as unknown as GetVectorsCommandOutput + const keys = validateVectorKeys(input.keys, MAX_GET_VECTOR_KEYS) const wantData = input.returnData === true const wantMeta = input.returnMetadata === true @@ -272,8 +594,7 @@ export class PgVectorStore implements VectorStore { async deleteVectors(input: DeleteVectorsInput): Promise { const bucket = input.vectorBucketName! const index = input.indexName! - const keys = input.keys ?? [] - if (keys.length === 0) return {} as DeleteVectorsOutput + const keys = validateVectorKeys(input.keys, MAX_DELETE_VECTOR_KEYS) return handlePgVectorError( async () => { @@ -296,8 +617,8 @@ export class PgVectorStore implements VectorStore { } const wantMeta = input.returnMetadata === true - const wantDistance = input.returnDistance !== false - const topK = validatePositiveInt('topK', input.topK ?? 10, MAX_TOP_K) + const wantDistance = input.returnDistance === true + const topK = validatePositiveInt('topK', input.topK ?? 10, MAX_QUERY_TOP_K) return handlePgVectorError( async () => { @@ -317,11 +638,10 @@ export class PgVectorStore implements VectorStore { let whereClause = '' if (input.filter) { - const translated = translateFilter(input.filter as unknown as S3VectorFilter) - // The translator emits $N placeholders in left-to-right document - // order matching translated.params. Knex.raw uses `?` positionally, - // so we strip the numbering and let knex consume params in order. - whereClause = ' WHERE ' + translated.sql.replace(/\$\d+/g, '?') + const translated = translateFilterForKnex(input.filter as unknown as S3VectorFilter) + // Knex.raw uses `?` positionally, so repeated translated $N + // placeholders must be expanded into repeated bindings. + whereClause = ' WHERE ' + translated.sql params.push(...translated.params) } @@ -331,10 +651,11 @@ export class PgVectorStore implements VectorStore { // Inline the table name — it's a fixed prefix + hex hash (safe). // Avoids mixing `??` with `?::halfvec` casts, which trips up knex's // placeholder counter. - const sql = `SELECT ${cols.join(', ')} FROM ${qualifiedTable(bucket, index)}${whereClause} + const table = tableName(bucket, index) + const sql = `SELECT ${cols.join(', ')} FROM ${qualifiedTableName(table)}${whereClause} ORDER BY embedding ${distanceOp} ?::halfvec ASC LIMIT ?` - const result = await this.db().raw(sql, params) + const result = await this.queryVectorsRaw(this.db(), table, sql, params, topK) const rows = result.rows as Array<{ key: string distance?: number @@ -368,8 +689,9 @@ export class PgVectorStore implements VectorStore { const index = input.indexName! const wantData = input.returnData === true const wantMeta = input.returnMetadata === true - const maxResults = validatePositiveInt('maxResults', input.maxResults ?? 100, MAX_LIST_RESULTS) + const maxResults = validatePositiveInt('maxResults', input.maxResults ?? 500, MAX_LIST_RESULTS) const cursor = input.nextToken + const segment = validateListSegment(input) return handlePgVectorError( async () => { @@ -378,13 +700,18 @@ export class PgVectorStore implements VectorStore { if (wantMeta) cols.push('metadata') const params: unknown[] = [qualifiedTable(bucket, index)] - let whereClause = '' + const whereClauses: string[] = [] if (cursor) { - whereClause = ' WHERE key > ?' + whereClauses.push('key > ?') params.push(cursor) } - params.push(maxResults) + if (segment) { + whereClauses.push('mod(abs(hashtext(key)::bigint), ?::bigint) = ?::bigint') + params.push(segment.segmentCount, segment.segmentIndex) + } + params.push(maxResults + 1) + const whereClause = whereClauses.length > 0 ? ` WHERE ${whereClauses.join(' AND ')}` : '' const sql = `SELECT ${cols.join(', ')} FROM ??${whereClause} ORDER BY key ASC LIMIT ?` const result = await this.db().raw(sql, params) @@ -393,9 +720,9 @@ export class PgVectorStore implements VectorStore { embedding?: string metadata?: DocumentType }> - const nextToken = rows.length === maxResults ? rows[rows.length - 1].key : undefined + const { pageRows, nextToken } = paginateNPlusOne(rows, maxResults, (row) => row.key) return { - vectors: rows.map((r) => ({ + vectors: pageRows.map((r) => ({ key: r.key, data: wantData && r.embedding ? { float32: parseVectorLiteral(r.embedding) } : undefined, diff --git a/src/storage/protocols/vector/adapter/s3-vector.test.ts b/src/storage/protocols/vector/adapter/s3-vector.test.ts new file mode 100644 index 000000000..e85609b9d --- /dev/null +++ b/src/storage/protocols/vector/adapter/s3-vector.test.ts @@ -0,0 +1,46 @@ +import { PutVectorsCommand, S3VectorsClient, ValidationException } from '@aws-sdk/client-s3vectors' +import { ErrorCode } from '@internal/errors' +import { type Mock, vi } from 'vitest' +import { S3Vector } from './s3-vector' + +describe('S3Vector', () => { + let send: Mock + + beforeEach(() => { + send = vi.fn() + }) + + function createStore() { + return new S3Vector({ send } as unknown as S3VectorsClient) + } + + it('maps S3Vectors validation failures to invalid parameter errors', async () => { + const upstreamError = new ValidationException({ + message: + "Invalid record for key '5797803-0': Filterable metadata must have at most 2048 bytes", + $metadata: { httpStatusCode: 400 }, + fieldList: [ + { + path: 'vectors[0].metadata', + message: 'Filterable metadata must have at most 2048 bytes', + }, + ], + }) + send.mockRejectedValueOnce(upstreamError) + + await expect( + createStore().putVectors({ + vectorBucketName: 'bucket', + indexName: 'index', + vectors: [{ key: '5797803-0', data: { float32: [1, 2, 3] } }], + }) + ).rejects.toMatchObject({ + code: ErrorCode.InvalidParameter, + httpStatusCode: 400, + message: upstreamError.message, + originalError: upstreamError, + }) + + expect(send).toHaveBeenCalledWith(expect.any(PutVectorsCommand)) + }) +}) diff --git a/src/storage/protocols/vector/adapter/s3-vector.ts b/src/storage/protocols/vector/adapter/s3-vector.ts index 87053b1a9..4dfb3ccb8 100644 --- a/src/storage/protocols/vector/adapter/s3-vector.ts +++ b/src/storage/protocols/vector/adapter/s3-vector.ts @@ -24,11 +24,14 @@ import { QueryVectorsInput, QueryVectorsOutput, S3VectorsClient, + ValidationException, } from '@aws-sdk/client-s3vectors' import { ERRORS } from '@internal/errors' import { getConfig } from '../../../../config' export interface VectorStore { + maxDimensions?: number + transactionalIndexOperations?: boolean createVectorIndex(command: CreateIndexCommandInput): Promise deleteVectorIndex(param: DeleteIndexCommandInput): Promise putVectors(command: PutVectorsInput): Promise @@ -124,6 +127,10 @@ export class S3Vector implements VectorStore { throw ERRORS.S3VectorNotFoundException(resource.type, e.message) } + if (e instanceof ValidationException) { + throw ERRORS.InvalidParameter(resource.name, { error: e, message: e.message }) + } + throw e } } diff --git a/src/storage/protocols/vector/errors.ts b/src/storage/protocols/vector/errors.ts new file mode 100644 index 000000000..1418e1799 --- /dev/null +++ b/src/storage/protocols/vector/errors.ts @@ -0,0 +1,21 @@ +import { ErrorCode } from '@internal/errors/codes' + +interface ErrorWithCode { + code?: unknown +} + +export function isVectorResourceNotFoundError(error: unknown): boolean { + return ( + typeof error === 'object' && + error !== null && + (error as ErrorWithCode).code === ErrorCode.S3VectorNotFoundException + ) +} + +export function isVectorResourceConflictError(error: unknown): boolean { + return ( + typeof error === 'object' && + error !== null && + (error as ErrorWithCode).code === ErrorCode.S3VectorConflictException + ) +} diff --git a/src/storage/protocols/vector/knex.ts b/src/storage/protocols/vector/knex.ts index 445a1b78b..9102b1487 100644 --- a/src/storage/protocols/vector/knex.ts +++ b/src/storage/protocols/vector/knex.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from 'node:async_hooks' import { ListVectorBucketsInput } from '@aws-sdk/client-s3vectors' import { wait } from '@internal/concurrency' import { ERRORS } from '@internal/errors' @@ -7,6 +8,7 @@ import { VectorBucket } from '@storage/schemas' import { VectorIndex } from '@storage/schemas/vector' import { Knex } from 'knex' import { DatabaseError } from 'pg' +import { paginateNPlusOne } from './pagination' type DBVectorIndex = VectorIndex & { id: string; created_at: Date; updated_at: Date } export type VectorLockResourceType = 'bucket' | 'index' | 'global' @@ -61,8 +63,26 @@ export interface VectorMetadataDB { findVectorIndexForBucket(vectorBucketName: string, indexName: string): Promise } +const vectorTransactionStorage = new AsyncLocalStorage<{ root: Knex; transaction: Knex }>() + +export function createVectorTransactionKnexResolver(knex: Knex): { + resolve: () => Knex + root: () => Knex +} { + return { + resolve: () => { + const transaction = vectorTransactionStorage.getStore() + return transaction?.root === knex ? transaction.transaction : knex + }, + root: () => knex, + } +} + export class KnexVectorMetadataDB implements VectorMetadataDB { - constructor(protected readonly knex: Knex) {} + constructor( + protected readonly knex: Knex, + private readonly rootKnex: Knex = knex + ) {} lockResource(resourceType: VectorLockResourceType, resourceId: string): Promise { const lockId = hashStringToInt(`vector:${resourceType}:${resourceId}`) @@ -101,14 +121,15 @@ export class KnexVectorMetadataDB implements VectorMetadataDB { const maxResults = param.maxResults ? Math.min(param.maxResults, 500) : 500 const result = await query.orderBy('id', 'asc').limit(maxResults + 1) - - const hasMore = result.length > maxResults - - const buckets = result.slice(0, maxResults) + const { pageRows: buckets, nextToken } = paginateNPlusOne( + result, + maxResults, + (bucket) => bucket.id + ) return { vectorBuckets: buckets, - nextToken: hasMore ? buckets[buckets.length - 1].id : undefined, + nextToken, } } @@ -179,13 +200,15 @@ export class KnexVectorMetadataDB implements VectorMetadataDB { } const result = await query.limit(maxResults + 1) - const hasMore = result.length > maxResults - - const indexes = result.slice(0, maxResults) + const { pageRows: indexes, nextToken } = paginateNPlusOne( + result, + maxResults, + (index) => index.name + ) return { indexes, - nextToken: hasMore ? indexes[indexes.length - 1].name : undefined, + nextToken, } } @@ -205,7 +228,7 @@ export class KnexVectorMetadataDB implements VectorMetadataDB { async createVectorIndex(data: CreateVectorIndexParams) { try { - return await this.knex + const rows = await this.knex .withSchema('storage') .table('vector_indexes') .insert({ @@ -216,6 +239,9 @@ export class KnexVectorMetadataDB implements VectorMetadataDB { distance_metric: data.distanceMetric, metadata_configuration: data.metadataConfiguration, }) + .returning('*') + + return rows[0] } catch (e) { if (e instanceof Error && e instanceof DatabaseError) { if (e.code === '23505') { @@ -237,9 +263,11 @@ export class KnexVectorMetadataDB implements VectorMetadataDB { while (attempt < maxRetries) { try { return await this.knex.transaction(async (trx) => { - const trxDb = new KnexVectorMetadataDB(trx) - const result = await fn(trxDb) - return result + const transaction = trx as unknown as Knex + const trxDb = new KnexVectorMetadataDB(transaction, this.rootKnex) + return vectorTransactionStorage.run({ root: this.rootKnex, transaction }, async () => + fn(trxDb) + ) }, config) } catch (error) { attempt++ diff --git a/src/storage/protocols/vector/limits.ts b/src/storage/protocols/vector/limits.ts new file mode 100644 index 000000000..25af4306f --- /dev/null +++ b/src/storage/protocols/vector/limits.ts @@ -0,0 +1,63 @@ +import type { PutVectorsInput } from '@aws-sdk/client-s3vectors' +import { ERRORS } from '@internal/errors' + +export const MIN_VECTOR_DIMENSIONS = 1 +export const MAX_QUERY_TOP_K = 100 +export const MAX_LIST_RESULTS = 1_000 +export const MAX_SEGMENT_COUNT = 16 +export const MAX_PUT_VECTORS = 500 +export const MAX_GET_VECTOR_KEYS = 100 +export const MAX_DELETE_VECTOR_KEYS = 500 +export const MAX_VECTOR_KEY_LENGTH = 1_024 + +export function validateVectorKey(key: unknown, parameterName: string): void { + if (typeof key !== 'string' || key.length < 1 || key.length > MAX_VECTOR_KEY_LENGTH) { + throw ERRORS.InvalidParameter(parameterName, { + message: `${parameterName} must be between 1 and ${MAX_VECTOR_KEY_LENGTH} characters`, + }) + } +} + +export function validateVectorKeys(keys: string[] | undefined, max: number): string[] { + if (!Array.isArray(keys) || keys.length < 1 || keys.length > max) { + throw ERRORS.InvalidParameter('keys', { + message: `keys must contain between 1 and ${max} entries`, + }) + } + + for (const key of keys) { + validateVectorKey(key, 'keys') + } + + return keys +} + +export function validatePutVectors( + vectors: PutVectorsInput['vectors'] +): NonNullable { + if (!Array.isArray(vectors) || vectors.length < 1 || vectors.length > MAX_PUT_VECTORS) { + throw ERRORS.InvalidParameter('vectors', { + message: `vectors must contain between 1 and ${MAX_PUT_VECTORS} entries`, + }) + } + + const seenKeys = new Set() + + for (const vector of vectors) { + if (!vector || vector.key === undefined) { + throw ERRORS.MissingParameter('vectors.key') + } + + validateVectorKey(vector.key, 'vectors.key') + + if (seenKeys.has(vector.key)) { + throw ERRORS.InvalidParameter('vectors', { + message: 'Request must not contain duplicate keys', + }) + } + + seenKeys.add(vector.key) + } + + return vectors +} diff --git a/src/storage/protocols/vector/pagination.ts b/src/storage/protocols/vector/pagination.ts new file mode 100644 index 000000000..7112403c3 --- /dev/null +++ b/src/storage/protocols/vector/pagination.ts @@ -0,0 +1,14 @@ +export function paginateNPlusOne( + rows: T[], + maxResults: number, + getNextToken: (row: T) => string +): { pageRows: T[]; nextToken?: string } { + const pageRows = rows.slice(0, maxResults) + const hasMore = rows.length > maxResults + + return { + pageRows, + nextToken: + hasMore && pageRows.length > 0 ? getNextToken(pageRows[pageRows.length - 1]) : undefined, + } +} diff --git a/src/storage/protocols/vector/vector-store.test.ts b/src/storage/protocols/vector/vector-store.test.ts index 47947c446..3dd7c6144 100644 --- a/src/storage/protocols/vector/vector-store.test.ts +++ b/src/storage/protocols/vector/vector-store.test.ts @@ -1,5 +1,4 @@ import { - ConflictException, CreateIndexCommandOutput, DeleteIndexCommandOutput, DeleteVectorsOutput, @@ -9,12 +8,15 @@ import { QueryVectorsOutput, } from '@aws-sdk/client-s3vectors' import { ERRORS } from '@internal/errors' +import { logSchema } from '@internal/monitoring' import { Sharder } from '@internal/sharding' import { VectorBucket } from '@storage/schemas' +import type { Knex } from 'knex' import { type Mocked, vi } from 'vitest' import { type VectorStore } from './adapter/s3-vector' import { - type KnexVectorMetadataDB, + createVectorTransactionKnexResolver, + KnexVectorMetadataDB, type VectorLockResourceType, type VectorMetadataDB, } from './knex' @@ -32,6 +34,12 @@ function createMockVectorStore(): Mocked { } } +function createTransactionalMockVectorStore(): Mocked { + return Object.assign(createMockVectorStore(), { + transactionalIndexOperations: true, + }) +} + function createMockSharder(): Mocked { return { createShard: vi.fn(), @@ -75,6 +83,38 @@ function createVectorBucketRecord(bucketName: string): VectorBucket { } as unknown as VectorBucket } +function createVectorIndexRecord( + overrides: Partial>> = {} +): Awaited> { + return { + bucket_id: 'bucket-a', + created_at: new Date(), + data_type: 'float32', + dimension: 4, + distance_metric: 'cosine', + id: 'index-id-a', + metadata_configuration: undefined, + name: 'index-a', + updated_at: new Date(), + ...overrides, + } as Awaited> +} + +function metadataWithJsonByteLength(key: string, byteLength: number): Record { + const emptyByteLength = Buffer.byteLength(JSON.stringify({ [key]: '' }), 'utf8') + const valueByteLength = byteLength - emptyByteLength + + if (valueByteLength < 0) { + throw new Error(`Cannot build ${byteLength}-byte metadata for key "${key}"`) + } + + return { [key]: 'x'.repeat(valueByteLength) } +} + +function metadataWithKeyCount(count: number): Record { + return Object.fromEntries(Array.from({ length: count }, (_, i) => [`key-${i}`, 'value'])) +} + function createDeterministicVectorDb(options: { bucketCount: number existingBuckets?: string[] @@ -138,10 +178,7 @@ function createDeterministicVectorDb(options: { await options.onCreateVectorBucket?.(bucketName) if (state.existingBuckets.has(bucketName)) { - throw new ConflictException({ - message: `vector bucket "${bucketName}" already exists`, - $metadata: {}, - }) + throw ERRORS.S3VectorConflictException('vector bucket', bucketName) } state.existingBuckets.add(bucketName) @@ -189,6 +226,23 @@ function createDeterministicVectorDb(options: { } describe('VectorStoreManager bucket lifecycle', () => { + it('exposes the active metadata transaction through a scoped knex resolver', async () => { + const trx = { isTransaction: true } as unknown as Knex + const rootKnex = { + transaction: async (fn: (transaction: typeof trx) => Promise) => fn(trx), + } as unknown as Knex + const db = new KnexVectorMetadataDB(rootKnex) + const resolver = createVectorTransactionKnexResolver(rootKnex) + const resolved: unknown[] = [] + + await db.withTransaction(async () => { + resolved.push(resolver.resolve()) + }) + resolved.push(resolver.resolve()) + + expect(resolved).toEqual([trx, rootKnex]) + }) + it('serializes concurrent creates for the final bucket slot', async () => { const releaseFirstCreate = Promise.withResolvers() const firstCreateStarted = Promise.withResolvers() @@ -227,7 +281,7 @@ describe('VectorStoreManager bucket lifecycle', () => { ]) }) - it('keeps createBucket idempotent for an existing bucket even when at capacity', async () => { + it('returns conflict for an existing bucket even when at capacity', async () => { const db = createDeterministicVectorDb({ bucketCount: 1, existingBuckets: ['bucket-a'], @@ -239,12 +293,31 @@ describe('VectorStoreManager bucket lifecycle', () => { maxIndexCount: Infinity, }) - await expect(manager.createBucket('bucket-a')).resolves.toBeUndefined() + await expect(manager.createBucket('bucket-a')).rejects.toMatchObject({ + code: 'ConflictException', + }) await expect(manager.createBucket('bucket-b')).rejects.toMatchObject({ code: 'S3VectorMaxBucketsExceeded', }) }) + it('returns conflict for an existing bucket when capacity is available', async () => { + const db = createDeterministicVectorDb({ + bucketCount: 1, + existingBuckets: ['bucket-a'], + }) + + const manager = new VectorStoreManager(createMockVectorStore(), db, createMockSharder(), { + tenantId: 'test-tenant', + maxBucketCount: 2, + maxIndexCount: Infinity, + }) + + await expect(manager.createBucket('bucket-a')).rejects.toMatchObject({ + code: 'ConflictException', + }) + }) + it('shares the bucket-count lock between delete and create so capacity is observed after delete commits', async () => { const releaseDelete = Promise.withResolvers() const deleteReachedRemoval = Promise.withResolvers() @@ -276,6 +349,61 @@ describe('VectorStoreManager bucket lifecycle', () => { await expect(createPromise).resolves.toBeUndefined() }) + it('serializes deleting and recreating the same bucket through the global count lock', async () => { + const releaseDelete = Promise.withResolvers() + const deleteReachedRemoval = Promise.withResolvers() + const events: string[] = [] + + const db = createDeterministicVectorDb({ + bucketCount: 1, + existingBuckets: ['bucket-a'], + onLockResource: (resourceType, resourceId) => { + events.push(`lock:${resourceType}:${resourceId}`) + }, + onDeleteVectorBucket: async (bucketName) => { + events.push(`delete:start:${bucketName}`) + deleteReachedRemoval.resolve() + await releaseDelete.promise + events.push(`delete:end:${bucketName}`) + }, + onCreateVectorBucket: (bucketName) => { + events.push(`create:${bucketName}`) + }, + }) + + const manager = new VectorStoreManager(createMockVectorStore(), db, createMockSharder(), { + tenantId: 'test-tenant', + maxBucketCount: 1, + maxIndexCount: Infinity, + }) + + const deletePromise = manager.deleteBucket('bucket-a') + await deleteReachedRemoval.promise + + const recreatePromise = manager.createBucket('bucket-a') + await Promise.resolve() + + expect(events).toEqual([ + 'lock:bucket:bucket-a', + `lock:global:${VECTOR_BUCKET_COUNT_LOCK}`, + 'delete:start:bucket-a', + `lock:global:${VECTOR_BUCKET_COUNT_LOCK}`, + ]) + + releaseDelete.resolve() + + await expect(deletePromise).resolves.toBeUndefined() + await expect(recreatePromise).resolves.toBeUndefined() + expect(events).toEqual([ + 'lock:bucket:bucket-a', + `lock:global:${VECTOR_BUCKET_COUNT_LOCK}`, + 'delete:start:bucket-a', + `lock:global:${VECTOR_BUCKET_COUNT_LOCK}`, + 'delete:end:bucket-a', + 'create:bucket-a', + ]) + }) + it('does not block unrelated creates while delete waits on the target bucket lock', async () => { const releaseBucketLock = Promise.withResolvers() const deleteWaitingOnBucketLock = Promise.withResolvers() @@ -365,4 +493,1890 @@ describe('VectorStoreManager bucket lifecycle', () => { expect(sharder.reserve).not.toHaveBeenCalled() expect(vectorStore.createVectorIndex).not.toHaveBeenCalled() }) + + it('rejects dimensions above the backend limit before metadata or shard work', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + maxDimensions: 4_000, + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4096, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: 'dimension must be an integer in [1, 4000] for this vector backend, got: 4096', + }) + + expect(db.findVectorBucket).not.toHaveBeenCalled() + expect(db.withTransaction).not.toHaveBeenCalled() + expect(db.createVectorIndex).not.toHaveBeenCalled() + expect(sharder.reserve).not.toHaveBeenCalled() + expect(vectorStore.createVectorIndex).not.toHaveBeenCalled() + }) + + it('creates the physical vector index inside the metadata transaction', async () => { + const callOrder: string[] = [] + let inMetadataTransaction = false + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createTransactionalMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + callOrder.push('metadata:start') + inMetadataTransaction = true + try { + return await fn(db as unknown as KnexVectorMetadataDB) + } finally { + inMetadataTransaction = false + callOrder.push('metadata:end') + } + }) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockImplementation(async () => { + callOrder.push('confirm') + }) + vectorStore.createVectorIndex.mockImplementation(async () => { + callOrder.push(`physical:${inMetadataTransaction ? 'inside' : 'outside'}`) + return {} as CreateIndexCommandOutput + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + + expect(callOrder).toEqual(['metadata:start', 'physical:inside', 'confirm', 'metadata:end']) + expect(vectorStore.createVectorIndex).toHaveBeenCalledWith({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + }) + + it('rejects duplicate vector index metadata before reserving a shard or creating the physical index', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockRejectedValue( + ERRORS.S3VectorConflictException('vector index', 'index-a') + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toMatchObject({ code: 'ConflictException' }) + + expect(sharder.reserve).not.toHaveBeenCalled() + expect(vectorStore.createVectorIndex).not.toHaveBeenCalled() + expect(sharder.cancel).not.toHaveBeenCalled() + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + }) + + it('returns conflict for an existing index even when the bucket is at index capacity', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(1) + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: 1, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toMatchObject({ code: 'ConflictException' }) + + expect(sharder.reserve).not.toHaveBeenCalled() + expect(vectorStore.createVectorIndex).not.toHaveBeenCalled() + }) + + it('serializes concurrent creates for the same index and leaves the duplicate without cleanup work', async () => { + const firstMetadataInserted = Promise.withResolvers() + const releaseFirstMetadata = Promise.withResolvers() + const physicalCreateStarted = Promise.withResolvers() + const releasePhysicalCreate = Promise.withResolvers() + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + const indexes = new Set() + const bucketLockWaiters: Array<() => void> = [] + let bucketLockHeld = false + + async function acquireBucketLock() { + if (!bucketLockHeld) { + bucketLockHeld = true + return + } + + await new Promise((resolve) => { + bucketLockWaiters.push(resolve) + }) + } + + function releaseBucketLock() { + const next = bucketLockWaiters.shift() + if (next) { + next() + return + } + + bucketLockHeld = false + } + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockImplementation(async () => indexes.size) + db.createVectorIndex.mockImplementation(async (data) => { + const key = `${data.vectorBucketName}/${data.indexName}` + if (indexes.has(key)) { + throw ERRORS.S3VectorConflictException('vector index', data.indexName) + } + + indexes.add(key) + firstMetadataInserted.resolve() + await releaseFirstMetadata.promise + return {} as Awaited> + }) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + let holdsBucketLock = false + const tx: Partial = { + ...db, + lockResource: async (resourceType, resourceId) => { + if (resourceType === 'bucket' && resourceId === 'bucket-a') { + await acquireBucketLock() + holdsBucketLock = true + } + }, + } + + try { + return await fn(tx as KnexVectorMetadataDB) + } finally { + if (holdsBucketLock) { + releaseBucketLock() + } + } + }) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockResolvedValue() + vectorStore.createVectorIndex.mockImplementation(async () => { + physicalCreateStarted.resolve() + await releasePhysicalCreate.promise + return {} as CreateIndexCommandOutput + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + const command = { + dataType: 'float32' as const, + dimension: 4, + distanceMetric: 'cosine' as const, + indexName: 'index-a', + vectorBucketName: 'bucket-a', + } + + const firstCreate = manager.createVectorIndex(command) + await firstMetadataInserted.promise + + const duplicateCreate = manager.createVectorIndex(command) + releaseFirstMetadata.resolve() + await physicalCreateStarted.promise + releasePhysicalCreate.resolve() + + const results = await Promise.allSettled([firstCreate, duplicateCreate]) + + expect(results).toEqual([ + { status: 'fulfilled', value: undefined }, + { + status: 'rejected', + reason: expect.objectContaining({ code: 'ConflictException' }), + }, + ]) + expect(db.createVectorIndex).toHaveBeenCalledTimes(2) + expect(sharder.reserve).toHaveBeenCalledTimes(1) + expect(sharder.confirm).toHaveBeenCalledTimes(1) + expect(vectorStore.createVectorIndex).toHaveBeenCalledTimes(1) + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.cancel).not.toHaveBeenCalled() + }) + + it('cleans up the physical index and cancels the shard reservation when physical creation fails before commit', async () => { + const createError = new Error('DDL failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.cancel.mockResolvedValue() + vectorStore.createVectorIndex.mockRejectedValue(createError) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(createError) + + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(sharder.cancel).toHaveBeenCalledWith('reservation-a') + expect(sharder.confirm).not.toHaveBeenCalled() + }) + + it('deletes the physical index, metadata row, and shard reservation when shard confirmation fails after creation', async () => { + const confirmError = new Error('confirm failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockRejectedValue(confirmError) + sharder.cancel.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(confirmError) + + expect(vectorStore.createVectorIndex).toHaveBeenCalledWith({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(sharder.cancel).toHaveBeenCalledWith('reservation-a') + }) + + it('does not delete the physical index or free a failed create reservation after the same index was recreated', async () => { + const confirmError = new Error('confirm failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createTransactionalMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockRejectedValue(confirmError) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(confirmError) + + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).not.toHaveBeenCalled() + expect(sharder.cancel).toHaveBeenCalledWith('reservation-a') + }) + + it('does not replay nontransactional physical creation when the metadata transaction callback is retried', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + transactionalIndexOperations: false, + }) + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue(createVectorIndexRecord()) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + await fn(db as unknown as KnexVectorMetadataDB) + await fn(db as unknown as KnexVectorMetadataDB) + }) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + + expect(db.createVectorIndex).toHaveBeenCalledTimes(2) + expect(vectorStore.createVectorIndex).toHaveBeenCalledTimes(1) + expect(sharder.confirm).toHaveBeenCalledTimes(1) + }) + + it('waits for failed index creation cleanup to delete the physical index before freeing the shard', async () => { + const confirmError = new Error('confirm failed') + const physicalDeleteStarted = Promise.withResolvers() + const allowPhysicalDelete = Promise.withResolvers() + const callOrder: string[] = [] + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockRejectedValue(confirmError) + sharder.cancel.mockImplementation(async () => { + callOrder.push('cancel') + }) + vectorStore.deleteVectorIndex.mockImplementation(async () => { + callOrder.push('physical:start') + physicalDeleteStarted.resolve() + await allowPhysicalDelete.promise + callOrder.push('physical:end') + return {} as DeleteIndexCommandOutput + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + const createResult = manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + + await physicalDeleteStarted.promise + + expect(callOrder).toEqual(['physical:start']) + + allowPhysicalDelete.resolve() + + await expect(createResult).rejects.toBe(confirmError) + expect(callOrder).toEqual(['physical:start', 'physical:end', 'cancel']) + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('confirms the shard reservation when nontransactional physical creation finds an existing index', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + transactionalIndexOperations: false, + }) + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue(createVectorIndexRecord()) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockResolvedValue() + vectorStore.createVectorIndex.mockRejectedValue( + ERRORS.S3VectorConflictException('vector-index', 'index-a') + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).resolves.toBeUndefined() + + expect(sharder.confirm).toHaveBeenCalledWith('reservation-a', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).not.toHaveBeenCalled() + expect(sharder.cancel).not.toHaveBeenCalled() + }) + + it('does not delete an existing physical index when nontransactional conflict adoption fails confirmation', async () => { + const confirmError = new Error('confirm failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + transactionalIndexOperations: false, + }) + const index = createVectorIndexRecord({ id: 'created-index-id' }) + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue(index) + db.findVectorIndexForBucket.mockResolvedValue(index) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockRejectedValue(confirmError) + sharder.cancel.mockResolvedValue() + vectorStore.createVectorIndex.mockRejectedValue( + ERRORS.S3VectorConflictException('vector-index', 'index-a') + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(confirmError) + + expect(sharder.confirm).toHaveBeenCalledWith('reservation-a', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(sharder.cancel).toHaveBeenCalledWith('reservation-a') + }) + + it('rejects filters that reference non-filterable metadata keys before querying the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + createVectorIndexRecord({ + metadata_configuration: { + nonFilterableMetadataKeys: ['private-note'], + }, + }) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.queryVectors({ + filter: { + $and: [{ 'private-note': 'hidden' }], + } as never, + indexName: 'index-a', + queryVector: { float32: [1, 0] }, + topK: 1, + vectorBucketName: 'bucket-a', + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + + expect(vectorStore.queryVectors).not.toHaveBeenCalled() + }) + + it('rejects sibling non-filterable keys when a logical filter is present', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + createVectorIndexRecord({ + metadata_configuration: { + nonFilterableMetadataKeys: ['private-note'], + }, + }) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.queryVectors({ + filter: { + $or: [{ category: 'public' }], + 'private-note': 'hidden', + } as never, + indexName: 'index-a', + queryVector: { float32: [1, 0] }, + topK: 1, + vectorBucketName: 'bucket-a', + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: 'Metadata key "private-note" is configured as non-filterable', + }) + + expect(vectorStore.queryVectors).not.toHaveBeenCalled() + }) + + it('confirms the shard reservation when transactional physical creation finds an existing index', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + transactionalIndexOperations: true, + }) + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.findVectorIndexForBucket.mockRejectedValue( + ERRORS.S3VectorNotFoundException('vector-index', 'index-a') + ) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockResolvedValue() + vectorStore.createVectorIndex.mockRejectedValue( + ERRORS.S3VectorConflictException('vector-index', 'index-a') + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).resolves.toBeUndefined() + + expect(sharder.confirm).toHaveBeenCalledWith('reservation-a', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).not.toHaveBeenCalled() + expect(sharder.cancel).not.toHaveBeenCalled() + }) + + it('cleans up the physical index and confirmed shard when metadata commit fails after creation', async () => { + const commitError = new Error('commit failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createTransactionalMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.findVectorIndexForBucket.mockRejectedValue( + ERRORS.S3VectorNotFoundException('vector-index', 'index-a') + ) + db.withTransaction + .mockImplementationOnce(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + await fn(db as unknown as KnexVectorMetadataDB) + throw commitError + }) + .mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.cancel.mockResolvedValue() + sharder.freeByResource.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(commitError) + + expect(vectorStore.createVectorIndex).toHaveBeenCalledWith({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(sharder.confirm).toHaveBeenCalledWith('reservation-a', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + expect(sharder.cancel).not.toHaveBeenCalled() + }) + + it('does not clean up a confirmed create when post-failure recheck sees committed metadata', async () => { + const commitError = new Error('commit failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createTransactionalMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + db.withTransaction + .mockImplementationOnce(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + await fn(db as unknown as KnexVectorMetadataDB) + throw commitError + }) + .mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.confirm.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(commitError) + + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).not.toHaveBeenCalled() + expect(sharder.cancel).not.toHaveBeenCalled() + }) + + it('does not log cleanup failure when rollback cleanup finds metadata already gone', async () => { + const commitError = new Error('commit failed') + const cleanupLogSpy = vi.spyOn(logSchema, 'error').mockImplementation(() => undefined) + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createTransactionalMockVectorStore() + + db.findVectorBucket.mockResolvedValue(createVectorBucketRecord('bucket-a')) + db.countIndexes.mockResolvedValue(0) + db.createVectorIndex.mockResolvedValue( + {} as Awaited> + ) + db.findVectorIndexForBucket.mockRejectedValue( + ERRORS.S3VectorNotFoundException('vector-index', 'index-a') + ) + db.withTransaction + .mockImplementationOnce(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + await fn(db as unknown as KnexVectorMetadataDB) + throw commitError + }) + .mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.reserve.mockResolvedValue({ + leaseExpiresAt: '', + reservationId: 'reservation-a', + shardId: '1', + shardKey: 'shard-a', + slotNo: 0, + }) + sharder.freeByResource.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + try { + await expect( + manager.createVectorIndex({ + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(commitError) + + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(sharder.freeByResource).toHaveBeenCalledWith('1', { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + const loggedCleanupFailure = cleanupLogSpy.mock.calls.some( + ([, message]) => message === 'Vector index creation cleanup failed' + ) + expect(loggedCleanupFailure).toBe(false) + } finally { + cleanupLogSpy.mockRestore() + } + }) + + it('deletes metadata, physical index, and shard allocation in one bucket-locked transaction', async () => { + const callOrder: string[] = [] + let inMetadataTransaction = false + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + callOrder.push('metadata:start') + inMetadataTransaction = true + try { + const result = await fn(db as unknown as KnexVectorMetadataDB) + callOrder.push('metadata:commit') + return result + } finally { + inMetadataTransaction = false + callOrder.push('metadata:end') + } + }) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + sharder.freeByResource.mockImplementation(async () => { + callOrder.push(`free:${inMetadataTransaction ? 'inside' : 'outside'}`) + }) + vectorStore.deleteVectorIndex.mockImplementation(async () => { + callOrder.push(`physical:${inMetadataTransaction ? 'inside' : 'outside'}`) + return {} as DeleteIndexCommandOutput + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + + expect(callOrder).toEqual([ + 'metadata:start', + 'physical:inside', + 'free:inside', + 'metadata:commit', + 'metadata:end', + ]) + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('deletes transactional physical vector index before metadata delete commits', async () => { + const callOrder: string[] = [] + let inMetadataTransaction = false + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = Object.assign(createMockVectorStore(), { + transactionalIndexOperations: true, + }) + + db.findVectorIndexForBucket.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + callOrder.push('metadata:start') + inMetadataTransaction = true + try { + const result = await fn(db as unknown as KnexVectorMetadataDB) + callOrder.push('metadata:commit') + return result + } finally { + inMetadataTransaction = false + callOrder.push('metadata:end') + } + }) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + sharder.freeByResource.mockImplementation(async () => { + callOrder.push(`free:${inMetadataTransaction ? 'inside' : 'outside'}`) + }) + vectorStore.deleteVectorIndex.mockImplementation(async () => { + callOrder.push(`physical:${inMetadataTransaction ? 'inside' : 'outside'}`) + return {} as DeleteIndexCommandOutput + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + + expect(callOrder).toEqual([ + 'metadata:start', + 'physical:inside', + 'free:inside', + 'metadata:commit', + 'metadata:end', + ]) + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('removes metadata and frees the shard when the physical vector index was already deleted', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket + .mockResolvedValueOnce( + {} as Awaited> + ) + .mockRejectedValueOnce(ERRORS.S3VectorNotFoundException('vector-index', 'index-a')) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + sharder.freeByResource.mockResolvedValue() + vectorStore.deleteVectorIndex.mockRejectedValue( + ERRORS.S3VectorNotFoundException('vector-index', 'index-a') + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).resolves.toBeUndefined() + + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('returns not found before physical delete or shard free when metadata is missing', async () => { + const missingIndex = ERRORS.S3VectorNotFoundException('vector-index', 'index-a') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockRejectedValue(missingIndex) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(missingIndex) + + expect(db.deleteVectorIndex).not.toHaveBeenCalled() + expect(vectorStore.deleteVectorIndex).not.toHaveBeenCalled() + expect(sharder.freeByResource).not.toHaveBeenCalled() + }) + + it('matches master sequencing by doing physical delete and shard free before commit can fail', async () => { + const commitError = new Error('commit failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + {} as Awaited> + ) + db.deleteVectorIndex.mockResolvedValue() + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => { + await fn(db as unknown as KnexVectorMetadataDB) + throw commitError + }) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + sharder.freeByResource.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(commitError) + + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('rejects oversized filterable metadata before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [ + { + key: 'vec-a', + data: { float32: [1, 0, 0, 0] }, + metadata: { large: 'x'.repeat(2_050) }, + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: "Invalid record for key 'vec-a': Filterable metadata must have at most 2048 bytes", + }) + + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('rejects a PutVectors batch with duplicate keys before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [ + { key: 'dup', data: { float32: [1, 0, 0, 0] } }, + { key: 'other', data: { float32: [0, 1, 0, 0] } }, + { key: 'dup', data: { float32: [0, 0, 1, 0] } }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: 'Request must not contain duplicate keys', + }) + + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('rejects PutVectors keys above the S3Vectors length limit before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [{ key: 'x'.repeat(1025), data: { float32: [1, 0, 0, 0] } }], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('rejects PutVectors batches above the S3Vectors count limit before metadata lookup', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: Array.from({ length: 501 }, (_, i) => ({ + key: `vec-${i}`, + data: { float32: [1, 0, 0, 0] }, + })), + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + + expect(db.findVectorIndexForBucket).not.toHaveBeenCalled() + expect(sharder.findShardByResourceId).not.toHaveBeenCalled() + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('rejects PutVectors entries without keys before metadata lookup', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [{ data: { float32: [1, 0, 0, 0] } } as never], + }) + ).rejects.toMatchObject({ + code: 'MissingParameter', + }) + + expect(db.findVectorIndexForBucket).not.toHaveBeenCalled() + expect(sharder.findShardByResourceId).not.toHaveBeenCalled() + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('rejects GetVectors keys above the S3Vectors length limit before metadata lookup', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.getVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + keys: ['x'.repeat(1025)], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + + expect(db.findVectorIndexForBucket).not.toHaveBeenCalled() + expect(sharder.findShardByResourceId).not.toHaveBeenCalled() + expect(vectorStore.getVectors).not.toHaveBeenCalled() + }) + + it('rejects DeleteVectors keys above the S3Vectors length limit before metadata lookup', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + keys: ['x'.repeat(1025)], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + + expect(db.findVectorIndexForBucket).not.toHaveBeenCalled() + expect(sharder.findShardByResourceId).not.toHaveBeenCalled() + expect(vectorStore.deleteVectors).not.toHaveBeenCalled() + }) + + it('allows filterable metadata exactly at the 2 KB boundary', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + const metadata = metadataWithJsonByteLength('boundary', 2_048) + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + ).resolves.toBeUndefined() + + expect(vectorStore.putVectors).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + }) + + it('rejects total metadata above 40 KB even when it is non-filterable', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + createVectorIndexRecord({ + metadata_configuration: { + nonFilterableMetadataKeys: ['large'], + }, + }) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [ + { + key: 'vec-a', + data: { float32: [1, 0, 0, 0] }, + metadata: metadataWithJsonByteLength('large', 40 * 1_024 + 1), + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: "Invalid record for key 'vec-a': Total metadata must have at most 40960 bytes", + }) + + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('allows total metadata exactly at the 40 KB boundary', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + const metadata = metadataWithJsonByteLength('large', 40 * 1_024) + + db.findVectorIndexForBucket.mockResolvedValue( + createVectorIndexRecord({ + metadata_configuration: { + nonFilterableMetadataKeys: ['large'], + }, + }) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + ).resolves.toBeUndefined() + + expect(vectorStore.putVectors).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + }) + + it('rejects metadata with more than 50 keys before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [ + { + key: 'vec-a', + data: { float32: [1, 0, 0, 0] }, + metadata: metadataWithKeyCount(51), + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: "Invalid record for key 'vec-a': Metadata must have at most 50 keys", + }) + + expect(vectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('allows metadata with exactly 50 keys before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + const metadata = metadataWithKeyCount(50) + + db.findVectorIndexForBucket.mockResolvedValue(createVectorIndexRecord()) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + ).resolves.toBeUndefined() + + expect(vectorStore.putVectors).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + vectors: [{ key: 'vec-a', data: { float32: [1, 0, 0, 0] }, metadata }], + }) + }) + + it('allows oversized non-filterable metadata before calling the vector backend', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + createVectorIndexRecord({ + metadata_configuration: { + nonFilterableMetadataKeys: ['large'], + }, + }) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.putVectors({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + vectors: [ + { + key: 'vec-a', + data: { float32: [1, 0, 0, 0] }, + metadata: { large: 'x'.repeat(2_050), category: 'small' }, + }, + ], + }) + ).resolves.toBeUndefined() + + expect(vectorStore.putVectors).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + vectors: [ + { + key: 'vec-a', + data: { float32: [1, 0, 0, 0] }, + metadata: { large: 'x'.repeat(2_050), category: 'small' }, + }, + ], + }) + }) + + it('deletes the physical vector index and frees the shard in the metadata delete transaction', async () => { + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + {} as Awaited> + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + sharder.freeByResource.mockResolvedValue() + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).resolves.toBeUndefined() + + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledWith({ + indexName: 'test-tenant-index-a', + vectorBucketName: 'shard-a', + }) + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) + + it('keeps the metadata row and shard allocated when physical delete fails so retry starts from the index row', async () => { + const deleteError = new Error('delete failed') + const db = createMockVectorDb() + const sharder = createMockSharder() + const vectorStore = createMockVectorStore() + + db.findVectorIndexForBucket.mockResolvedValue( + {} as Awaited> + ) + db.withTransaction.mockImplementation(async (fn: (db: KnexVectorMetadataDB) => unknown) => + fn(db as unknown as KnexVectorMetadataDB) + ) + sharder.findShardByResourceId.mockResolvedValue({ + capacity: 1, + created_at: new Date().toISOString(), + id: 1, + kind: 'vector', + next_slot: 1, + shard_key: 'shard-a', + status: 'active', + }) + vectorStore.deleteVectorIndex.mockRejectedValueOnce(deleteError).mockResolvedValueOnce({ + $metadata: {}, + } as DeleteIndexCommandOutput) + + const manager = new VectorStoreManager(vectorStore, db, sharder, { + tenantId: 'test-tenant', + maxBucketCount: Infinity, + maxIndexCount: Infinity, + }) + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).rejects.toBe(deleteError) + + expect(db.withTransaction).toHaveBeenCalledTimes(1) + expect(db.deleteVectorIndex).toHaveBeenCalledWith('bucket-a', 'index-a') + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledTimes(1) + expect(sharder.freeByResource).not.toHaveBeenCalled() + + await expect( + manager.deleteIndex({ + indexName: 'index-a', + vectorBucketName: 'bucket-a', + }) + ).resolves.toBeUndefined() + + expect(db.withTransaction).toHaveBeenCalledTimes(2) + expect(db.deleteVectorIndex).toHaveBeenCalledTimes(2) + expect(vectorStore.deleteVectorIndex).toHaveBeenCalledTimes(2) + expect(sharder.freeByResource).toHaveBeenCalledWith(1, { + bucketName: 'bucket-a', + kind: 'vector', + logicalName: 'index-a', + tenantId: 'test-tenant', + }) + }) }) diff --git a/src/storage/protocols/vector/vector-store.ts b/src/storage/protocols/vector/vector-store.ts index 60161cfe6..e2cd898a3 100644 --- a/src/storage/protocols/vector/vector-store.ts +++ b/src/storage/protocols/vector/vector-store.ts @@ -1,5 +1,4 @@ import { - ConflictException, CreateIndexInput, DeleteIndexInput, DeleteVectorsInput, @@ -16,11 +15,17 @@ import { QueryVectorsInput, } from '@aws-sdk/client-s3vectors' import { ERRORS } from '@internal/errors' -import { ErrorCode } from '@internal/errors/codes' import { logger, logSchema } from '@internal/monitoring' import { Sharder } from '@internal/sharding/sharder' import { VectorStore } from './adapter/s3-vector' +import { isVectorResourceConflictError, isVectorResourceNotFoundError } from './errors' import { VectorMetadataDB } from './knex' +import { + MAX_DELETE_VECTOR_KEYS, + MAX_GET_VECTOR_KEYS, + validatePutVectors, + validateVectorKeys, +} from './limits' interface VectorStoreConfig { tenantId: string @@ -28,7 +33,156 @@ interface VectorStoreConfig { maxIndexCount: number } +type VectorShardReservation = Awaited> +type VectorIndexMetadata = Awaited> +type VectorShardResource = { + kind: 'vector' + tenantId: string + bucketName: string + logicalName: string +} + export const VECTOR_BUCKET_COUNT_LOCK = '__vector_bucket_count__' +const MAX_FILTERABLE_METADATA_BYTES = 2_048 +const MAX_TOTAL_METADATA_BYTES = 40 * 1_024 +const MAX_METADATA_KEYS = 50 + +function getNonFilterableMetadataKeys(metadataConfiguration: unknown): ReadonlySet { + if ( + typeof metadataConfiguration !== 'object' || + metadataConfiguration === null || + !('nonFilterableMetadataKeys' in metadataConfiguration) + ) { + return new Set() + } + + const keys = (metadataConfiguration as { nonFilterableMetadataKeys?: unknown }) + .nonFilterableMetadataKeys + if (!Array.isArray(keys)) { + return new Set() + } + + return new Set(keys.filter((key): key is string => typeof key === 'string')) +} + +function isMetadataObject(metadata: unknown): metadata is Record { + if (typeof metadata !== 'object' || metadata === null || Array.isArray(metadata)) { + return false + } + + return true +} + +function getJsonByteLength(metadata: Record): number { + return Buffer.byteLength(JSON.stringify(metadata), 'utf8') +} + +function getFilterableMetadataByteLength( + metadata: unknown, + nonFilterableKeys: ReadonlySet +): number { + if (!isMetadataObject(metadata)) { + return 0 + } + + const filterableMetadata = Object.fromEntries( + Object.entries(metadata).filter(([key]) => !nonFilterableKeys.has(key)) + ) + return getJsonByteLength(filterableMetadata) +} + +function validateMetadataLimits( + vectors: PutVectorsInput['vectors'], + metadataConfiguration: unknown +): void { + const nonFilterableKeys = getNonFilterableMetadataKeys(metadataConfiguration) + + for (const vector of vectors ?? []) { + if (isMetadataObject(vector.metadata)) { + if (Object.keys(vector.metadata).length > MAX_METADATA_KEYS) { + throw ERRORS.InvalidParameter('vectors.metadata', { + message: `Invalid record for key '${vector.key ?? ''}': Metadata must have at most ${MAX_METADATA_KEYS} keys`, + }) + } + + if (getJsonByteLength(vector.metadata) > MAX_TOTAL_METADATA_BYTES) { + throw ERRORS.InvalidParameter('vectors.metadata', { + message: `Invalid record for key '${vector.key ?? ''}': Total metadata must have at most ${MAX_TOTAL_METADATA_BYTES} bytes`, + }) + } + } + + if ( + getFilterableMetadataByteLength(vector.metadata, nonFilterableKeys) > + MAX_FILTERABLE_METADATA_BYTES + ) { + throw ERRORS.InvalidParameter('vectors.metadata', { + message: `Invalid record for key '${vector.key ?? ''}': Filterable metadata must have at most ${MAX_FILTERABLE_METADATA_BYTES} bytes`, + }) + } + } +} + +function validateCreateIndexDimensionForStore( + dimension: CreateIndexInput['dimension'], + maxDimensions: number | undefined +): void { + if (maxDimensions === undefined || dimension === undefined) { + return + } + + if (!Number.isInteger(dimension) || dimension < 1 || dimension > maxDimensions) { + throw ERRORS.InvalidParameter('dimension', { + message: `dimension must be an integer in [1, ${maxDimensions}] for this vector backend, got: ${dimension}`, + }) + } +} + +function collectFilterMetadataKeys(filter: unknown, keys = new Set()): ReadonlySet { + if (typeof filter !== 'object' || filter === null || Array.isArray(filter)) { + return keys + } + + const record = filter as Record + for (const op of ['$and', '$or'] as const) { + const children = record[op] + if (Array.isArray(children)) { + for (const child of children) { + collectFilterMetadataKeys(child, keys) + } + } + } + + for (const key of Object.keys(record)) { + if (!key.startsWith('$')) { + keys.add(key) + } + } + + return keys +} + +function validateFilterableMetadataKeys( + filter: QueryVectorsInput['filter'], + metadataConfiguration: unknown +): void { + if (!filter) { + return + } + + const nonFilterableKeys = getNonFilterableMetadataKeys(metadataConfiguration) + if (nonFilterableKeys.size === 0) { + return + } + + for (const key of collectFilterMetadataKeys(filter)) { + if (nonFilterableKeys.has(key)) { + throw ERRORS.InvalidParameter('filter', { + message: `Metadata key "${key}" is configured as non-filterable`, + }) + } + } +} export class VectorStoreManager { constructor( @@ -42,6 +196,23 @@ export class VectorStoreManager { return `${this.config.tenantId}-${name}` } + private physicalCreateIndexInput(command: CreateIndexInput, shardKey: string): CreateIndexInput { + const input = { + ...command, + indexName: this.getIndexName(command.indexName!), + vectorBucketName: shardKey, + } + + if ( + input.metadataConfiguration?.nonFilterableMetadataKeys && + input.metadataConfiguration.nonFilterableMetadataKeys.length === 0 + ) { + delete input.metadataConfiguration + } + + return input + } + async createBucket(bucketName: string): Promise { await this.db.withTransaction(async (tnx) => { await tnx.lockResource('global', VECTOR_BUCKET_COUNT_LOCK) @@ -50,9 +221,9 @@ export class VectorStoreManager { if (bucketCount >= this.config.maxBucketCount) { try { await tnx.findVectorBucket(bucketName) - return + throw ERRORS.S3VectorConflictException('vector bucket', bucketName) } catch (e) { - if ((e as { code?: string }).code !== ErrorCode.S3VectorNotFoundException) { + if (!isVectorResourceNotFoundError(e)) { throw e } } @@ -60,14 +231,7 @@ export class VectorStoreManager { throw ERRORS.S3VectorMaxBucketsExceeded(this.config.maxBucketCount) } - try { - await tnx.createVectorBucket(bucketName) - } catch (e) { - if (e instanceof ConflictException) { - return - } - throw e - } + await tnx.createVectorBucket(bucketName) }) } @@ -133,14 +297,24 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } + validateCreateIndexDimensionForStore(command.dimension, this.vectorStore.maxDimensions) + await this.db.findVectorBucket(command.vectorBucketName) - const createIndexInput = { - ...command, - indexName: this.getIndexName(command.indexName), + const shardResource = { + kind: 'vector' as const, + bucketName: command.vectorBucketName, + tenantId: this.config.tenantId, + logicalName: command.indexName, + } + + if (!this.vectorStore.transactionalIndexOperations) { + return this.createVectorIndexWithNonTransactionalPhysicalCreate(command, shardResource) } - let shardReservation: { reservationId: string; shardKey: string; shardId: string } | undefined + let shardReservation: VectorShardReservation | undefined + let physicalIndexNeedsCleanup = false + let shardConfirmed = false try { await this.db.withTransaction(async (tx) => { @@ -150,64 +324,56 @@ export class VectorStoreManager { const indexCount = await tx.countIndexes(command.vectorBucketName!) if (indexCount >= this.config.maxIndexCount) { + try { + await tx.findVectorIndexForBucket(command.vectorBucketName!, command.indexName!) + throw ERRORS.S3VectorConflictException('vector index', command.indexName!) + } catch (e) { + if (!isVectorResourceNotFoundError(e)) { + throw e + } + } + throw ERRORS.S3VectorMaxIndexesExceeded(this.config.maxIndexCount) } await tx.createVectorIndex({ - dataType: createIndexInput.dataType!, - dimension: createIndexInput.dimension!, - distanceMetric: createIndexInput.distanceMetric!, + dataType: command.dataType!, + dimension: command.dimension!, + distanceMetric: command.distanceMetric!, indexName: command.indexName!, - metadataConfiguration: createIndexInput.metadataConfiguration, + metadataConfiguration: command.metadataConfiguration, vectorBucketName: command.vectorBucketName!, }) - shardReservation = await this.sharding.reserve({ - kind: 'vector', - bucketName: command.vectorBucketName!, - tenantId: this.config.tenantId, - logicalName: command.indexName!, - }) + shardReservation = await this.sharding.reserve(shardResource) if (!shardReservation) { throw ERRORS.NoAvailableShard() } - try { - if ( - createIndexInput.metadataConfiguration && - createIndexInput.metadataConfiguration.nonFilterableMetadataKeys && - createIndexInput.metadataConfiguration.nonFilterableMetadataKeys.length === 0 - ) { - delete createIndexInput.metadataConfiguration - } + const reservation = shardReservation + try { + physicalIndexNeedsCleanup = true await this.vectorStore.createVectorIndex({ - ...createIndexInput, - vectorBucketName: shardReservation.shardKey, + ...this.physicalCreateIndexInput(command, reservation.shardKey), }) - await this.sharding.confirm(shardReservation.reservationId, { - kind: 'vector', - bucketName: command.vectorBucketName!, - tenantId: this.config.tenantId, - logicalName: command.indexName!, - }) + await this.sharding.confirm(reservation.reservationId, shardResource) + shardConfirmed = true } catch (e) { + if (isVectorResourceConflictError(e)) { + physicalIndexNeedsCleanup = false + await this.sharding.confirm(reservation.reservationId, shardResource) + shardConfirmed = true + return + } + logSchema.error(logger, 'Vector index creation failed', { type: 'vector', error: e, project: this.config.tenantId, }) - if (e instanceof ConflictException) { - await this.sharding.confirm(shardReservation.reservationId, { - kind: 'vector', - bucketName: command.vectorBucketName!, - tenantId: this.config.tenantId, - logicalName: command.indexName!, - }) - return - } throw e } @@ -219,12 +385,257 @@ export class VectorStoreManager { project: this.config.tenantId, }) if (shardReservation) { - await this.sharding.cancel(shardReservation.reservationId) + await this.cleanupFailedVectorIndexCreate( + command, + shardReservation, + physicalIndexNeedsCleanup, + shardConfirmed, + shardResource + ) + } + throw error + } + } + + private async createVectorIndexWithNonTransactionalPhysicalCreate( + command: CreateIndexInput, + shardResource: VectorShardResource + ): Promise { + let shardReservation: VectorShardReservation | undefined + let createdIndex: VectorIndexMetadata | undefined + let physicalIndexNeedsCleanup = false + let shardConfirmed = false + + try { + await this.db.withTransaction(async (tx) => { + await tx.lockResource('bucket', command.vectorBucketName!) + await tx.findVectorBucket(command.vectorBucketName!) + + const indexCount = await tx.countIndexes(command.vectorBucketName!) + + if (indexCount >= this.config.maxIndexCount) { + try { + await tx.findVectorIndexForBucket(command.vectorBucketName!, command.indexName!) + throw ERRORS.S3VectorConflictException('vector index', command.indexName!) + } catch (e) { + if (!isVectorResourceNotFoundError(e)) { + throw e + } + } + + throw ERRORS.S3VectorMaxIndexesExceeded(this.config.maxIndexCount) + } + + createdIndex = await tx.createVectorIndex({ + dataType: command.dataType!, + dimension: command.dimension!, + distanceMetric: command.distanceMetric!, + indexName: command.indexName!, + metadataConfiguration: command.metadataConfiguration, + vectorBucketName: command.vectorBucketName!, + }) + + shardReservation = await this.sharding.reserve(shardResource) + + if (!shardReservation) { + throw ERRORS.NoAvailableShard() + } + }) + + const reservation = shardReservation + if (!reservation) { + throw ERRORS.NoAvailableShard() + } + + try { + physicalIndexNeedsCleanup = true + await this.vectorStore.createVectorIndex( + this.physicalCreateIndexInput(command, reservation.shardKey) + ) + } catch (e) { + if (isVectorResourceConflictError(e)) { + physicalIndexNeedsCleanup = false + await this.sharding.confirm(reservation.reservationId, shardResource) + shardConfirmed = true + return + } + + logSchema.error(logger, 'Vector index creation failed', { + type: 'vector', + error: e, + project: this.config.tenantId, + }) + + throw e + } + + await this.sharding.confirm(reservation.reservationId, shardResource) + shardConfirmed = true + } catch (error) { + logSchema.error(logger, 'Create vector index transaction failed', { + type: 'vector', + error, + project: this.config.tenantId, + }) + if (shardReservation) { + await this.cleanupFailedCommittedVectorIndexCreate( + command, + shardReservation, + physicalIndexNeedsCleanup, + shardConfirmed, + shardResource, + createdIndex?.id + ) } throw error } } + protected async cleanupFailedVectorIndexCreate( + command: CreateIndexInput, + shardReservation: VectorShardReservation, + physicalIndexNeedsCleanup: boolean, + shardConfirmed: boolean, + shardResource: VectorShardResource + ) { + const cleanupErrors: unknown[] = [] + const runCleanup = async (cleanup: () => Promise) => { + try { + await cleanup() + } catch (error) { + cleanupErrors.push(error) + } + } + + await runCleanup(() => + this.db.withTransaction(async (tx) => { + await tx.lockResource('bucket', command.vectorBucketName!) + + if (await this.findVectorIndexMetadata(tx, command.vectorBucketName!, command.indexName!)) { + if (!shardConfirmed) { + await this.sharding.cancel(shardReservation.reservationId) + } + return + } + + // Keep cleanup serial so the shard slot is not reusable before physical teardown finishes. + if (physicalIndexNeedsCleanup) { + await this.deletePhysicalVectorIndexIfExists( + shardReservation.shardKey, + this.getIndexName(command.indexName!) + ) + } + + await this.sharding.freeByResource(shardReservation.shardId, shardResource) + + if (!shardConfirmed) { + await this.sharding.cancel(shardReservation.reservationId) + } + }) + ) + + if (cleanupErrors.length > 0) { + logSchema.error(logger, 'Vector index creation cleanup failed', { + type: 'vector', + error: cleanupErrors, + project: this.config.tenantId, + }) + } + } + + protected async cleanupFailedCommittedVectorIndexCreate( + command: CreateIndexInput, + shardReservation: VectorShardReservation, + physicalIndexNeedsCleanup: boolean, + shardConfirmed: boolean, + shardResource: VectorShardResource, + createdIndexId: string | undefined + ) { + const cleanupErrors: unknown[] = [] + const runCleanup = async (cleanup: () => Promise) => { + try { + await cleanup() + } catch (error) { + cleanupErrors.push(error) + } + } + + await runCleanup(() => + this.db.withTransaction(async (tx) => { + await tx.lockResource('bucket', command.vectorBucketName!) + + const currentIndex = await this.findVectorIndexMetadata( + tx, + command.vectorBucketName!, + command.indexName! + ) + if (currentIndex) { + if (createdIndexId && currentIndex.id !== createdIndexId) { + if (!shardConfirmed) { + await this.sharding.cancel(shardReservation.reservationId) + } + return + } + + await tx.deleteVectorIndex(command.vectorBucketName!, command.indexName!) + } + + // Keep cleanup serial so the shard slot is not reusable before physical teardown finishes. + if (physicalIndexNeedsCleanup) { + await this.deletePhysicalVectorIndexIfExists( + shardReservation.shardKey, + this.getIndexName(command.indexName!) + ) + } + + await this.sharding.freeByResource(shardReservation.shardId, shardResource) + + if (!shardConfirmed) { + await this.sharding.cancel(shardReservation.reservationId) + } + }) + ) + + if (cleanupErrors.length > 0) { + logSchema.error(logger, 'Vector index creation cleanup failed', { + type: 'vector', + error: cleanupErrors, + project: this.config.tenantId, + }) + } + } + + private async deletePhysicalVectorIndexIfExists( + vectorBucketName: string, + indexName: string + ): Promise { + try { + await this.vectorStore.deleteVectorIndex({ + vectorBucketName, + indexName, + }) + } catch (e) { + if (!isVectorResourceNotFoundError(e)) { + throw e + } + } + } + + private async findVectorIndexMetadata( + db: Pick, + vectorBucketName: string, + indexName: string + ): Promise { + try { + return await db.findVectorIndexForBucket(vectorBucketName, indexName) + } catch (e) { + if (!isVectorResourceNotFoundError(e)) { + throw e + } + return undefined + } + } + async deleteIndex(command: DeleteIndexInput): Promise { if (!command.indexName) { throw ERRORS.MissingParameter('indexName') @@ -234,35 +645,27 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } - await this.db.findVectorIndexForBucket(command.vectorBucketName, command.indexName) - const vectorIndexName = this.getIndexName(command.indexName) + const shardResource = { + kind: 'vector' as const, + tenantId: this.config.tenantId, + logicalName: command.indexName, + bucketName: command.vectorBucketName, + } await this.db.withTransaction(async (tx) => { - const shard = await this.sharding.findShardByResourceId({ - kind: 'vector', - tenantId: this.config.tenantId, - logicalName: command.indexName!, - bucketName: command.vectorBucketName!, - }) + await tx.lockResource('bucket', command.vectorBucketName!) + await tx.findVectorIndexForBucket(command.vectorBucketName!, command.indexName!) + + const shard = await this.sharding.findShardByResourceId(shardResource) if (!shard) { throw ERRORS.NoAvailableShard() } await tx.deleteVectorIndex(command.vectorBucketName!, command.indexName!) - - await this.sharding.freeByResource(shard.id, { - kind: 'vector', - tenantId: this.config.tenantId, - bucketName: command.vectorBucketName!, - logicalName: command.indexName!, - }) - - await this.vectorStore.deleteVectorIndex({ - vectorBucketName: shard.shard_key, - indexName: vectorIndexName, - }) + await this.deletePhysicalVectorIndexIfExists(shard.shard_key, vectorIndexName) + await this.sharding.freeByResource(shard.id, shardResource) }) } @@ -322,7 +725,9 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } - const [shard] = await Promise.all([ + const vectors = validatePutVectors(command.vectors) + + const [shard, index] = await Promise.all([ this.sharding.findShardByResourceId({ kind: 'vector', tenantId: this.config.tenantId, @@ -336,8 +741,11 @@ export class VectorStoreManager { throw ERRORS.NoAvailableShard() } + validateMetadataLimits(vectors, index.metadata_configuration) + const putVectorsInput = { ...command, + vectors, vectorBucketName: shard.shard_key, indexName: this.getIndexName(command.indexName), } @@ -353,6 +761,8 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } + validateVectorKeys(command.keys, MAX_DELETE_VECTOR_KEYS) + const [shard] = await Promise.all([ this.sharding.findShardByResourceId({ kind: 'vector', @@ -422,7 +832,7 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } - const [shard] = await Promise.all([ + const [shard, index] = await Promise.all([ this.sharding.findShardByResourceId({ kind: 'vector', tenantId: this.config.tenantId, @@ -436,6 +846,8 @@ export class VectorStoreManager { throw ERRORS.NoAvailableShard() } + validateFilterableMetadataKeys(command.filter, index.metadata_configuration) + const queryInput = { ...command, vectorBucketName: shard.shard_key, @@ -453,6 +865,8 @@ export class VectorStoreManager { throw ERRORS.MissingParameter('vectorBucketName') } + validateVectorKeys(command.keys, MAX_GET_VECTOR_KEYS) + const [shard] = await Promise.all([ this.sharding.findShardByResourceId({ kind: 'vector', diff --git a/src/test/pgvector-adapter.test.ts b/src/test/pgvector-adapter.test.ts index a247d147d..b0f742aef 100644 --- a/src/test/pgvector-adapter.test.ts +++ b/src/test/pgvector-adapter.test.ts @@ -1,12 +1,57 @@ import { PgVectorStore } from '@storage/protocols/vector' import Knex from 'knex' -import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' const TEST_DATABASE_URL = process.env.VECTOR_DATABASE_URL || process.env.DATABASE_URL || 'postgresql://postgres:postgres@127.0.0.1/postgres' +function isTableAccessMethodLookup(sql: unknown): boolean { + const text = String(sql) + return text.includes('FROM pg_class') && text.includes('pg_am') && text.includes('relam') +} + +function serializedRowsFromRawCall(call: unknown[]): Array<{ + key: string + embedding: string + metadata: Record +}> { + const params = call[1] as unknown[] + return JSON.parse(params[1] as string) +} + +function quoteIdent(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"` +} + +function qualifiedStorageVectorTable(table: string): string { + return `${quoteIdent('storage_vectors')}.${quoteIdent(table)}` +} + +function createManualUpsertDb( + db: Knex.Knex, + afterRaw: (callNo: number) => Promise +): Knex.Knex { + return { + transaction: async (fn: (trx: Knex.Knex) => Promise) => + db.transaction(async (trx) => { + let callNo = 0 + const wrappedTrx = { + raw: async (sql: string, bindings?: unknown[]) => { + callNo += 1 + const result = + bindings === undefined ? await trx.raw(sql) : await trx.raw(sql, bindings) + await afterRaw(callNo) + return result + }, + } as unknown as Knex.Knex + + return fn(wrappedTrx) + }), + } as unknown as Knex.Knex +} + // pgvector availability is probed in beforeAll. We can't probe at module load // because the project builds to CJS, where top-level await is unsupported. // Tests guard at the top of each `it` body and skip when unavailable. @@ -49,12 +94,1098 @@ describe('PgVectorStore (real pgvector)', () => { afterAll(async () => { if (!pgvectorAvailable) return + + await store + .deleteVectorIndex({ vectorBucketName: bucket, indexName: index }) + .catch(() => undefined) + await knex.destroy() + }) + + it('rejects maxResults=0 before querying Postgres', async () => { + const raw = vi.fn() + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.listVectors({ + vectorBucketName: bucket, + indexName: index, + maxResults: 0, + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects maxResults above the pgvector list limit before querying Postgres', async () => { + const raw = vi.fn() + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.listVectors({ + vectorBucketName: bucket, + indexName: index, + maxResults: 1001, + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('uses the S3Vectors default page size when maxResults is omitted', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.listVectors({ + vectorBucketName: bucket, + indexName: index, + }) + ).resolves.toEqual({ + vectors: [], + nextToken: undefined, + }) + + expect(raw).toHaveBeenCalledTimes(1) + expect(raw.mock.calls[0]?.[1]).toEqual([expect.any(String), 501]) + }) + + it('rejects GetVectors requests above the S3Vectors key limit before querying Postgres', async () => { + const raw = vi.fn() + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.getVectors({ + vectorBucketName: bucket, + indexName: index, + keys: Array.from({ length: 101 }, (_, i) => `key-${i}`), + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects GetVectors keys above the S3Vectors key length before querying Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.getVectors({ + vectorBucketName: bucket, + indexName: index, + keys: ['x'.repeat(1025)], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects DeleteVectors requests above the S3Vectors key limit before querying Postgres', async () => { + const raw = vi.fn() + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.deleteVectors({ + vectorBucketName: bucket, + indexName: index, + keys: Array.from({ length: 501 }, (_, i) => `key-${i}`), + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects DeleteVectors keys above the S3Vectors key length before querying Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.deleteVectors({ + vectorBucketName: bucket, + indexName: index, + keys: ['x'.repeat(1025)], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it.each([ + -1, 0, 4001, + ])('rejects invalid dimension=%s before creating pgvector tables', async (dimension) => { + const raw = vi.fn() + const transaction = vi.fn() + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.createVectorIndex({ + vectorBucketName: bucket, + indexName: index, + dataType: 'float32', + dimension, + distanceMetric: 'cosine', + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + expect(transaction).not.toHaveBeenCalled() + }) + + it('rejects PutVectors keys above the S3Vectors key length before writing to Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { + key: 'x'.repeat(1025), + data: { float32: [1, 0] }, + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects PutVectors requests above the S3Vectors count limit before writing to Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: Array.from({ length: 501 }, (_, i) => ({ + key: `vec-${i}`, + data: { float32: [1, 0] }, + })), + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects nested metadata objects before writing to Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { + key: 'nested-metadata', + data: { float32: [1, 0] }, + metadata: { + nested: { value: 'not supported' }, + } as never, + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('allows list-valued metadata before writing to Postgres', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { + key: 'list-metadata', + data: { float32: [1, 0] }, + metadata: { + tags: ['cats', 'docs', 2026, true], + } as never, + }, + ], + }) + ).resolves.toEqual({}) + + const insertCall = raw.mock.calls.find(([sql]) => !isTableAccessMethodLookup(sql)) + expect(insertCall).toBeDefined() + expect(serializedRowsFromRawCall(insertCall!).at(0)?.metadata).toEqual({ + tags: ['cats', 'docs', 2026, true], + }) + }) + + it('rejects nested metadata arrays before writing to Postgres', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { + key: 'nested-list-metadata', + data: { float32: [1, 0] }, + metadata: { + tags: [['nested']], + } as never, + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it.each([ + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + ])('rejects non-finite metadata number %s before writing to Postgres', async (score) => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { + key: 'non-finite-metadata', + data: { float32: [1, 0] }, + metadata: { score } as never, + }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects duplicate put keys before regular upsert', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { key: 'dup', data: { float32: [1, 0] }, metadata: { version: 1 } }, + { key: 'other', data: { float32: [0, 1] }, metadata: { version: 1 } }, + { key: 'dup', data: { float32: [0.5, 0.5] }, metadata: { version: 2 } }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: 'Request must not contain duplicate keys', + }) + + expect(raw).not.toHaveBeenCalled() + }) + + it('rejects topK above the pgvector query limit before querying Postgres', async () => { + const raw = vi.fn() + const localStore = new PgVectorStore({ raw } as unknown as Knex.Knex) + + await expect( + localStore.queryVectors({ + vectorBucketName: bucket, + indexName: index, + queryVector: { float32: [1, 0] }, + topK: 101, + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + }) + expect(raw).not.toHaveBeenCalled() + }) + + it('allows topK at the S3Vectors query limit', async () => { + const localBucket = `bucket-top-k-${Date.now()}-${Math.random()}` + const localIndex = `index-top-k-${Date.now()}-${Math.random()}` + const raw = vi.fn(async (sql: string, _params?: unknown[]) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_cosine_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.queryVectors({ + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 100, + }) + ).resolves.toMatchObject({ + vectors: [], + }) + + const queryCall = raw.mock.calls.find(([sql]) => String(sql).includes('ORDER BY embedding')) + expect(queryCall).toBeDefined() + expect(queryCall![1]).toEqual(expect.arrayContaining([100])) + expect(transaction).toHaveBeenCalledTimes(1) + const efSearchCall = raw.mock.calls.find(([sql]) => + String(sql).includes("set_config('hnsw.ef_search'") + ) + expect(efSearchCall?.[1]).toEqual(['100']) + }) + + it('sets hnsw ef_search to the default floor for standard scans below the default', async () => { + const localBucket = `bucket-top-k-floor-${Date.now()}-${Math.random()}` + const localIndex = `index-top-k-floor-${Date.now()}-${Math.random()}` + const raw = vi.fn(async (sql: string, _params?: unknown[]) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_cosine_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.queryVectors({ + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 10, + }) + ).resolves.toMatchObject({ + vectors: [], + }) + + expect(transaction).toHaveBeenCalledTimes(1) + const efSearchCall = raw.mock.calls.find(([sql]) => + String(sql).includes("set_config('hnsw.ef_search'") + ) + expect(efSearchCall?.[1]).toEqual(['40']) + const queryCall = raw.mock.calls.find(([sql]) => String(sql).includes('ORDER BY embedding')) + expect(queryCall?.[1]).toEqual(expect.arrayContaining([10])) + }) + + it('omits distance values when queryVectors has returnDistance=false', async () => { + const localBucket = `bucket-no-distance-${Date.now()}-${Math.random()}` + const localIndex = `index-no-distance-${Date.now()}-${Math.random()}` + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_l2_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + if (text.includes('ORDER BY embedding')) { + return { rows: [{ key: 'a', distance: 123, metadata: { group: 'test' } }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + const result = await localStore.queryVectors({ + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 1, + returnDistance: false, + returnMetadata: true, + }) + + expect(result.distanceMetric).toBe('euclidean') + expect(result.vectors).toEqual([ + { + key: 'a', + distance: undefined, + metadata: { group: 'test' }, + }, + ]) + const queryCall = raw.mock.calls.find(([sql]) => String(sql).includes('ORDER BY embedding')) + expect(String(queryCall?.[0])).not.toContain('AS distance') + }) + + it('omits distance values when queryVectors does not request returnDistance', async () => { + const localBucket = `bucket-default-distance-${Date.now()}-${Math.random()}` + const localIndex = `index-default-distance-${Date.now()}-${Math.random()}` + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_l2_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + if (text.includes('ORDER BY embedding')) { + return { rows: [{ key: 'a', distance: 123, metadata: { group: 'test' } }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + const result = await localStore.queryVectors({ + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 1, + returnMetadata: true, + }) + + expect(result.distanceMetric).toBe('euclidean') + expect(result.vectors).toEqual([ + { + key: 'a', + distance: undefined, + metadata: { group: 'test' }, + }, + ]) + const queryCall = raw.mock.calls.find(([sql]) => String(sql).includes('ORDER BY embedding')) + expect(String(queryCall?.[0])).not.toContain('AS distance') + }) + + it('reuses the cached distance metric for repeated queries on the same index', async () => { + const localBucket = `bucket-cache-${Date.now()}-${Math.random()}` + const localIndex = `index-cache-${Date.now()}-${Math.random()}` + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_l2_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + const command = { + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 1, + } + + await localStore.queryVectors(command) + await localStore.queryVectors(command) + + const metricLookups = raw.mock.calls.filter(([sql]) => String(sql).includes('FROM pg_index')) + const vectorQueries = raw.mock.calls.filter(([sql]) => + String(sql).includes('ORDER BY embedding') + ) + expect(metricLookups).toHaveLength(1) + expect(vectorQueries).toHaveLength(2) + expect(vectorQueries.every(([sql]) => String(sql).includes('<->'))).toBe(true) + }) + + it('uses exact scan semantics and retries the OrioleDB probe after a transient probe failure', async () => { + const localBucket = `bucket-probe-failure-${Date.now()}-${Math.random()}` + const localIndex = `index-probe-failure-${Date.now()}-${Math.random()}` + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_cosine_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + const lookupCalls = raw.mock.calls.filter(([callSql]) => + isTableAccessMethodLookup(callSql) + ).length + if (lookupCalls === 1) { + throw new Error('temporary probe failure') + } + return { rows: [{ amname: 'orioledb' }] } + } + if (text.includes('ORDER BY embedding')) { + throw new Error('query should run through exact scan transaction') + } + return { rows: [] } + }) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + const command = { + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 1, + } + + await localStore.queryVectors(command) + await localStore.queryVectors(command) + + const capabilityLookups = raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql)) + expect(capabilityLookups).toHaveLength(2) + expect(transaction).toHaveBeenCalledTimes(2) + const exactScanSettings = trxRaw.mock.calls.filter(([sql]) => + String(sql).includes('enable_indexscan') + ) + expect(exactScanSettings).toHaveLength(2) + expect(exactScanSettings.every(([sql]) => String(sql).includes('enable_bitmapscan'))).toBe(true) + }) + + it('treats an empty table capability probe as unknown and does not cache it', async () => { + const localBucket = `bucket-empty-probe-${Date.now()}-${Math.random()}` + const localIndex = `index-empty-probe-${Date.now()}-${Math.random()}` + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (text.includes('FROM pg_index')) { + return { rows: [{ opcname: 'halfvec_cosine_ops' }] } + } + if (isTableAccessMethodLookup(text)) { + return { rows: [] } + } + if (text.includes('ORDER BY embedding')) { + throw new Error('query should run through exact scan transaction') + } + return { rows: [] } + }) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + const command = { + vectorBucketName: localBucket, + indexName: localIndex, + queryVector: { float32: [1, 0] }, + topK: 1, + } + + await localStore.queryVectors(command) + await localStore.queryVectors(command) + + const capabilityLookups = raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql)) + expect(capabilityLookups).toHaveLength(2) + expect(transaction).toHaveBeenCalledTimes(2) + const exactScanSettings = trxRaw.mock.calls.filter(([sql]) => + String(sql).includes('enable_indexscan') + ) + expect(exactScanSettings).toHaveLength(2) + expect(exactScanSettings.every(([sql]) => String(sql).includes('enable_bitmapscan'))).toBe(true) + }) + + it('does not classify heap pools as bridged from an unexpected self-updated tuple error', async () => { + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + if (text.includes('ON CONFLICT (key) DO UPDATE')) { + throw new Error('unexpected self-updated tuple') + } + return { rows: [] } + }) + const transaction = vi.fn() + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'a', data: { float32: [1, 0] } }], + }) + ).rejects.toThrow('unexpected self-updated tuple') + + expect(transaction).not.toHaveBeenCalled() + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(2) + }) + + it('does not re-probe table capability for unrelated regular upsert failures', async () => { + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'heap' }] } + } + if (text.includes('ON CONFLICT (key) DO UPDATE')) { + throw new Error('different vector dimensions 2 and 3') + } + return { rows: [] } + }) + const transaction = vi.fn() + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'a', data: { float32: [1, 0] } }], + }) + ).rejects.toThrow('different vector dimensions') + + expect(transaction).not.toHaveBeenCalled() + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(1) + }) + + it('uses actual table access method for OrioleDB bridged handling', async () => { + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (isTableAccessMethodLookup(text)) { + return { rows: [{ amname: 'orioledb' }] } + } + if (text.includes('SHOW default_table_access_method')) { + return { rows: [{ default_table_access_method: 'heap' }] } + } + if (text.includes('ON CONFLICT (key) DO UPDATE')) { + throw new Error('should not use regular ON CONFLICT upsert') + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'a', data: { float32: [1, 0] } }], + }) + + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(1) + expect( + raw.mock.calls.some(([sql]) => String(sql).includes('SHOW default_table_access_method')) + ).toBe(false) + expect(transaction).toHaveBeenCalledTimes(1) + expect(trxRaw).toHaveBeenCalledTimes(3) + expect(String(trxRaw.mock.calls[1][0])).toContain('ON CONFLICT (key) DO NOTHING') + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'b', data: { float32: [0, 1] } }], + }) + + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(1) + expect(transaction).toHaveBeenCalledTimes(2) + }) + + it('rejects duplicate put keys before OrioleDB manual upsert', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'orioledb' }] } + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { key: 'dup', data: { float32: [1, 0] }, metadata: { version: 1 } }, + { key: 'other', data: { float32: [0, 1] }, metadata: { version: 1 } }, + { key: 'dup', data: { float32: [0.5, 0.5] }, metadata: { version: 2 } }, + ], + }) + ).rejects.toMatchObject({ + code: 'InvalidParameter', + message: 'Request must not contain duplicate keys', + }) + + expect(raw).not.toHaveBeenCalled() + expect(transaction).not.toHaveBeenCalled() + expect(trxRaw).not.toHaveBeenCalled() + }) + + it('runs a final update after OrioleDB manual inserts to preserve concurrent upsert semantics', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'orioledb' }] } + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'race-key', data: { float32: [1, 0] }, metadata: { writer: 'late' } }], + }) + + expect(trxRaw).toHaveBeenCalledTimes(3) + expect(String(trxRaw.mock.calls[0][0])).toContain('UPDATE') + expect(String(trxRaw.mock.calls[1][0])).toContain('ON CONFLICT (key) DO NOTHING') + expect(String(trxRaw.mock.calls[2][0])).toContain('UPDATE') + }) + + it('sorts OrioleDB manual upsert rows by key before writing', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'orioledb' }] } + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { key: 'z', data: { float32: [1, 0] } }, + { key: 'a', data: { float32: [0, 1] } }, + { key: 'm', data: { float32: [0.5, 0.5] } }, + ], + }) + + const serializedRows = (trxRaw.mock.calls[0][1] as unknown[])[0] as string + expect(JSON.parse(serializedRows).map((row: { key: string }) => row.key)).toEqual([ + 'a', + 'm', + 'z', + ]) + }) + + it('retries OrioleDB manual upsert when Postgres reports a deadlock', async () => { + const raw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'orioledb' }] } + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const deadlock = Object.assign(new Error('deadlock detected'), { code: '40P01' }) + const transaction = vi + .fn() + .mockRejectedValueOnce(deadlock) + .mockImplementationOnce(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'retry-key', data: { float32: [1, 0] } }], + }) + ).resolves.toEqual({}) + + expect(transaction).toHaveBeenCalledTimes(2) + expect(trxRaw).toHaveBeenCalledTimes(3) + }) + + it('preserves manual upsert semantics when two real transactions both miss a new key', async (ctx) => { + if (!pgvectorAvailable) return ctx.skip() + + const manualStore = store as unknown as { + putVectorsManually(db: Knex.Knex, table: string, serializedRows: string): Promise + } + const manualTable = `manual_race_${Date.now()}_${Math.random().toString(16).slice(2)}` + const qualified = qualifiedStorageVectorTable(manualTable) + const firstWriterRows = JSON.stringify([ + { + key: 'race-key', + embedding: '[1,0]', + metadata: { writer: 'first' }, + }, + ]) + const secondWriterRows = JSON.stringify([ + { + key: 'race-key', + embedding: '[0,1]', + metadata: { writer: 'second' }, + }, + ]) + const secondWriterMissedInitialUpdate = Promise.withResolvers() + const allowSecondWriterToContinue = Promise.withResolvers() + + await knex.raw( + `CREATE TABLE storage_vectors.?? + ( + key text PRIMARY KEY, + embedding halfvec(2) NOT NULL, + metadata jsonb NOT NULL DEFAULT '{}'::jsonb + )`, + [manualTable] + ) + try { - await store.deleteVectorIndex({ vectorBucketName: bucket, indexName: index }) - } catch { - /* swallow cleanup errors */ + const firstWriter = manualStore.putVectorsManually( + createManualUpsertDb(knex, async (callNo) => { + if (callNo === 1) { + await secondWriterMissedInitialUpdate.promise + } + }), + qualified, + firstWriterRows + ) + const secondWriter = manualStore.putVectorsManually( + createManualUpsertDb(knex, async (callNo) => { + if (callNo === 1) { + secondWriterMissedInitialUpdate.resolve() + await allowSecondWriterToContinue.promise + } + }), + qualified, + secondWriterRows + ) + + await secondWriterMissedInitialUpdate.promise + await firstWriter + allowSecondWriterToContinue.resolve() + await secondWriter + + const result = await knex.raw( + `SELECT embedding::text AS embedding, metadata + FROM storage_vectors.?? + WHERE key = ?`, + [manualTable, 'race-key'] + ) + + expect(result.rows).toHaveLength(1) + expect(result.rows[0]).toMatchObject({ + embedding: '[0,1]', + metadata: { writer: 'second' }, + }) + } finally { + allowSecondWriterToContinue.resolve() + secondWriterMissedInitialUpdate.resolve() + await knex.raw('DROP TABLE IF EXISTS storage_vectors.??', [manualTable]) } - await knex.destroy() + }) + + it('falls back when bridged HNSW indexes reject ON CONFLICT DO UPDATE', async () => { + const raw = vi.fn(async (sql: string) => { + const text = String(sql) + if (isTableAccessMethodLookup(text)) { + const lookupCalls = raw.mock.calls.filter(([callSql]) => + isTableAccessMethodLookup(callSql) + ).length + if (lookupCalls === 1) { + throw new Error('temporary probe failure') + } + return { rows: [{ amname: 'orioledb' }] } + } + if (text.includes('ON CONFLICT (key) DO UPDATE')) { + throw new Error('unexpected self-updated tuple') + } + return { rows: [] } + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + await expect( + localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [ + { key: 'a', data: { float32: [1, 0] }, metadata: { category: 'cats' } }, + { key: 'b', data: { float32: [0, 1] }, metadata: { category: 'dogs' } }, + ], + }) + ).resolves.toEqual({}) + + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(2) + expect(transaction).toHaveBeenCalledTimes(1) + expect(trxRaw).toHaveBeenCalledTimes(3) + expect(String(trxRaw.mock.calls[1][0])).toContain('ON CONFLICT (key) DO NOTHING') + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'c', data: { float32: [1, 1] } }], + }) + + expect(raw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(2) + expect(transaction).toHaveBeenCalledTimes(2) + }) + + it('does not let an older failed capability probe delete a newer cached probe', async () => { + const isCapabilityLookup = (sql: unknown) => + String(sql).includes('JOIN pg_am am ON am.oid = c.relam') + const firstLookup = Promise.withResolvers<{ rows: Array<{ amname: string }> }>() + let lookupCount = 0 + const raw = vi.fn((sql: string) => { + if (String(sql).includes('FROM pg_index')) { + return Promise.resolve({ rows: [{ opcname: 'halfvec_cosine_ops' }] }) + } + if (isCapabilityLookup(sql)) { + lookupCount += 1 + if (lookupCount === 1) { + return firstLookup.promise + } + return Promise.resolve({ rows: [{ amname: 'orioledb' }] }) + } + return Promise.resolve({ rows: [] }) + }) + const trxRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof trxRaw }) => Promise) => + fn({ raw: trxRaw }) + ) + const localStore = new PgVectorStore({ raw, transaction } as unknown as Knex.Knex) + + const firstQuery = localStore.queryVectors({ + vectorBucketName: bucket, + indexName: index, + queryVector: { float32: [1, 0] }, + topK: 1, + }) + await new Promise((resolve) => setImmediate(resolve)) + expect(raw.mock.calls.filter(([sql]) => isCapabilityLookup(sql))).toHaveLength(1) + + await localStore.deleteVectorIndex({ vectorBucketName: bucket, indexName: index }) + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'a', data: { float32: [1, 0] } }], + }) + expect(raw.mock.calls.filter(([sql]) => isCapabilityLookup(sql))).toHaveLength(2) + + firstLookup.reject(new Error('temporary probe failure')) + await expect(firstQuery).resolves.toMatchObject({ vectors: [] }) + + await localStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'b', data: { float32: [0, 1] } }], + }) + + expect(raw.mock.calls.filter(([sql]) => isCapabilityLookup(sql))).toHaveLength(2) + }) + + it('creates an index through an active transaction resolver using a savepoint', async () => { + const raw = vi.fn().mockResolvedValue({ rows: [] }) + const transaction = vi.fn(async (fn: (trx: { raw: typeof raw }) => Promise) => + fn({ raw }) + ) + const localStore = new PgVectorStore({ + resolve: () => + ({ + isTransaction: true, + transaction, + }) as unknown as Knex.Knex, + }) + + await expect( + localStore.createVectorIndex({ + vectorBucketName: bucket, + indexName: index, + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + }) + ).resolves.toEqual({ $metadata: {} }) + + expect(transaction).toHaveBeenCalledTimes(1) + expect(raw).toHaveBeenCalledTimes(2) + expect(String(raw.mock.calls[0][0])).toContain('CREATE TABLE') + expect(String(raw.mock.calls[1][0])).toContain('CREATE INDEX') + }) + + it('invalidates a root capability cache after transaction-scoped index creation', async () => { + const rootRaw = vi.fn(async (sql: string) => { + if (isTableAccessMethodLookup(sql)) { + return { rows: [{ amname: 'heap' }] } + } + return { rows: [] } + }) + const rootDb = { raw: rootRaw } as unknown as Knex.Knex + const rootStore = new PgVectorStore(rootDb) + + await rootStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'a', data: { float32: [1, 0] } }], + }) + expect(rootRaw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(1) + + const transactionRaw = vi.fn().mockResolvedValue({ rows: [] }) + const transactionDb = { + isTransaction: true, + transaction: vi.fn(async (fn: (trx: { raw: typeof transactionRaw }) => Promise) => + fn({ raw: transactionRaw }) + ), + } as unknown as Knex.Knex + const transactionStore = new PgVectorStore({ + resolve: () => transactionDb, + root: () => rootDb, + } as never) + + await transactionStore.createVectorIndex({ + vectorBucketName: bucket, + indexName: index, + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + }) + + await rootStore.putVectors({ + vectorBucketName: bucket, + indexName: index, + vectors: [{ key: 'b', data: { float32: [0, 1] } }], + }) + + expect(rootRaw.mock.calls.filter(([sql]) => isTableAccessMethodLookup(sql))).toHaveLength(2) }) it('creates an index, puts vectors, queries, filters, fetches, lists, deletes', async (ctx) => { @@ -86,6 +1217,16 @@ describe('PgVectorStore (real pgvector)', () => { data: { float32: [0, 0, 1, 0] }, metadata: { category: 'cats', score: 9 }, }, + { + key: 'd', + data: { float32: [0, 0, 0, 1] }, + metadata: { category: ['cats', 'birds'], score: 7 }, + }, + { + key: 'e', + data: { float32: [0.5, 0.5, 0, 0] }, + metadata: { category: 'numeric-list', score: [1, 5, 9] }, + }, ], }) @@ -109,8 +1250,40 @@ describe('PgVectorStore (real pgvector)', () => { filter: { category: 'cats' } as never, returnMetadata: true, }) - const keys = (filtered.vectors ?? []).map((v) => v.key).sort() - expect(keys).toEqual(['a', 'c']) + function sortedVectorKeys(vectors: Array<{ key?: string }> | undefined): string[] { + return (vectors ?? []) + .map((v) => { + if (v.key === undefined) { + throw new Error('QueryVectors result is missing key') + } + return v.key + }) + .sort() + } + + const keys = sortedVectorKeys(filtered.vectors) + expect(keys).toEqual(['a', 'c', 'd']) + + async function filteredKeys(filter: unknown): Promise { + const result = await store.queryVectors({ + vectorBucketName: bucket, + indexName: index, + queryVector: { float32: [1, 0, 0, 0] }, + topK: 5, + filter: filter as never, + }) + + return sortedVectorKeys(result.vectors) + } + + await expect(filteredKeys({ category: { $ne: 'cats' } })).resolves.toEqual(['b', 'e']) + await expect(filteredKeys({ category: { $in: ['birds'] } })).resolves.toEqual(['d']) + await expect(filteredKeys({ category: { $nin: ['cats', 'birds'] } })).resolves.toEqual([ + 'b', + 'e', + ]) + await expect(filteredKeys({ score: { $eq: 5 } })).resolves.toEqual(['a', 'e']) + await expect(filteredKeys({ score: { $gt: 3 } })).resolves.toEqual(['a', 'c', 'd']) const fetched = await store.getVectors({ vectorBucketName: bucket, @@ -127,7 +1300,7 @@ describe('PgVectorStore (real pgvector)', () => { indexName: index, maxResults: 100, }) - expect((list.vectors ?? []).map((v) => v.key).sort()).toEqual(['a', 'b', 'c']) + expect((list.vectors ?? []).map((v) => v.key).sort()).toEqual(['a', 'b', 'c', 'd', 'e']) const firstPage = await store.listVectors({ vectorBucketName: bucket, @@ -156,7 +1329,7 @@ describe('PgVectorStore (real pgvector)', () => { keys: ['a'], }) expect(afterDel.vectors).toHaveLength(0) - }, 20_000) + }) it('ranks results by L2 distance when distanceMetric is euclidean', async (ctx) => { if (!pgvectorAvailable) return ctx.skip() @@ -194,7 +1367,164 @@ describe('PgVectorStore (real pgvector)', () => { } finally { await store.deleteVectorIndex({ vectorBucketName: bucket, indexName: euclIndex }) } - }, 20_000) + }) + + it('finds the nearest vector inserted after HNSW index creation', async (ctx) => { + if (!pgvectorAvailable) return ctx.skip() + const exactScanIndex = `tenant-a-exact-scan-${Date.now()}` + const distractors = Array.from({ length: 128 }, (_, i) => ({ + key: `far-${i.toString().padStart(3, '0')}`, + data: { float32: [100 + i, -100 - i] }, + })) + + await store.createVectorIndex({ + vectorBucketName: bucket, + indexName: exactScanIndex, + dataType: 'float32', + dimension: 2, + distanceMetric: 'euclidean', + }) + try { + await store.putVectors({ + vectorBucketName: bucket, + indexName: exactScanIndex, + vectors: [ + ...distractors, + { + key: 'true-nearest-after-index', + data: { float32: [0.001, 0] }, + }, + ], + }) + + const result = await store.queryVectors({ + vectorBucketName: bucket, + indexName: exactScanIndex, + queryVector: { float32: [0, 0] }, + topK: 1, + returnDistance: true, + }) + + expect(result.distanceMetric).toBe('euclidean') + expect(result.vectors?.map((v) => v.key)).toEqual(['true-nearest-after-index']) + expect(result.vectors?.[0].distance).toBeCloseTo(0.001, 4) + } finally { + await store.deleteVectorIndex({ vectorBucketName: bucket, indexName: exactScanIndex }) + } + }) + + it('returns topK rows above the default HNSW ef_search on standard index scans', async (ctx) => { + if (!pgvectorAvailable) return ctx.skip() + + const forcedKnex = Knex({ + client: 'pg', + connection: { connectionString: TEST_DATABASE_URL, connectionTimeoutMillis: 5_000 }, + pool: { min: 0, max: 1 }, + }) + const forcedStore = new PgVectorStore(forcedKnex) + const topKIndex = `tenant-a-topk-hnsw-${Date.now()}` + const vectors = Array.from({ length: 150 }, (_, i) => ({ + key: `vec-${i.toString().padStart(3, '0')}`, + data: { float32: [i + 1, 0] }, + })) + + try { + await forcedKnex.raw('SET default_table_access_method = heap') + await forcedStore.createVectorIndex({ + vectorBucketName: bucket, + indexName: topKIndex, + dataType: 'float32', + dimension: 2, + distanceMetric: 'euclidean', + }) + await forcedKnex.raw('SET enable_seqscan = off') + await forcedKnex.raw('SET hnsw.ef_search = 40') + await forcedKnex.raw('SET hnsw.iterative_scan = off') + await forcedStore.putVectors({ + vectorBucketName: bucket, + indexName: topKIndex, + vectors, + }) + + const result = await forcedStore.queryVectors({ + vectorBucketName: bucket, + indexName: topKIndex, + queryVector: { float32: [0, 0] }, + topK: 100, + }) + + expect(result.vectors).toHaveLength(100) + } finally { + try { + await forcedStore.deleteVectorIndex({ vectorBucketName: bucket, indexName: topKIndex }) + } finally { + await forcedKnex.destroy() + } + } + }) + + it('returns an empty list from an empty index with maxResults=1', async (ctx) => { + if (!pgvectorAvailable) return ctx.skip() + const emptyIndex = `tenant-a-empty-${Date.now()}` + await store.createVectorIndex({ + vectorBucketName: bucket, + indexName: emptyIndex, + dataType: 'float32', + dimension: 2, + distanceMetric: 'cosine', + }) + try { + const result = await store.listVectors({ + vectorBucketName: bucket, + indexName: emptyIndex, + maxResults: 1, + }) + + expect(result.vectors).toEqual([]) + expect(result.nextToken).toBeUndefined() + } finally { + await store.deleteVectorIndex({ vectorBucketName: bucket, indexName: emptyIndex }) + } + }) + + it('returns an empty page when nextToken references a deleted vector with no later keys', async (ctx) => { + if (!pgvectorAvailable) return ctx.skip() + const deletedCursorIndex = `tenant-a-deleted-cursor-${Date.now()}` + await store.createVectorIndex({ + vectorBucketName: bucket, + indexName: deletedCursorIndex, + dataType: 'float32', + dimension: 2, + distanceMetric: 'cosine', + }) + try { + await store.putVectors({ + vectorBucketName: bucket, + indexName: deletedCursorIndex, + vectors: [{ key: 'deleted-cursor', data: { float32: [1, 0] } }], + }) + await store.deleteVectors({ + vectorBucketName: bucket, + indexName: deletedCursorIndex, + keys: ['deleted-cursor'], + }) + + const result = await store.listVectors({ + vectorBucketName: bucket, + indexName: deletedCursorIndex, + maxResults: 1, + nextToken: 'deleted-cursor', + }) + + expect(result.vectors).toEqual([]) + expect(result.nextToken).toBeUndefined() + } finally { + await store.deleteVectorIndex({ + vectorBucketName: bucket, + indexName: deletedCursorIndex, + }) + } + }) it('returns NotFound when querying a missing index', async (ctx) => { if (!pgvectorAvailable) return ctx.skip() diff --git a/src/test/sharding.test.ts b/src/test/sharding.test.ts index f8e89e348..c260fc384 100644 --- a/src/test/sharding.test.ts +++ b/src/test/sharding.test.ts @@ -191,6 +191,28 @@ describe('Sharding System', () => { expect(resv.status).toBe('cancelled') }) + it('should not cancel a confirmed reservation', async () => { + const reservation = await catalog.reserve({ + kind: 'vector', + tenantId: 'tenant-1', + bucketName: 'bucket-1', + logicalName: 'index-1', + }) + + await catalog.confirm(reservation.reservationId, { + kind: 'vector', + tenantId: 'tenant-1', + bucketName: 'bucket-1', + logicalName: 'index-1', + }) + + await catalog.cancel(reservation.reservationId) + + const resv = await db('shard_reservation').where({ id: reservation.reservationId }).first() + + expect(resv.status).toBe('confirmed') + }) + it('should throw error when confirming non-existent reservation', async () => { await expect( catalog.confirm(randomUUID(), { diff --git a/src/test/vectors-pgvector.test.ts b/src/test/vectors-pgvector.test.ts index 3869f28d6..a9189e5cd 100644 --- a/src/test/vectors-pgvector.test.ts +++ b/src/test/vectors-pgvector.test.ts @@ -1,5 +1,11 @@ +import { createHash } from 'node:crypto' import { BucketScopedSingleShard } from '@internal/sharding' -import { KnexVectorMetadataDB, PgVectorStore, VectorStoreManager } from '@storage/protocols/vector' +import { + createVectorTransactionKnexResolver, + KnexVectorMetadataDB, + PgVectorStore, + VectorStoreManager, +} from '@storage/protocols/vector' import Knex from 'knex' import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' @@ -19,6 +25,18 @@ const TEST_DATABASE_URL = process.env.DATABASE_URL || 'postgresql://postgres:postgres@127.0.0.1/postgres' +function quoteIdent(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"` +} + +function pgvectorTableName(vectorBucketName: string, indexName: string): string { + const hash = createHash('sha256') + .update(`${vectorBucketName}\x00${indexName}`) + .digest('hex') + .slice(0, 24) + return `vector_${hash}` +} + let pgvectorAvailable = false const tenantId = 'pgvector-it-tenant' @@ -69,16 +87,11 @@ describe('Vectors via VectorStoreManager + real pgvector', () => { afterAll(async () => { if (!pgvectorAvailable) return - // Best-effort cleanup; the per-test cleanups already drop indexes/buckets. - try { - await metadataDb - .withTransaction(async (tx) => { - await tx.deleteVectorBucket(bucketName) - }) - .catch(() => undefined) - } catch { - /* swallow */ - } + await metadataDb + .withTransaction(async (tx) => { + await tx.deleteVectorBucket(bucketName) + }) + .catch(() => undefined) await knex.destroy() }) @@ -86,6 +99,57 @@ describe('Vectors via VectorStoreManager + real pgvector', () => { if (!pgvectorAvailable) return ctx.skip() }) + it('persists metadata when transactional createIndex adopts an existing physical table', async () => { + const bucket = `pgvector-it-adopt-${Date.now()}` + const indexName = `it-adopt-${Date.now()}` + const physicalBucket = `pgvector__${bucket}` + const physicalIndex = `${tenantId}-${indexName}` + const physicalTable = pgvectorTableName(physicalBucket, physicalIndex) + const qualifiedTable = `storage_vectors.${quoteIdent(physicalTable)}` + const shard = new BucketScopedSingleShard({ + keyPrefix: 'pgvector__', + capacity: Number.MAX_SAFE_INTEGER, + }) + const transactionalManager = new VectorStoreManager( + new PgVectorStore(createVectorTransactionKnexResolver(knex)), + metadataDb, + shard, + { + tenantId, + maxBucketCount: Infinity, + maxIndexCount: Infinity, + } + ) + + try { + await knex('storage.buckets_vectors').insert({ id: bucket }).onConflict('id').ignore() + await knex.raw(`CREATE TABLE ${qualifiedTable} + ( + key text PRIMARY KEY, + embedding halfvec(4) NOT NULL, + metadata jsonb NOT NULL DEFAULT '{}'::jsonb + )`) + + await transactionalManager.createVectorIndex({ + vectorBucketName: bucket, + indexName, + dataType: 'float32', + dimension: 4, + distanceMetric: 'cosine', + }) + + const metadataRows = await knex('storage.vector_indexes') + .where({ bucket_id: bucket, name: indexName }) + .select('name', 'bucket_id') + + expect(metadataRows).toEqual([{ name: indexName, bucket_id: bucket }]) + } finally { + await knex.raw(`DROP TABLE IF EXISTS ${qualifiedTable}`) + await knex('storage.vector_indexes').where({ bucket_id: bucket, name: indexName }).del() + await knex('storage.buckets_vectors').where({ id: bucket }).del() + } + }) + it('createBucket → createVectorIndex → putVectors → queryVectors → getVectors → deleteVectors → deleteIndex', async () => { const indexName = `it-index-${Date.now()}` await manager.createBucket(bucketName) @@ -228,7 +292,7 @@ describe('Vectors via VectorStoreManager + real pgvector', () => { await manager.putVectors({ vectorBucketName: bucketName, indexName, - vectors: Array.from({ length: 5 }, (_, i) => ({ + vectors: Array.from({ length: 4 }, (_, i) => ({ key: `k-${i.toString().padStart(2, '0')}`, data: { float32: [i, 0] }, metadata: { i }, @@ -250,11 +314,64 @@ describe('Vectors via VectorStoreManager + real pgvector', () => { nextToken: firstPage.nextToken, }) expect(secondPage.vectors).toHaveLength(2) - expect(secondPage.vectors?.[0].key).not.toBe(firstPage.vectors?.[0].key) + expect(secondPage.vectors?.map((vector) => vector.key)).toEqual(['k-02', 'k-03']) + expect(secondPage.nextToken).toBeUndefined() await manager.deleteIndex({ vectorBucketName: bucketName, indexName }) }, 30_000) + it('listVectors returns disjoint keysets for parallel segments', async () => { + const segmentBucketName = `pgvector-it-segments-${Date.now()}` + const indexName = `it-segments-${Date.now()}` + const insertedKeys = Array.from({ length: 32 }, (_, i) => `k-${i.toString().padStart(2, '0')}`) + + await manager.createBucket(segmentBucketName) + await manager.createVectorIndex({ + vectorBucketName: segmentBucketName, + indexName, + dataType: 'float32', + dimension: 2, + distanceMetric: 'euclidean', + }) + + try { + await manager.putVectors({ + vectorBucketName: segmentBucketName, + indexName, + vectors: insertedKeys.map((key, i) => ({ + key, + data: { float32: [i, 0] }, + })), + }) + + const segments = await Promise.all( + [0, 1].map((segmentIndex) => + manager.listVectors({ + vectorBucketName: segmentBucketName, + indexName, + maxResults: 1000, + segmentCount: 2, + segmentIndex, + }) + ) + ) + const segmentKeys = segments.map( + (segment) => segment.vectors?.map((vector) => vector.key) ?? [] + ) + const combinedKeys = segmentKeys.flat() + + expect(segmentKeys[0].length).toBeGreaterThan(0) + expect(segmentKeys[1].length).toBeGreaterThan(0) + expect(new Set(combinedKeys).size).toBe(combinedKeys.length) + expect(combinedKeys.sort()).toEqual(insertedKeys) + } finally { + await manager + .deleteIndex({ vectorBucketName: segmentBucketName, indexName }) + .catch(() => undefined) + await manager.deleteBucket(segmentBucketName).catch(() => undefined) + } + }, 30_000) + it('rejects creating an index inside a missing bucket', async () => { await expect( manager.createVectorIndex({ diff --git a/src/test/vectors.test.ts b/src/test/vectors.test.ts index b5b07cbc6..62fec7d29 100644 --- a/src/test/vectors.test.ts +++ b/src/test/vectors.test.ts @@ -7,6 +7,7 @@ import { QueryVectorsOutput, } from '@aws-sdk/client-s3vectors' import { signJWT } from '@internal/auth' +import { ERRORS } from '@internal/errors' import { SingleShard } from '@internal/sharding' import { KnexVectorMetadataDB, VectorStore, VectorStoreManager } from '@storage/protocols/vector' import { FastifyInstance } from 'fastify' @@ -166,6 +167,30 @@ describe('Vectors API', () => { ) }) + it('should allow route validation for S3Vectors create index requests with 4096 dimensions', async () => { + const request = { + ...validCreateIndexRequest, + dimension: 4096, + indexName: 'test-index-4096', + } + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/CreateIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: request, + }) + + expect(response.statusCode).toBe(200) + expect(mockVectorStore.createVectorIndex).toHaveBeenCalledWith({ + ...request, + vectorBucketName: vectorBucketS3, + indexName: `${tenantId}-test-index-4096`, + }) + }) + it('should require authentication with service role', async () => { const response = await appInstance.inject({ method: 'POST', @@ -270,6 +295,47 @@ describe('Vectors API', () => { // Vector service not called when validation fails }) + it('should reject numeric string dimension without coercing it', async () => { + const invalidRequest = { + ...validCreateIndexRequest, + dimension: '1536', + } + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/CreateIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: invalidRequest, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.createVectorIndex).not.toHaveBeenCalled() + }) + + it.each([ + ['fractional', 3.5], + ['too large', 4097], + ])('should reject %s dimension before creating index metadata', async (_label, dimension) => { + const invalidRequest = { + ...validCreateIndexRequest, + dimension, + } + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/CreateIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: invalidRequest, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.createVectorIndex).not.toHaveBeenCalled() + }) + it('should validate metadataConfiguration structure', async () => { const invalidRequest = { ...validCreateIndexRequest, @@ -292,6 +358,34 @@ describe('Vectors API', () => { // Vector service not called when validation fails }) + it.each([ + ['empty nonFilterableMetadataKeys', []], + ['duplicate nonFilterableMetadataKeys', ['key1', 'key1']], + ['too many nonFilterableMetadataKeys', Array.from({ length: 11 }, (_, i) => `key-${i}`)], + ['empty nonFilterableMetadataKeys item', ['']], + ['too long nonFilterableMetadataKeys item', ['x'.repeat(64)]], + ['numeric nonFilterableMetadataKeys item', [123]], + ])('should validate metadataConfiguration %s', async (_label, nonFilterableMetadataKeys) => { + const invalidRequest = { + ...validCreateIndexRequest, + metadataConfiguration: { + nonFilterableMetadataKeys, + }, + } + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/CreateIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: invalidRequest, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.createVectorIndex).not.toHaveBeenCalled() + }) + it('should handle vector service not configured', async () => { mergeConfig({ vectorEnabled: false }) @@ -459,7 +553,28 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) - it('should handle duplicate bucket creation gracefully', async () => { + it('should reject numeric bucket names without coercing them', async () => { + const numericBucketName = Date.now() + + try { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/CreateVectorBucket', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName: numericBucketName, + }, + }) + + expect(response.statusCode).toBe(400) + } finally { + await s3Vector.deleteBucket(String(numericBucketName)).catch(() => undefined) + } + }) + + it('should return conflict for duplicate bucket creation', async () => { // First creation const newVectorBucketName = `test-bucket-${Date.now()}-dup` const response1 = await appInstance.inject({ @@ -576,6 +691,36 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject numeric bucket names without coercing and deleting them', async () => { + const numericBucketName = Date.now() + await s3Vector.createBucket(String(numericBucketName)) + + try { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/DeleteVectorBucket', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName: numericBucketName, + }, + }) + + expect(response.statusCode).toBe(400) + + const bucketRecord = await storageTest.database.connection.pool + .acquire() + .table('storage.buckets_vectors') + .where({ id: String(numericBucketName) }) + .first() + + expect(bucketRecord).toBeDefined() + } finally { + await s3Vector.deleteBucket(String(numericBucketName)) + } + }) + it('should handle non-existent bucket', async () => { const response = await appInstance.inject({ method: 'POST', @@ -644,6 +789,23 @@ describe('Vectors API', () => { } }) + it('should reject numeric and stringified controls without coercing them', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectorBuckets', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + maxResults: '1', + nextToken: 123, + prefix: 456, + }, + }) + + expect(response.statusCode).toBe(400) + }) + it('should support pagination with nextToken', async () => { const response1 = await appInstance.inject({ method: 'POST', @@ -907,6 +1069,28 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject numeric bucket names without coercing them', async () => { + const numericBucketName = Date.now() + await s3Vector.createBucket(String(numericBucketName)) + + try { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/GetVectorBucket', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName: numericBucketName, + }, + }) + + expect(response.statusCode).toBe(400) + } finally { + await s3Vector.deleteBucket(String(numericBucketName)) + } + }) + it('should handle non-existent bucket', async () => { const response = await appInstance.inject({ method: 'POST', @@ -1004,6 +1188,51 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject numeric index names without coercing and deleting them', async () => { + const numericIndexName = Date.now() + await s3Vector.createVectorIndex({ + dataType: 'float32', + dimension: 1536, + distanceMetric: 'cosine', + indexName: String(numericIndexName), + vectorBucketName, + }) + + try { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/DeleteIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + indexName: numericIndexName, + vectorBucketName, + }, + }) + + expect(response.statusCode).toBe(400) + + const indexRecord = await storageTest.database.connection.pool + .acquire() + .table('storage.vector_indexes') + .where({ + name: String(numericIndexName), + bucket_id: vectorBucketName, + }) + .first() + + expect(indexRecord).toBeDefined() + } finally { + await s3Vector + .deleteIndex({ + indexName: String(numericIndexName), + vectorBucketName, + }) + .catch(() => undefined) + } + }) + it('should validate indexName pattern', async () => { const response = await appInstance.inject({ method: 'POST', @@ -1116,6 +1345,24 @@ describe('Vectors API', () => { expect(body.indexes.length).toBeLessThanOrEqual(1) }) + it('should reject numeric and stringified controls without coercing them', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/ListIndexes', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + maxResults: '1', + nextToken: 123, + prefix: 456, + }, + }) + + expect(response.statusCode).toBe(400) + }) + it('supports pagination with nextToken for indexes', async () => { const page1Response = await appInstance.inject({ method: 'POST', @@ -1465,6 +1712,43 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject numeric index names without coercing them', async () => { + const numericIndexName = Date.now() + await s3Vector.createVectorIndex({ + dataType: 'float32', + dimension: 1536, + distanceMetric: 'cosine', + indexName: String(numericIndexName), + vectorBucketName, + metadataConfiguration: { + nonFilterableMetadataKeys: ['key1'], + }, + }) + + try { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/GetIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + indexName: numericIndexName, + vectorBucketName, + }, + }) + + expect(response.statusCode).toBe(400) + } finally { + await s3Vector + .deleteIndex({ + indexName: String(numericIndexName), + vectorBucketName, + }) + .catch(() => undefined) + } + }) + it('should validate indexName pattern', async () => { const response = await appInstance.inject({ method: 'POST', @@ -1540,10 +1824,13 @@ describe('Vectors API', () => { float32: [1.0, 2.0, 3.0], }, metadata: { + active: true, category: 'test', + score: 0.75, }, }, { + key: 'vec2', data: { float32: [4.0, 5.0, 6.0], }, @@ -1564,14 +1851,60 @@ describe('Vectors API', () => { float32: [1.0, 2.0, 3.0], }, metadata: { + active: true, category: 'test', + score: 0.75, }, }, { + key: 'vec2', data: { float32: [4.0, 5.0, 6.0], }, - key: undefined, + }, + ], + vectorBucketName: vectorBucketS3, + }) + }) + + it('should accept list-valued metadata', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: 'vec-list-metadata', + data: { + float32: [1.0, 2.0, 3.0], + }, + metadata: { + active: true, + tags: ['docs', 'search', 2026, false], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(200) + expect(mockVectorStore.putVectors).toHaveBeenCalledWith({ + indexName: `${tenantId}-${indexName}`, + vectors: [ + { + key: 'vec-list-metadata', + data: { + float32: [1.0, 2.0, 3.0], + }, + metadata: { + active: true, + tags: ['docs', 'search', 2026, false], + }, }, ], vectorBucketName: vectorBucketS3, @@ -1637,13 +1970,7 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) - it('should validate maxItems limit', async () => { - const tooManyVectors = Array.from({ length: 501 }, (_, i) => ({ - data: { - float32: [1.0, 2.0, 3.0], - }, - })) - + it('should return InvalidRequest for route validation failures', async () => { const response = await appInstance.inject({ method: 'POST', url: '/vector/PutVectors', @@ -1653,14 +1980,26 @@ describe('Vectors API', () => { payload: { vectorBucketName, indexName, - vector: tooManyVectors, + vectors: [ + { + key: 'invalid-data', + data: { + float32: ['1.0', 2.0, 3.0], + }, + }, + ], }, }) expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() }) - it('should handle non-existent index', async () => { + it('should reject empty float32 data before calling the vector store', async () => { const response = await appInstance.inject({ method: 'POST', url: '/vector/PutVectors', @@ -1669,39 +2008,288 @@ describe('Vectors API', () => { }, payload: { vectorBucketName, - indexName: 'non-existent-index', + indexName, vectors: [ { + key: 'empty-data', data: { - float32: [1.0, 2.0, 3.0], + float32: [], }, }, ], }, }) - expect(response.statusCode).toBe(404) + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('float32') + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() }) - }) - - describe('POST /vector/QueryVectors', () => { - let indexName: string - beforeEach(async () => { - indexName = `test-index-${Date.now()}` - await appInstance.inject({ + it('should reject missing vector keys before calling the vector store', async () => { + const response = await appInstance.inject({ method: 'POST', - url: '/vector/CreateIndex', + url: '/vector/PutVectors', headers: { authorization: `Bearer ${serviceToken}`, }, payload: { - dataType: 'float32', - dimension: 3, - distanceMetric: 'cosine', - indexName, vectorBucketName, - }, + indexName, + vectors: [ + { + data: { + float32: [1.0, 2.0, 3.0], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('key') + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should reject empty vector keys before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: '', + data: { + float32: [1.0, 2.0, 3.0], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('key') + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should reject vector keys above the S3Vectors length limit before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: 'x'.repeat(1025), + data: { + float32: [1.0, 2.0, 3.0], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should reject nested metadata objects before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: 'nested-metadata', + data: { + float32: [1.0, 2.0, 3.0], + }, + metadata: { + nested: { + value: 'not supported', + }, + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should reject oversized filterable metadata before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: 'vec-large-metadata', + data: { + float32: [1.0, 2.0, 3.0], + }, + metadata: { + large: 'x'.repeat(2_050), + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidParameter', + message: + "Invalid record for key 'vec-large-metadata': Filterable metadata must have at most 2048 bytes", + }) + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should validate maxItems limit', async () => { + const tooManyVectors = Array.from({ length: 501 }, (_, i) => ({ + key: `vec-${i}`, + data: { + float32: [1.0, 2.0, 3.0], + }, + })) + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: tooManyVectors, + }, + }) + + expect(response.statusCode).toBe(400) + expect(response.body).toContain('must NOT have more than 500 items') + expect(mockVectorStore.putVectors).not.toHaveBeenCalled() + }) + + it('should handle non-existent index', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName: 'non-existent-index', + vectors: [ + { + key: 'vec1', + data: { + float32: [1.0, 2.0, 3.0], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(404) + }) + + it('should return validation failures from the vector store as HTTP 400', async () => { + const message = + "Invalid record for key '5797803-0': Filterable metadata must have at most 2048 bytes" + mockVectorStore.putVectors.mockRejectedValueOnce( + ERRORS.InvalidParameter(indexName, { message }) + ) + + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/PutVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + vectors: [ + { + key: '5797803-0', + data: { + float32: [1.0, 2.0, 3.0], + }, + }, + ], + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidParameter', + error: 'InvalidParameter', + message, + }) + }) + }) + + describe('POST /vector/QueryVectors', () => { + let indexName: string + + beforeEach(async () => { + indexName = `test-index-${Date.now()}` + await appInstance.inject({ + method: 'POST', + url: '/vector/CreateIndex', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + dataType: 'float32', + dimension: 3, + distanceMetric: 'cosine', + indexName, + vectorBucketName, + }, }) mockVectorStore.queryVectors.mockResolvedValue({ @@ -1836,6 +2424,58 @@ describe('Vectors API', () => { expect(mockVectorStore.queryVectors).not.toHaveBeenCalled() }) + it('should reject topK above the documented maximum before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/QueryVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + queryVector: { + float32: [1.0, 2.0, 3.0], + }, + topK: 101, + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('topK') + expect(mockVectorStore.queryVectors).not.toHaveBeenCalled() + }) + + it('should reject fractional topK before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/QueryVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + queryVector: { + float32: [1.0, 2.0, 3.0], + }, + topK: 1.5, + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('topK') + expect(mockVectorStore.queryVectors).not.toHaveBeenCalled() + }) + it('should require authentication with service role', async () => { const response = await appInstance.inject({ method: 'POST', @@ -1895,6 +2535,32 @@ describe('Vectors API', () => { expect(body.message).toContain('float32') }) + it('should reject empty queryVector.float32 before calling the vector store', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/QueryVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + queryVector: { + float32: [], + }, + topK: 10, + }, + }) + + expect(response.statusCode).toBe(400) + expect(parseJsonBody(response.body)).toMatchObject({ + statusCode: '400', + code: 'InvalidRequest', + }) + expect(response.body).toContain('float32') + expect(mockVectorStore.queryVectors).not.toHaveBeenCalled() + }) + it('should handle non-existent index', async () => { const response = await appInstance.inject({ method: 'POST', @@ -1993,6 +2659,60 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject more than 500 vector keys', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/DeleteVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: Array.from({ length: 501 }, (_, i) => `vec-${i}`), + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.deleteVectors).not.toHaveBeenCalled() + }) + + it('should reject vector keys above the S3Vectors length limit', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/DeleteVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: ['x'.repeat(1025)], + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.deleteVectors).not.toHaveBeenCalled() + }) + + it('should reject numeric vector keys without coercing them', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/DeleteVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: [123], + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.deleteVectors).not.toHaveBeenCalled() + }) + it('should handle non-existent index', async () => { const response = await appInstance.inject({ method: 'POST', @@ -2184,7 +2904,7 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) - it('should validate maxResults range', async () => { + it('should allow the S3Vectors maxResults upper bound', async () => { const response = await appInstance.inject({ method: 'POST', url: '/vector/ListVectors', @@ -2194,11 +2914,55 @@ describe('Vectors API', () => { payload: { vectorBucketName, indexName, - maxResults: 501, + maxResults: 1000, + }, + }) + + expect(response.statusCode).toBe(200) + expect(mockVectorStore.listVectors).toHaveBeenCalledWith( + expect.objectContaining({ + maxResults: 1000, + }) + ) + }) + + it('should reject maxResults above the S3Vectors limit', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + maxResults: 1001, }, }) expect(response.statusCode).toBe(400) + expect(mockVectorStore.listVectors).not.toHaveBeenCalled() + }) + + it('should reject numeric and boolean strings without coercing them', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + maxResults: '1000', + returnData: 'true', + segmentCount: '2', + segmentIndex: '1', + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.listVectors).not.toHaveBeenCalled() }) it('should validate segmentIndex range', async () => { @@ -2219,6 +2983,58 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should require segmentCount and segmentIndex together', async () => { + const segmentCountOnly = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + segmentCount: 4, + }, + }) + + expect(segmentCountOnly.statusCode).toBe(400) + + const segmentIndexOnly = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + segmentIndex: 0, + }, + }) + + expect(segmentIndexOnly.statusCode).toBe(400) + expect(mockVectorStore.listVectors).not.toHaveBeenCalled() + }) + + it('should reject segmentIndex greater than or equal to segmentCount', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/ListVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + segmentCount: 4, + segmentIndex: 4, + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.listVectors).not.toHaveBeenCalled() + }) + it('should handle non-existent index', async () => { const response = await appInstance.inject({ method: 'POST', @@ -2359,6 +3175,61 @@ describe('Vectors API', () => { expect(response.statusCode).toBe(400) }) + it('should reject more than 100 vector keys', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/GetVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: Array.from({ length: 101 }, (_, i) => `vec-${i}`), + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.getVectors).not.toHaveBeenCalled() + }) + + it('should reject vector keys above the S3Vectors length limit', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/GetVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: ['x'.repeat(1025)], + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.getVectors).not.toHaveBeenCalled() + }) + + it('should reject numeric vector keys and boolean strings without coercing them', async () => { + const response = await appInstance.inject({ + method: 'POST', + url: '/vector/GetVectors', + headers: { + authorization: `Bearer ${serviceToken}`, + }, + payload: { + vectorBucketName, + indexName, + keys: [123], + returnData: 'true', + }, + }) + + expect(response.statusCode).toBe(400) + expect(mockVectorStore.getVectors).not.toHaveBeenCalled() + }) + it('should handle non-existent index', async () => { const response = await appInstance.inject({ method: 'POST', diff --git a/src/test/webhooks.test.ts b/src/test/webhooks.test.ts index 8023d2c37..c64586a64 100644 --- a/src/test/webhooks.test.ts +++ b/src/test/webhooks.test.ts @@ -449,10 +449,11 @@ describe('Webhooks', () => { }, }) + const objectCreatedAt = new Date(Date.now() - 1_000) const objects = await Promise.all([ - createObject(pg, emptyTestBucketName), - createObject(pg, emptyTestBucketName), - createObject(pg, emptyTestBucketName), + createObject(pg, emptyTestBucketName, objectCreatedAt), + createObject(pg, emptyTestBucketName, objectCreatedAt), + createObject(pg, emptyTestBucketName, objectCreatedAt), ]) const response = await appInstance.inject({ @@ -678,9 +679,10 @@ describe('Webhooks', () => { }) }) -async function createObject(pg: TenantConnection, bucketId: string) { +async function createObject(pg: TenantConnection, bucketId: string, createdAt?: Date) { const objectName = randomUUID() const tnx = await pg.transaction() + const createdAtIso = createdAt?.toISOString() const [data] = await tnx .from('objects') @@ -689,6 +691,9 @@ async function createObject(pg: TenantConnection, bucketId: string) { name: objectName.toString(), bucket_id: bucketId, version: randomUUID(), + created_at: createdAtIso, + updated_at: createdAtIso, + last_accessed_at: createdAtIso, metadata: { cacheControl: 'no-cache', contentLength: 3746,