feat: add analytics bucket sharding

Author: fenos

migrations/multitenant/0022-iceberg-catalog-sharding.sql

@@ -0,0 +1,62 @@
+DO $$
+    DECLARE
+        iceberg_shards text[] = COALESCE(current_setting('storage.iceberg_shards', true), '[]::text[]')::text[];
+        iceberg_default_shard text = COALESCE(current_setting('storage.iceberg_default_shard', true), '')::text;
+        i_shard_key text;
+    BEGIN
+
+        ALTER TABLE iceberg_namespaces ADD COLUMN IF NOT EXISTS metadata JSONB NOT NULL DEFAULT '{}';
+
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS remote_table_id TEXT NULL;
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS shard_key TEXT NULL;
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS shard_id bigint NULL;
+
+        -- Only allow deleting namespaces if empty
+        ALTER TABLE iceberg_tables DROP CONSTRAINT IF EXISTS iceberg_tables_namespace_id_fkey;
+
+        ALTER TABLE iceberg_tables
+            ADD CONSTRAINT iceberg_tables_namespace_id_fkey
+                FOREIGN KEY (namespace_id)
+                    REFERENCES iceberg_namespaces(id) ON DELETE RESTRICT;
+
+
+        FOREACH i_shard_key IN ARRAY iceberg_shards
+            LOOP
+                INSERT INTO shard (kind, shard_key, capacity) VALUES ('iceberg-table', i_shard_key, 1000)
+                ON CONFLICT (kind, shard_key) DO NOTHING;
+            END LOOP;
+
+
+        IF array_length(iceberg_shards, 1) = 0 THEN
+            RETURN;
+        END IF;
+
+        UPDATE iceberg_tables
+        SET shard_id = (
+            SELECT id FROM shard WHERE kind = 'iceberg-table' AND shard_key = iceberg_default_shard LIMIT 1
+        ), shard_key = iceberg_default_shard
+        WHERE shard_id IS NULL;
+
+        WITH all_iceberg_tables as (
+            SELECT t.id, t.tenant_id, t.name, t.shard_key, t.shard_id, t.namespace_id, row_number() OVER () as seq_num
+            FROM iceberg_tables t
+        ),
+        set_shard_reservation AS (
+             INSERT INTO shard_reservation (tenant_id, shard_id, kind, resource_id, lease_expires_at, slot_no, status)
+                 SELECT it.tenant_id, s.id, 'iceberg-table', (it.namespace_id || '/' || it.name), now() + interval '5 minutes', it.seq_num - 1, 'confirmed'
+                 FROM all_iceberg_tables it
+                          JOIN shard s ON s.kind = 'iceberg-table' AND s.shard_key = it.shard_key
+        ),
+        shard_slot AS (
+         INSERT INTO shard_slots (tenant_id, shard_id, resource_id, slot_no)
+             SELECT it.tenant_id, it.shard_id, it.id, it.seq_num - 1
+             FROM all_iceberg_tables it
+             RETURNING slot_no
+        )
+        UPDATE shard
+        SET next_slot = (SELECT COALESCE((SELECT MAX(slot_no) + 1 FROM shard_slot), next_slot))
+        WHERE shard.kind = 'iceberg-table'
+          AND shard.shard_key = iceberg_default_shard;
+
+    END
+$$;

migrations/tenant/0046-iceberg-table-metadata.sql

@@ -0,0 +1,16 @@
+DO $$
+    DECLARE
+        is_multitenant bool = COALESCE(current_setting('storage.multitenant', true), 'false')::boolean;
+    BEGIN
+
+        IF is_multitenant THEN
+            RETURN;
+        END IF;
+
+        ALTER TABLE iceberg_namespaces ADD COLUMN IF NOT EXISTS metadata JSONB NOT NULL DEFAULT '{}';
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS remote_table_id TEXT NULL;
+
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS shard_key TEXT NULL;
+        ALTER TABLE iceberg_tables ADD COLUMN IF NOT EXISTS shard_id TEXT NULL;
+END
+$$;

src/config.ts

