diff --git a/src/context/chartdb-context/chartdb-provider.tsx b/src/context/chartdb-context/chartdb-provider.tsx index 7a0b04a8..99a3b985 100644 --- a/src/context/chartdb-context/chartdb-provider.tsx +++ b/src/context/chartdb-context/chartdb-provider.tsx @@ -6,7 +6,10 @@ import type { ChartDBContext, ChartDBEvent } from './chartdb-context'; import { chartDBContext } from './chartdb-context'; import { DatabaseType } from '@/lib/domain/database-type'; import type { DBField } from '@/lib/domain/db-field'; -import type { DBIndex } from '@/lib/domain/db-index'; +import { + getTableIndexesWithPrimaryKey, + type DBIndex, +} from '@/lib/domain/db-index'; import type { DBRelationship } from '@/lib/domain/db-relationship'; import { useStorage } from '@/hooks/use-storage'; import { useRedoUndoStack } from '@/hooks/use-redo-undo-stack'; @@ -348,6 +351,11 @@ export const ChartDBProvider: React.FC< order: tables.length, ...attributes, }; + + table.indexes = getTableIndexesWithPrimaryKey({ + table, + }); + await addTable(table); return table; @@ -639,17 +647,30 @@ export const ChartDBProvider: React.FC< options = { updateHistory: true } ) => { const prevField = getField(tableId, fieldId); + + const updateTableFn = (table: DBTable) => { + const updatedTable: DBTable = { + ...table, + fields: table.fields.map((f) => + f.id === fieldId ? { ...f, ...field } : f + ), + } satisfies DBTable; + + updatedTable.indexes = getTableIndexesWithPrimaryKey({ + table: updatedTable, + }); + + return updatedTable; + }; + setTables((tables) => - tables.map((table) => - table.id === tableId - ? { - ...table, - fields: table.fields.map((f) => - f.id === fieldId ? { ...f, ...field } : f - ), - } - : table - ) + tables.map((table) => { + if (table.id === tableId) { + return updateTableFn(table); + } + + return table; + }) ); const table = await db.getTable({ diagramId, id: tableId }); @@ -664,10 +685,7 @@ export const ChartDBProvider: React.FC< db.updateTable({ id: tableId, attributes: { - ...table, - fields: table.fields.map((f) => - f.id === fieldId ? { ...f, ...field } : f - ), + ...updateTableFn(table), }, }), ]); @@ -694,19 +712,29 @@ export const ChartDBProvider: React.FC< fieldId: string, options = { updateHistory: true } ) => { + const updateTableFn = (table: DBTable) => { + const updatedTable: DBTable = { + ...table, + fields: table.fields.filter((f) => f.id !== fieldId), + } satisfies DBTable; + + updatedTable.indexes = getTableIndexesWithPrimaryKey({ + table: updatedTable, + }); + + return updatedTable; + }; + const fields = getTable(tableId)?.fields ?? []; const prevField = getField(tableId, fieldId); setTables((tables) => - tables.map((table) => - table.id === tableId - ? { - ...table, - fields: table.fields.filter( - (f) => f.id !== fieldId - ), - } - : table - ) + tables.map((table) => { + if (table.id === tableId) { + return updateTableFn(table); + } + + return table; + }) ); events.emit({ @@ -730,8 +758,7 @@ export const ChartDBProvider: React.FC< db.updateTable({ id: tableId, attributes: { - ...table, - fields: table.fields.filter((f) => f.id !== fieldId), + ...updateTableFn(table), }, }), ]); diff --git a/src/lib/data/export-metadata/__tests__/export-sql-dbml.test.ts b/src/lib/data/export-metadata/__tests__/export-sql-dbml.test.ts index 47f9e4e4..28f28f55 100644 --- a/src/lib/data/export-metadata/__tests__/export-sql-dbml.test.ts +++ b/src/lib/data/export-metadata/__tests__/export-sql-dbml.test.ts @@ -124,6 +124,96 @@ describe('DBML Export - SQL Generation Tests', () => { ); }); + it('should not create duplicate index for composite primary key', () => { + const tableId = testId(); + const field1Id = testId(); + const field2Id = testId(); + const field3Id = testId(); + + const diagram: Diagram = createDiagram({ + id: testId(), + name: 'Landlord System', + databaseType: DatabaseType.POSTGRESQL, + tables: [ + createTable({ + id: tableId, + name: 'users_master_table', + schema: 'landlord', + fields: [ + createField({ + id: field1Id, + name: 'master_user_id', + type: { id: 'bigint', name: 'bigint' }, + primaryKey: true, + nullable: false, + unique: false, + }), + createField({ + id: field2Id, + name: 'tenant_id', + type: { id: 'bigint', name: 'bigint' }, + primaryKey: true, + nullable: false, + unique: false, + }), + createField({ + id: field3Id, + name: 'tenant_user_id', + type: { id: 'bigint', name: 'bigint' }, + primaryKey: true, + nullable: false, + unique: false, + }), + createField({ + id: testId(), + name: 'enabled', + type: { id: 'boolean', name: 'boolean' }, + primaryKey: false, + nullable: true, + unique: false, + }), + ], + indexes: [ + { + id: testId(), + name: 'idx_users_master_table_master_user_id_tenant_id_tenant_user_id', + unique: false, + fieldIds: [field1Id, field2Id, field3Id], + createdAt: testTime, + }, + { + id: testId(), + name: 'index_1', + unique: true, + fieldIds: [field2Id, field3Id], + createdAt: testTime, + }, + ], + }), + ], + relationships: [], + }); + + const sql = exportBaseSQL({ + diagram, + targetDatabaseType: DatabaseType.POSTGRESQL, + isDBMLFlow: true, + }); + + // Should contain composite primary key constraint + expect(sql).toContain( + 'PRIMARY KEY (master_user_id, tenant_id, tenant_user_id)' + ); + + // Should NOT contain the duplicate index for the primary key fields + expect(sql).not.toContain( + 'CREATE INDEX idx_users_master_table_master_user_id_tenant_id_tenant_user_id' + ); + + // Should still contain the unique index on subset of fields + expect(sql).toContain('CREATE UNIQUE INDEX index_1'); + }); + it('should handle single primary keys inline', () => { const diagram: Diagram = createDiagram({ id: testId(), diff --git a/src/lib/data/export-metadata/export-per-type/mssql.ts b/src/lib/data/export-metadata/export-per-type/mssql.ts index 5a52c347..071a6ab2 100644 --- a/src/lib/data/export-metadata/export-per-type/mssql.ts +++ b/src/lib/data/export-metadata/export-per-type/mssql.ts @@ -178,7 +178,15 @@ export function exportMSSQL({ }) .join(',\n')}${ table.fields.filter((f) => f.primaryKey).length > 0 - ? `,\n PRIMARY KEY (${table.fields + ? `,\n ${(() => { + // Find PK index to get the constraint name + const pkIndex = table.indexes.find( + (idx) => idx.isPrimaryKey + ); + return pkIndex?.name + ? `CONSTRAINT [${pkIndex.name}] ` + : ''; + })()}PRIMARY KEY (${table.fields .filter((f) => f.primaryKey) .map((f) => `[${f.name}]`) .join(', ')})` diff --git a/src/lib/data/export-metadata/export-per-type/mysql.ts b/src/lib/data/export-metadata/export-per-type/mysql.ts index 53601525..c8a75ab6 100644 --- a/src/lib/data/export-metadata/export-per-type/mysql.ts +++ b/src/lib/data/export-metadata/export-per-type/mysql.ts @@ -313,7 +313,15 @@ export function exportMySQL({ .join(',\n')}${ // Add PRIMARY KEY as table constraint primaryKeyFields.length > 0 - ? `,\n PRIMARY KEY (${primaryKeyFields + ? `,\n ${(() => { + // Find PK index to get the constraint name + const pkIndex = table.indexes.find( + (idx) => idx.isPrimaryKey + ); + return pkIndex?.name + ? `CONSTRAINT \`${pkIndex.name}\` ` + : ''; + })()}PRIMARY KEY (${primaryKeyFields .map((f) => `\`${f.name}\``) .join(', ')})` : '' diff --git a/src/lib/data/export-metadata/export-per-type/postgresql.ts b/src/lib/data/export-metadata/export-per-type/postgresql.ts index a745bd5d..dbe1ffc5 100644 --- a/src/lib/data/export-metadata/export-per-type/postgresql.ts +++ b/src/lib/data/export-metadata/export-per-type/postgresql.ts @@ -325,7 +325,15 @@ export function exportPostgreSQL({ }) .join(',\n')}${ primaryKeyFields.length > 0 - ? `,\n PRIMARY KEY (${primaryKeyFields + ? `,\n ${(() => { + // Find PK index to get the constraint name + const pkIndex = table.indexes.find( + (idx) => idx.isPrimaryKey + ); + return pkIndex?.name + ? `CONSTRAINT "${pkIndex.name}" ` + : ''; + })()}PRIMARY KEY (${primaryKeyFields .map((f) => `"${f.name}"`) .join(', ')})` : '' diff --git a/src/lib/data/export-metadata/export-sql-script.ts b/src/lib/data/export-metadata/export-sql-script.ts index f152918e..007bb888 100644 --- a/src/lib/data/export-metadata/export-sql-script.ts +++ b/src/lib/data/export-metadata/export-sql-script.ts @@ -313,21 +313,33 @@ export const exportBaseSQL = ({ } } - // Handle PRIMARY KEY constraint - only add inline if not composite - if (field.primaryKey && !hasCompositePrimaryKey) { + // Handle PRIMARY KEY constraint - only add inline if no PK index with custom name + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + if (field.primaryKey && !hasCompositePrimaryKey && !pkIndex?.name) { sqlScript += ' PRIMARY KEY'; } - // Add a comma after each field except the last one (or before composite primary key) - if (index < table.fields.length - 1 || hasCompositePrimaryKey) { + // Add a comma after each field except the last one (or before PK constraint) + const needsPKConstraint = + hasCompositePrimaryKey || + (primaryKeyFields.length === 1 && pkIndex?.name); + if (index < table.fields.length - 1 || needsPKConstraint) { sqlScript += ',\n'; } }); - // Add composite primary key constraint if needed - if (hasCompositePrimaryKey) { + // Add primary key constraint if needed (for composite PKs or single PK with custom name) + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + if ( + hasCompositePrimaryKey || + (primaryKeyFields.length === 1 && pkIndex?.name) + ) { const pkFieldNames = primaryKeyFields.map((f) => f.name).join(', '); - sqlScript += `\n PRIMARY KEY (${pkFieldNames})`; + if (pkIndex?.name) { + sqlScript += `\n CONSTRAINT ${pkIndex.name} PRIMARY KEY (${pkFieldNames})`; + } else { + sqlScript += `\n PRIMARY KEY (${pkFieldNames})`; + } } sqlScript += '\n);\n'; @@ -349,12 +361,33 @@ export const exportBaseSQL = ({ // Generate SQL for indexes table.indexes.forEach((index) => { - const fieldNames = index.fieldIds - .map( - (fieldId) => - table.fields.find((field) => field.id === fieldId)?.name + // Skip the primary key index (it's already handled as a constraint) + if (index.isPrimaryKey) { + return; + } + + // Get the fields for this index + const indexFields = index.fieldIds + .map((fieldId) => table.fields.find((f) => f.id === fieldId)) + .filter( + (field): field is NonNullable => + field !== undefined + ); + + // Skip if this index exactly matches the primary key fields + // This prevents creating redundant indexes for composite primary keys + if ( + primaryKeyFields.length > 0 && + primaryKeyFields.length === indexFields.length && + primaryKeyFields.every((pk) => + indexFields.some((field) => field.id === pk.id) ) - .filter(Boolean) + ) { + return; // Skip this index as it's redundant with the primary key + } + + const fieldNames = indexFields + .map((field) => field.name) .join(', '); if (fieldNames) { diff --git a/src/lib/dbml/dbml-export/__tests__/composite-pk-export.test.ts b/src/lib/dbml/dbml-export/__tests__/composite-pk-export.test.ts new file mode 100644 index 00000000..d78fa169 --- /dev/null +++ b/src/lib/dbml/dbml-export/__tests__/composite-pk-export.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect } from 'vitest'; +import { generateDBMLFromDiagram } from '../dbml-export'; +import { DatabaseType } from '@/lib/domain/database-type'; +import type { Diagram } from '@/lib/domain/diagram'; +import { generateId } from '@/lib/utils'; + +describe('Composite Primary Key Name Export', () => { + it('should export composite primary key with name in DBML', () => { + const diagram: Diagram = { + id: generateId(), + name: 'Test', + databaseType: DatabaseType.POSTGRESQL, + createdAt: new Date(), + updatedAt: new Date(), + tables: [ + { + id: generateId(), + name: 'users_master_table', + schema: 'landlord', + x: 0, + y: 0, + color: '#FFF', + isView: false, + createdAt: Date.now(), + fields: [ + { + id: generateId(), + name: 'master_user_id', + type: { id: 'bigint', name: 'bigint' }, + nullable: false, + primaryKey: true, + unique: false, + createdAt: Date.now(), + }, + { + id: generateId(), + name: 'tenant_id', + type: { id: 'bigint', name: 'bigint' }, + nullable: false, + primaryKey: true, + unique: false, + createdAt: Date.now(), + }, + { + id: generateId(), + name: 'tenant_user_id', + type: { id: 'bigint', name: 'bigint' }, + nullable: false, + primaryKey: true, + unique: false, + createdAt: Date.now(), + }, + { + id: generateId(), + name: 'enabled', + type: { id: 'boolean', name: 'boolean' }, + nullable: true, + primaryKey: false, + unique: false, + createdAt: Date.now(), + }, + ], + indexes: [ + { + id: generateId(), + name: 'users_master_table_index_1', + unique: true, + fieldIds: ['dummy1', 'dummy2'], // Will be replaced + createdAt: Date.now(), + }, + ], + }, + ], + relationships: [], + }; + + // Fix field IDs in the index and add PK index + const table = diagram.tables![0]; + const masterUserIdField = table.fields.find( + (f) => f.name === 'master_user_id' + ); + const tenantIdField = table.fields.find((f) => f.name === 'tenant_id'); + const tenantUserIdField = table.fields.find( + (f) => f.name === 'tenant_user_id' + ); + table.indexes[0].fieldIds = [tenantIdField!.id, tenantUserIdField!.id]; + + // Add the PK index with name + table.indexes.push({ + id: generateId(), + name: 'moshe', + unique: true, + isPrimaryKey: true, + fieldIds: [ + masterUserIdField!.id, + tenantIdField!.id, + tenantUserIdField!.id, + ], + createdAt: Date.now(), + }); + + const result = generateDBMLFromDiagram(diagram); + + // Check that the DBML contains the composite PK with name + expect(result.standardDbml).toContain( + '(master_user_id, tenant_id, tenant_user_id) [pk, name: "moshe"]' + ); + + // Check that the unique index is also present + expect(result.standardDbml).toContain( + '(tenant_id, tenant_user_id) [unique, name: "users_master_table_index_1"]' + ); + }); +}); diff --git a/src/lib/dbml/dbml-export/__tests__/dbml-export-issue-fix.test.ts b/src/lib/dbml/dbml-export/__tests__/dbml-export-issue-fix.test.ts index 7b4f4815..213cc063 100644 --- a/src/lib/dbml/dbml-export/__tests__/dbml-export-issue-fix.test.ts +++ b/src/lib/dbml/dbml-export/__tests__/dbml-export-issue-fix.test.ts @@ -1383,12 +1383,9 @@ Ref "fk_0_table_2_id_fk":"table_1"."id" < "table_2"."id" const result = generateDBMLFromDiagram(diagram); // Check that the inline DBML has proper indentation + // Note: indexes on primary key fields should be filtered out expect(result.inlineDbml).toContain(`Table "table_1" { "id" bigint [pk, not null] - - Indexes { - id [name: "index_1"] - } }`); expect(result.inlineDbml).toContain(`Table "table_2" { diff --git a/src/lib/dbml/dbml-export/dbml-export.ts b/src/lib/dbml/dbml-export/dbml-export.ts index f49a54d0..2a70efbf 100644 --- a/src/lib/dbml/dbml-export/dbml-export.ts +++ b/src/lib/dbml/dbml-export/dbml-export.ts @@ -605,6 +605,45 @@ const fixTableBracketSyntax = (dbml: string): string => { ); }; +// Restore composite primary key names in the DBML +const restoreCompositePKNames = (dbml: string, tables: DBTable[]): string => { + if (!tables || tables.length === 0) return dbml; + + let result = dbml; + + tables.forEach((table) => { + // Check if this table has a PK index with a name + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + if (pkIndex?.name) { + const primaryKeyFields = table.fields.filter((f) => f.primaryKey); + if (primaryKeyFields.length >= 1) { + // Build the column list for the composite PK + const columnList = primaryKeyFields + .map((f) => f.name) + .join(', '); + + // Build the table identifier pattern + const tableIdentifier = table.schema + ? `"${table.schema.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}"\\."${table.name.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}"` + : `"${table.name.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}"`; + + // Pattern to match the composite PK index line + // Match patterns like: (col1, col2, col3) [pk] + const pkPattern = new RegExp( + `(Table ${tableIdentifier} \\{[^}]*?Indexes \\{[^}]*?)(\\(${columnList.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}\\) \\[pk\\])`, + 'gs' + ); + + // Replace with the named version + const replacement = `$1(${columnList}) [pk, name: "${pkIndex.name}"]`; + result = result.replace(pkPattern, replacement); + } + } + }); + + return result; +}; + // Restore schema information that may have been stripped by the DBML importer const restoreTableSchemas = (dbml: string, tables: DBTable[]): string => { if (!tables || tables.length === 0) return dbml; @@ -870,14 +909,16 @@ export function generateDBMLFromDiagram(diagram: Diagram): DBMLExportResult { ...table, name: safeTableName, fields: processedFields, - indexes: (table.indexes || []).map((index) => ({ - ...index, - name: index.name - ? /[^\w]/.test(index.name) - ? `"${index.name.replace(/"/g, '\\"')}"` - : index.name - : `idx_${Math.random().toString(36).substring(2, 8)}`, - })), + indexes: (table.indexes || []) + .filter((index) => !index.isPrimaryKey) // Filter out PK indexes as they're handled separately + .map((index) => ({ + ...index, + name: index.name + ? /[^\w]/.test(index.name) + ? `"${index.name.replace(/"/g, '\\"')}"` + : index.name + : `idx_${Math.random().toString(36).substring(2, 8)}`, + })), }; }; @@ -939,6 +980,9 @@ export function generateDBMLFromDiagram(diagram: Diagram): DBMLExportResult { // Restore schema information that may have been stripped by DBML importer standard = restoreTableSchemas(standard, uniqueTables); + // Restore composite primary key names + standard = restoreCompositePKNames(standard, uniqueTables); + // Prepend Enum DBML to the standard output if (enumsDBML) { standard = enumsDBML + '\n\n' + standard; diff --git a/src/lib/dbml/dbml-import/__tests__/composite-pk-name.test.ts b/src/lib/dbml/dbml-import/__tests__/composite-pk-name.test.ts new file mode 100644 index 00000000..a2fe3e5d --- /dev/null +++ b/src/lib/dbml/dbml-import/__tests__/composite-pk-name.test.ts @@ -0,0 +1,190 @@ +import { describe, it, expect } from 'vitest'; +import { importDBMLToDiagram } from '../dbml-import'; +import { exportPostgreSQL } from '@/lib/data/export-metadata/export-per-type/postgresql'; +import { exportMySQL } from '@/lib/data/export-metadata/export-per-type/mysql'; +import { exportMSSQL } from '@/lib/data/export-metadata/export-per-type/mssql'; +import { DatabaseType } from '@/lib/domain/database-type'; + +describe('Composite Primary Key with Name', () => { + it('should preserve composite primary key name in DBML import and SQL export', async () => { + const dbmlContent = ` +Table "landlord"."users_master_table" { + "master_user_id" bigint [not null] + "tenant_id" bigint [not null] + "tenant_user_id" bigint [not null] + "enabled" boolean + + Indexes { + (master_user_id, tenant_id, tenant_user_id) [pk, name: "idx_users_master_table_master_user_id_tenant_id_tenant_user_id"] + (tenant_id, tenant_user_id) [unique, name: "index_1"] + } +} +`; + + // Import DBML + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.POSTGRESQL, + }); + + // Check that the composite PK name was captured + expect(diagram.tables).toBeDefined(); + const table = diagram.tables![0]; + + // Check for the PK index + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + expect(pkIndex).toBeDefined(); + expect(pkIndex!.name).toBe( + 'idx_users_master_table_master_user_id_tenant_id_tenant_user_id' + ); + + // Check that fields are marked as primary keys + const pkFields = table.fields.filter((f) => f.primaryKey); + expect(pkFields).toHaveLength(3); + expect(pkFields.map((f) => f.name)).toEqual([ + 'master_user_id', + 'tenant_id', + 'tenant_user_id', + ]); + + // Check that we have both the PK index and the unique index + expect(table.indexes).toHaveLength(2); + const uniqueIndex = table.indexes.find((idx) => !idx.isPrimaryKey); + expect(uniqueIndex!.name).toBe('index_1'); + expect(uniqueIndex!.unique).toBe(true); + }); + + it('should export composite primary key with CONSTRAINT name in PostgreSQL', async () => { + const dbmlContent = ` +Table "users" { + "id" bigint [not null] + "tenant_id" bigint [not null] + + Indexes { + (id, tenant_id) [pk, name: "pk_users_composite"] + } +} +`; + + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.POSTGRESQL, + }); + + const sqlScript = exportPostgreSQL({ diagram }); + + // Check that the SQL contains the named constraint + expect(sqlScript).toContain( + 'CONSTRAINT "pk_users_composite" PRIMARY KEY ("id", "tenant_id")' + ); + expect(sqlScript).not.toContain('PRIMARY KEY ("id", "tenant_id"),'); // Should not have unnamed PK + }); + + it('should export composite primary key with CONSTRAINT name in MySQL', async () => { + const dbmlContent = ` +Table "orders" { + "order_id" int [not null] + "product_id" int [not null] + + Indexes { + (order_id, product_id) [pk, name: "orders_order_product_pk"] + } +} +`; + + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.MYSQL, + }); + + const sqlScript = exportMySQL({ diagram }); + + // Check that the SQL contains the named constraint + expect(sqlScript).toContain( + 'CONSTRAINT `orders_order_product_pk` PRIMARY KEY (`order_id`, `product_id`)' + ); + }); + + it('should export composite primary key with CONSTRAINT name in MSSQL', async () => { + const dbmlContent = ` +Table "products" { + "category_id" int [not null] + "product_id" int [not null] + + Indexes { + (category_id, product_id) [pk, name: "pk_products"] + } +} +`; + + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.SQL_SERVER, + }); + + const sqlScript = exportMSSQL({ diagram }); + + // Check that the SQL contains the named constraint + expect(sqlScript).toContain( + 'CONSTRAINT [pk_products] PRIMARY KEY ([category_id], [product_id])' + ); + }); + + it('should merge duplicate PK index with name', async () => { + const dbmlContent = ` +Table "test" { + "a" int [not null] + "b" int [not null] + + Indexes { + (a, b) [pk] + (a, b) [name: "test_pk_name"] + } +} +`; + + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.POSTGRESQL, + }); + + expect(diagram.tables).toBeDefined(); + const table = diagram.tables![0]; + + // Should capture the name from the duplicate index + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + expect(pkIndex).toBeDefined(); + expect(pkIndex!.name).toBe('test_pk_name'); + + // Should only have the PK index + expect(table.indexes).toHaveLength(1); + + // Fields should be marked as primary keys + expect(table.fields.filter((f) => f.primaryKey)).toHaveLength(2); + }); + + it('should handle composite PK without name', async () => { + const dbmlContent = ` +Table "simple" { + "x" int [not null] + "y" int [not null] + + Indexes { + (x, y) [pk] + } +} +`; + + const diagram = await importDBMLToDiagram(dbmlContent, { + databaseType: DatabaseType.POSTGRESQL, + }); + + expect(diagram.tables).toBeDefined(); + const table = diagram.tables![0]; + + // PK index should not exist for composite PK without name + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + expect(pkIndex).toBeDefined(); + + const sqlScript = exportPostgreSQL({ diagram }); + + // Should have unnamed PRIMARY KEY + expect(sqlScript).toContain('PRIMARY KEY ("x", "y")'); + expect(sqlScript).toContain('CONSTRAINT'); + }); +}); diff --git a/src/lib/dbml/dbml-import/__tests__/dbml-import-fantasy-examples.test.ts b/src/lib/dbml/dbml-import/__tests__/dbml-import-fantasy-examples.test.ts index 7c5a3292..c81be0a1 100644 --- a/src/lib/dbml/dbml-import/__tests__/dbml-import-fantasy-examples.test.ts +++ b/src/lib/dbml/dbml-import/__tests__/dbml-import-fantasy-examples.test.ts @@ -177,7 +177,7 @@ Table ranks { expect(wizardsTable?.fields).toHaveLength(11); // Check indexes - expect(wizardsTable?.indexes).toHaveLength(2); + expect(wizardsTable?.indexes).toHaveLength(3); const emailIndex = wizardsTable?.indexes.find((idx) => idx.name.includes('email') ); @@ -920,7 +920,7 @@ Note dragon_note { expect(hoardsTable).toBeDefined(); // Verify all indexes are imported correctly - expect(hoardsTable?.indexes).toHaveLength(3); // Should have 3 indexes as defined in DBML + expect(hoardsTable?.indexes).toHaveLength(4); // 3 from DBML + 1 implicit PK index // Verify named indexes const uniqueDragonIndex = hoardsTable?.indexes.find( @@ -1119,7 +1119,7 @@ Table "public_3"."comments" { ).toBe('timestamp'); // Check posts indexes thoroughly - expect(postsTable?.indexes).toHaveLength(2); + expect(postsTable?.indexes).toHaveLength(3); // Index 1: Composite unique index on (content, user_id) const compositeIndex = postsTable?.indexes.find( @@ -1154,7 +1154,7 @@ Table "public_3"."comments" { // Check comments table expect(commentsTable?.fields).toHaveLength(5); - expect(commentsTable?.indexes).toHaveLength(1); + expect(commentsTable?.indexes).toHaveLength(2); // Index: Unique index on id const idIndex = commentsTable?.indexes.find( diff --git a/src/lib/dbml/dbml-import/dbml-import.ts b/src/lib/dbml/dbml-import/dbml-import.ts index 8d590f5a..10098b47 100644 --- a/src/lib/dbml/dbml-import/dbml-import.ts +++ b/src/lib/dbml/dbml-import/dbml-import.ts @@ -9,7 +9,7 @@ import { findDataTypeDataById } from '@/lib/data/data-types/data-types'; import { defaultTableColor } from '@/lib/colors'; import { DatabaseType } from '@/lib/domain/database-type'; import type Field from '@dbml/core/types/model_structure/field'; -import type { DBIndex } from '@/lib/domain'; +import { getTableIndexesWithPrimaryKey, type DBIndex } from '@/lib/domain'; import { DBCustomTypeKind, type DBCustomType, @@ -100,6 +100,7 @@ interface DBMLIndex { columns: (string | DBMLIndexColumn)[]; unique?: boolean; name?: string; + pk?: boolean; // Primary key index flag } interface DBMLTable { @@ -387,15 +388,19 @@ export const importDBMLToDiagram = async ( ); } - // Generate a consistent index name + // For PK indexes, only use the name if explicitly provided + // For regular indexes, generate a default name if needed const indexName = dbmlIndex.name || - `idx_${table.name}_${indexColumns.join('_')}`; + (!dbmlIndex.pk + ? `idx_${table.name}_${indexColumns.join('_')}` + : undefined); return { columns: indexColumns, unique: dbmlIndex.unique || false, name: indexName, + pk: Boolean(dbmlIndex.pk) || false, }; }) || [], }); @@ -484,29 +489,126 @@ export const importDBMLToDiagram = async ( }; }); - // Convert DBML indexes to ChartDB indexes - const indexes: DBIndex[] = - table.indexes?.map((dbmlIndex) => { - const fieldIds = dbmlIndex.columns.map((columnName) => { - const field = fields.find((f) => f.name === columnName); - if (!field) { - throw new Error( - `Index references non-existent column: ${columnName}` - ); - } - return field.id; - }); + // Process composite primary keys from indexes with [pk] attribute + let compositePKFields: string[] = []; + let compositePKIndexName: string | undefined; - return { - id: generateId(), - name: - dbmlIndex.name || - `idx_${table.name}_${(dbmlIndex.columns as string[]).join('_')}`, - fieldIds, - unique: dbmlIndex.unique || false, - createdAt: Date.now(), - }; - }) || []; + // Find PK indexes and mark fields as primary keys + table.indexes?.forEach((dbmlIndex) => { + if (dbmlIndex.pk) { + // Extract column names from the columns array + compositePKFields = dbmlIndex.columns.map((col) => + typeof col === 'string' ? col : col.value + ); + // Only store the name if it was explicitly provided (not undefined) + if (dbmlIndex.name) { + compositePKIndexName = dbmlIndex.name; + } + // Mark fields as primary keys + dbmlIndex.columns.forEach((col) => { + const columnName = + typeof col === 'string' ? col : col.value; + const field = fields.find((f) => f.name === columnName); + if (field) { + field.primaryKey = true; + } + }); + } + }); + + // If we found a PK without a name, look for a duplicate index with just a name + if (compositePKFields.length > 0 && !compositePKIndexName) { + table.indexes?.forEach((dbmlIndex) => { + if ( + !dbmlIndex.pk && + dbmlIndex.name && + dbmlIndex.columns.length === compositePKFields.length + ) { + // Check if columns match + const indexColumns = dbmlIndex.columns.map((col) => + typeof col === 'string' ? col : col.value + ); + if ( + indexColumns.every( + (col, i) => col === compositePKFields[i] + ) + ) { + compositePKIndexName = dbmlIndex.name; + } + } + }); + } + + // Convert DBML indexes to ChartDB indexes (excluding PK indexes and their duplicates) + const indexes: DBIndex[] = + table.indexes + ?.filter((dbmlIndex) => { + // Skip PK indexes - we'll handle them separately + if (dbmlIndex.pk) return false; + + // Skip duplicate indexes that match the composite PK + // (when user has both [pk] and [name: "..."] on same fields) + if ( + compositePKFields.length > 0 && + dbmlIndex.columns.length === + compositePKFields.length && + dbmlIndex.columns.every((col, i) => { + const colName = + typeof col === 'string' ? col : col.value; + return colName === compositePKFields[i]; + }) + ) { + return false; + } + + return true; + }) + .map((dbmlIndex) => { + const fieldIds = dbmlIndex.columns.map((columnName) => { + const field = fields.find( + (f) => f.name === columnName + ); + if (!field) { + throw new Error( + `Index references non-existent column: ${columnName}` + ); + } + return field.id; + }); + + return { + id: generateId(), + name: + dbmlIndex.name || + `idx_${table.name}_${(dbmlIndex.columns as string[]).join('_')}`, + fieldIds, + unique: dbmlIndex.unique || false, + createdAt: Date.now(), + }; + }) || []; + + // Add PK as an index if it exists and has a name + // Only create the PK index if there's an explicit name for it + if (compositePKFields.length >= 1 && compositePKIndexName) { + const pkFieldIds = compositePKFields.map((columnName) => { + const field = fields.find((f) => f.name === columnName); + if (!field) { + throw new Error( + `PK references non-existent column: ${columnName}` + ); + } + return field.id; + }); + + indexes.push({ + id: generateId(), + name: compositePKIndexName, + fieldIds: pkFieldIds, + unique: true, + isPrimaryKey: true, + createdAt: Date.now(), + }); + } // Extract table note/comment let tableComment: string | undefined; @@ -521,7 +623,7 @@ export const importDBMLToDiagram = async ( } } - return { + const tableToReturn: DBTable = { id: generateId(), name: table.name.replace(/['"]/g, ''), schema: @@ -540,6 +642,13 @@ export const importDBMLToDiagram = async ( createdAt: Date.now(), comments: tableComment, } satisfies DBTable; + + return { + ...tableToReturn, + indexes: getTableIndexesWithPrimaryKey({ + table: tableToReturn, + }), + }; }); // Create relationships using the refs diff --git a/src/lib/domain/__tests__/composite-pk-metadata-import.test.ts b/src/lib/domain/__tests__/composite-pk-metadata-import.test.ts new file mode 100644 index 00000000..252ff579 --- /dev/null +++ b/src/lib/domain/__tests__/composite-pk-metadata-import.test.ts @@ -0,0 +1,192 @@ +import { describe, it, expect } from 'vitest'; +import { createTablesFromMetadata } from '../db-table'; +import { DatabaseType } from '../database-type'; +import type { DatabaseMetadata } from '@/lib/data/import-metadata/metadata-types/database-metadata'; + +describe('Composite Primary Key Name from Metadata Import', () => { + it('should capture composite primary key name from metadata indexes', () => { + const metadata: DatabaseMetadata = { + database_name: 'test_db', + version: '', + fk_info: [], + pk_info: [ + { + schema: 'landlord', + table: 'users_master_table', + column: 'master_user_id', + pk_def: 'PRIMARY KEY (master_user_id, tenant_id, tenant_user_id)', + }, + { + schema: 'landlord', + table: 'users_master_table', + column: 'tenant_id', + pk_def: 'PRIMARY KEY (master_user_id, tenant_id, tenant_user_id)', + }, + { + schema: 'landlord', + table: 'users_master_table', + column: 'tenant_user_id', + pk_def: 'PRIMARY KEY (master_user_id, tenant_id, tenant_user_id)', + }, + ], + columns: [ + { + schema: 'landlord', + table: 'users_master_table', + name: 'master_user_id', + ordinal_position: 1, + type: 'bigint', + character_maximum_length: null, + precision: null, + nullable: false, + default: '', + collation: '', + comment: '', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'tenant_id', + ordinal_position: 2, + type: 'bigint', + character_maximum_length: null, + precision: null, + nullable: false, + default: '', + collation: '', + comment: '', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'tenant_user_id', + ordinal_position: 3, + type: 'bigint', + character_maximum_length: null, + precision: null, + nullable: false, + default: '', + collation: '', + comment: '', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'enabled', + ordinal_position: 4, + type: 'boolean', + character_maximum_length: null, + precision: null, + nullable: true, + default: '', + collation: '', + comment: '', + }, + ], + indexes: [ + // The composite PK index named "moshe" + { + schema: 'landlord', + table: 'users_master_table', + name: 'moshe', + column: 'master_user_id', + index_type: 'btree', + cardinality: 0, + size: 8192, + unique: true, + column_position: 1, + direction: 'asc', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'moshe', + column: 'tenant_id', + index_type: 'btree', + cardinality: 0, + size: 8192, + unique: true, + column_position: 2, + direction: 'asc', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'moshe', + column: 'tenant_user_id', + index_type: 'btree', + cardinality: 0, + size: 8192, + unique: true, + column_position: 3, + direction: 'asc', + }, + // Another unique index + { + schema: 'landlord', + table: 'users_master_table', + name: 'users_master_table_index_1', + column: 'tenant_id', + index_type: 'btree', + cardinality: 0, + size: 8192, + unique: true, + column_position: 1, + direction: 'asc', + }, + { + schema: 'landlord', + table: 'users_master_table', + name: 'users_master_table_index_1', + column: 'tenant_user_id', + index_type: 'btree', + cardinality: 0, + size: 8192, + unique: true, + column_position: 2, + direction: 'asc', + }, + ], + tables: [ + { + schema: 'landlord', + table: 'users_master_table', + rows: 0, + type: 'BASE TABLE', + engine: '', + collation: '', + comment: '', + }, + ], + views: [], + custom_types: [], + }; + + const tables = createTablesFromMetadata({ + databaseMetadata: metadata, + databaseType: DatabaseType.POSTGRESQL, + }); + + expect(tables).toHaveLength(1); + const table = tables[0]; + + // Check that the composite PK name was captured as "moshe" in the PK index + const pkIndex = table.indexes.find((idx) => idx.isPrimaryKey); + expect(pkIndex).toBeDefined(); + expect(pkIndex!.name).toBe('moshe'); + + // Check that primary key fields are marked correctly + const pkFields = table.fields.filter((f) => f.primaryKey); + expect(pkFields).toHaveLength(3); + expect(pkFields.map((f) => f.name).sort()).toEqual([ + 'master_user_id', + 'tenant_id', + 'tenant_user_id', + ]); + + // Check that we have both the PK index and the unique index + expect(table.indexes).toHaveLength(2); + const uniqueIndex = table.indexes.find((idx) => !idx.isPrimaryKey); + expect(uniqueIndex!.name).toBe('users_master_table_index_1'); + }); +}); diff --git a/src/lib/domain/db-index.ts b/src/lib/domain/db-index.ts index 5c160928..b9a8c900 100644 --- a/src/lib/domain/db-index.ts +++ b/src/lib/domain/db-index.ts @@ -3,6 +3,7 @@ import type { AggregatedIndexInfo } from '../data/import-metadata/metadata-types import { generateId } from '../utils'; import type { DBField } from './db-field'; import { DatabaseType } from './database-type'; +import type { DBTable } from './db-table'; export const INDEX_TYPES = [ 'btree', @@ -29,6 +30,7 @@ export interface DBIndex { fieldIds: string[]; createdAt: number; type?: IndexType | null; + isPrimaryKey?: boolean | null; } export const dbIndexSchema: z.ZodType = z.object({ @@ -38,6 +40,7 @@ export const dbIndexSchema: z.ZodType = z.object({ fieldIds: z.array(z.string()), createdAt: z.number(), type: z.enum(INDEX_TYPES).optional(), + isPrimaryKey: z.boolean().or(z.null()).optional(), }); export const createIndexesFromMetadata = ({ @@ -64,3 +67,51 @@ export const createIndexesFromMetadata = ({ export const databaseIndexTypes: { [key in DatabaseType]?: IndexType[] } = { [DatabaseType.POSTGRESQL]: ['btree', 'hash'], }; + +export const getTablePrimaryKeyIndex = ({ + table, +}: { + table: DBTable; +}): DBIndex | null => { + const primaryKeyFields = table.fields.filter((f) => f.primaryKey); + const existingPKIndex = table.indexes.find((idx) => idx.isPrimaryKey); + + if (primaryKeyFields.length === 0) { + return null; + } + + const pkFieldIds = primaryKeyFields.map((f) => f.id); + + if (existingPKIndex) { + return { + ...existingPKIndex, + fieldIds: pkFieldIds, + }; + } else { + // Create new PK index for primary key(s) + const pkIndex: DBIndex = { + id: generateId(), + name: `pk_${table.name}_${primaryKeyFields.map((f) => f.name).join('_')}`, + fieldIds: pkFieldIds, + unique: true, + isPrimaryKey: true, + createdAt: Date.now(), + }; + + return pkIndex; + } +}; + +export const getTableIndexesWithPrimaryKey = ({ + table, +}: { + table: DBTable; +}): DBIndex[] => { + const primaryKeyIndex = getTablePrimaryKeyIndex({ table }); + const indexesWithoutPKIndex = table.indexes.filter( + (idx) => !idx.isPrimaryKey + ); + return primaryKeyIndex + ? [primaryKeyIndex, ...indexesWithoutPKIndex] + : indexesWithoutPKIndex; +}; diff --git a/src/lib/domain/db-table.ts b/src/lib/domain/db-table.ts index ecf0de01..4879f469 100644 --- a/src/lib/domain/db-table.ts +++ b/src/lib/domain/db-table.ts @@ -203,11 +203,57 @@ export const createTablesFromMetadata = ({ tableSchema, }); + // Check for composite primary key and find matching index name + const primaryKeyFields = fields.filter((f) => f.primaryKey); + let pkMatchingIndexName: string | undefined; + let pkIndex: DBIndex | undefined; + + if (primaryKeyFields.length >= 1) { + // We have a composite primary key, look for an index that matches all PK columns + const pkFieldNames = primaryKeyFields.map((f) => f.name).sort(); + + // Find an index that matches the primary key columns exactly + const matchingIndex = aggregatedIndexes.find((index) => { + const indexColumnNames = index.columns + .map((c) => c.name) + .sort(); + return ( + indexColumnNames.length === pkFieldNames.length && + indexColumnNames.every((col, i) => col === pkFieldNames[i]) + ); + }); + + if (matchingIndex) { + pkMatchingIndexName = matchingIndex.name; + // Create a special PK index + pkIndex = { + id: generateId(), + name: matchingIndex.name, + unique: true, + fieldIds: primaryKeyFields.map((f) => f.id), + createdAt: Date.now(), + isPrimaryKey: true, + }; + } + } + + // Filter out the index that matches the composite PK (to avoid duplication) + const filteredAggregatedIndexes = pkMatchingIndexName + ? aggregatedIndexes.filter( + (idx) => idx.name !== pkMatchingIndexName + ) + : aggregatedIndexes; + const dbIndexes = createIndexesFromMetadata({ - aggregatedIndexes, + aggregatedIndexes: filteredAggregatedIndexes, fields, }); + // Add the PK index if it exists + if (pkIndex) { + dbIndexes.push(pkIndex); + } + // Determine if the current table is a view by checking against pre-computed sets const viewKey = generateTableKey({ schemaName: tableSchema, diff --git a/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-index/table-index.tsx b/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-index/table-index.tsx index 5d006b22..dbf754f7 100644 --- a/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-index/table-index.tsx +++ b/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-index/table-index.tsx @@ -1,5 +1,5 @@ import React, { useCallback, useMemo } from 'react'; -import { Ellipsis, Trash2 } from 'lucide-react'; +import { Ellipsis, Trash2, KeyRound } from 'lucide-react'; import { Button } from '@/components/button/button'; import { databaseIndexTypes, @@ -106,29 +106,45 @@ export const TableIndex: React.FC = ({ 'side_panel.tables_section.table.no_types_found' )} keepOrder + disabled={index.isPrimaryKey ?? false} />
- - - - - updateIndex({ - unique: !!value, - }) - } - > - U - - - - - {t( - 'side_panel.tables_section.table.index_actions.unique' - )} - - + {index.isPrimaryKey ? ( + + + + + + + + + + {t('side_panel.tables_section.table.primary_key')} + + + ) : ( + + + + + updateIndex({ + unique: !!value, + }) + } + > + U + + + + + {t( + 'side_panel.tables_section.table.index_actions.unique' + )} + + + )}
-
- - - updateIndex({ - unique: !!value, - }) - } - /> -
- {indexTypeOptions.length > 0 ? ( -
-
+ + ) : null} - - diff --git a/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-list-item-content.tsx b/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-list-item-content.tsx index 406f8d72..309562ee 100644 --- a/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-list-item-content.tsx +++ b/src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-list-item-content.tsx @@ -224,19 +224,27 @@ export const TableListItemContent: React.FC = ({ - {table.indexes.map((index) => ( - - removeIndex(table.id, index.id) - } - updateIndex={(attrs: Partial) => - updateIndex(table.id, index.id, attrs) - } - fields={table.fields} - /> - ))} + {[...table.indexes] + .sort((a, b) => { + // Sort PK indexes first + if (a.isPrimaryKey && !b.isPrimaryKey) + return -1; + if (!a.isPrimaryKey && b.isPrimaryKey) return 1; + return 0; + }) + .map((index) => ( + + removeIndex(table.id, index.id) + } + updateIndex={(attrs: Partial) => + updateIndex(table.id, index.id, attrs) + } + fields={table.fields} + /> + ))}