From e962a7d1e4f8bf22b1355306573edbb72d82eb2c Mon Sep 17 00:00:00 2001
From: Dmitry Leontev <dmitry.leontev@cube.dev>
Date: Thu, 9 Feb 2023 15:53:04 +0300
Subject: [PATCH 1/4] feat(api-gateway, server-core, query-orchestrator):
 sql-runner - support running queries against extDb

---
 packages/cubejs-api-gateway/src/gateway.ts    | 22 ++++++++++++++-----
 .../cubejs-api-gateway/test/index.test.ts     | 22 ++++++++++++++++++-
 packages/cubejs-api-gateway/test/mocks.ts     |  8 +++++++
 .../src/CubeStoreDriver.ts                    | 10 ++++++++-
 .../src/orchestrator/QueryCache.ts            |  8 +++++--
 .../src/orchestrator/QueryOrchestrator.ts     |  4 ++--
 .../test/unit/QueryOrchestrator.test.js       | 11 +++++++---
 .../src/core/CompilerApi.js                   |  8 +++++++
 .../src/core/OrchestratorApi.ts               |  5 ++++-
 .../test/unit/index.test.ts                   | 22 +++++++++++++++++++
 10 files changed, 104 insertions(+), 16 deletions(-)

diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts
index 4c8ff0784b8be..31891ccb53486 100644
--- a/packages/cubejs-api-gateway/src/gateway.ts
+++ b/packages/cubejs-api-gateway/src/gateway.ts
@@ -1187,8 +1187,17 @@ class ApiGateway {
           query.query = query.query.slice(0, -1);
         }
 
-        const driver = await orchestratorApi
-          .driverFactory(query.dataSource || 'default');
+        const driver = !query.external
+          ? await orchestratorApi.driverFactory(query.dataSource || 'default')
+          : await orchestratorApi.options?.externalDriverFactory();
+
+        if (!driver) {
+          throw new UserError(
+            `A driver for query not found. Please check ${
+              query.external ? 'externalDriverFactory' : 'driverFactory'
+            } config setting`
+          );
+        }
 
         driver.wrapQueryWithLimit(query);
       }
