extend redshift schema queries with external tables metadata

Author: KSDaemon

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

@@ -4,11 +4,15 @@
  * @fileoverview The `RedshiftDriver` and related types declaration.
  */
 
-import { getEnv } from '@cubejs-backend/shared';
+import { assertDataSource, getEnv } from '@cubejs-backend/shared';
 import { PostgresDriver, PostgresDriverConfiguration } from '@cubejs-backend/postgres-driver';
 import {
+  DatabaseStructure,
   DownloadTableCSVData,
   DriverCapabilities,
+  QueryColumnsResult,
+  QuerySchemasResult,
+  QueryTablesResult,
   StreamOptions,
   StreamTableDataWithTypes,
   UnloadOptions
@@ -40,10 +44,14 @@ export interface RedshiftDriverConfiguration extends PostgresDriverConfiguration
   exportBucket?: RedshiftDriverExportAWS;
 }
 
+const IGNORED_SCHEMAS = ['pg_catalog', 'pg_internal', 'information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA'];
+
 /**
  * Redshift driver class.
  */
 export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration> {
+  private readonly dbName: string;
+
   /**
    * Returns default concurrency value.
    */
@@ -55,7 +63,7 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
    * Class constructor.
    */
   public constructor(
-    options: RedshiftDriverConfiguration & {
+    config: RedshiftDriverConfiguration & {
       /**
        * Data source name.
        */
@@ -73,7 +81,15 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
       testConnectionTimeout?: number,
     } = {}
   ) {
-    super(options);
+    super(config);
+
+    const dataSource =
+      config.dataSource ||
+      assertDataSource('default');
+
+    // We need a DB name for querying external tables.
+    // It's not possible to get it later from the pool
+    this.dbName = getEnv('dbName', { dataSource });
   }
 
   protected primaryKeysQuery() {
@@ -84,6 +100,104 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
     return null;
   }
 
+  /**
+   * @override
+   */
+  protected informationSchemaQuery() {
+    return `
+      SELECT columns.column_name as ${this.quoteIdentifier('column_name')},
+             columns.table_name as ${this.quoteIdentifier('table_name')},
+             columns.table_schema as ${this.quoteIdentifier('table_schema')},
+             columns.data_type as ${this.quoteIdentifier('data_type')}
+      FROM information_schema.columns
+      WHERE columns.table_schema NOT IN (${IGNORED_SCHEMAS.map(s => `'${s}'`).join(',')})
+   `;
+  }
+
+  /**
+   * In Redshift external tables are not shown in regular Postgres information_schema,
+   * so it needs to be queried separately.
+   * @override
+   */
+  public async tablesSchema(): Promise<DatabaseStructure> {
+    const query = this.informationSchemaQuery();
+    const tablesSchema = await this.query(query, []).then(data => data.reduce<DatabaseStructure>(this.informationColumnsSchemaReducer, {}));
+
+    const allSchemas = await this.getSchemas();
+    const externalSchemas = allSchemas.filter(s => !tablesSchema[s.schema_name]).map(s => s.schema_name);
+
+    for (const externalSchema of externalSchemas) {
+      tablesSchema[externalSchema] = {};
+      const tablesRes = await this.tablesForExternalSchema(externalSchema);
+      const tables = tablesRes.map(t => t.table_name);
+      for (const tableName of tables) {
+        const columnRes = await this.columnsForExternalTable(externalSchema, tableName);
+        tablesSchema[externalSchema][tableName] = columnRes.map(def => ({
+          name: def.column_name,
+          type: def.data_type,
+          attributes: []
+        }));
+      }
+    }
+
+    return tablesSchema;
+  }
+
+  // eslint-disable-next-line camelcase
+  private async tablesForExternalSchema(schemaName: string): Promise<{ table_name: string }[]> {
+    return this.query(`SHOW TABLES FROM SCHEMA ${this.dbName}.${schemaName}`, []);
+  }
+
+  private async columnsForExternalTable(schemaName: string, tableName: string): Promise<QueryColumnsResult[]> {
+    return this.query(`SHOW COLUMNS FROM TABLE ${this.dbName}.${schemaName}.${tableName}`, []);
+  }
+
+  /**
+   * @override
+   */
+  protected getSchemasQuery() {
+    return `
+      SELECT table_schema as ${this.quoteIdentifier('schema_name')}
+      FROM information_schema.tables
+      WHERE table_schema NOT IN (${IGNORED_SCHEMAS.map(s => `'${s}'`).join(',')})
+      GROUP BY table_schema
+    `;
+  }
+
+  /**
+   * From the Redshift docs:
+   * SHOW SCHEMAS FROM DATABASE database_name [LIKE 'filter_pattern'] [LIMIT row_limit ]
+   * It returns regular schemas (queryable from information_schema) and external ones.
+   * @override
+   */
+  public async getSchemas(): Promise<QuerySchemasResult[]> {
+    const schemas = await this.query<QuerySchemasResult>(`SHOW SCHEMAS FROM DATABASE ${this.dbName}`, []);
+
+    return schemas.filter(s => !IGNORED_SCHEMAS.includes(s.schema_name));
+  }
+
+  public async getColumnsForSpecificTables(tables: QueryTablesResult[]): Promise<QueryColumnsResult[]> {
+    const columns = await super.getColumnsForSpecificTables(tables);
+
+    // We might request the external tables, their descriptions won't be returned
+    // by the super.getColumnsForSpecificTables(). Need to request them separately.
+    const missedTables = tables.filter(table => !columns.some(column => column.schema_name === table.schema_name && column.table_name === table.table_name));
+
+    for (const table of missedTables) {
+      const columnRes = await this.columnsForExternalTable(table.schema_name, table.table_name);
+      columnRes.forEach(c => {
+        columns.push({
+          schema_name: c.schema_name,
+          table_name: c.table_name,
+          column_name: c.column_name,
+          data_type: c.data_type,
+        });
+      });
+    }
+
+    return columns;
+  }
+
   /**
    * @override
    */