feat: Add ER diagram visualization for database schemas

- Right-click a schema node → "Show ER Diagram" opens an interactive
  React Flow-based diagram showing tables, columns, PK/FK markers,
  and foreign key relationships with column-level connections
- Dagre layout engine for automatic positioning that minimizes
  edge crossings
- Per-table color coding with colored relationship edges
- Focus mode: click a table or use the search field to highlight
  it and its connected tables, dimming everything else
- Foreign key queries added for PostgreSQL, MySQL, MSSQL, and SQLite
- Fixed pre-existing PG driver notice listener leak

Author: jcjaskula-aws

package.json

@@ -57,6 +57,9 @@
     "source-map-support": "^0.5.11",
     "typescript": "~4.8.3",
     "@vscode/vsce": "^2.19.0",
-    "ts-node": "^9.1.1"
+    "ts-node": "^9.1.1",
+    "react-flow-renderer": "^10.3.17",
+    "dagre": "^0.8.5",
+    "@types/dagre": "^0.7.52"
   }
 }

packages/base-driver/src/index.ts

@@ -115,6 +115,24 @@ export default abstract class AbstractDriver<ConnectionType extends any, DriverO
     return Promise.resolve([]);
   }
 
+  public async getERDiagramData(schema: NSDatabase.ISchema): Promise<{ tables: NSDatabase.ITable[], columns: { [tableName: string]: NSDatabase.IColumn[] }, foreignKeys: NSDatabase.IForeignKey[] }> {
+    const tables: NSDatabase.ITable[] = await this.queryResults(this.queries.fetchTables(schema));
+    const columns: { [tableName: string]: NSDatabase.IColumn[] } = {};
+    for (const table of tables) {
+      const cols: NSDatabase.IColumn[] = await this.queryResults(this.queries.fetchColumns(table));
+      columns[table.label] = cols;
+    }
+    let foreignKeys: NSDatabase.IForeignKey[] = [];
+    if (this.queries.fetchForeignKeys) {
+      try {
+        foreignKeys = await this.queryResults(this.queries.fetchForeignKeys(schema));
+      } catch (e) {
+        this.log.warn('Failed to fetch foreign keys: %O', e);
+      }
+    }
+    return { tables, columns, foreignKeys };
+  }
+
   public async toAbsolutePath(fsPath: string) {
     if (!path.isAbsolute(fsPath) && /\$\{workspaceFolder:(.+)}/g.test(fsPath)) {
       const workspaceName = fsPath.match(/\$\{workspaceFolder:(.+)}/)[1];

packages/base-driver/src/lib/factory.test.ts

@@ -17,5 +17,44 @@ describe(`query generator`, () => {
 
     const expected = `SELECT * FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = 'table' AND TABLE_CATALOG = 'catalog' AND TABLE_SCHEMA = 'schema' AND something = 'smtg'`;
     expect(result).toBe(expected);
-  })
+  });
+
+  it(`should generate foreign key queries with schema and database params`, () => {
+    const fetchForeignKeys = queryFactory<{ schema: string; database: string }>`
+    SELECT
+      tc.constraint_name AS "constraintName",
+      tc.table_schema AS "sourceTableSchema",
+      tc.table_name AS "sourceTableName",
+      kcu.column_name AS "sourceColumnName",
+      ccu.table_schema AS "targetTableSchema",
+      ccu.table_name AS "targetTableName",
+      ccu.column_name AS "targetColumnName"
+    FROM information_schema.table_constraints AS tc
+    WHERE
+      tc.constraint_type = 'FOREIGN KEY'
+      AND tc.table_schema = '${p => p.schema}'
+      AND tc.table_catalog = '${p => p.database}'
+    `;
+    const result = fetchForeignKeys({ schema: 'public', database: 'mydb' });
+
+    expect(result).toContain("tc.table_schema = 'public'");
+    expect(result).toContain("tc.table_catalog = 'mydb'");
+    expect(result).toContain("tc.constraint_type = 'FOREIGN KEY'");
+    expect(result).toContain('"constraintName"');
+    expect(result).toContain('"sourceTableName"');
+    expect(result).toContain('"targetTableName"');
+  });
+
+  it(`should handle queries with no dynamic params`, () => {
+    const staticQuery = queryFactory`SELECT 1 AS result`;
+    expect(staticQuery({})).toBe('SELECT 1 AS result');
+  });
+
+  it(`should expose raw query template`, () => {
+    const query = queryFactory<{ name: string }>`
+    SELECT * FROM tables WHERE name = '${p => p.name}'
+    `;
+    expect(query.raw).toBeDefined();
+    expect(query.raw).toContain('${');
+  });
 });

packages/driver.mssql/src/ls/queries.ts

