diff --git a/src/lib/data/export-metadata/export-sql-script.ts b/src/lib/data/export-metadata/export-sql-script.ts index af8c4e34..6c464ea6 100644 --- a/src/lib/data/export-metadata/export-sql-script.ts +++ b/src/lib/data/export-metadata/export-sql-script.ts @@ -84,7 +84,55 @@ export const exportBaseSQL = ({ schemas.forEach((schema) => { sqlScript += `CREATE SCHEMA IF NOT EXISTS ${schema};\n`; }); - sqlScript += '\n'; + if (schemas.size > 0) sqlScript += '\n'; // Add newline only if schemas were added + + // Add CREATE TYPE statements for ENUMs and COMPOSITE types from diagram.customTypes + if (diagram.customTypes && diagram.customTypes.length > 0) { + diagram.customTypes.forEach((customType) => { + const typeNameWithSchema = customType.schema + ? `${customType.schema}.${customType.name}` + : customType.name; + + if ( + customType.kind === 'enum' && + customType.values && + customType.values.length > 0 + ) { + // For PostgreSQL, generate CREATE TYPE ... AS ENUM + // For other DBs, this might need adjustment or be omitted if not supported directly + // or if we rely on the DBML generator to create Enums separately (as currently done) + // For now, let's assume PostgreSQL-style for demonstration if isDBMLFlow is false. + // If isDBMLFlow is true, we let TableDBML.tsx handle Enum syntax directly. + if ( + targetDatabaseType === DatabaseType.POSTGRESQL && + !isDBMLFlow + ) { + const enumValues = customType.values + .map((v) => `'${v.replace(/'/g, "''")}'`) + .join(', '); + sqlScript += `CREATE TYPE ${typeNameWithSchema} AS ENUM (${enumValues});\n`; + } + } else if ( + customType.kind === 'composite' && + customType.fields && + customType.fields.length > 0 + ) { + // For PostgreSQL, generate CREATE TYPE ... AS (...) + // This is crucial for composite types to be recognized by the DBML importer + if ( + targetDatabaseType === DatabaseType.POSTGRESQL || + isDBMLFlow + ) { + // Assume other DBs might not support this or DBML flow needs it + const compositeFields = customType.fields + .map((f) => `${f.field} ${simplifyDataType(f.type)}`) + .join(',\n '); + sqlScript += `CREATE TYPE ${typeNameWithSchema} AS (\n ${compositeFields}\n);\n`; + } + } + }); + sqlScript += '\n'; // Add a newline if custom types were processed + } // Add CREATE SEQUENCE statements const sequences = new Set(); @@ -119,8 +167,45 @@ export const exportBaseSQL = ({ let typeName = simplifyDataType(field.type.name); // Handle ENUM type - if (typeName.toLowerCase() === 'enum') { - // Map enum to TEXT for broader compatibility, especially with DBML importer + // If we are generating SQL for DBML flow, and we ALREADY generated CREATE TYPE for enums (e.g., for PG), + // then we should use the enum type name. Otherwise, map to text. + // However, the current TableDBML.tsx generates its own Enum blocks, so for DBML flow, + // converting to TEXT here might still be the safest bet to avoid conflicts if SQL enums aren't perfectly parsed. + // Let's adjust: if it's a known custom enum type, use its name for PG, otherwise TEXT. + const customEnumType = diagram.customTypes?.find( + (ct) => + ct.name === field.type.name && + ct.kind === 'enum' && + (ct.schema ? ct.schema === table.schema : true) + ); + + if ( + customEnumType && + targetDatabaseType === DatabaseType.POSTGRESQL && + !isDBMLFlow + ) { + typeName = customEnumType.schema + ? `${customEnumType.schema}.${customEnumType.name}` + : customEnumType.name; + } else if (typeName.toLowerCase() === 'enum') { + // Fallback for non-PG or if custom type not found, or for DBML flow if not handled by CREATE TYPE above + typeName = 'text'; + } + + // Check if the field type is a known composite custom type + const customCompositeType = diagram.customTypes?.find( + (ct) => + ct.name === field.type.name && + ct.kind === 'composite' && + (ct.schema ? ct.schema === table.schema : true) + ); + + if (customCompositeType) { + typeName = customCompositeType.schema + ? `${customCompositeType.schema}.${customCompositeType.name}` + : customCompositeType.name; + } else if (typeName.toLowerCase() === 'user-defined') { + // If it's 'user-defined' but not a known composite, fallback to TEXT typeName = 'text'; } @@ -129,11 +214,6 @@ export const exportBaseSQL = ({ typeName = 'text[]'; } - // Temp fix for 'user-defined' to be text - if (typeName.toLowerCase() === 'user-defined') { - typeName = 'text'; - } - sqlScript += ` ${field.name} ${typeName}`; // Add size for character types diff --git a/src/pages/editor-page/side-panel/tables-section/table-dbml/table-dbml.tsx b/src/pages/editor-page/side-panel/tables-section/table-dbml/table-dbml.tsx index 8402da9f..4b8fc0ca 100644 --- a/src/pages/editor-page/side-panel/tables-section/table-dbml/table-dbml.tsx +++ b/src/pages/editor-page/side-panel/tables-section/table-dbml/table-dbml.tsx @@ -12,11 +12,41 @@ import { setupDBMLLanguage } from '@/components/code-snippet/languages/dbml-lang import { DatabaseType } from '@/lib/domain/database-type'; import { ArrowLeftRight } from 'lucide-react'; import { type DBField } from '@/lib/domain/db-field'; +import type { DBCustomType } from '@/lib/domain/db-custom-type'; +import { DBCustomTypeKind } from '@/lib/domain/db-custom-type'; export interface TableDBMLProps { filteredTables: DBTable[]; } +// Use DBCustomType for generating Enum DBML +const generateEnumsDBML = (customTypes: DBCustomType[] | undefined): string => { + if (!customTypes || customTypes.length === 0) { + return ''; + } + + // Filter for enum types and map them + return customTypes + .filter((ct) => ct.kind === DBCustomTypeKind.enum) + .map((enumDef) => { + const enumIdentifier = enumDef.schema + ? `"${enumDef.schema}"."${enumDef.name.replace(/"/g, '\\"')}"` + : `"${enumDef.name.replace(/"/g, '\\"')}"`; + + const valuesString = (enumDef.values || []) // Ensure values array exists + .map((valueName) => { + // valueName is a string as per DBCustomType + const valLine = ` "${valueName.replace(/"/g, '\\"')}"`; + // If you have notes per enum value, you'd need to adjust DBCustomType + // For now, assuming no notes per value in DBCustomType + return valLine; + }) + .join('\n'); + return `Enum ${enumIdentifier} {\n${valuesString}\n}\n`; + }) + .join('\n'); +}; + const getEditorTheme = (theme: EffectiveTheme) => { return theme === 'dark' ? 'dbml-dark' : 'dbml-light'; }; @@ -141,6 +171,19 @@ const sanitizeSQLforDBML = (sql: string): string => { } ); + // Comment out self-referencing foreign keys to prevent "Two endpoints are the same" error + // Example: ALTER TABLE public.class ADD CONSTRAINT ... FOREIGN KEY (class_id) REFERENCES public.class (class_id); + const lines = sanitized.split('\n'); + const processedLines = lines.map((line) => { + const selfRefFKPattern = + /ALTER\s+TABLE\s+(?:\S+\.)?(\S+)\s+ADD\s+CONSTRAINT\s+\S+\s+FOREIGN\s+KEY\s*\([^)]+\)\s+REFERENCES\s+(?:\S+\.)?\1\s*\([^)]+\)\s*;/i; + if (selfRefFKPattern.test(line)) { + return `-- ${line}`; // Comment out the line + } + return line; + }); + sanitized = processedLines.join('\n'); + // Replace any remaining problematic characters sanitized = sanitized.replace(/\?\?/g, '__'); @@ -287,7 +330,7 @@ export const TableDBML: React.FC = ({ filteredTables }) => { const { effectiveTheme } = useTheme(); const { toast } = useToast(); const [dbmlFormat, setDbmlFormat] = useState<'inline' | 'standard'>( - 'standard' + 'inline' ); // --- Effect for handling empty field name warnings --- @@ -439,6 +482,9 @@ export const TableDBML: React.FC = ({ filteredTables }) => { let inline = ''; let baseScript = ''; // Define baseScript outside try + // Use finalDiagramForExport.customTypes which should be DBCustomType[] + const enumsDBML = generateEnumsDBML(finalDiagramForExport.customTypes); + try { baseScript = exportBaseSQL({ diagram: finalDiagramForExport, // Use final diagram @@ -467,6 +513,9 @@ export const TableDBML: React.FC = ({ filteredTables }) => { ) ); + // Prepend Enum DBML to the standard output + standard = enumsDBML + '\n' + standard; + inline = normalizeCharTypeFormat(convertToInlineRefs(standard)); } catch (error: unknown) { console.error( @@ -495,6 +544,15 @@ export const TableDBML: React.FC = ({ filteredTables }) => { variant: 'destructive', }); } + + // If an error occurred, still prepend enums if they exist, or they'll be lost. + // The error message will then follow. + if (standard.startsWith('// Error generating DBML:')) { + standard = enumsDBML + standard; + } + if (inline.startsWith('// Error generating DBML:')) { + inline = enumsDBML + inline; + } } return { standardDbml: standard, inlineDbml: inline }; }, [currentDiagram, filteredTables, toast]); // Keep toast dependency for now, although direct call is removed