bunsane 0.1.4 → 0.2.0

package/database/DatabaseHelper.ts

@@ -1,5 +1,7 @@
-import db from "database";
-import { logger as MainLogger } from "core/Logger";
+import db from "./index";
+import { logger as MainLogger } from "../core/Logger";
+import { getMetadataStorage } from "../core/metadata";
+import { ensureMultipleJSONBPathIndexes } from "./IndexingStrategy";
 const logger = MainLogger.child({ scope: "DatabaseHelper" });
 
 const BUNSANE_RELATION_TYPED_COLUMN = process.env.BUNSANE_RELATION_TYPED_COLUMN === 'true' || false;
@@ -21,10 +23,15 @@ const retryWithBackoff = async (fn: () => Promise<void>, maxRetries: number = 3,
         try {
             await fn();
             return;
-        } catch (error) {
+        } catch (error: any) {
+            const isDeadlock = error?.code === '40P01' || error?.message?.includes('deadlock');
             if (i === maxRetries - 1) throw error;
             const delay = baseDelay * Math.pow(2, i);
-            logger.warn(`Operation failed, retrying in ${delay}ms: ${error}`);
+            if (isDeadlock) {
+                logger.warn(`Deadlock detected, retrying in ${delay}ms (attempt ${i + 1}/${maxRetries})`);
+            } else {
+                logger.warn(`Operation failed, retrying in ${delay}ms: ${error}`);
+            }
             await sleep(delay);
         }
     }