@@ -188,6 +188,33 @@ FETCH NEXT ${p => p.limit || 100} ROWS ONLY
 `;
 
 
+export const fetchForeignKeys: IBaseQueries['fetchForeignKeys'] = queryFactory`
+SELECT
+  FK.CONSTRAINT_NAME AS "constraintName",
+  FK.TABLE_SCHEMA AS "sourceTableSchema",
+  FK.TABLE_NAME AS "sourceTableName",
+  CU.COLUMN_NAME AS "sourceColumnName",
+  PK.TABLE_SCHEMA AS "targetTableSchema",
+  PK.TABLE_NAME AS "targetTableName",
+  PT.COLUMN_NAME AS "targetColumnName"
+FROM
+  ${p => p.database ? `${escapeTableName({ database: p.database, schema: "INFORMATION_SCHEMA", label: "REFERENTIAL_CONSTRAINTS" })}` : 'INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS'} AS C
+  JOIN ${p => p.database ? `${escapeTableName({ database: p.database, schema: "INFORMATION_SCHEMA", label: "TABLE_CONSTRAINTS" })}` : 'INFORMATION_SCHEMA.TABLE_CONSTRAINTS'} AS FK
+    ON C.CONSTRAINT_NAME = FK.CONSTRAINT_NAME AND C.CONSTRAINT_SCHEMA = FK.CONSTRAINT_SCHEMA
+  JOIN ${p => p.database ? `${escapeTableName({ database: p.database, schema: "INFORMATION_SCHEMA", label: "TABLE_CONSTRAINTS" })}` : 'INFORMATION_SCHEMA.TABLE_CONSTRAINTS'} AS PK
+    ON C.UNIQUE_CONSTRAINT_NAME = PK.CONSTRAINT_NAME AND C.UNIQUE_CONSTRAINT_SCHEMA = PK.CONSTRAINT_SCHEMA
+  JOIN ${p => p.database ? `${escapeTableName({ database: p.database, schema: "INFORMATION_SCHEMA", label: "KEY_COLUMN_USAGE" })}` : 'INFORMATION_SCHEMA.KEY_COLUMN_USAGE'} AS CU
+    ON C.CONSTRAINT_NAME = CU.CONSTRAINT_NAME AND C.CONSTRAINT_SCHEMA = CU.CONSTRAINT_SCHEMA
+  JOIN ${p => p.database ? `${escapeTableName({ database: p.database, schema: "INFORMATION_SCHEMA", label: "KEY_COLUMN_USAGE" })}` : 'INFORMATION_SCHEMA.KEY_COLUMN_USAGE'} AS PT
+    ON C.UNIQUE_CONSTRAINT_NAME = PT.CONSTRAINT_NAME AND C.UNIQUE_CONSTRAINT_SCHEMA = PT.CONSTRAINT_SCHEMA
+WHERE
+  FK.TABLE_SCHEMA = '${p => p.schema}'
+  AND FK.TABLE_CATALOG = '${p => p.database}'
+ORDER BY
+  FK.TABLE_NAME,
+  CU.COLUMN_NAME
+`;
+
 // export default {
 //   fetchFunctions: `
 // SELECT

packages/driver.mysql/src/ls/queries.ts

@@ -182,6 +182,25 @@ ORDER BY
 LIMIT ${p => p.limit || 100}
 
 `;
+export const fetchForeignKeys: IBaseQueries['fetchForeignKeys'] = queryFactory`
+SELECT
+  KCU.CONSTRAINT_NAME AS "constraintName",
+  KCU.TABLE_SCHEMA AS "sourceTableSchema",
+  KCU.TABLE_NAME AS "sourceTableName",
+  KCU.COLUMN_NAME AS "sourceColumnName",
+  KCU.REFERENCED_TABLE_SCHEMA AS "targetTableSchema",
+  KCU.REFERENCED_TABLE_NAME AS "targetTableName",
+  KCU.REFERENCED_COLUMN_NAME AS "targetColumnName"
+FROM
+  INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS KCU
+WHERE
+  KCU.REFERENCED_TABLE_NAME IS NOT NULL
+  AND KCU.TABLE_SCHEMA = '${p => p.schema}'
+ORDER BY
+  KCU.TABLE_NAME,
+  KCU.COLUMN_NAME
+`;
+
 // export default {
 //   fetchFunctions: `
 // SELECT

packages/driver.pg/src/ls/driver.ts

