@dbml/connector 5.0.0 → 5.1.0

package/src/connectors/mysqlConnector.ts

@@ -1,6 +1,8 @@
 import { createConnection, Connection } from 'mysql2/promise';
 import { flatten } from 'lodash';
-import { DefaultInfo, DatabaseSchema, DefaultType, Field, IndexColumn, FieldsDictionary, Table, Enum, IndexesDictionary, TableConstraintsDictionary, Index, TableConstraint, Ref } from './types';
+import {
+  CheckConstraintDictionary, DefaultInfo, DatabaseSchema, DefaultType, Field, IndexColumn, FieldsDictionary, Table, Enum, IndexesDictionary, TableConstraintsDictionary, Index, Ref,
+} from './types';
 
 const NUMBER_REGEX = '^-?[0-9]+(.[0-9]+)?$';
 
@@ -221,7 +223,7 @@ async function generateEnums (client: Connection, schemaName: string): Promise<E
 }
 
 /**
- * Mysql is automatically create index for primary keys, foreign keys, unique constraint. -> Ignore
+ * Mysql automatically creates index for primary keys, foreign keys, unique constraint. -> Ignore
  */
 async function generateIndexes (client: Connection, schemaName: string): Promise<IndexesDictionary> {
   const query = `
@@ -321,9 +323,9 @@ async function generatePrimaryAndUniqueConstraint (client: Connection, schemaNam
   const rows: any = queryResponse[0];
 
   const inlineConstraintList = rows.filter((constraint: Record<string, any>) => constraint.columnCount === 1);
-  const outOfConstraintList = rows.filter((constraint: Record<string, any>) => constraint.columnCount > 1);
+  const outOfLineConstraintList = rows.filter((constraint: Record<string, any>) => constraint.columnCount > 1);
 
-  const compositeTableConstraintMap = outOfConstraintList.reduce((acc: any, row: Record<string, any>) => {
+  const compositeTableConstraintMap = outOfLineConstraintList.reduce((acc: any, row: Record<string, any>) => {
     const {
       tableName, constraintName, columnNames, constraintType,
     } = row;
@@ -355,12 +357,12 @@ async function generatePrimaryAndUniqueConstraint (client: Connection, schemaNam
     return acc;
   }, {});
 
-  const compositeConstraintMap: any = {};
+  const compositeConstraintMap: IndexesDictionary = {};
   Object.keys(compositeTableConstraintMap).forEach((tableName) => {
     compositeConstraintMap[tableName] = flatten(Object.values(compositeTableConstraintMap[tableName]));
   });
 
-  const constraintMap = inlineConstraintList.reduce((acc: any, row: Record<string, any>) => {
+  const constraintMap = inlineConstraintList.reduce((acc: TableConstraintsDictionary, row: Record<string, any>) => {
     const { tableName, columnNames, constraintType } = row;
 
     const key = tableName;
@@ -372,7 +374,7 @@ async function generatePrimaryAndUniqueConstraint (client: Connection, schemaNam
 
     columnList.forEach((columnName: string) => {
       if (!acc[key][columnName]) {
-        acc[key][columnName] = {};
+        acc[key][columnName] = { checks: [] };
       }
 
       if (constraintType === 'PRIMARY KEY') {
@@ -386,7 +388,58 @@ async function generatePrimaryAndUniqueConstraint (client: Connection, schemaNam
     return acc;
   }, {});
 
-  return { compositeConstraintMap, constraintMap };
+  return { compositeConstraintMap, primaryAndUniqueConstraintMap: constraintMap };
+}
+
+function extractColumnsFromCheckConstraint (checkClause: string): string[] {
+  return [...checkClause.matchAll(/`(?:``|[^`])*`/g)].map((col) => col.toString().slice(1, -1));
+}
+
+async function generateCheckConstraints (client: Connection, schemaName: string): Promise<{ checks: CheckConstraintDictionary, tableConstraints: TableConstraintsDictionary }> {
+  const query = `
+    select
+      cc.constraint_name as constraintName,
+      tc.table_name as tableName,
+      cc.check_clause as checkClause
+    from information_schema.check_constraints cc
+    join information_schema.table_constraints tc
+      on cc.constraint_name = tc.constraint_name
+    where cc.constraint_schema = ?
+  `;
+
+  const queryResponse = await client.query(query, [schemaName]);
+  const rows: any = queryResponse[0];
+  const checks: CheckConstraintDictionary = {};
+  const tableConstraints: TableConstraintsDictionary = {};
+  rows.forEach((row: any) => {
+    const {
+      constraintName,
+      tableName,
+      checkClause,
+    } = row;
+    const key = tableName;
+    const cols = extractColumnsFromCheckConstraint(checkClause);
+    if (cols.length === 1) {
+      const col = cols[0];
+      if (!tableConstraints[key]) tableConstraints[key] = {};
+      if (!tableConstraints[key][col]) tableConstraints[key][col] = { checks: [] };
+      tableConstraints[key][col].checks.push({
+        name: constraintName,
+        // The check definition has the form: `(expr)`
+        // The expression starts at 1 and ends at -1
+        expression: checkClause.slice(1, -1),
+      });
+      return;
+    }
+    if (!checks[key]) checks[key] = [];
+    checks[key].push({
+      name: constraintName,
+      // The check definition has the form: `(expr)`
+      // The expression starts at 1 and ends at -1
+      expression: checkClause.slice(1, -1),
+    });
+  });
+  return { checks, tableConstraints };
 }
 
 async function generateForeignKeys (client: Connection, schemaName: string): Promise<Ref[]> {
@@ -453,7 +506,7 @@ async function generateForeignKeys (client: Connection, schemaName: string): Pro
   return foreignKeyList;
 }
 
