feat: primary and foreign keys driver queries

packages/cubejs-base-driver/src/BaseDriver.ts

@@ -40,6 +40,8 @@ import {
   QueryTablesResult,
   QueryColumnsResult,
   TableMemoryData,
+  PrimaryKeysQueryResult,
+  ForeignKeysQueryResult,
 } from './driver.interface';
 
 const sortByKeys = (unordered: any) => {
@@ -169,7 +171,7 @@ export abstract class BaseDriver implements DriverInterface {
     const query = `
       SELECT table_schema as ${this.quoteIdentifier('schema_name')},
             table_name as ${this.quoteIdentifier('table_name')}
-      FROM information_schema.tables
+      FROM information_schema.tables as columns
       WHERE table_schema IN (${schemasPlaceholders})
     `;
     return query;
@@ -181,18 +183,65 @@ export abstract class BaseDriver implements DriverInterface {
              columns.table_name as ${this.quoteIdentifier('table_name')},
              columns.table_schema as ${this.quoteIdentifier('schema_name')},
              columns.data_type as ${this.quoteIdentifier('data_type')}
-      FROM information_schema.columns
+      FROM information_schema.columns as columns
       WHERE ${conditionString}
     `;
+
     return query;
   }
 
+  protected primaryKeysQuery(_?: string): string | null {
+    return null;
+  }
+
+  protected foreignKeysQuery(_?: string): string | null {
+    return null;
+  }
+
+  protected async primaryKeys(conditionString?: string, params?: string[]): Promise<PrimaryKeysQueryResult[]> {
+    const query = this.primaryKeysQuery(conditionString);
+
+    if (!query) {
+      return [];
+    }
+
+    try {
+      return (await this.query<PrimaryKeysQueryResult>(query, params));
+    } catch (error: any) {
+      if (this.logger) {
+        this.logger('Primary Keys Query failed. Primary Keys will be defined by heuristics', {
+          error: (error.stack || error).toString()
+        });
+      }
+      return [];
+    }
+  }
+
+  protected async foreignKeys(conditionString?: string, params?: string[]): Promise<ForeignKeysQueryResult[]> {
+    const query = this.foreignKeysQuery(conditionString);
+
+    if (!query) {
+      return [];
+    }
+
+    try {
+      return (await this.query<ForeignKeysQueryResult>(query, params));
+    } catch (error: any) {
+      if (this.logger) {
+        this.logger('Foreign Keys Query failed. Joins will be defined by heuristics', {
+          error: (error.stack || error).toString()
+        });
+      }
+      return [];
+    }
+  }
+
   protected getColumnNameForSchemaName() {
-    return 'table_schema';
+    return 'columns.table_schema';
   }
 
   protected getColumnNameForTableName() {
-    return 'table_name';
+    return 'columns.table_name';
   }
 
   protected getSslOptions(dataSource: string): TLSConnectionOptions | undefined {
@@ -316,7 +365,11 @@ export abstract class BaseDriver implements DriverInterface {
     let schema = (result[i.table_schema] || {});
     const tables = (schema[i.table_name] || []);
 
-    tables.push({ name: i.column_name, type: i.data_type, attributes: i.key_type ? ['primaryKey'] : [] });
+    tables.push({
+      name: i.column_name,
+      type: i.data_type,
+      attributes: i.key_type ? ['primaryKey'] : []
+    });
 
     tables.sort();
     schema[i.table_name] = tables;
@@ -332,6 +385,41 @@ export abstract class BaseDriver implements DriverInterface {
     return this.query(query).then(data => reduce(this.informationColumnsSchemaReducer, {}, data));
   }
 
+  // Extended version of tablesSchema containing primary and foreign keys
+  public async tablesSchemaV2() {
+    const query = this.informationSchemaQuery();
+
+    const tablesSchema = await this.query(query).then(data => reduce(this.informationColumnsSchemaReducer, {}, data));
+    const [primaryKeys, foreignKeys] = await Promise.all([this.primaryKeys(), this.foreignKeys()]);
+
+    for (const pk of primaryKeys) {
+      if (Array.isArray(tablesSchema?.[pk.table_schema]?.[pk.table_name])) {
+        tablesSchema[pk.table_schema][pk.table_name] = tablesSchema[pk.table_schema][pk.table_name].map((it: any) => {
+          if (it.name === pk.column_name) {
+            it.attributes = ['primaryKey'];
+          }
+          return it;
+        });
+      }
+    }
+
+    for (const foreignKey of foreignKeys) {
+      if (Array.isArray(tablesSchema?.[foreignKey.table_schema]?.[foreignKey.table_name])) {
+        tablesSchema[foreignKey.table_schema][foreignKey.table_name] = tablesSchema[foreignKey.table_schema][foreignKey.table_name].map((it: any) => {
+          if (it.name === foreignKey.column_name) {
+            it.foreignKeys = [...(it.foreignKeys || []), {
+              targetTable: foreignKey.target_table,
+              targetColumn: foreignKey.target_column
+            }];
+          }
+          return it;
+        });
+      }
+    }
+
+    return tablesSchema;
+  }
+
   public async createSchemaIfNotExists(schemaName: string): Promise<void> {
     const schemas = await this.query(
       `SELECT schema_name FROM information_schema.schemata WHERE schema_name = ${this.param(0)}`,
@@ -356,7 +444,7 @@ export abstract class BaseDriver implements DriverInterface {
     return this.query<QueryTablesResult>(query, schemaNames);
   }
 
-  public getColumnsForSpecificTables(tables: QueryTablesResult[]) {
+  public async getColumnsForSpecificTables(tables: QueryTablesResult[]) {
     const groupedBySchema: Record<string, string[]> = {};
     tables.forEach((t) => {
       if (!groupedBySchema[t.schema_name]) {
@@ -381,8 +469,26 @@ export abstract class BaseDriver implements DriverInterface {
     const conditionString = conditions.join(' OR ');
 
     const query = this.getColumnsForSpecificTablesQuery(conditionString);
+
+    const [primaryKeys, foreignKeys] = await Promise.all([
+      this.primaryKeys(conditionString, parameters),
+      this.foreignKeys(conditionString, parameters)
+    ]);
+
+    const columns = await this.query<QueryColumnsResult>(query, parameters);
+
+    for (const column of columns) {
+      if (primaryKeys.some(pk => pk.table_schema === column.schema_name && pk.table_name === column.table_name && pk.column_name === column.column_name)) {
+        column.attributes = ['primaryKey'];
+      }
+
+      column.foreign_keys = foreignKeys.filter(fk => fk.table_schema === column.schema_name && fk.table_name === column.table_name && fk.column_name === column.column_name).map(fk => ({
+        target_table: fk.target_table,
+        target_column: fk.target_column
+      }));
+    }
 
-    return this.query<QueryColumnsResult>(query, parameters);
+    return columns;
   }
 
   public getTablesQuery(schemaName: string) {

packages/cubejs-base-driver/src/driver.interface.ts

@@ -9,12 +9,21 @@ export interface TableColumn {
   attributes?: string[]
 }
 
+export type ForeignKey = {
+  // eslint-disable-next-line camelcase
+  target_table: string;
+  // eslint-disable-next-line camelcase
+  target_column: string;
+};
+
 export interface TableColumnQueryResult {
   // eslint-disable-next-line camelcase
   column_name: string;
   // eslint-disable-next-line camelcase
   data_type: GenericDataBaseType;
   attributes?: string[]
+  // eslint-disable-next-line camelcase
+  foreign_keys?: ForeignKey[]
 }
 
 export type TableStructure = TableColumn[];
@@ -181,6 +190,33 @@ export type QueryTablesResult = { schema_name: string, table_name: string };
 // eslint-disable-next-line camelcase
 export type QueryColumnsResult = { schema_name: string, table_name: string } & TableColumnQueryResult;
 
+export type PrimaryKeysQueryResult = {
+  // eslint-disable-next-line camelcase
+  table_schema: string
+  // eslint-disable-next-line camelcase
+  table_name: string
+  // eslint-disable-next-line camelcase
+  column_name: string
+};
+
+export type ForeignKeysQueryResult = {
+  // eslint-disable-next-line camelcase
+  table_schema: string
+  // eslint-disable-next-line camelcase
+  table_name: string
+  // eslint-disable-next-line camelcase
+  column_name: string
+  // eslint-disable-next-line camelcase
+  target_table: string
+  // eslint-disable-next-line camelcase
+  target_column: string
+};
+
+export type TableKeysFilter = {
+  tableSchema: string,
+  tableName: string[]
+};
+
 export interface DriverInterface {
   createSchemaIfNotExists(schemaName: string): Promise<void>;
   uploadTableWithIndexes(

packages/cubejs-mysql-driver/src/MySqlDriver.ts

@@ -178,6 +178,41 @@ export class MySqlDriver extends BaseDriver implements DriverInterface {
     });
   }
 
+  protected primaryKeysQuery(conditionString?: string): string | null {
+    return `SELECT
+      TABLE_SCHEMA as ${this.quoteIdentifier('table_schema')}, 
+      TABLE_NAME as ${this.quoteIdentifier('table_name')},
+      COLUMN_NAME as ${this.quoteIdentifier('column_name')}
+  FROM
+      information_schema.KEY_COLUMN_USAGE
+  WHERE
+      CONSTRAINT_NAME = 'PRIMARY'
+      AND TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
+      ${conditionString ? ` AND (${conditionString})` : ''}
+  ORDER BY
+      TABLE_SCHEMA,
+      TABLE_NAME,
+      ORDINAL_POSITION;`;
+  }
+
+  protected foreignKeysQuery(conditionString?: string): string | null {
+    return `SELECT
+        tc.table_schema as ${this.quoteIdentifier('table_schema')},
+        tc.table_name as ${this.quoteIdentifier('table_name')},
+        kcu.column_name as ${this.quoteIdentifier('column_name')},
+        columns.table_name as ${this.quoteIdentifier('target_table')},
+        columns.column_name as ${this.quoteIdentifier('target_column')}
+    FROM
+        information_schema.table_constraints AS tc
+    JOIN information_schema.key_column_usage AS kcu
+        ON tc.constraint_name = kcu.constraint_name
+    JOIN information_schema.key_column_usage AS columns
+        ON columns.constraint_name = tc.constraint_name
+    WHERE
+        columns.table_name NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
+        AND tc.constraint_type = 'FOREIGN KEY'${conditionString ? ` AND (${conditionString})` : ''};`;
+  }
+
   public readOnly() {
     return !!this.config.readOnly;
   }

packages/cubejs-postgres-driver/src/PostgresDriver.ts

@@ -142,6 +142,38 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
     this.enabled = true;
   }
 
+  protected primaryKeysQuery(conditionString?: string): string | null {
+    return `SELECT 
+      columns.table_schema as ${this.quoteIdentifier('table_schema')},
+      columns.table_name as ${this.quoteIdentifier('table_name')}, 
+      columns.column_name as ${this.quoteIdentifier('column_name')}
+    FROM information_schema.table_constraints tc
+    JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_name)
+    JOIN information_schema.columns AS columns ON columns.table_schema = tc.constraint_schema
+      AND tc.table_name = columns.table_name AND ccu.column_name = columns.column_name
+    WHERE constraint_type = 'PRIMARY KEY' AND columns.table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA')${conditionString ? ` AND (${conditionString})` : ''}`;
+  }
+
+  protected foreignKeysQuery(conditionString?: string): string | null {
+    return `SELECT
+        tc.table_schema as ${this.quoteIdentifier('table_schema')},
+        tc.table_name as ${this.quoteIdentifier('table_name')},
+        kcu.column_name as ${this.quoteIdentifier('column_name')},
+        columns.table_name as ${this.quoteIdentifier('target_table')},
+        columns.column_name as ${this.quoteIdentifier('target_column')}
+      FROM
+        information_schema.table_constraints AS tc
+      JOIN information_schema.key_column_usage AS kcu
+        ON tc.constraint_name = kcu.constraint_name
+      JOIN information_schema.constraint_column_usage AS columns
+        ON columns.constraint_name = tc.constraint_name
+      WHERE
+         constraint_type = 'FOREIGN KEY'
+         AND ${this.getColumnNameForSchemaName()} NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA')
+         ${conditionString ? ` AND (${conditionString})` : ''}
+    `;
+  }
+
   /**
    * The easiest way how to add additional configuration from env variables, because
    * you cannot call method in RedshiftDriver.constructor before super.

packages/cubejs-redshift-driver/src/RedshiftDriver.ts

@@ -72,6 +72,14 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
     super(options);
   }
 
+  protected primaryKeysQuery() {
+    return null;
+  }
+
+  protected foreignKeysQuery() {
+    return null;
+  }
+
   /**
    * @override
    */