feat(server-core, query-orchestrator): sql-runner - supp… (#6142)

Author: wolfeeeg

packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts

@@ -81,7 +81,15 @@ export class CubeStoreDriver extends BaseDriver implements DriverInterface {
   }
 
   public informationSchemaQuery() {
-    return `${super.informationSchemaQuery()} AND columns.table_schema = '${this.config.database}'`;
+    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 as columns
+      WHERE columns.table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA')
+      AND columns.table_schema = '${this.config.database}'
+   `;
   }
 
   public createTableSqlWithOptions(tableName, columns, options: CreateTableOptions) {

packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts

@@ -428,6 +428,7 @@ export class QueryCache {
       useCsvQuery,
       persistent,
       aliasNameToMember,
+      tablesSchema,
     }: {
       cacheKey: CacheKey,
       dataSource: string,
@@ -438,6 +439,7 @@ export class QueryCache {
       useCsvQuery?: boolean,
       persistent?: boolean,
       aliasNameToMember?: { [alias: string]: string },
+      tablesSchema?: boolean,
     }
   ) {
     const queue = external
@@ -451,6 +453,7 @@ export class QueryCache {
       requestId,
       inlineTables,
       useCsvQuery,
+      tablesSchema,
     };
 
     const opt = {
@@ -530,6 +533,10 @@ export class QueryCache {
         `SQL_QUERY_EXT_${this.redisPrefix}`,
         this.options.externalDriverFactory,
         (client, q) => {
+          if (q.tablesSchema) {
+            return client.tablesSchema();
+          }
+
           this.logger('Executing SQL', {
             ...q
           });
@@ -972,9 +979,4 @@ export class QueryCache {
   public async testConnection() {
     return this.cacheDriver.testConnection();
   }
-  
-  public async fetchSchema(dataSource: string) {
-    const queue = await this.getQueue(dataSource);
-    return queue.executeQueryInQueue('query', [`Fetch schema for ${dataSource}`, []], { tablesSchema: true });
-  }
 }

packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts

@@ -482,7 +482,12 @@ export class QueryOrchestrator {
     return this.preAggregations.updateRefreshEndReached();
   }
 
-  public async fetchSchema(dataSource: string) {
-    return this.queryCache.fetchSchema(dataSource);
+  public async fetchSchema(dataSource?: string, external?: boolean) {
+    return this.queryCache.queryWithRetryAndRelease('', [], {
+      cacheKey: [`Fetch schema for ${dataSource}`, []],
+      dataSource,
+      external,
+      tablesSchema: true
+    });
   }
 }

packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js

@@ -128,8 +128,8 @@ class MockDriver {
 }
 
 class ExternalMockDriver extends MockDriver {
-  constructor() {
-    super();
+  constructor({ schemaData } = {}) {
+    super({ schemaData });
     this.indexes = [];
     this.csvFiles = [];
   }
@@ -223,7 +223,7 @@ describe('QueryOrchestrator', () => {
     const csvMockDriverLocal = new MockDriver({ csvImport: 'true' });
     const mockDriverUnloadWithoutTempTableSupportLocal = new MockDriverUnloadWithoutTempTableSupport();
     const streamingSourceMockDriverLocal = new StreamingSourceMockDriver();
-    const externalMockDriverLocal = new ExternalMockDriver();
+    const externalMockDriverLocal = new ExternalMockDriver({ schemaData });
 
     const redisPrefix = `ORCHESTRATOR_TEST_${testCount++}`;
     const driverFactory = (dataSource) => {
@@ -1679,4 +1679,9 @@ describe('QueryOrchestrator', () => {
       expect(data['Foo.query']).toMatch(/orders_d/);
     }));
   });
+
+  test('fetch table schema from external db', async () => {
+    const schema = await queryOrchestrator.fetchSchema(undefined, true);
+    expect(schema).toEqual(schemaData);
+  });
 });