@@ -1250,7 +1259,8 @@ class ApiGateway {
 
   protected async dbSchema({ query, context, res }: {
     query: {
-      dataSource: string;
+      dataSource?: string;
+      external?: boolean;
     };
     context?: RequestContext;
     res: ResponseResultFn;
@@ -1263,9 +1273,9 @@ class ApiGateway {
         );
       }
 
-      if (!query.dataSource) {
+      if (!query.dataSource && !query.external) {
         throw new UserError(
-          'A user\'s query must contain dataSource.'
+          'A user\'s query must contain dataSource or be an external'
         );
       }
 
@@ -1273,7 +1283,7 @@ class ApiGateway {
    
       const schema = await orchestratorApi
         .getQueryOrchestrator()
-        .fetchSchema(query.dataSource);
+        .fetchSchema(query.dataSource, !!query.external);
 
       res({ data: schema });
     } catch (e) {
diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts
index 90451e8cc5893..12f3c0fa86848 100644
--- a/packages/cubejs-api-gateway/test/index.test.ts
+++ b/packages/cubejs-api-gateway/test/index.test.ts
@@ -633,6 +633,26 @@ describe('API Gateway', () => {
           }
         }]
       },
+      {
+        route: 'sql-runner',
+        scope: ['sql-runner'],
+        method: 'post',
+        successBody: {
+          query: {
+            query: 'SELECT * FROM sql-runner; ',
+            external: true,
+            limit: 10000,
+            resultFilter: {
+              objectLimit: 2,
+              stringLimit: 2,
+              objectTypes: ['Buffer'],
+              limit: 1,
+              offset: 1,
+            },
+          }
+        },
+        successResult: { data: [{ string: 'st', number: 1, buffer: '[4', bufferTwo: 'Placeholder', object: '{"' }] },
+      },
       { route: 'data-sources', scope: ['sql-runner'], successResult: { dataSources: [{ dataSource: 'default', dbType: 'postgres' }] } },
       {
         route: 'db-schema',
@@ -672,7 +692,7 @@ describe('API Gateway', () => {
           {
             result: {
               status: 400,
-              error: 'A user\'s query must contain dataSource.'
+              error: 'A user\'s query must contain dataSource or be an external'
             },
             body: { query: {} }
           }
diff --git a/packages/cubejs-api-gateway/test/mocks.ts b/packages/cubejs-api-gateway/test/mocks.ts
index 91382d9cc808b..d6ff85df8b01b 100644
--- a/packages/cubejs-api-gateway/test/mocks.ts
+++ b/packages/cubejs-api-gateway/test/mocks.ts
@@ -181,6 +181,14 @@ export class AdapterApiMock {
 
   public $testOrchestratorConnectionsDone: boolean = false;
 
+  public options = {
+    externalDriverFactory: () => ({
+      wrapQueryWithLimit(query: { query: string; limit: number }) {
+        query.query = `SELECT * FROM (${query.query}) AS t LIMIT ${query.limit}`;
+      },
+    })
+  };
+
   public async testConnection() {
     this.$testConnectionsDone = true;
 
diff --git a/packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts b/packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts
index 71f780a2c2e55..1f2019be3882a 100644
--- a/packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts
+++ b/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) {
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
index bcf9eca4defa1..0c099d65ecff8 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
@@ -531,6 +531,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
           });
@@ -974,8 +978,8 @@ export class QueryCache {
     return this.cacheDriver.testConnection();
   }
   
-  public async fetchSchema(dataSource: string) {
-    const queue = await this.getQueue(dataSource);
+  public async fetchSchema(dataSource?: string, isExternalQuery?: boolean) {
+    const queue = !isExternalQuery ? await this.getQueue(dataSource) : await this.getExternalQueue();
     return queue.executeQueryInQueue('query', [`Fetch schema for ${dataSource}`, []], { tablesSchema: true });
   }
 }
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
index e0e70151b80f0..0194ff5b96019 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
@@ -482,7 +482,7 @@ export class QueryOrchestrator {
     return this.preAggregations.updateRefreshEndReached();
   }
 
-  public async fetchSchema(dataSource: string) {
-    return this.queryCache.fetchSchema(dataSource);
+  public async fetchSchema(dataSource?: string, isExternalQuery?: boolean) {
+    return this.queryCache.fetchSchema(dataSource, isExternalQuery);
   }
 }
diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js b/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js
index 2d5cbea39b95c..ad26caa822159 100644
--- a/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js
+++ b/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);
+  });
 });
diff --git a/packages/cubejs-server-core/src/core/CompilerApi.js b/packages/cubejs-server-core/src/core/CompilerApi.js
index 90050144a0237..da8ae19e99006 100644
--- a/packages/cubejs-server-core/src/core/CompilerApi.js
+++ b/packages/cubejs-server-core/src/core/CompilerApi.js
@@ -240,6 +240,14 @@ export class CompilerApi {
       })
     );
 
+    if (orchestratorApi.options?.externalDriverFactory) {
+      dataSources.push({
+        dataSource: 'externalDb',
+        dbType: this.options.externalDbType || '',
+        isExternal: true,
+      });
+    }
+
     return {
       dataSources: dataSources.filter((source) => source),
     };
diff --git a/packages/cubejs-server-core/src/core/OrchestratorApi.ts b/packages/cubejs-server-core/src/core/OrchestratorApi.ts
index 38836b99be819..d6ea74777c478 100644
--- a/packages/cubejs-server-core/src/core/OrchestratorApi.ts
+++ b/packages/cubejs-server-core/src/core/OrchestratorApi.ts
@@ -130,7 +130,10 @@ export class OrchestratorApi {
         return res;
       }
 
-      data.dbType = await extractDbType(data);
+      if (!query.external) {
+        data.dbType = await extractDbType(data);
+      }
+    
       data.extDbType = extractExternalDbType(data);
 
       return data;
