From d79b0c205d6007144484c59cdaf3e191417fb75d Mon Sep 17 00:00:00 2001
From: 1ilit <lilitmousheghyan@gmail.com>
Date: Thu, 27 Jun 2024 11:12:09 +0300
Subject: [PATCH] Return types after import

---
 src/components/EditorHeader/Modal/Modal.jsx |  8 ++++++--
 src/components/Workspace.jsx                |  2 +-
 src/data/databases.js                       | 22 +++++++++++++--------
 src/utils/importSQL/index.js                |  2 +-
 src/utils/importSQL/postgres.js             |  3 ++-
 5 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/components/EditorHeader/Modal/Modal.jsx b/src/components/EditorHeader/Modal/Modal.jsx
index 5f253be..8b634c1 100644
--- a/src/components/EditorHeader/Modal/Modal.jsx
+++ b/src/components/EditorHeader/Modal/Modal.jsx
@@ -138,14 +138,18 @@ export default function Modal({
       return;
     }
 
-    const d = importSQL(ast, database === DB.GENERIC ? importDb : database, database);
+    const d = importSQL(
+      ast,
+      database === DB.GENERIC ? importDb : database,
+      database,
+    );
     if (importSource.overwrite) {
       setTables(d.tables);
       setRelationships(d.relationships);
       setTransform((prev) => ({ ...prev, pan: { x: 0, y: 0 } }));
       setNotes([]);
       setAreas([]);
-      setTypes([]);
+      setTypes(d.types ?? []);
       setUndoStack([]);
       setRedoStack([]);
     } else {
diff --git a/src/components/Workspace.jsx b/src/components/Workspace.jsx
index df38e03..cdd4e24 100644
--- a/src/components/Workspace.jsx
+++ b/src/components/Workspace.jsx
@@ -355,7 +355,7 @@ export default function WorkSpace() {
         okButtonProps={{ disabled: selectedDb === "" }}
       >
         <div className="grid grid-cols-3 gap-4 place-content-center">
-          {databases.map((x) => (
+          {Object.values(databases).map((x) => (
             <div
               key={x.name}
               onClick={() => setSelectedDb(x.label)}
diff --git a/src/data/databases.js b/src/data/databases.js
index 70996e3..78e2e8d 100644
--- a/src/data/databases.js
+++ b/src/data/databases.js
@@ -6,36 +6,42 @@ import mssqlImage from "../assets/mssql-icon.png";
 import i18n from "../i18n/i18n";
 import { DB } from "./constants";
 
-export const databases = [
-  {
+export const databases = {
+  [DB.MYSQL]: {
     name: "MySQL",
     label: DB.MYSQL,
     image: mysqlImage,
+    hasTypes: false,
   },
-  {
+  [DB.POSTGRES]: {
     name: "PostgreSQL",
     label: DB.POSTGRES,
     image: postgresImage,
+    hasTypes: true,
   },
-  {
+  [DB.SQLITE]: {
     name: "SQLite",
     label: DB.SQLITE,
     image: sqliteImage,
+    hasTypes: false,
   },
-  {
+  [DB.MARIADB]: {
     name: "MariaDB",
     label: DB.MARIADB,
     image: mariadbImage,
+    hasTypes: false,
   },
-  {
+  [DB.MSSQL]: {
     name: "MSSQL",
     label: DB.MSSQL,
     image: mssqlImage,
+    hasTypes: false,
   },
-  {
+  [DB.GENERIC]: {
     name: i18n.t("generic"),
     label: DB.GENERIC,
     image: null,
     description: i18n.t("generic_description"),
+    hasTypes: true,
   },
-];
+};
diff --git a/src/utils/importSQL/index.js b/src/utils/importSQL/index.js
index 59a9b00..4bf1951 100644
--- a/src/utils/importSQL/index.js
+++ b/src/utils/importSQL/index.js
@@ -10,7 +10,7 @@ import { fromPostgres } from "./postgres";
 import { fromSQLite } from "./sqlite";
 
 export function importSQL(ast, toDb = DB.MYSQL, diagramDb = DB.GENERIC) {
-  let diagram = { tables: [], relationships: [] };
+  let diagram;
   switch (toDb) {
     case DB.SQLITE:
       diagram = fromSQLite(ast, diagramDb);
diff --git a/src/utils/importSQL/postgres.js b/src/utils/importSQL/postgres.js
index 61703a3..089c303 100644
--- a/src/utils/importSQL/postgres.js
+++ b/src/utils/importSQL/postgres.js
@@ -20,6 +20,7 @@ const affinity = {
 export function fromPostgres(ast, diagramDb = DB.GENERIC) {
   const tables = [];
   const relationships = [];
+  const types = [];
 
   ast.forEach((e) => {
     if (e.type === "create") {
@@ -241,5 +242,5 @@ export function fromPostgres(ast, diagramDb = DB.GENERIC) {
 
   relationships.forEach((r, i) => (r.id = i));
 
-  return { tables, relationships };
+  return { tables, relationships, types };
 }