@@ -178,6 +178,7 @@ type StorageConfigType = {
 
   icebergEnabled: boolean
   icebergWarehouse: string
+  icebergShards: string[]
   icebergCatalogUrl: string
   icebergCatalogAuthType: IcebergCatalogAuthType
   icebergCatalogToken?: string
@@ -515,6 +516,7 @@ export function getConfig(options?: { reload?: boolean }): StorageConfigType {
 
     icebergEnabled: getOptionalConfigFromEnv('ICEBERG_ENABLED') === 'true',
     icebergWarehouse: getOptionalConfigFromEnv('ICEBERG_WAREHOUSE') || '',
+    icebergShards: getOptionalConfigFromEnv('ICEBERG_SHARDS')?.trim().split(',') || [],
     icebergCatalogUrl:
       getOptionalConfigFromEnv('ICEBERG_CATALOG_URL') ||
       `https://s3tables.ap-southeast-1.amazonaws.com/iceberg/v1`,

src/http/plugins/iceberg.ts

@@ -5,6 +5,7 @@ import { getTenantConfig, multitenantKnex } from '@internal/database'
 import { getCatalogAuthStrategy, TenantAwareRestCatalog } from '@storage/protocols/iceberg/catalog'
 import { getConfig } from '../../config'
 import { ICEBERG_BUCKET_RESERVED_SUFFIX } from '@storage/limits'
+import { KnexShardStoreFactory, ShardCatalog, SingleShard } from '@internal/sharding'
 
 declare module 'fastify' {
   interface FastifyRequest {
@@ -48,8 +49,13 @@ export const icebergRestCatalog = fastifyPlugin(async function (fastify: Fastify
       tenantId: req.tenantId,
       limits: limits,
       restCatalogUrl: icebergCatalogUrl,
-      warehouse: icebergWarehouse,
       auth: catalogAuthType,
+      sharding: isMultitenant
+        ? new ShardCatalog(new KnexShardStoreFactory(multitenantKnex))
+        : new SingleShard({
+            shardKey: icebergWarehouse,
+            capacity: 10000,
+          }),
       metastore: new KnexMetastore(isMultitenant ? multitenantKnex : req.db.pool.acquire(), {
         multiTenant: isMultitenant,
         schema: isMultitenant ? 'public' : 'storage',

src/http/routes/iceberg/namespace.ts

@@ -10,6 +10,7 @@ const createNamespaceSchema = {
     type: 'object',
     properties: {
       namespace: { type: 'string', examples: ['namespace'] },
+      properties: { type: 'object', additionalProperties: { type: 'string' } },
     },
     required: ['namespace'],
   },
@@ -104,6 +105,7 @@ export default async function routes(fastify: FastifyInstance) {
       const result = await request.icebergCatalog.createNamespace({
         namespace: [request.body.namespace],
         warehouse: request.params.prefix,
+        properties: request.body.properties,
       })
 
       return response.send(result)
@@ -140,7 +142,7 @@ export default async function routes(fastify: FastifyInstance) {
       config: {
         operation: { type: ROUTE_OPERATIONS.ICEBERG_NAMESPACE_EXISTS },
       },
-      schema: { ...listNamespaceSchema, tags: ['iceberg'] },
+      schema: { ...loadNamespaceSchema, tags: ['iceberg'] },
     },
     async (request, response) => {
       if (!request.icebergCatalog) {

src/http/routes/s3/router.ts

@@ -283,7 +283,7 @@ export class Router<Context = unknown, S extends Schema = Schema> {
       const headerValue = headerParts[1]
 
       const matchHeaderName = received[headerName] !== undefined
-      const matchHeaderValue = headerValue ? received[headerName].startsWith(headerValue) : true
+      const matchHeaderValue = headerValue ? received[headerName]?.startsWith(headerValue) : true
 
       return matchHeaderName && matchHeaderValue
     })

src/internal/database/migrations/migrate.ts

@@ -30,6 +30,7 @@ const {
   dbInstallRoles,
   dbRefreshMigrationHashesOnMismatch,
   dbMigrationFreezeAt,
+  icebergShards,
 } = getConfig()
 
 /**
@@ -640,14 +641,19 @@ function runMigrations({
       const migrationsToRun = filterMigrations(intendedMigrations, appliedMigrations)
       const completedMigrations = []
 
+      const icebergShardVar = `{${icebergShards.map((s) => `"${s}"`).join(',')}}`
+      const icebergDefaultShard = icebergShards.length > 0 ? icebergShards[0] : ''
+
       if (migrationsToRun.length > 0) {
         await client.query(SQL`SELECT 
           set_config('storage.install_roles', ${dbInstallRoles}, false),
           set_config('storage.multitenant', ${isMultitenant ? 'true' : 'false'}, false),
           set_config('storage.anon_role', ${dbAnonRole}, false),
           set_config('storage.authenticated_role', ${dbAuthenticatedRole}, false),
           set_config('storage.service_role', ${dbServiceRole}, false),
-          set_config('storage.super_user', ${dbSuperUser}, false)
+          set_config('storage.super_user', ${dbSuperUser}, false),
+          set_config('storage.iceberg_default_shard', ${icebergDefaultShard}, false),
+          set_config('storage.iceberg_shards', ${icebergShardVar}, false);
         `)
       }
 

src/internal/database/migrations/types.ts

@@ -44,4 +44,5 @@ export const DBMigration = {
   'fix-object-level': 43,
   'vector-bucket-type': 44,
   'vector-buckets': 45,
+  'iceberg-table-metadata': 46,
 }

src/internal/errors/codes.ts

@@ -43,14 +43,16 @@ export enum ErrorCode {
   NotSupported = 'NotSupported',
   IcebergError = 'IcebergError',
   IcebergMaximumResourceLimit = 'IcebergMaximumResourceLimit',
+  IcebergResourceNotEmpty = 'IcebergResourceNotEmpty',
   NoSuchCatalog = 'NoSuchCatalog',
 
   S3VectorConflictException = 'ConflictException',
   S3VectorNotFoundException = 'NotFoundException',
   S3VectorBucketNotEmpty = 'VectorBucketNotEmpty',
   S3VectorMaxBucketsExceeded = 'S3VectorMaxBucketsExceeded',
   S3VectorMaxIndexesExceeded = 'S3VectorMaxIndexesExceeded',
-  S3VectorNoAvailableShard = 'S3VectorNoAvailableShard',
+  NoAvailableShard = 'NoAvailableShard',
+  ShardNotFound = 'ShardNotFound',
 }
 
 export const ERRORS = {
@@ -69,6 +71,13 @@ export const ERRORS = {
       message: `The maximum number of this resource ${limit} is reached`,
       originalError: e,
     }),
+  IcebergResourceNotEmpty: (resource: string, name: string, e?: Error) =>
+    new StorageBackendError({
+      code: ErrorCode.IcebergResourceNotEmpty,
+      httpStatusCode: 400,
+      message: `The resource ${resource}: ${name} is not empty`,
+      originalError: e,
+    }),
   FeatureNotEnabled: (resource: string, feature: string, e?: Error) =>
     new StorageBackendError({
       code: ErrorCode.InvalidRequest,
@@ -472,11 +481,18 @@ export const ERRORS = {
       message: `Maximum number of indexes exceeded. Max allowed is ${maxIndexes}. Contact support to increase your limit.`,
     })
   },
-  S3VectorNoAvailableShard() {
+  NoAvailableShard() {
     return new StorageBackendError({
-      code: ErrorCode.S3VectorNoAvailableShard,
+      code: ErrorCode.NoAvailableShard,
       httpStatusCode: 500,
-      message: `No available shards are available to host the vector index. Please try again later.`,
+      message: `No available shards are available to host the resource. Please try again later.`,
+    })
+  },
+  ShardNotFound(shardId: string) {
+    return new StorageBackendError({
+      code: ErrorCode.ShardNotFound,
+      httpStatusCode: 404,
+      message: `Shard not found: ${shardId}`,
     })
   },
 }

src/internal/sharding/knex.ts

@@ -10,9 +10,18 @@ import {
 } from './store'
 import { hashStringToInt } from '@internal/hashing'
 
-export class KnexShardStoreFactory implements ShardStoreFactory {
+export class KnexShardStoreFactory implements ShardStoreFactory<Knex.Transaction> {
   constructor(private knex: Knex) {}
+
+  withExistingTransaction(tnx: Knex.Transaction): ShardStoreFactory {
+    return new KnexShardStoreFactory(tnx)
+  }
   async withTransaction<T>(fn: (store: ShardStore) => Promise<T>): Promise<T> {
+    if (this.knex.isTransaction) {
+      // Already in a transaction, use current connection
+      return fn(new KnexShardStore(this.knex))
+    }
+
     try {
       return await this.knex.transaction(async (trx) => {
         return fn(new KnexShardStore(trx))
@@ -33,6 +42,11 @@ class KnexShardStore implements ShardStore {
     return this.db.raw<T>(sql, params as any)
   }
 
+  async findShardById(shardId: number): Promise<ShardRow | null> {
+    const shard = await this.db<ShardRow>('shard').select('*').where({ id: shardId }).first()
+    return shard ?? null
+  }
+
   async advisoryLockByString(key: string): Promise<void> {
     const id = hashStringToInt(key)
     await this.q(`SELECT pg_advisory_xact_lock(?::bigint)`, [id])