fix: Use last rollup table types to avoid type guessing for unionWithSourceData lambda queries (#6337)

Author: paveltiunov

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

@@ -362,7 +362,7 @@ export abstract class BaseDriver implements DriverInterface {
     return value;
   }
 
-  public async tableColumnTypes(table: string) {
+  public async tableColumnTypes(table: string): Promise<TableStructure> {
     const [schema, name] = table.split('.');
 
     const columns = await this.query<TableColumnQueryResult>(

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

@@ -12,7 +12,7 @@ import {
   StreamTableData,
   StreamingSourceTableData,
   QueryOptions,
-  ExternalDriverCompatibilities,
+  ExternalDriverCompatibilities, TableStructure, TableColumnQueryResult,
 } from '@cubejs-backend/base-driver';
 import { getEnv } from '@cubejs-backend/shared';
 import { format as formatSql, escape } from 'sqlstring';
@@ -163,6 +163,22 @@ export class CubeStoreDriver extends BaseDriver implements DriverInterface {
     );
   }
 
+  public async tableColumnTypes(table: string): Promise<TableStructure> {
+    const [schema, name] = table.split('.');
+
+    const columns = await this.query<TableColumnQueryResult>(
+      `SELECT column_name as ${this.quoteIdentifier('column_name')},
+             table_name as ${this.quoteIdentifier('table_name')},
+             table_schema as ${this.quoteIdentifier('table_schema')},
+             data_type as ${this.quoteIdentifier('data_type')}
+      FROM information_schema.columns
+      WHERE table_name = ${this.param(0)} AND table_schema = ${this.param(1)}`,
+      [name, schema]
+    );
+
+    return columns.map(c => ({ name: c.column_name, type: this.toGenericType(c.data_type) }));
+  }
+
   public quoteIdentifier(identifier: string): string {
     return `\`${identifier}\``;
   }

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

@@ -22,8 +22,8 @@ import {
   DriverInterface,
   InlineTable,
   SaveCancelFn,
-  StreamOptions,
-  UnloadOptions
+  StreamOptions, TableStructure,
+  UnloadOptions,
 } from '@cubejs-backend/base-driver';
 import { CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
 import { PreAggTableToTempTable, Query, QueryBody, QueryCache, QueryTuple, QueryWithParams } from './QueryCache';
@@ -263,6 +263,8 @@ class PreAggregationLoadCache {
 
   private tables: { [redisKey: string]: TableCacheEntry[] };
 
+  private tableColumnTypes: { [cacheKey: string]: { [tableName: string]: TableStructure } };
+
   // TODO this is in memory cache structure as well however it depends on
   // data source only and load cache is per data source for now.
   // Make it per data source key in case load cache scope is broaden.
@@ -290,10 +292,11 @@ class PreAggregationLoadCache {
     this.tablePrefixes = options.tablePrefixes;
     this.versionEntries = {};
     this.tables = {};
+    this.tableColumnTypes = {};
   }
 
   protected async tablesFromCache(preAggregation, forceRenew?) {
-    let tables = forceRenew ? null : await this.queryCache.getCacheDriver().get(this.tablesRedisKey(preAggregation));
+    let tables = forceRenew ? null : await this.queryCache.getCacheDriver().get(this.tablesCachePrefixKey(preAggregation));
     if (!tables) {
       tables = await this.preAggregations.getLoadCacheQueue(this.dataSource).executeInQueue(
         'query',
@@ -315,7 +318,7 @@ class PreAggregationLoadCache {
 
     const newTables = await this.fetchTablesNoCache(preAggregation);
     await this.queryCache.getCacheDriver().set(
-      this.tablesRedisKey(preAggregation),
+      this.tablesCachePrefixKey(preAggregation),
       newTables,
       this.preAggregations.options.preAggregationsSchemaCacheExpire || 60 * 60
     );
@@ -332,12 +335,12 @@ class PreAggregationLoadCache {
     return client.getTablesQuery(preAggregation.preAggregationsSchema);
   }
 
-  public tablesRedisKey(preAggregation: PreAggregationDescription) {
+  public tablesCachePrefixKey(preAggregation: PreAggregationDescription) {
     return this.queryCache.getKey('SQL_PRE_AGGREGATIONS_TABLES', `${preAggregation.dataSource}${preAggregation.preAggregationsSchema}${preAggregation.external ? '_EXT' : ''}`);
   }
 
   protected async getTablesQuery(preAggregation) {
-    const redisKey = this.tablesRedisKey(preAggregation);
+    const redisKey = this.tablesCachePrefixKey(preAggregation);
     if (!this.tables[redisKey]) {
       const tables = this.preAggregations.options.skipExternalCacheAndQueue && preAggregation.external ?
         await this.fetchTablesNoCache(preAggregation) :
@@ -350,6 +353,22 @@ class PreAggregationLoadCache {
     return this.tables[redisKey];
   }
 
+  protected async getTableColumnTypes(preAggregation: PreAggregationDescription, tableName: string): Promise<TableStructure> {
+    const prefixKey = this.tablesCachePrefixKey(preAggregation);
+    if (!this.tableColumnTypes[prefixKey]?.[tableName]) {
+      if (!this.preAggregations.options.skipExternalCacheAndQueue && preAggregation.external) {
+        throw new Error(`Lambda union with source data feature is supported only by external rollups stored in Cube Store but was invoked for '${preAggregation.preAggregationId}'`);
+      }
+      const client = await this.externalDriverFactory();
+      const columnTypes = await client.tableColumnTypes(tableName);
+      if (!this.tableColumnTypes[prefixKey]) {
+        this.tableColumnTypes[prefixKey] = {};
+      }
+      this.tableColumnTypes[prefixKey][tableName] = columnTypes;
+    }
+    return this.tableColumnTypes[prefixKey][tableName];
+  }
+
   private async calculateVersionEntries(preAggregation): Promise<VersionEntriesObj> {
     let versionEntries = tablesToVersionEntries(
       preAggregation.preAggregationsSchema,
@@ -394,7 +413,7 @@ class PreAggregationLoadCache {
     if (this.tablePrefixes && !this.tablePrefixes.find(p => preAggregation.tableName.split('.')[1].startsWith(p))) {
       throw new Error(`Load cache tries to load table ${preAggregation.tableName} outside of tablePrefixes filter: ${this.tablePrefixes.join(', ')}`);
     }
-    const redisKey = this.tablesRedisKey(preAggregation);
+    const redisKey = this.tablesCachePrefixKey(preAggregation);
     if (!this.versionEntries[redisKey]) {
       this.versionEntries[redisKey] = this.calculateVersionEntries(preAggregation).catch(e => {
         delete this.versionEntries[redisKey];
@@ -450,6 +469,7 @@ class PreAggregationLoadCache {
   protected async reset(preAggregation) {
     await this.tablesFromCache(preAggregation, true);
     this.tables = {};
+    this.tableColumnTypes = {};
     this.queryStageState = undefined;
     this.versionEntries = {};
   }
@@ -1655,8 +1675,9 @@ export class PreAggregationPartitionRangeLoader {
 
       if (this.preAggregation.rollupLambdaId) {
         if (this.lambdaQuery && loadResults.length > 0) {
-          const { buildRangeEnd } = loadResults[loadResults.length - 1];
-          lambdaTable = await this.downloadLambdaTable(buildRangeEnd);
+          const { buildRangeEnd, targetTableName } = loadResults[loadResults.length - 1];
+          const lambdaTypes = await this.loadCache.getTableColumnTypes(this.preAggregation, targetTableName);
+          lambdaTable = await this.downloadLambdaTable(buildRangeEnd, lambdaTypes);
         }
         const rollupLambdaResults = this.preAggregationsTablesToTempTables.filter(tempTableResult => tempTableResult[1].rollupLambdaId === this.preAggregation.rollupLambdaId);
         const filteredResults = loadResults.filter(
@@ -1708,7 +1729,7 @@ export class PreAggregationPartitionRangeLoader {
   /**
    * Downloads the lambda table from the source DB.
    */
-  private async downloadLambdaTable(fromDate: string): Promise<InlineTable> {
+  private async downloadLambdaTable(fromDate: string, lambdaTypes: TableStructure): Promise<InlineTable> {
     const { sqlAndParams, cacheKeyQueries } = this.lambdaQuery;
     const [query, params] = sqlAndParams;
     const values = params.map((p) => {
@@ -1733,6 +1754,7 @@ export class PreAggregationPartitionRangeLoader {
         dataSource: this.dataSource,
         external: false,
         useCsvQuery: true,
+        lambdaTypes,
       }
     );
     if (data.rowCount === this.options.maxSourceRowLimit) {
@@ -2393,7 +2415,7 @@ export class PreAggregations {
       preAggregations.map(
         async preAggregation => {
           const { dataSource, preAggregationsSchema } = preAggregation;
-          const cacheKey = getLoadCacheByDataSource(dataSource, preAggregationsSchema).tablesRedisKey(preAggregation);
+          const cacheKey = getLoadCacheByDataSource(dataSource, preAggregationsSchema).tablesCachePrefixKey(preAggregation);
           if (!firstByCacheKey[cacheKey]) {
             firstByCacheKey[cacheKey] = getLoadCacheByDataSource(dataSource, preAggregationsSchema).getVersionEntries(preAggregation);
             const res = await firstByCacheKey[cacheKey];

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

@@ -1,8 +1,9 @@
 import csvWriter from 'csv-write-stream';
 import LRUCache from 'lru-cache';
+import { pipeline } from 'stream';
 import { MaybeCancelablePromise, streamToArray } from '@cubejs-backend/shared';
 import { CubeStoreCacheDriver, CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
-import { BaseDriver, InlineTables, CacheDriverInterface } from '@cubejs-backend/base-driver';
+import { BaseDriver, InlineTables, CacheDriverInterface, TableStructure } from '@cubejs-backend/base-driver';
 
 import { QueryQueue } from './QueryQueue';
 import { ContinueWaitError } from './ContinueWaitError';
@@ -246,6 +247,7 @@ export class QueryCache {
           persistent: queryBody.persistent,
           dataSource: queryBody.dataSource,
           useCsvQuery: queryBody.useCsvQuery,
+          lambdaTypes: queryBody.lambdaTypes,
           aliasNameToMember: queryBody.aliasNameToMember,
         });
       } else {
@@ -426,6 +428,7 @@ export class QueryCache {
       requestId,
       inlineTables,
       useCsvQuery,
+      lambdaTypes,
       persistent,
       aliasNameToMember,
       tablesSchema,
@@ -437,6 +440,7 @@ export class QueryCache {
       requestId?: string,
       inlineTables?: InlineTables,
       useCsvQuery?: boolean,
+      lambdaTypes?: TableStructure,
       persistent?: boolean,
       aliasNameToMember?: { [alias: string]: string },
       tablesSchema?: boolean,
@@ -453,6 +457,7 @@ export class QueryCache {
       requestId,
       inlineTables,
       useCsvQuery,
+      lambdaTypes,
       tablesSchema,
     };
 
@@ -504,24 +509,41 @@ export class QueryCache {
   }
 
   private async csvQuery(client, q) {
-    const tableData = await client.downloadQueryResults(q.query, q.values, q);
-    const headers = tableData.types.map(c => c.name);
+    const headers = q.lambdaTypes.map(c => c.name);
     const writer = csvWriter({
       headers,
       sendHeaders: false,
     });
-    tableData.rows.forEach(
-      row => writer.write(row)
-    );
-    writer.end();
-    const lines = await streamToArray(writer);
-    if (tableData.release) {
-      await tableData.release();
+    let tableData;
+    try {
+      if (client.stream) {
+        tableData = await client.stream(q.query, q.values, q);
+        const errors = [];
+        await pipeline(tableData.rowStream, writer, (err) => {
+          if (err) {
+            errors.push(err);
+          }
+        });
+        if (errors.length > 0) {
+          throw new Error(`Lambda query errors ${errors.join(', ')}`);
+        }
+      } else {
+        tableData = await client.downloadQueryResults(q.query, q.values, q);
+        tableData.rows.forEach(
+          row => writer.write(row)
+        );
+        writer.end();
+      }
+    } finally {
+      if (tableData?.release) {
+        await tableData.release();
+      }
     }
+    const lines = await streamToArray(writer);
     const rowCount = lines.length;
     const csvRows = lines.join('');
     return {
-      types: tableData.types,
+      types: q.lambdaTypes,
       csvRows,
       rowCount,
     };
@@ -708,6 +730,7 @@ export class QueryCache {
       external?: boolean,
       dataSource: string,
       useCsvQuery?: boolean,
+      lambdaTypes?: TableStructure,
       persistent?: boolean,
     }
   ) {
@@ -741,6 +764,7 @@ export class QueryCache {
               requestId: options.requestId,
               dataSource: options.dataSource,
               useCsvQuery: options.useCsvQuery,
+              lambdaTypes: options.lambdaTypes,
               persistent: options.persistent,
             }
           ),
@@ -815,6 +839,7 @@ export class QueryCache {
       forceNoCache?: boolean,
       useInMemory?: boolean,
       useCsvQuery?: boolean,
+      lambdaTypes?: TableStructure,
       persistent?: boolean,
     }
   ) {
@@ -831,6 +856,7 @@ export class QueryCache {
         persistent: options.persistent,
         dataSource: options.dataSource,
         useCsvQuery: options.useCsvQuery,
+        lambdaTypes: options.lambdaTypes,
       }).then(res => {
         const result = {
           time: (new Date()).getTime(),