-function combineIndexAndCompositeConstraint (userDefinedIndexMap: any, compositeConstraintMap: any) {
+function combineIndexAndCompositeConstraint (userDefinedIndexMap: IndexesDictionary, compositeConstraintMap: IndexesDictionary): IndexesDictionary {
   const indexMap = Object.assign(userDefinedIndexMap, {});
 
   Object.keys(compositeConstraintMap).forEach((tableName) => {
@@ -484,6 +537,7 @@ async function fetchSchemaJson (connection: string): Promise<DatabaseSchema> {
     generateIndexes(client, schemaName),
     generatePrimaryAndUniqueConstraint(client, schemaName),
     generateForeignKeys(client, schemaName),
+    generateCheckConstraints(client, schemaName),
   ]);
   client.end();
 
@@ -491,20 +545,31 @@ async function fetchSchemaJson (connection: string): Promise<DatabaseSchema> {
     { tableList, fieldMap },
     enumList,
     rawIndexMap,
-    { constraintMap, compositeConstraintMap },
+    { primaryAndUniqueConstraintMap, compositeConstraintMap },
     foreignKeyList,
+    { checks: tableCheckConstraintMap, tableConstraints: columnCheckConstraintMap },
   ] = result;
 
   // combine normal index and composite key
   const indexMap = combineIndexAndCompositeConstraint(rawIndexMap, compositeConstraintMap);
 
+  const tableConstraints: TableConstraintsDictionary = { ...primaryAndUniqueConstraintMap };
+  Object.keys(columnCheckConstraintMap).forEach((tableName) => {
+    if (!tableConstraints[tableName]) tableConstraints[tableName] = {};
+    Object.keys(columnCheckConstraintMap[tableName]).forEach((colName) => {
+      if (!tableConstraints[tableName][colName]) tableConstraints[tableName][colName] = { checks: [] };
+      tableConstraints[tableName][colName].checks.push(...columnCheckConstraintMap[tableName][colName].checks);
+    });
+  });
+
   return {
     tables: tableList,
     fields: fieldMap,
     refs: foreignKeyList,
     enums: enumList,
     indexes: indexMap,
-    tableConstraints: constraintMap,
+    tableConstraints,
+    checks: tableCheckConstraintMap,
   };
 }
 

package/src/connectors/postgresConnector.ts

@@ -2,6 +2,7 @@
 import { Client } from 'pg';
 import { buildSchemaQuery, parseConnectionString } from '../utils/parseSchema';
 import {
+  CheckConstraintDictionary,
   DatabaseSchema,
   DefaultInfo,
   DefaultType,
@@ -301,7 +302,7 @@ const generateRawRefs = async (client: Client, schemas: string[]): Promise<Ref[]
   return refs;
 };
 
-const generateIndexes = async (client: Client, schemas: string[]) => {
+const generateIndexesAndConstraints = async (client: Client, schemas: string[]) => {
   const indexListSql = `
     WITH user_tables AS (
       SELECT
@@ -348,8 +349,8 @@ const generateIndexes = async (client: Client, schemas: string[]) => {
       ii.is_unique,
       ii.is_primary,
       ii.index_type,
-      ii.columns,
-      ii.expressions,
+      ii.columns AS index_columns,
+      ii.expressions AS index_expressions,
       ii.constraint_type  -- Retained constraint type here
     FROM
       user_tables ut
@@ -364,38 +365,38 @@ const generateIndexes = async (client: Client, schemas: string[]) => {
     ;
   `;
   const indexListResult = await client.query(indexListSql);
-  const { outOfLineConstraints, inlineConstraints } = indexListResult.rows.reduce((acc, row) => {
-    const { constraint_type, columns } = row;
+  const { outOfLineIndexes, inlineIndexes } = indexListResult.rows.reduce((acc, row) => {
+    const { constraint_type, index_columns } = row;
 
-    if (columns === 'null' || columns.trim() === '') return acc;
+    if (index_columns === 'null' || index_columns.trim() === '') return acc;
 
-    const isSingleColumn = columns.split(',').length === 1;
-    const isInlineConstraint = isSingleColumn && (constraint_type === 'PRIMARY KEY' || constraint_type === 'UNIQUE');
-    if (isInlineConstraint) {
-      acc.inlineConstraints.push(row);
+    const isSingleColumn = index_columns.split(',').length === 1;
+    const isInlineIndex = isSingleColumn && (constraint_type === 'PRIMARY KEY' || constraint_type === 'UNIQUE');
+    if (isInlineIndex) {
+      acc.inlineIndexes.push(row);
     } else {
-      acc.outOfLineConstraints.push(row);
+      acc.outOfLineIndexes.push(row);
     }
     return acc;
-  }, { outOfLineConstraints: [], inlineConstraints: [] });
+  }, { outOfLineIndexes: [], inlineIndexes: [] });
 
-  const indexes = outOfLineConstraints.reduce((acc: IndexesDictionary, indexRow: Record<string, any>) => {
+  const indexes = outOfLineIndexes.reduce((acc: IndexesDictionary, indexRow: Record<string, any>) => {
     const {
       table_schema,
       table_name,
       index_name,
       index_type,
-      columns,
-      expressions,
+      index_columns,
+      index_expressions,
     } = indexRow;
-    const indexColumns = columns.split(',').map((column: string) => {
+    const indexColumns = index_columns.split(',').map((column: string) => {
       return {
         type: 'column',
         value: column.trim(),
       };
     });
 
-    const indexExpressions = expressions ? expressions.split(',').map((expression: string) => {
+    const indexExpressions = index_expressions ? index_expressions.split(',').map((expression: string) => {
       return {
         type: 'expression',
         value: expression,
@@ -421,32 +422,105 @@ const generateIndexes = async (client: Client, schemas: string[]) => {
     return acc;
   }, {});
 
-  const tableConstraints = inlineConstraints.reduce((acc: TableConstraintsDictionary, row: Record<string, any>) => {
+  const tableConstraints: TableConstraintsDictionary = {};
+
+  inlineIndexes.forEach((row: Record<string, any>) => {
     const {
       table_schema,
       table_name,
-      columns,
+      index_columns,
       constraint_type,
     } = row;
     const key = `${table_schema}.${table_name}`;
-    if (!acc[key]) acc[key] = {};
+    if (!tableConstraints[key]) tableConstraints[key] = {};
 
-    const columnNames = columns.split(',').map((column: string) => column.trim());
+    const columnNames = index_columns.split(',').map((column: string) => column.trim());
     columnNames.forEach((columnName: string) => {
-      if (!acc[key][columnName]) acc[key][columnName] = {};
+      if (!tableConstraints[key][columnName]) tableConstraints[key][columnName] = { checks: [] };
       if (constraint_type === 'PRIMARY KEY') {
-        acc[key][columnName].pk = true;
+        tableConstraints[key][columnName].pk = true;
       }
-      if (constraint_type === 'UNIQUE' && !acc[key][columnName].pk) {
-        acc[key][columnName].unique = true;
+      if (constraint_type === 'UNIQUE' && !tableConstraints[key][columnName].pk) {
+        tableConstraints[key][columnName].unique = true;
       }
     });
-    return acc;
   }, {});
 
+  const checkListSql = `
+    SELECT
+      n.nspname AS table_schema,
+      c.relname AS table_name,
+      MAX(a.attname) AS column_name,
+      con.conname AS check_name,
+      pg_get_constraintdef(con.oid) AS check_definition,
+      CASE
+        WHEN COUNT(*) = 1 THEN TRUE
+        ELSE FALSE
+      END AS is_column_constraint 
+    FROM
+      pg_catalog.pg_constraint AS con
+    JOIN
+      pg_catalog.pg_class AS c
+      ON con.conrelid = c.oid
+    LEFT JOIN
+      pg_catalog.pg_namespace AS n
+      ON c.relnamespace = n.oid
+    LEFT JOIN
+      pg_catalog.pg_attribute AS a
+      ON a.attrelid = c.oid AND a.attnum = ANY(con.conkey)
+    WHERE con.contype = 'c' -- 'c' for CHECK constraints
+    ${buildSchemaQuery('n.nspname', schemas)}
+    GROUP BY
+      con.conname,
+      n.nspname,
+      c.relname,
+      con.oid
+  `;
+  const checkListResult = await client.query(checkListSql);
+  const checks: CheckConstraintDictionary = {};
+  checkListResult.rows.forEach((row) => {
+    const {
+      table_schema,
+      table_name,
+      check_name,
+      check_definition,
+      column_name,
+      is_column_constraint,
+    } = row;
+    if (typeof table_schema !== 'string' || typeof table_name !== 'string' || typeof check_definition !== 'string') {
+      return;
+    }
+    if (!check_definition.match(/CHECK \(\(.*\)\)/)) {
+      return;
+    }
+    if (is_column_constraint && typeof column_name !== 'string') {
+      return;
+    }
+    const key = `${table_schema}.${table_name}`;
+    if (is_column_constraint) {
+      if (!tableConstraints[key]) tableConstraints[key] = {};
+      if (!tableConstraints[key][column_name]) tableConstraints[key][column_name] = { checks: [] };
+      tableConstraints[key][column_name].checks.push({
+        name: check_name || undefined,
+        // The check definition has the form: `CHECK ((expr))`
+        // The expression starts at 8 and ends at -2
+        expression: check_definition.slice(8, -2),
+      });
+      return;
+    }
+    if (!checks[key]) checks[key] = [];
+    checks[key].push({
+      name: check_name || undefined,
+      // The check definition has the form: `CHECK ((expr))`
+      // The expression starts at 8 and ends at -2
+      expression: check_definition.slice(8, -2),
+    });
+  });
+
   return {
     indexes,
     tableConstraints,
+    checks,
   };
 };
 
@@ -496,7 +570,7 @@ const fetchSchemaJson = async (connection: string): Promise<DatabaseSchema> => {
   const client = await getValidatedClient(connectionString);
 
   const tablesAndFieldsRes = generateTablesAndFields(client, schemas);
-  const indexesRes = generateIndexes(client, schemas);
+  const indexesRes = generateIndexesAndConstraints(client, schemas);
   const refsRes = generateRawRefs(client, schemas);
   const enumsRes = generateRawEnums(client, schemas);
 
@@ -509,7 +583,7 @@ const fetchSchemaJson = async (connection: string): Promise<DatabaseSchema> => {
   client.end();
 
   const { tables, fields } = res[0];
-  const { indexes, tableConstraints } = res[1];
+  const { indexes, tableConstraints, checks } = res[1];
   const refs = res[2];
   const enums = res[3];
 
@@ -520,9 +594,10 @@ const fetchSchemaJson = async (connection: string): Promise<DatabaseSchema> => {
     enums,
     indexes,
     tableConstraints,
+    checks,
   };
 };
 
 export {
-  fetchSchemaJson
+  fetchSchemaJson,
 };

package/src/connectors/snowflakeConnector.ts

@@ -43,12 +43,11 @@ const parseConnectionString = (connectionString: string): Record<string, string>
   let match;
 
   while ((match = regex.exec(connectionString)) !== null) {
-      params[match[1].trim()] = match[2].trim();
+    params[match[1].trim()] = match[2].trim();
   }
 
   return params;
-}
-
+};
 
 const connect = async (connection: Connection): Promise<void> => {
   return new Promise((resolve, reject) => {
@@ -61,13 +60,13 @@ const connect = async (connection: Connection): Promise<void> => {
       }
     });
   });
-}
+};
 
 const executeQuery = (connection: Connection, sqlText: string): Promise<any[]> => {
   return new Promise((resolve, reject) => {
     console.log('Executing query...');
     connection.execute({
-      sqlText: sqlText,
+      sqlText,
       complete: (err, stmt, rows) => {
         if (err) {
           reject(err);
@@ -77,14 +76,14 @@ const executeQuery = (connection: Connection, sqlText: string): Promise<any[]> =
             resolve(rows);
           }
         }
-      }
+      },
     });
   });
-}
+};
 
 const isLogLevel = (value: string): value is LogLevel => {
-  return ['ERROR', 'WARN', 'INFO', 'DEBUG', 'TRACE'].includes(value as LogLevel)
-}
+  return ['ERROR', 'WARN', 'INFO', 'DEBUG', 'TRACE'].includes(value as LogLevel);
+};
 
 const connectToSnowflake = async (config: Record<string, string>): Promise<Connection> => {
   const logLevel = isLogLevel(config.LOG_LEVEL) ? config.LOG_LEVEL : 'INFO';
@@ -266,14 +265,16 @@ const generateTablesAndFields = async (conn: Connection, schemas: string[]): Pro
 
 const createConstraintKeysMap = (keys: Record<string, string>[], schemas: string[], constraintType: 'primary' | 'unique'): Record<string, ConstraintRow> => {
   return keys.reduce((acc: Record<string, ConstraintRow>, row: Record<string, string>) => {
-    const { schema_name, table_name, column_name, constraint_name } = row;
+    const {
+      schema_name, table_name, column_name, constraint_name,
+    } = row;
     const selectedSchema = schemas.length > 0 ? schemas.includes(schema_name) : true;
     if (!selectedSchema) { return acc; }
 
     const key = `${schema_name}.${table_name}.${constraint_name}`;
 
     if (acc[key]) {
-      const columnNames = acc[key].columnNames;
+      const { columnNames } = acc[key];
       columnNames.push(column_name);
       return acc;
     }
@@ -291,8 +292,6 @@ const createConstraintKeysMap = (keys: Record<string, string>[], schemas: string
   }, {});
 };
 const generateIndexes = async (conn: Connection, databaseName: string, schemas: string[]) => {
-
-
   const getPrimaryKeysSql = `
     SHOW PRIMARY KEYS IN DATABASE ${databaseName};
   `;
@@ -309,7 +308,9 @@ const generateIndexes = async (conn: Connection, databaseName: string, schemas:
 
   const allConstraints: ConstraintRow[] = [Object.values(primaryKeysByConstraint), Object.values(uniqueKeysByConstraint)].flat();
   const { indexes, tableConstraints } = allConstraints.reduce((acc: GeneratedIndexes, row: ConstraintRow): GeneratedIndexes => {
-    const { schemaName, tableName, constraintName, columnNames, type, primary, unique } = row;
+    const {
+      schemaName, tableName, constraintName, columnNames, type, primary, unique,
+    } = row;
     const key = `${schemaName}.${tableName}`;
 
     if (columnNames.length < 2) {
@@ -317,6 +318,7 @@ const generateIndexes = async (conn: Connection, databaseName: string, schemas:
         [columnNames[0]]: {
           pk: primary,
           unique,
+          checks: [],
         },
       };
       return acc;
@@ -340,7 +342,7 @@ const generateIndexes = async (conn: Connection, databaseName: string, schemas:
     }
 
     return acc;
-  }, { indexes: {}, tableConstraints: {}});
+  }, { indexes: {}, tableConstraints: {} });
 
   return {
     indexes,
@@ -384,6 +386,8 @@ const fetchSchemaJson = async (connection: string): Promise<DatabaseSchema> => {
     enums: [],
     indexes,
     tableConstraints,
+    // Checks are not supported in snowflake
+    checks: {},
   };
 };
 

package/src/connectors/types.ts

@@ -1,5 +1,3 @@
-import { columns } from "mssql";
-
 interface NoteInfo {
   value: string;
 }
@@ -9,7 +7,6 @@ interface TypeInfo {
   schemaName: string | null;
 }
 
-
 type DefaultType = 'string' | 'number' | 'boolean' | 'expression';
 
 interface DefaultInfo {
@@ -38,12 +35,22 @@ interface FieldsDictionary {
   [key: string]: Field[]; // Represents fields within tables, indexed by schemaName.tableName
 }
 
+interface CheckConstraint {
+  name?: string;
+  expression: string;
+}
+
+interface CheckConstraintDictionary {
+  [tableName: string]: CheckConstraint[];
+}
+
 interface TableConstraint {
-  [fieldName: string]: {
-    pk?: boolean;
-    unique?: boolean;
+ [fieldName: string]: {
+   pk?: boolean;
+   unique?: boolean;
+   checks: CheckConstraint[];
   };
-}
+ }
 
 interface TableConstraintsDictionary {
   [tableIdentifier: string]: TableConstraint; // Represents constraints within tables, indexed by schemaName.tableName
@@ -97,6 +104,7 @@ interface DatabaseSchema {
   tableConstraints: TableConstraintsDictionary;
   refs: Ref[];
   indexes: IndexesDictionary; // Represents the indexes property
+  checks: CheckConstraintDictionary;
 }
 
 interface BigQueryCredentials {
@@ -108,59 +116,6 @@ interface BigQueryCredentials {
   datasets: string[],
 }
 
-// Currently, we parse the check constraints as enum values in mssql.
-// However, one check constraint can have multiple set of enum values.
-// So, we need to store the enum values in a dictionary.
-// Examples:
-// CREATE TABLE AddressInfo (
-//   A1 UNIQUEIDENTIFIER NOT NULL,
-//   A2 UNIQUEIDENTIFIER NOT NULL,
-//   A3 UNIQUEIDENTIFIER NOT NULL,
-//   A4 UNIQUEIDENTIFIER NOT NULL,
-//   CONSTRAINT CK_Address_Valid CHECK (
-//       A1 IN ('1111', '2222', '333') AND
-//       A2 IN ('1111', '2222') AND
-//       A3 IN ('2222', '1111') AND
-//       A4 IN ('2222', '3333')
-//   )
-// );
-// GO
-// In the above example, we have a check constraint CK_Address_Valid2 with 4 columns: A1, A2, A3, A4.
-// Each column has a different set of enum values. So, we need to store the enum values in a dictionary.
-// In the above example, the enum values dictionary will look like this:
-// [
-//   {
-//     columns : ['A1'],
-//     enumValues: [
-//       { name: '1111' },
-//       { name: '2222' },
-//       { name: '333' }
-//     ],
-//     constraint_name: 'CK_Address_Valid_A1'
-//   },
-//   {
-//     columns : ['A2, A3'], => In this case, these two columns have the same set of enum values.
-//     enumValues: [
-//       { name: '1111' },
-//       { name: '2222' }
-//     ],
-//     constraint_name: 'CK_Address_Valid_A2_A3'
-//   },
-//   {
-//     columns : ['A4'],
-//     enumValues: [
-//       { name: '2222' },
-//       { name: '3333' }
-//     ],
-//     constraint_name: 'CK_Address_Valid_A4'
-//   }
-// ]
-interface EnumValuesDict {
-  columns: string[],
-  enumValues: EnumValue[],
-  constraint_name: string,
-}
-
 export {
   NoteInfo,
   TypeInfo,
@@ -180,5 +135,5 @@ export {
   IndexesDictionary,
   DatabaseSchema,
   BigQueryCredentials,
-  EnumValuesDict,
+  CheckConstraintDictionary,
 };