diff --git a/packages/cubejs-server-core/test/unit/index.test.ts b/packages/cubejs-server-core/test/unit/index.test.ts
index c2849b806014b..e4f7f5b00bf72 100644
--- a/packages/cubejs-server-core/test/unit/index.test.ts
+++ b/packages/cubejs-server-core/test/unit/index.test.ts
@@ -422,6 +422,28 @@ describe('index.test', () => {
       dataSourcesSpy.mockClear();
     });
 
+    test.only('CompilerApi dataSources with externalDb', async () => {
+      const dataSources = await compilerApi.dataSources({
+        driverFactory: jest.fn(async () => true),
+        options: {
+          externalDriverFactory: jest.fn(async () => true),
+        }
+      });
+
+      expect(dataSources).toHaveProperty('dataSources');
+      expect(dataSources.dataSources).toEqual([{
+        dataSource: 'main',
+        dbType: 'mysql',
+      }, {
+        dataSource: 'externalDb',
+        dbType: '',
+        isExternal: true,
+      }]);
+      
+      expect(dataSourcesSpy).toHaveBeenCalled();
+      dataSourcesSpy.mockClear();
+    });
+
     test('CompilerApi dataSources with driverFactory error', async () => {
       const dataSources = await compilerApi.dataSources({
         driverFactory: jest.fn(async () => {

From 8144b4942dbb88bdd52994a1579f53695bf5acb1 Mon Sep 17 00:00:00 2001
From: Dmitry Leontev <dmitry.leontev@cube.dev>
Date: Thu, 16 Feb 2023 16:02:04 +0300
Subject: [PATCH 2/4] get rid api-gateway fixes

---
 packages/cubejs-api-gateway/src/gateway.ts    | 22 +++++--------------
 .../cubejs-api-gateway/test/index.test.ts     | 22 +------------------
 packages/cubejs-api-gateway/test/mocks.ts     |  8 -------
 .../test/unit/index.test.ts                   |  2 +-
 4 files changed, 8 insertions(+), 46 deletions(-)

diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts
index 31891ccb53486..4c8ff0784b8be 100644
--- a/packages/cubejs-api-gateway/src/gateway.ts
+++ b/packages/cubejs-api-gateway/src/gateway.ts
@@ -1187,17 +1187,8 @@ class ApiGateway {
           query.query = query.query.slice(0, -1);
         }
 
-        const driver = !query.external
-          ? await orchestratorApi.driverFactory(query.dataSource || 'default')
-          : await orchestratorApi.options?.externalDriverFactory();
-
-        if (!driver) {
-          throw new UserError(
-            `A driver for query not found. Please check ${
-              query.external ? 'externalDriverFactory' : 'driverFactory'
-            } config setting`
-          );
-        }
+        const driver = await orchestratorApi
+          .driverFactory(query.dataSource || 'default');
 
         driver.wrapQueryWithLimit(query);
       }
@@ -1259,8 +1250,7 @@ class ApiGateway {
 
   protected async dbSchema({ query, context, res }: {
     query: {
-      dataSource?: string;
-      external?: boolean;
+      dataSource: string;
     };
     context?: RequestContext;
     res: ResponseResultFn;
@@ -1273,9 +1263,9 @@ class ApiGateway {
         );
       }
 
-      if (!query.dataSource && !query.external) {
+      if (!query.dataSource) {
         throw new UserError(
-          'A user\'s query must contain dataSource or be an external'
+          'A user\'s query must contain dataSource.'
         );
       }
 
@@ -1283,7 +1273,7 @@ class ApiGateway {
    
       const schema = await orchestratorApi
         .getQueryOrchestrator()
-        .fetchSchema(query.dataSource, !!query.external);
+        .fetchSchema(query.dataSource);
 
       res({ data: schema });
     } catch (e) {
diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts
index 12f3c0fa86848..90451e8cc5893 100644
--- a/packages/cubejs-api-gateway/test/index.test.ts
+++ b/packages/cubejs-api-gateway/test/index.test.ts
@@ -633,26 +633,6 @@ describe('API Gateway', () => {
           }
         }]
       },
-      {
-        route: 'sql-runner',
-        scope: ['sql-runner'],
-        method: 'post',
-        successBody: {
-          query: {
-            query: 'SELECT * FROM sql-runner; ',
-            external: true,
-            limit: 10000,
-            resultFilter: {
-              objectLimit: 2,
-              stringLimit: 2,
-              objectTypes: ['Buffer'],
-              limit: 1,
-              offset: 1,
-            },
-          }
-        },
-        successResult: { data: [{ string: 'st', number: 1, buffer: '[4', bufferTwo: 'Placeholder', object: '{"' }] },
-      },
       { route: 'data-sources', scope: ['sql-runner'], successResult: { dataSources: [{ dataSource: 'default', dbType: 'postgres' }] } },
       {
         route: 'db-schema',
@@ -692,7 +672,7 @@ describe('API Gateway', () => {
           {
             result: {
               status: 400,
-              error: 'A user\'s query must contain dataSource or be an external'
+              error: 'A user\'s query must contain dataSource.'
             },
             body: { query: {} }
           }
diff --git a/packages/cubejs-api-gateway/test/mocks.ts b/packages/cubejs-api-gateway/test/mocks.ts
index d6ff85df8b01b..91382d9cc808b 100644
--- a/packages/cubejs-api-gateway/test/mocks.ts
+++ b/packages/cubejs-api-gateway/test/mocks.ts
@@ -181,14 +181,6 @@ export class AdapterApiMock {
 
   public $testOrchestratorConnectionsDone: boolean = false;
 
-  public options = {
-    externalDriverFactory: () => ({
-      wrapQueryWithLimit(query: { query: string; limit: number }) {
-        query.query = `SELECT * FROM (${query.query}) AS t LIMIT ${query.limit}`;
-      },
-    })
-  };
-
   public async testConnection() {
     this.$testConnectionsDone = true;
 
diff --git a/packages/cubejs-server-core/test/unit/index.test.ts b/packages/cubejs-server-core/test/unit/index.test.ts
index e4f7f5b00bf72..04e311c9debe6 100644
--- a/packages/cubejs-server-core/test/unit/index.test.ts
+++ b/packages/cubejs-server-core/test/unit/index.test.ts
@@ -422,7 +422,7 @@ describe('index.test', () => {
       dataSourcesSpy.mockClear();
     });
 
-    test.only('CompilerApi dataSources with externalDb', async () => {
+    test('CompilerApi dataSources with externalDb', async () => {
       const dataSources = await compilerApi.dataSources({
         driverFactory: jest.fn(async () => true),
         options: {

From 8aee542bbe7907a10f735385d82df375cd06e859 Mon Sep 17 00:00:00 2001
From: Dmitry Leontev <dmitry.leontev@cube.dev>
Date: Tue, 7 Mar 2023 13:32:29 +0300
Subject: [PATCH 3/4] fix cr

---
 .../src/orchestrator/QueryCache.ts                       | 8 +++-----
 .../src/orchestrator/QueryOrchestrator.ts                | 9 +++++++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
index 0c099d65ecff8..9869c01ae0468 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
@@ -429,6 +429,7 @@ export class QueryCache {
       useCsvQuery,
       persistent,
       aliasNameToMember,
+      tablesSchema,
     }: {
       cacheKey: CacheKey,
       dataSource: string,
@@ -439,6 +440,7 @@ export class QueryCache {
       useCsvQuery?: boolean,
       persistent?: boolean,
       aliasNameToMember?: { [alias: string]: string },
+      tablesSchema?: boolean,
     }
   ) {
     const queue = external
@@ -452,6 +454,7 @@ export class QueryCache {
       requestId,
       inlineTables,
       useCsvQuery,
+      tablesSchema,
     };
 
     const opt = {
@@ -977,9 +980,4 @@ export class QueryCache {
   public async testConnection() {
     return this.cacheDriver.testConnection();
   }
-  
-  public async fetchSchema(dataSource?: string, isExternalQuery?: boolean) {
-    const queue = !isExternalQuery ? await this.getQueue(dataSource) : await this.getExternalQueue();
-    return queue.executeQueryInQueue('query', [`Fetch schema for ${dataSource}`, []], { tablesSchema: true });
-  }
 }
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
index 0194ff5b96019..10ea4c8948436 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
@@ -482,7 +482,12 @@ export class QueryOrchestrator {
     return this.preAggregations.updateRefreshEndReached();
   }
 
-  public async fetchSchema(dataSource?: string, isExternalQuery?: boolean) {
-    return this.queryCache.fetchSchema(dataSource, isExternalQuery);
+  public async fetchSchema(dataSource?: string, external?: boolean) {
+    return this.queryCache.queryWithRetryAndRelease('', [], {
+      cacheKey: [`Fetch schema for ${dataSource}`, []],
+      dataSource,
+      external,
+      tablesSchema: true
+    });
   }
 }

From 7ce5942c6ad205fb65890605b6a3304554bf60dd Mon Sep 17 00:00:00 2001
From: Dmitry Leontev <dmitry.leontev@cube.dev>
Date: Fri, 17 Mar 2023 11:03:54 +0300
Subject: [PATCH 4/4] fix CR

---
 .../src/core/CompilerApi.js                   |  8 -------
 .../src/core/OrchestratorApi.ts               |  5 +----
 .../test/unit/index.test.ts                   | 22 -------------------
 3 files changed, 1 insertion(+), 34 deletions(-)

diff --git a/packages/cubejs-server-core/src/core/CompilerApi.js b/packages/cubejs-server-core/src/core/CompilerApi.js
index da8ae19e99006..90050144a0237 100644
--- a/packages/cubejs-server-core/src/core/CompilerApi.js
+++ b/packages/cubejs-server-core/src/core/CompilerApi.js
@@ -240,14 +240,6 @@ export class CompilerApi {
       })
     );
 
-    if (orchestratorApi.options?.externalDriverFactory) {
-      dataSources.push({
-        dataSource: 'externalDb',
-        dbType: this.options.externalDbType || '',
-        isExternal: true,
-      });
-    }
-
     return {
       dataSources: dataSources.filter((source) => source),
     };
diff --git a/packages/cubejs-server-core/src/core/OrchestratorApi.ts b/packages/cubejs-server-core/src/core/OrchestratorApi.ts
index d6ea74777c478..38836b99be819 100644
--- a/packages/cubejs-server-core/src/core/OrchestratorApi.ts
+++ b/packages/cubejs-server-core/src/core/OrchestratorApi.ts
@@ -130,10 +130,7 @@ export class OrchestratorApi {
         return res;
       }
 
-      if (!query.external) {
-        data.dbType = await extractDbType(data);
-      }
-    
+      data.dbType = await extractDbType(data);
       data.extDbType = extractExternalDbType(data);
 
       return data;
diff --git a/packages/cubejs-server-core/test/unit/index.test.ts b/packages/cubejs-server-core/test/unit/index.test.ts
index 04e311c9debe6..c2849b806014b 100644
--- a/packages/cubejs-server-core/test/unit/index.test.ts
+++ b/packages/cubejs-server-core/test/unit/index.test.ts
@@ -422,28 +422,6 @@ describe('index.test', () => {
       dataSourcesSpy.mockClear();
     });
 
-    test('CompilerApi dataSources with externalDb', async () => {
-      const dataSources = await compilerApi.dataSources({
-        driverFactory: jest.fn(async () => true),
-        options: {
-          externalDriverFactory: jest.fn(async () => true),
-        }
-      });
-
-      expect(dataSources).toHaveProperty('dataSources');
-      expect(dataSources.dataSources).toEqual([{
-        dataSource: 'main',
-        dbType: 'mysql',
-      }, {
-        dataSource: 'externalDb',
-        dbType: '',
-        isExternal: true,
-      }]);
-      
-      expect(dataSourcesSpy).toHaveBeenCalled();
-      dataSourcesSpy.mockClear();
-    });
-
     test('CompilerApi dataSources with driverFactory error', async () => {
       const dataSources = await compilerApi.dataSources({
         driverFactory: jest.fn(async () => {