feat: vector buckets (#774)

* feat: vector buckets
* feat: sharding of vector buckets

Author: fenos

.env.test.sample

@@ -22,4 +22,8 @@ AWS_DEFAULT_REGION=ap-southeast-1
 STORAGE_S3_ENDPOINT=http://127.0.0.1:9000
 STORAGE_S3_PROTOCOL=http
 STORAGE_S3_FORCE_PATH_STYLE=true
-REQUEST_X_FORWARDED_HOST_REGEXP=
+REQUEST_X_FORWARDED_HOST_REGEXP=
+
+VECTOR_ENABLED=true
+ICEBERG_ENABLED=true
+ICEBERG_BUCKET_DETECTION_MODE="BUCKET"

.eslintrc.js

@@ -1,18 +1,18 @@
 module.exports = {
-  ignorePatterns: ['src/test/assets/**', 'src/test/db/**', 'src/test/*.yaml'],
+  ignorePatterns: ['src/test/assets/**', 'src/test/db/**', 'src/test/*.yaml', 'src/**/**/*.md'],
   parser: '@typescript-eslint/parser',
   extends: ['plugin:@typescript-eslint/recommended', 'plugin:prettier/recommended'],
   parserOptions: {
     ecmaVersion: 2020, // Allows for the parsing of modern ECMAScript features
     sourceType: 'module', // Allows for the use of imports
-    "project": "./tsconfig.json",
+    project: './tsconfig.json',
   },
   rules: {
     '@typescript-eslint/no-floating-promises': 'error',
     '@typescript-eslint/no-explicit-any': 'warn',
     '@typescript-eslint/no-unused-vars': [
       'warn',
-      { 'argsIgnorePattern': '^_+$', 'varsIgnorePattern': '^_+$' } // allows intentionally unused variables named _
+      { argsIgnorePattern: '^_+$', varsIgnorePattern: '^_+$' }, // allows intentionally unused variables named _
     ],
     '@typescript-eslint/no-require-imports': 'warn',
   },

.github/workflows/ci.yml

@@ -86,6 +86,9 @@ jobs:
           MULTI_TENANT: false
           S3_PROTOCOL_ACCESS_KEY_ID: ${{ secrets.TENANT_ID }}
           S3_PROTOCOL_ACCESS_KEY_SECRET: ${{ secrets.SERVICE_KEY }}
+          VECTOR_S3_BUCKETS: supa-test-local-dev
+          VECTOR_ENABLED: true
+          ICEBERG_ENABLED: true
 
       - name: Upload coverage results to Coveralls
         uses: coverallsapp/github-action@master

.gitignore

@@ -8,4 +8,5 @@ static/api.json
 data/
 bin/
 coverage/
-.idea/
+.idea/
+src/scripts/*.py

migrations/multitenant/0020-vector-buckets-feature.sql

@@ -0,0 +1,3 @@
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS feature_vector_buckets boolean NOT NULL DEFAULT false;
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS feature_vector_buckets_max_buckets int NOT NULL DEFAULT 10;
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS feature_vector_buckets_max_indexes int NOT NULL DEFAULT 5;

migrations/multitenant/0021-sharding-resources.sql

@@ -0,0 +1,78 @@
+
+
+-- Main shards table.
+CREATE TABLE IF NOT EXISTS shard (
+    id           BIGSERIAL PRIMARY KEY,
+    kind         TEXT NOT NULL,
+    shard_key    TEXT NOT NULL,
+    capacity     INT  NOT NULL DEFAULT 10000,
+    next_slot    INT  NOT NULL DEFAULT 0,
+    status       TEXT NOT NULL DEFAULT 'active',    -- active|draining|disabled
+    created_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
+    UNIQUE (kind, shard_key)
+);
+
+-- Sparse slot rows: only the slots that have ever been used exist here.
+-- A "free" slot is a row with reservation_id NULL and resource_id NULL.
+CREATE TABLE IF NOT EXISTS shard_slots (
+    shard_id       BIGINT NOT NULL REFERENCES shard(id) ON DELETE CASCADE,
+    slot_no        INT    NOT NULL,
+    tenant_id      TEXT,
+    resource_id    TEXT,               -- set when confirmed
+    PRIMARY KEY (shard_id, slot_no)
+);
+
+-- Reservations with short leases
+CREATE TABLE IF NOT EXISTS shard_reservation (
+    id               UUID PRIMARY KEY default gen_random_uuid(),
+    kind             text NOT NULL,
+    tenant_id        TEXT,
+    resource_id      TEXT NOT NULL,               -- e.g. "vector::bucket::name"
+    shard_id         BIGINT NOT NULL,
+    slot_no          INT NOT NULL,
+    status           TEXT NOT NULL DEFAULT 'pending',  -- pending|confirmed|expired|cancelled
+    lease_expires_at TIMESTAMPTZ NOT NULL,
+    created_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
+    UNIQUE (kind, resource_id),
+    UNIQUE (shard_id, slot_no)
+);
+
+-- Fast “used count” per shard
+CREATE INDEX IF NOT EXISTS shard_slots_used_idx
+    ON shard_slots (shard_id)
+    WHERE resource_id IS NOT NULL;
+
+ALTER TABLE shard
+    ADD CONSTRAINT shard_capacity_not_less_than_minted
+        CHECK (capacity >= next_slot);
+
+
+-- Create index for counting slots by tenant
+CREATE INDEX IF NOT EXISTS shard_slots_tenant_id_idx
+    ON shard_slots (tenant_id);
+
+-- Create index for counting reservations by tenant
+CREATE INDEX IF NOT EXISTS shard_reservation_tenant_id_idx
+    ON shard_reservation (tenant_id);
+
+-- Create index for counting used slots by tenant
+CREATE INDEX IF NOT EXISTS shard_slots_tenant_resource_idx
+    ON shard_slots (tenant_id, shard_id)
+    WHERE resource_id IS NOT NULL;
+
+
+ALTER TABLE shard_reservation
+    ADD CONSTRAINT fk_shard_slot
+        FOREIGN KEY (shard_id, slot_no)
+            REFERENCES shard_slots(shard_id, slot_no)
+            ON DELETE RESTRICT;
+
+
+CREATE INDEX IF NOT EXISTS shard_slots_free_idx
+    ON shard_slots (shard_id, slot_no)
+    WHERE resource_id IS NULL;
+
+-- Add index for finding active reservations by slot
+CREATE INDEX IF NOT EXISTS shard_reservation_active_slot_idx
+    ON shard_reservation (shard_id, slot_no, lease_expires_at)
+    WHERE status = 'pending';

migrations/tenant/0044-vector-bucket-type.sql

@@ -0,0 +1,13 @@
+DO $$
+    DECLARE
+    BEGIN
+        IF NOT EXISTS (
+            SELECT 1
+            FROM pg_enum
+                     JOIN pg_type ON pg_enum.enumtypid = pg_type.oid
+            WHERE pg_type.typname = 'buckettype'
+              AND enumlabel = 'VECTOR'
+        ) THEN
+            ALTER TYPE storage.BucketType ADD VALUE 'VECTOR';
+        END IF;
+END$$;

migrations/tenant/0045-vector-buckets.sql

@@ -0,0 +1,34 @@
+DO $$
+    DECLARE
+        anon_role text = COALESCE(current_setting('storage.anon_role', true), 'anon');
+        authenticated_role text = COALESCE(current_setting('storage.authenticated_role', true), 'authenticated');
+        service_role text = COALESCE(current_setting('storage.service_role', true), 'service_role');
+    BEGIN
+        CREATE TABLE IF NOT EXISTS storage.buckets_vectors (
+            id text not null primary key,
+            type storage.BucketType NOT NULL default 'VECTOR',
+            created_at timestamptz NOT NULL default now(),
+            updated_at timestamptz NOT NULL default now()
+        );
+
+        CREATE TABLE IF NOT EXISTS storage.vector_indexes
+        (
+            id                     text             primary key default gen_random_uuid(),
+            name                   text COLLATE "C" NOT NULL,
+            bucket_id              text             NOT NULL references storage.buckets_vectors (id),
+            data_type              text             NOT NULL,
+            dimension              integer          NOT NULL,
+            distance_metric        text             NOT NULL,
+            metadata_configuration jsonb            NULL,
+            created_at             timestamptz      NOT NULL default now(),
+            updated_at             timestamptz      NOT NULL default now()
+        );
+
+        ALTER TABLE storage.buckets_vectors ENABLE ROW LEVEL SECURITY;
+        ALTER TABLE storage.vector_indexes ENABLE ROW LEVEL SECURITY;
+
+        EXECUTE 'GRANT SELECT ON TABLE storage.buckets_vectors TO ' || service_role || ', ' || authenticated_role || ', ' || anon_role;
+        EXECUTE 'GRANT SELECT ON TABLE storage.vector_indexes TO ' || service_role || ', ' || authenticated_role || ', ' || anon_role;
+
+        CREATE UNIQUE INDEX IF NOT EXISTS vector_indexes_name_bucket_id_idx ON storage.vector_indexes (name, bucket_id);
+END$$;