@@ -104,12 +104,15 @@ export default class PostgreSQL extends AbstractDriver<Pool, PoolConfig> impleme
   public query: (typeof AbstractDriver)['prototype']['query'] = (query, opt = {}) => {
     const messages = [];
     let cli : PoolClient;
+    let noticeHandler: (notice: any) => void;
     const { requestId } = opt;
     return this.open()
       .then(async (pool) => {
         cli = await pool.connect();
-        cli.on('notice', notice => messages.push(this.prepareMessage(`${notice.name.toUpperCase()}: ${notice.message}`)));
+        noticeHandler = notice => messages.push(this.prepareMessage(`${notice.name.toUpperCase()}: ${notice.message}`));
+        cli.on('notice', noticeHandler);
         const results = await cli.query({ text: query.toString(), rowMode: 'array' });
+        cli.removeListener('notice', noticeHandler);
         cli.release();
         return results;
       })
@@ -137,6 +140,7 @@ export default class PostgreSQL extends AbstractDriver<Pool, PoolConfig> impleme
         });
       })
       .catch(err => {
+        if (cli && noticeHandler) cli.removeListener('notice', noticeHandler);
         cli && cli.release();
         return [<NSDatabase.IResult>{
           connId: this.getId(),

packages/driver.pg/src/ls/queries.ts

@@ -225,6 +225,31 @@ ORDER BY
   schema_name;
 `;
 
+const fetchForeignKeys: IBaseQueries['fetchForeignKeys'] = queryFactory`
+SELECT
+  tc.constraint_name AS "constraintName",
+  tc.table_schema AS "sourceTableSchema",
+  tc.table_name AS "sourceTableName",
+  kcu.column_name AS "sourceColumnName",
+  ccu.table_schema AS "targetTableSchema",
+  ccu.table_name AS "targetTableName",
+  ccu.column_name AS "targetColumnName"
+FROM information_schema.table_constraints AS tc
+JOIN information_schema.key_column_usage AS kcu
+  ON tc.constraint_name = kcu.constraint_name
+  AND tc.table_schema = kcu.table_schema
+JOIN information_schema.constraint_column_usage AS ccu
+  ON ccu.constraint_name = tc.constraint_name
+  AND ccu.table_schema = tc.table_schema
+WHERE
+  tc.constraint_type = 'FOREIGN KEY'
+  AND tc.table_schema = '${p => p.schema}'
+  AND tc.table_catalog = '${p => p.database}'
+ORDER BY
+  tc.table_name,
+  kcu.column_name
+`;
+
 export default {
   describeTable,
   countRecords,
@@ -238,4 +263,5 @@ export default {
   fetchMaterializedViews,
   searchTables,
   searchColumns,
+  fetchForeignKeys,
 };

packages/driver.sqlite/src/ls/queries.ts

@@ -79,6 +79,22 @@ ORDER BY C.name ASC,
 LIMIT ${p => p.limit || 100}
 `;
 
+const fetchForeignKeys: IBaseQueries['fetchForeignKeys'] = queryFactory`
+SELECT
+  '' AS "constraintName",
+  '' AS "sourceTableSchema",
+  m.name AS "sourceTableName",
+  fk."from" AS "sourceColumnName",
+  '' AS "targetTableSchema",
+  fk."table" AS "targetTableName",
+  fk."to" AS "targetColumnName"
+FROM sqlite_master AS m
+JOIN pragma_foreign_key_list(m.name) AS fk ON 1 = 1
+WHERE m.type = 'table'
+  AND m.name NOT LIKE 'sqlite_%'
+ORDER BY m.name, fk.seq
+`;
+
 export default {
   describeTable,
   countRecords,
@@ -87,7 +103,8 @@ export default {
   fetchTables,
   fetchViews,
   searchTables,
-  searchColumns
+  searchColumns,
+  fetchForeignKeys,
 }
 // export default {
 //   listFks: `PRAGMA foreign_key_list(\':table\');`

packages/extension/build.json

@@ -16,6 +16,7 @@
       "entries": {
         "Settings": "../plugins/connection-manager/webview/ui/screens/Settings/index.tsx",
         "Results": "../plugins/connection-manager/webview/ui/screens/Results/index.tsx",
+        "ERDiagram": "../plugins/connection-manager/webview/ui/screens/ERDiagram/index.tsx",
         "theme": "../plugins/connection-manager/webview/ui/sass/theme.scss"
       },
       "type": "webview"

packages/extension/package.json

@@ -353,6 +353,15 @@
                 "title": "Open Docs",
                 "command": "sqltools.openDocs",
                 "category": "SQLTools"
+            },
+            {
+                "title": "Show ER Diagram",
+                "command": "sqltools.showERDiagram",
+                "category": "SQLTools Connection",
+                "icon": {
+                    "light": "icons/show-light.svg",
+                    "dark": "icons/show-dark.svg"
+                }
             }
         ],
         "keybindings": [
@@ -1046,6 +1055,10 @@
                     "command": "sqltools.refreshTree",
                     "when": "false"
                 },
+                {
+                    "command": "sqltools.showERDiagram",
+                    "when": "false"
+                },
                 {
                     "command": "sqltools.generateInsertQuery",
                     "when": "false"
@@ -1147,6 +1160,11 @@
                     "when": "view == sqltoolsViewConnectionExplorer && viewItem =~ /^connection\\.(table|view|materializedView)$/",
                     "group": "navigation@1"
                 },
+                {
+                    "command": "sqltools.showERDiagram",
+                    "when": "view == sqltoolsViewConnectionExplorer && viewItem == connection.schema",
+                    "group": "navigation@1"
+                },
                 {
                     "command": "sqltools.describeTable",
                     "when": "view == sqltoolsViewConnectionExplorer && viewItem =~ /^connection\\.(table|view)$/",