feat: add table selection for large database imports (#776)

* feat: add table selection UI for large database imports (>50 tables)

* some changes

* some changes

* some changes

* fix

* fix

---------

Co-authored-by: Guy Ben-Aharon <baguy3@gmail.com>

src/dialogs/create-diagram-dialog/create-diagram-dialog.tsx

@@ -15,9 +15,13 @@ import type { DatabaseEdition } from '@/lib/domain/database-edition';
 import { SelectDatabase } from './select-database/select-database';
 import { CreateDiagramDialogStep } from './create-diagram-dialog-step';
 import { ImportDatabase } from '../common/import-database/import-database';
+import { SelectTables } from '../common/select-tables/select-tables';
 import { useTranslation } from 'react-i18next';
 import type { BaseDialogProps } from '../common/base-dialog-props';
 import { sqlImportToDiagram } from '@/lib/data/sql-import';
+import type { SelectedTable } from '@/lib/data/import-metadata/filter-metadata';
+import { filterMetadataByTables } from '@/lib/data/import-metadata/filter-metadata';
+import { MAX_TABLES_WITHOUT_SHOWING_FILTER } from '../common/select-tables/constants';
 
 export interface CreateDiagramDialogProps extends BaseDialogProps {}
 
@@ -42,6 +46,8 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
     const { listDiagrams, addDiagram } = useStorage();
     const [diagramNumber, setDiagramNumber] = useState<number>(1);
     const navigate = useNavigate();
+    const [parsedMetadata, setParsedMetadata] = useState<DatabaseMetadata>();
+    const [isParsingMetadata, setIsParsingMetadata] = useState(false);
 
     useEffect(() => {
         setDatabaseEdition(undefined);
@@ -62,49 +68,72 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
         setDatabaseEdition(undefined);
         setScriptResult('');
         setImportMethod('query');
+        setParsedMetadata(undefined);
     }, [dialog.open]);
 
     const hasExistingDiagram = (diagramId ?? '').trim().length !== 0;
 
-    const importNewDiagram = useCallback(async () => {
-        let diagram: Diagram | undefined;
+    const importNewDiagram = useCallback(
+        async ({
+            selectedTables,
+            databaseMetadata,
+        }: {
+            selectedTables?: SelectedTable[];
+            databaseMetadata?: DatabaseMetadata;
+        } = {}) => {
+            let diagram: Diagram | undefined;
 
-        if (importMethod === 'ddl') {
-            diagram = await sqlImportToDiagram({
-                sqlContent: scriptResult,
-                sourceDatabaseType: databaseType,
-                targetDatabaseType: databaseType,
+            if (importMethod === 'ddl') {
+                diagram = await sqlImportToDiagram({
+                    sqlContent: scriptResult,
+                    sourceDatabaseType: databaseType,
+                    targetDatabaseType: databaseType,
+                });
+            } else {
+                let metadata: DatabaseMetadata | undefined = databaseMetadata;
+
+                if (!metadata) {
+                    metadata = loadDatabaseMetadata(scriptResult);
+                }
+
+                if (selectedTables && selectedTables.length > 0) {
+                    metadata = filterMetadataByTables({
+                        metadata,
+                        selectedTables,
+                    });
+                }
+
+                diagram = await loadFromDatabaseMetadata({
+                    databaseType,
+                    databaseMetadata: metadata,
+                    diagramNumber,
+                    databaseEdition:
+                        databaseEdition?.trim().length === 0
+                            ? undefined
+                            : databaseEdition,
+                });
+            }
+
+            await addDiagram({ diagram });
+            await updateConfig({
+                config: { defaultDiagramId: diagram.id },
             });
-        } else {
-            const databaseMetadata: DatabaseMetadata =
-                loadDatabaseMetadata(scriptResult);
 
-            diagram = await loadFromDatabaseMetadata({
-                databaseType,
-                databaseMetadata,
-                diagramNumber,
-                databaseEdition:
-                    databaseEdition?.trim().length === 0
-                        ? undefined
-                        : databaseEdition,
-            });
-        }
-
-        await addDiagram({ diagram });
-        await updateConfig({ config: { defaultDiagramId: diagram.id } });
-        closeCreateDiagramDialog();
-        navigate(`/diagrams/${diagram.id}`);
-    }, [
-        importMethod,
-        databaseType,
-        addDiagram,
-        databaseEdition,
-        closeCreateDiagramDialog,
-        navigate,
-        updateConfig,
-        scriptResult,
-        diagramNumber,
-    ]);
+            closeCreateDiagramDialog();
+            navigate(`/diagrams/${diagram.id}`);
+        },
+        [
+            importMethod,
+            databaseType,
+            addDiagram,
+            databaseEdition,
+            closeCreateDiagramDialog,
+            navigate,
+            updateConfig,
+            scriptResult,
+            diagramNumber,
+        ]
+    );
 
     const createEmptyDiagram = useCallback(async () => {
         const diagram: Diagram = {
@@ -138,10 +167,56 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
         openImportDBMLDialog,
     ]);
 
+    const importNewDiagramOrFilterTables = useCallback(async () => {
+        try {
+            setIsParsingMetadata(true);
+
+            if (importMethod === 'ddl') {
+                await importNewDiagram();
+            } else {
+                // Parse metadata asynchronously to avoid blocking the UI
+                const metadata = await new Promise<DatabaseMetadata>(
+                    (resolve, reject) => {
+                        setTimeout(() => {
+                            try {
+                                const result =
+                                    loadDatabaseMetadata(scriptResult);
+                                resolve(result);
+                            } catch (err) {
+                                reject(err);
+                            }
+                        }, 0);
+                    }
+                );
+
+                const totalTablesAndViews =
+                    metadata.tables.length + (metadata.views?.length || 0);
+
+                setParsedMetadata(metadata);
+
+                // Check if it's a large database that needs table selection
+                if (totalTablesAndViews > MAX_TABLES_WITHOUT_SHOWING_FILTER) {
+                    setStep(CreateDiagramDialogStep.SELECT_TABLES);
+                } else {
+                    await importNewDiagram({
+                        databaseMetadata: metadata,
+                    });
+                }
+            }
+        } finally {
+            setIsParsingMetadata(false);
+        }
+    }, [importMethod, scriptResult, importNewDiagram]);
+
     return (
         <Dialog
             {...dialog}
             onOpenChange={(open) => {
+                // Don't allow closing while parsing metadata
+                if (isParsingMetadata) {
+                    return;
+                }
+
                 if (!hasExistingDiagram) {
                     return;
                 }
@@ -154,6 +229,8 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
             <DialogContent
                 className="flex max-h-dvh w-full flex-col md:max-w-[900px]"
                 showClose={hasExistingDiagram}
+                onInteractOutside={(e) => e.preventDefault()}
+                onEscapeKeyDown={(e) => e.preventDefault()}
             >
                 {step === CreateDiagramDialogStep.SELECT_DATABASE ? (
                     <SelectDatabase
@@ -165,9 +242,9 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
                             setStep(CreateDiagramDialogStep.IMPORT_DATABASE)
                         }
                     />
-                ) : (
+                ) : step === CreateDiagramDialogStep.IMPORT_DATABASE ? (
                     <ImportDatabase
-                        onImport={importNewDiagram}
+                        onImport={importNewDiagramOrFilterTables}
                         onCreateEmptyDiagram={createEmptyDiagram}
                         databaseEdition={databaseEdition}
                         databaseType={databaseType}
@@ -180,8 +257,18 @@ export const CreateDiagramDialog: React.FC<CreateDiagramDialogProps> = ({
                         title={t('new_diagram_dialog.import_database.title')}
                         importMethod={importMethod}
                         setImportMethod={setImportMethod}
+                        keepDialogAfterImport={true}
                     />
-                )}
+                ) : step === CreateDiagramDialogStep.SELECT_TABLES ? (
+                    <SelectTables
+                        isLoading={isParsingMetadata || !parsedMetadata}
+                        databaseMetadata={parsedMetadata}
+                        onImport={importNewDiagram}
+                        onBack={() =>
+                            setStep(CreateDiagramDialogStep.IMPORT_DATABASE)
+                        }
+                    />
+                ) : null}
             </DialogContent>
         </Dialog>
     );