@@ -68,6 +75,7 @@ export const PrepareDatabase = async () => {
     }
     try {
         await CreateEntityComponentTable();
+        await PopulateComponentIds();
     } catch (error) {
         logger.error(`Failed to create entity component table: ${error}`);
         throw error;
@@ -103,32 +111,156 @@ export const CreateEntityTable = async () => {
 }
 
 export const CreateComponentTable = async () => {
+    const partitionStrategy = process.env.BUNSANE_PARTITION_STRATEGY === 'hash' ? 'hash' : 'list'; // Default to list (LIST+Direct is the recommended strategy)
+
+    // Check if the table already exists and what partitioning strategy it uses
+    const existingStrategy = await GetPartitionStrategy();
+    const tableExists = await db.unsafe(`
+        SELECT 1 FROM information_schema.tables
+        WHERE table_name = 'components'
+        AND table_schema = 'public'
+    `);
+
+    // If the table exists but has a different partitioning strategy, we need to recreate it
+    if (tableExists.length > 0 && existingStrategy !== partitionStrategy) {
+        logger.info(`Partitioning strategy changed from ${existingStrategy} to ${partitionStrategy}. Recreating components table...`);
+
+        // Drop the existing table and all its partitions
+        await db.unsafe(`DROP TABLE IF EXISTS components CASCADE`);
+
+        // Also clean up any orphaned partition tables
+        const orphanedPartitions = await db.unsafe(`
+            SELECT tablename
+            FROM pg_tables
+            WHERE tablename LIKE 'components_%'
+            AND schemaname = 'public'
+        `);
+
+        for (const partition of orphanedPartitions) {
+            await db.unsafe(`DROP TABLE IF EXISTS ${partition.tablename} CASCADE`);
+        }
+    }
+
+    if (partitionStrategy === 'hash') {
+        // Clean up any existing LIST partition tables before creating HASH partitions
+        await cleanupOldListPartitions();
+        await CreateHashPartitionedComponentTable();
+    } else {
+        // Original LIST partitioning
+        await db`CREATE TABLE IF NOT EXISTS components (
+            id UUID,
+            entity_id UUID REFERENCES entities(id) ON DELETE CASCADE,
+            type_id varchar(64) NOT NULL,
+            name varchar(128),
+            data jsonb,
+            created_at TIMESTAMP DEFAULT NOW(),
+            updated_at TIMESTAMP DEFAULT NOW(),
+            deleted_at TIMESTAMP,
+            PRIMARY KEY (id, type_id),
+            UNIQUE(entity_id, type_id)
+        ) PARTITION BY LIST (type_id);`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_entity_id ON components (entity_id)`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_type_id ON components (type_id)`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_data_gin ON components USING GIN (data)`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_entity_type_deleted ON components (entity_id, type_id, deleted_at)`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_type_deleted ON components (type_id, deleted_at) WHERE deleted_at IS NULL`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_deleted_entity ON components (deleted_at, entity_id) WHERE deleted_at IS NULL`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_entity_created_desc ON components (entity_id, created_at DESC)`;
+        await db`CREATE INDEX IF NOT EXISTS idx_components_type_entity_created ON components (type_id, entity_id, created_at DESC)`;
+    }
+}
+
+const cleanupOldListPartitions = async () => {
+    try {
+        logger.info(`Cleaning up old LIST partition tables before creating HASH partitions`);
+        
+        // Get all existing LIST partition tables
+        const existingPartitions = await db.unsafe(`
+            SELECT tablename 
+            FROM pg_tables 
+            WHERE tablename LIKE 'components_%' 
+            AND schemaname = 'public'
+            AND tablename != 'components'
+        `);
+        
+        for (const row of existingPartitions) {
+            const tableName = row.tablename;
+            logger.trace(`Dropping old LIST partition table: ${tableName}`);
+            await db.unsafe(`DROP TABLE IF EXISTS ${tableName} CASCADE`);
+        }
+        
+        // Drop the main components table if it exists (to recreate with HASH partitioning)
+        const mainTableExists = await db.unsafe(`
+            SELECT 1 FROM information_schema.tables 
+            WHERE table_name = 'components' 
+            AND table_schema = 'public'
+        `);
+        
+        if (mainTableExists.length > 0) {
+            logger.trace(`Dropping existing components table for HASH partitioning`);
+            await db.unsafe(`DROP TABLE IF EXISTS components CASCADE`);
+        }
+        
+        logger.info(`Cleaned up ${existingPartitions.length} old partition tables`);
+    } catch (error) {
+        logger.warn(`Could not clean up old LIST partitions: ${error}`);
+        // Continue anyway - the table creation might still work
+    }
+}
+
+export const CreateHashPartitionedComponentTable = async (partitionCount: number = 16) => {
     await db`CREATE TABLE IF NOT EXISTS components (
         id UUID,
-        entity_id UUID REFERENCES entities(id),
+        entity_id UUID REFERENCES entities(id) ON DELETE CASCADE,
         type_id varchar(64) NOT NULL,
         name varchar(128),
         data jsonb,
         created_at TIMESTAMP DEFAULT NOW(),
         updated_at TIMESTAMP DEFAULT NOW(),
         deleted_at TIMESTAMP,
-        PRIMARY KEY (id, type_id, entity_id)
-    ) PARTITION BY LIST (type_id);`;
-    await db`CREATE INDEX IF NOT EXISTS idx_components_entity_id ON components (entity_id);`
-    await db`CREATE INDEX IF NOT EXISTS idx_components_type_id ON components (type_id);`
-    await db`CREATE INDEX IF NOT EXISTS idx_components_data_gin ON components USING GIN (data);`
-
-    // Phase 2A: Add composite indexes for sorting optimization
-    await db`CREATE INDEX IF NOT EXISTS idx_components_entity_type_deleted ON components (entity_id, type_id, deleted_at);`
-    await db`CREATE INDEX IF NOT EXISTS idx_components_type_deleted ON components (type_id, deleted_at) WHERE deleted_at IS NULL;`
-    await db`CREATE INDEX IF NOT EXISTS idx_components_deleted_entity ON components (deleted_at, entity_id) WHERE deleted_at IS NULL;`
-}
+        PRIMARY KEY (id, type_id),
+        UNIQUE(entity_id, type_id)
+    ) PARTITION BY HASH (type_id);`;
 
+    // Create hash partitions
+    for (let i = 0; i < partitionCount; i++) {
+        await db.unsafe(`CREATE TABLE IF NOT EXISTS components_p${i}
+            PARTITION OF components
+            FOR VALUES WITH (MODULUS ${partitionCount}, REMAINDER ${i});`);
+    }
+
+    await db`CREATE INDEX IF NOT EXISTS idx_components_entity_id ON components (entity_id)`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_type_id ON components (type_id)`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_data_gin ON components USING GIN (data)`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_entity_type_deleted ON components (entity_id, type_id, deleted_at)`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_type_deleted ON components (type_id, deleted_at) WHERE deleted_at IS NULL`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_deleted_entity ON components (deleted_at, entity_id) WHERE deleted_at IS NULL`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_entity_created_desc ON components (entity_id, created_at DESC)`;
+    await db`CREATE INDEX IF NOT EXISTS idx_components_type_entity_created ON components (type_id, entity_id, created_at DESC)`;
+}
 export const UpdateComponentIndexes = async (table_name: string, indexedProperties: string[]) => {
     try {
         table_name = validateIdentifier(table_name);
         indexedProperties = indexedProperties.map(prop => validateIdentifier(prop));
         logger.trace(`Updating indexes for component table: ${table_name}`);
+
+        // Check if this is a hash partitioned table
+        const partitionStrategy = await GetPartitionStrategy();
+        if (partitionStrategy === 'hash' && table_name !== 'components') {
+            // For hash partitioning, indexes should be on the parent table
+            logger.trace(`Redirecting index update to parent table 'components' for hash partitioning`);
+            table_name = 'components';
+        }
+
+        // Check if table is partitioned
+        const partitionCheck = await db.unsafe(`
+            SELECT relkind
+            FROM pg_class
+            WHERE relname = '${table_name}' AND relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
+        `);
+        const isPartitioned = partitionCheck.length > 0 && partitionCheck[0].relkind === 'p';
+        const useConcurrently = !isPartitioned && !process.env.USE_PGLITE; // Cannot use CONCURRENTLY on partitioned tables or PGlite
+
         const indexes_list = await db.unsafe(`
             SELECT indexname 
             FROM pg_indexes 
@@ -144,7 +276,16 @@ export const UpdateComponentIndexes = async (table_name: string, indexedProperti
                 if (!existingIndexes.includes(indexName)) {
                     logger.trace(`Creating missing index ${indexName} for property ${prop}`);
                     await retryWithBackoff(async () => {
-                        await db.unsafe(`CREATE INDEX CONCURRENTLY IF NOT EXISTS ${indexName} ON ${table_name} USING GIN ((data->'${prop}'))`);
+                        try {
+                            await db.unsafe(`CREATE INDEX${useConcurrently ? ' CONCURRENTLY' : ''} IF NOT EXISTS ${indexName} ON ${table_name} USING GIN ((data->'${prop}'))`);
+                        } catch (error: any) {
+                            // Check if the error is about duplicate key (index already exists)
+                            if (error.message && error.message.includes('duplicate key value violates unique constraint "pg_class_relname_nsp_index"')) {
+                                logger.trace(`Index ${indexName} already exists (confirmed by error), skipping creation`);
+                                return;
+                            }
+                            throw error;
+                        }
                     });
                     addedIndexes.add(indexName);
                 } else {
@@ -160,7 +301,16 @@ export const UpdateComponentIndexes = async (table_name: string, indexedProperti
                 const prop = match[1];
                 if (!indexedProperties.includes(prop) && !addedIndexes.has(index)) {
                     await retryWithBackoff(async () => {
-                        await db.unsafe(`DROP INDEX CONCURRENTLY IF EXISTS ${index}`);
+                        try {
+                            await db.unsafe(`DROP INDEX${useConcurrently ? ' CONCURRENTLY' : ''} IF EXISTS ${index}`);
+                        } catch (error: any) {
+                            // Check if the error is about relation does not exist
+                            if (error.message && (error.message.includes('does not exist') || error.message.includes('not found'))) {
+                                logger.trace(`Index ${index} does not exist, skipping drop`);
+                                return;
+                            }
+                            throw error;
+                        }
                     });
                     logger.info(`Dropped obsolete index ${index} for property ${prop}`);
                 }
@@ -176,12 +326,34 @@ export const UpdateComponentIndexes = async (table_name: string, indexedProperti
 export const CreateComponentPartitionTable = async (comp_name: string, type_id: string) => {
     try {
         comp_name = validateIdentifier(comp_name);
-        // // type_id is a value, not identifier, so no validation
-        // if (indexedProperties) {
-        //     indexedProperties = indexedProperties.map(prop => validateIdentifier(prop));
-        // }
         logger.trace(`Attempt adding partition table for component: ${comp_name}`);
-        // const table_name = `components_${comp_name.toLowerCase().replace(/\s+/g, '_')}`;
+
+        // Check partitioning strategy
+        const partitionStrategy = await GetPartitionStrategy();
+        logger.trace(`Current partition strategy: ${partitionStrategy}`);
+
+        if (partitionStrategy === 'hash') {
+            // For HASH partitioning, partitions are pre-created and data is automatically distributed
+            // We just need to ensure indexes are created for this component type
+            logger.info(`Component ${comp_name} will use existing hash partitions`);
+            
+            // For hash partitioning, indexes are created at the parent table level
+            // But we can still create component-specific indexes if needed
+            const storage = getMetadataStorage();
+            const componentId = storage.getComponentId(comp_name);
+            const indexedFields = storage.getIndexedFields(componentId);
+            
+            if (indexedFields.length > 0) {
+                logger.trace(`Ensuring specialized indexes for ${comp_name} on hash partitions`);
+                // For hash partitioning, indexes on parent table should suffice
+                // But we can add component-specific logic here if needed
+                logger.trace(`Hash partitioning handles indexes at parent table level`);
+            }
+            
+            return;
+        }
+
+        // Original LIST partitioning logic
         const table_name = GenerateTableName(comp_name);
         logger.trace(`Checking for existing partition table: ${table_name}`);
         const existingPartition = await db.unsafe(`SELECT 1 FROM information_schema.tables 
@@ -202,17 +374,22 @@ export const CreateComponentPartitionTable = async (comp_name: string, type_id:
         });
         logger.trace(`Successfully created partition table: ${table_name}`);
 
-        // TODO: Not sure if this is needed here or should be handled separately
-        // if (BUNSANE_RELATION_TYPED_COLUMN && indexedProperties?.includes('value')) {
-        //     logger.trace(`Adding typed FK column for ${table_name}`);
-        //     await retryWithBackoff(async () => {
-        //         await db.unsafe(`ALTER TABLE ${table_name} ADD COLUMN IF NOT EXISTS fk_id UUID GENERATED ALWAYS AS ((data->>'value')::UUID) STORED`);
-        //     });
-        //     await retryWithBackoff(async () => {
-        //         await db.unsafe(`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_${table_name}_fk_id ON ${table_name} (fk_id)`);
-        //     });
-        //     logger.trace(`Added fk_id column and index for ${table_name}`);
-        // }
+        // Automatically create indexes based on component metadata
+        const storage = getMetadataStorage();
+        const componentId = storage.getComponentId(comp_name);
+        const indexedFields = storage.getIndexedFields(componentId);
+        
+        if (indexedFields.length > 0) {
+            logger.trace(`Creating ${indexedFields.length} specialized indexes for ${comp_name}`);
+            const indexDefinitions = indexedFields.map(field => ({
+                tableName: table_name,
+                field: field.propertyKey,
+                indexType: field.indexType,
+                isDateField: field.isDateField
+            }));
+            await ensureMultipleJSONBPathIndexes(table_name, indexDefinitions);
+            logger.trace(`Created specialized indexes for ${comp_name}`);
+        }
         
     } catch (error) {
         logger.error(`Failed to create component partition table for ${comp_name}: ${error}`);
@@ -223,6 +400,18 @@ export const CreateComponentPartitionTable = async (comp_name: string, type_id:
 export const DeleteComponentPartitionTable = async (comp_name: string) => {
     try {
         comp_name = validateIdentifier(comp_name);
+        
+        // Check partitioning strategy
+        const partitionStrategy = await GetPartitionStrategy();
+        
+        if (partitionStrategy === 'hash') {
+            // For HASH partitioning, partitions are managed automatically
+            // No individual partition tables to delete
+            logger.info(`Component ${comp_name} uses hash partitions - no individual table to delete`);
+            return;
+        }
+        
+        // Original LIST partitioning logic
         const table_name = `components_${comp_name.toLowerCase().replace(/\s+/g, '_')}`;
 
         const existingPartition = await db.unsafe(`
@@ -249,17 +438,204 @@ export const DeleteComponentPartitionTable = async (comp_name: string) => {
 
 export const CreateEntityComponentTable = async () => {
     await db`CREATE TABLE IF NOT EXISTS entity_components (
-        entity_id UUID REFERENCES entities(id),
+        entity_id UUID REFERENCES entities(id) ON DELETE CASCADE,
         type_id VARCHAR(64) NOT NULL,
+        component_id UUID,
+        created_at TIMESTAMP DEFAULT NOW(),
+        updated_at TIMESTAMP DEFAULT NOW(),
         deleted_at TIMESTAMP,
         UNIQUE(entity_id, type_id)
     );`;
-    await db`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entity_components_entity_id ON entity_components (entity_id);`
-    await db`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entity_components_type_id ON entity_components (type_id);`
-    await db`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entity_components_type_entity ON entity_components (type_id, entity_id);`
+    const concurrently = process.env.USE_PGLITE ? '' : ' CONCURRENTLY';
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_entity_id ON entity_components (entity_id)`);
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_type_id ON entity_components (type_id)`);
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_type_entity ON entity_components (type_id, entity_id)`);
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_type_entity_deleted ON entity_components (type_id, entity_id, deleted_at)`);
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_deleted_type ON entity_components (deleted_at, type_id) WHERE deleted_at IS NULL`);
+    await db.unsafe(`CREATE INDEX${concurrently} IF NOT EXISTS idx_entity_components_component_id ON entity_components (component_id)`);
+    
+    // Add component_id column if it doesn't exist (for backward compatibility)
+    try {
+        await db`ALTER TABLE entity_components ADD COLUMN IF NOT EXISTS component_id UUID`;
+        logger.info(`Added component_id column to entity_components table`);
+    } catch (error) {
+        logger.warn(`Could not add component_id column to entity_components table: ${error}`);
+    }
+}
+
+export const PopulateComponentIds = async () => {
+    try {
+        // Populate component_id for existing rows that don't have it set
+        await db`UPDATE entity_components 
+                 SET component_id = c.id 
+                 FROM components c 
+                 WHERE entity_components.entity_id = c.entity_id 
+                 AND entity_components.type_id = c.type_id 
+                 AND entity_components.component_id IS NULL`;
+        
+        logger.info(`Populated component_id for existing entity_components rows`);
+    } catch (error) {
+        logger.warn(`Could not populate component_id for existing rows: ${error}`);
+    }
+}
 
-    await db`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entity_components_type_entity_deleted ON entity_components (type_id, entity_id, deleted_at);`
-    await db`CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entity_components_deleted_type ON entity_components (deleted_at, type_id) WHERE deleted_at IS NULL;`
+export const EnsureDatabaseMigrations = async () => {
+    logger.trace(`Checking for database migrations...`);
+    
+    try {
+        // First, ensure the table exists and has the basic structure
+        await CreateEntityComponentTable();
+        
+        // Check if entity_components table has component_id column
+        const columnCheck = await db`SELECT column_name FROM information_schema.columns 
+                                     WHERE table_name = 'entity_components' 
+                                     AND column_name = 'component_id' 
+                                     AND table_schema = 'public'`;
+        
+        if (columnCheck.length === 0) {
+            logger.info(`entity_components table missing component_id column, adding it...`);
+            // Add the column
+            await db`ALTER TABLE entity_components ADD COLUMN component_id UUID`;
+            logger.info(`Added component_id column to entity_components table`);
+            
+            // Wait a bit for the column to be available
+            await new Promise(resolve => setTimeout(resolve, 500));
+            
+            // Populate existing data
+            await PopulateComponentIds();
+        } else {
+            logger.trace(`entity_components table already has component_id column`);
+        }
+    } catch (error) {
+        logger.error(`Failed during database migration: ${error}`);
+        // Try to add the column anyway in case the check failed
+        try {
+            await db`ALTER TABLE entity_components ADD COLUMN IF NOT EXISTS component_id UUID`;
+            logger.info(`Attempted to add component_id column as fallback`);
+        } catch (fallbackError) {
+            logger.error(`Fallback column addition also failed: ${fallbackError}`);
+        }
+    }
+}
+
+export const AnalyzeAllComponentTables = async (): Promise<void> => {
+    try {
+        logger.trace(`Analyzing all component tables`);
+
+        // Check partitioning strategy
+        const partitionStrategy = await GetPartitionStrategy();
+        
+        let tablePattern: string;
+        if (partitionStrategy === 'hash') {
+            // For hash partitioning, analyze the hash partition tables
+            tablePattern = 'components_p%';
+        } else {
+            // For list partitioning, analyze the component-specific partition tables
+            tablePattern = 'components_%';
+        }
+
+        // Get all component partition tables
+        const tables = await db.unsafe(`
+            SELECT tablename
+            FROM pg_tables
+            WHERE tablename LIKE '${tablePattern}' AND schemaname = 'public'
+        `);
+
+        for (const row of tables) {
+            logger.trace(`Running ANALYZE on table ${row.tablename}`);
+            await db.unsafe(`ANALYZE ${row.tablename}`);
+            logger.trace(`Completed ANALYZE on table ${row.tablename}`);
+        }
+
+        logger.info(`Completed ANALYZE on ${tables.length} component tables`);
+    } catch (error) {
+        logger.error(`Failed to analyze component tables: ${error}`);
+        throw error;
+    }
+}
+
+export const GetPartitionStrategy = async (): Promise<'list' | 'hash' | null> => {
+    try {
+        const result = await db.unsafe(`
+            SELECT 
+                CASE 
+                    WHEN partstrat = 'l' THEN 'list'
+                    WHEN partstrat = 'h' THEN 'hash'
+                    ELSE NULL
+                END as strategy
+            FROM pg_partitioned_table 
+            WHERE partrelid = (SELECT oid FROM pg_class WHERE relname = 'components' AND relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public'))
+        `);
+        return result.length > 0 ? result[0].strategy : null;
+    } catch (error) {
+        logger.warn(`Could not determine partition strategy: ${error}`);
+        return null;
+    }
+}
+
+export const BenchmarkPartitionCounts = async (partitionCounts: number[] = [8, 16, 32]) => {
+    const results: Array<{partitionCount: number, planningTime: number, executionTime: number}> = [];
+    
+    for (const count of partitionCounts) {
+        logger.info(`Benchmarking with ${count} partitions`);
+        
+        // Create temporary hash partitioned table
+        const tempTableName = `components_benchmark_${count}`;
+        await db.unsafe(`CREATE TABLE ${tempTableName} (
+            id UUID,
+            entity_id UUID,
+            type_id varchar(64) NOT NULL,
+            name varchar(128),
+            data jsonb,
+            created_at TIMESTAMP DEFAULT NOW(),
+            updated_at TIMESTAMP DEFAULT NOW(),
+            deleted_at TIMESTAMP,
+            PRIMARY KEY (id, type_id),
+            UNIQUE(entity_id, type_id)
+        ) PARTITION BY HASH (type_id);`);
+        
+        // Create partitions
+        for (let i = 0; i < count; i++) {
+            await db.unsafe(`CREATE TABLE ${tempTableName}_p${i}
+                PARTITION OF ${tempTableName}
+                FOR VALUES WITH (MODULUS ${count}, REMAINDER ${i});`);
+        }
+        
+        // Copy sample data (limit to avoid long benchmark)
+        await db.unsafe(`INSERT INTO ${tempTableName} (id, entity_id, type_id, name, data, created_at, updated_at, deleted_at)
+            SELECT id, entity_id, type_id, name, data, created_at, updated_at, deleted_at
+            FROM components 
+            TABLESAMPLE BERNOULLI(10) -- Sample 10% of data
+            LIMIT 10000;`);
+        
+        // Create indexes
+        await db.unsafe(`CREATE INDEX idx_${tempTableName}_type_id ON ${tempTableName} (type_id)`);
+        await db.unsafe(`ANALYZE ${tempTableName}`);
+        
+        // Run benchmark query
+        const explainResult = await db.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) 
+            SELECT DISTINCT ec.entity_id as id 
+            FROM entity_components ec 
+            WHERE ec.type_id = (SELECT type_id FROM ${tempTableName} LIMIT 1) 
+            AND ec.deleted_at IS NULL`);
+        
+        const plan = explainResult[0]['QUERY PLAN'] ? JSON.parse(explainResult[0]['QUERY PLAN']) : explainResult[0];
+        const planningTime = plan.Planning ? plan.Planning.Time : 0;
+        const executionTime = plan.Execution ? plan.Execution.Time : 0;
+        
+        results.push({
+            partitionCount: count,
+            planningTime,
+            executionTime
+        });
+        
+        // Clean up
+        await db.unsafe(`DROP TABLE ${tempTableName} CASCADE;`);
+        
+        logger.info(`Partition count ${count}: planning=${planningTime}ms, execution=${executionTime}ms`);
+    }
+    
+    return results;
 }
 
 export const GenerateTableName = (name: string) => `components_${name.toLowerCase().replace(/\s+/g, '_')}`;