feat: strategy without creating temp table (cube-js#5299)

* wip

* fix

* wip

* fix

* wip

* wip

* wip

* added types

* wip

* wip

* refactor

* one more refactor

* refactor databricks

* fix

* update types

* capabilities fix

* fix

* fix

* refactor

* add ESCAPING

Author: frolovdev

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

@@ -3,7 +3,25 @@ import fs from 'fs';
 import { getEnv, isFilePath, isSslKey, isSslCert } from '@cubejs-backend/shared';
 
 import { cancelCombinator } from './utils';
-import { CreateTableIndex, DownloadQueryResultsOptions, DownloadQueryResultsResult, DownloadTableCSVData, DownloadTableData, DownloadTableMemoryData, DriverInterface, ExternalDriverCompatibilities, IndexesSQL, isDownloadTableMemoryData, QueryOptions, Row, Rows, TableColumn, TableColumnQueryResult, TableQueryResult, TableStructure } from './driver.interface';
+import {
+  CreateTableIndex,
+  DownloadQueryResultsOptions,
+  DownloadQueryResultsResult,
+  DownloadTableCSVData,
+  DownloadTableData,
+  DownloadTableMemoryData,
+  DriverInterface,
+  ExternalDriverCompatibilities,
+  IndexesSQL,
+  isDownloadTableMemoryData,
+  QueryOptions,
+  Row,
+  TableColumn,
+  TableColumnQueryResult,
+  TableQueryResult,
+  TableStructure,
+  DriverCapabilities
+} from './driver.interface';
 
 const sortByKeys = (unordered) => {
   const ordered = {};
@@ -365,7 +383,7 @@ export abstract class BaseDriver implements DriverInterface {
     // override, if it's needed
   }
 
-  public capabilities(): ExternalDriverCompatibilities {
+  public capabilities(): DriverCapabilities {
     return {};
   }
 

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

@@ -97,6 +97,11 @@ export interface ExternalDriverCompatibilities {
   csvImport?: true,
   streamImport?: true,
 }
+
+export interface DriverCapabilities extends ExternalDriverCompatibilities {
+  unloadWithoutTempTable?: true,
+}
+
 export type StreamOptions = {
   highWaterMark: number;
 };
@@ -117,8 +122,14 @@ export type CreateTableIndex = {
   columns: string[]
 };
 
+type UnloadQuery = {
+  sql: string,
+  params: unknown[]
+};
+
 export type UnloadOptions = {
   maxFileSize: number,
+  query?: UnloadQuery;
 };
 
 export type QueryOptions = {
@@ -159,5 +170,5 @@ export interface DriverInterface {
   // Shutdown the driver
   release(): Promise<void>
 
-  capabilities(): ExternalDriverCompatibilities;
+  capabilities(): DriverCapabilities;
 }

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

@@ -8,8 +8,10 @@ import {
   BaseDriver,
   DownloadTableCSVData,
   DownloadTableMemoryData, DriverInterface, IndexesSQL, CreateTableIndex,
-  StreamTableData, ExternalDriverCompatibilities,
-  StreamingSourceTableData, QueryOptions,
+  StreamTableData,
+  DriverCapabilities,
+  StreamingSourceTableData,
+  QueryOptions,
 } from '@cubejs-backend/base-driver';
 import { getEnv } from '@cubejs-backend/shared';
 import { format as formatSql } from 'sqlstring';
@@ -364,7 +366,7 @@ export class CubeStoreDriver extends BaseDriver implements DriverInterface {
     return CubeStoreQuery;
   }
 
-  public capabilities(): ExternalDriverCompatibilities {
+  public capabilities(): DriverCapabilities {
     return {
       csvImport: true,
       streamImport: true,

packages/cubejs-databricks-jdbc-driver/src/DatabricksDriver.ts

@@ -10,9 +10,7 @@ import {
   SASProtocol,
   generateBlobSASQueryParameters,
 } from '@azure/storage-blob';
-import {
-  DownloadTableCSVData,
-} from '@cubejs-backend/base-driver';
+import { DriverCapabilities, UnloadOptions, } from '@cubejs-backend/base-driver';
 import {
   JDBCDriver,
   JDBCDriverConfiguration,
@@ -194,6 +192,23 @@ export class DatabricksDriver extends JDBCDriver {
     return result;
   }
 
+  private async queryColumnTypes(sql: string, params: unknown[]) {
+    const result = [];
+    // eslint-disable-next-line camelcase
+    const response = await this.query<{col_name: string; data_type: string}>(`DESCRIBE QUERY ${sql}`, params);
+
+    for (const column of response) {
+      // Databricks describe additional info by default after empty line.
+      if (column.col_name === '') {
+        break;
+      }
+
+      result.push({ name: column.col_name, type: this.toGenericType(column.data_type) });
+    }
+
+    return result;
+  }
+
   public async getTablesQuery(schemaName: string) {
     const response = await this.query(`SHOW TABLES IN ${this.quoteIdentifier(schemaName)}`, []);
 
@@ -248,21 +263,22 @@ export class DatabricksDriver extends JDBCDriver {
     return this.config.exportBucket !== undefined;
   }
 
-  /**
-   * Saves pre-aggs table to the bucket and returns links to download
-   * results.
-   */
-  public async unload(
-    tableName: string,
-  ): Promise<DownloadTableCSVData> {
-    const types = await this.tableColumnTypes(tableName);
-    const columns = types.map(t => t.name).join(', ');
+  public async unload(tableName: string, options: UnloadOptions) {
+    if (!['azure', 's3'].includes(this.config.bucketType as string)) {
+      throw new Error(`Unsupported export bucket type: ${
+        this.config.bucketType
+      }`);
+    }
+
+    const types = options.query ?
+      await this.unloadWithSql(tableName, options.query.sql, options.query.params) :
+      await this.unloadWithTable(tableName);
+
     const pathname = `${this.config.exportBucket}/${tableName}.csv`;
     const csvFile = await this.getCsvFiles(
-      tableName,
-      columns,
       pathname,
     );
+
     return {
       csvFile,
       types,
@@ -271,21 +287,41 @@ export class DatabricksDriver extends JDBCDriver {
   }
 
   /**
-   * Unload table to bucket using Databricks JDBC query and returns (async)
-   * csv files signed URLs array.
+   * Create table with query and unload it to bucket
+   */
+  private async unloadWithSql(tableName: string, sql: string, params: unknown[]) {
+    const types = await this.queryColumnTypes(sql, params);
+
+    await this.createExternalTableFromSql(tableName, sql, params);
+    
+    return types;
+  }
+
+  /**
+   * Create table from preaggregation table with location and unload it to bucket
+   */
+  private async unloadWithTable(tableName: string) {
+    const types = await this.tableColumnTypes(tableName);
+    const columns = types.map(t => t.name).join(', ');
+
+    await this.createExternalTableFromTable(tableName, columns);
+    
+    return types;
+  }
+
+  /**
+   * return csv files signed URLs array.
    */
   private async getCsvFiles(
-    table: string,
-    columns: string,
     pathname: string,
   ): Promise<string[]> {
     let res;
     switch (this.config.bucketType) {
       case 'azure':
-        res = await this.getAzureCsvFiles(table, columns, pathname);
+        res = await this.getSignedAzureUrls(pathname);
         break;
       case 's3':
-        res = await this.getS3CsvFiles(table, columns, pathname);
+        res = await this.getSignedS3Urls(pathname);
         break;
       default:
         throw new Error(`Unsupported export bucket type: ${
@@ -295,19 +331,6 @@ export class DatabricksDriver extends JDBCDriver {
     return res;
   }
 
-  /**
-   * Saves specified table to the Azure blob storage and returns (async)
-   * csv files signed URLs array.
-   */
-  private async getAzureCsvFiles(
-    table: string,
-    columns: string,
-    pathname: string,
-  ): Promise<string[]> {
-    await this.createExternalTable(table, columns);
-    return this.getSignedAzureUrls(pathname);
-  }
-
   /**
    * Returns Azure signed URLs of unloaded scv files.
    */
@@ -360,19 +383,6 @@ export class DatabricksDriver extends JDBCDriver {
     return csvFile;
   }
 
-  /**
-   * Saves specified table to the S3 bucket and returns (async) csv files
-   * signed URLs array.
-   */
-  private async getS3CsvFiles(
-    table: string,
-    columns: string,
-    pathname: string,
-  ): Promise<string[]> {
-    await this.createExternalTable(table, columns);
-    return this.getSignedS3Urls(pathname);
-  }
-
   /**
    * Returns S3 signed URLs of unloaded scv files.
    */
@@ -429,7 +439,19 @@ export class DatabricksDriver extends JDBCDriver {
    * `fs.s3a.access.key <aws-access-key>`
    * `fs.s3a.secret.key <aws-secret-key>`
    */
-  private async createExternalTable(table: string, columns: string,) {
+  private async createExternalTableFromSql(table: string, sql: string, params: unknown[]) {
+    await this.query(
+      `
+      CREATE TABLE ${table}_csv_export
+      USING CSV LOCATION '${this.config.exportBucketMountDir || this.config.exportBucket}/${table}.csv'
+      OPTIONS (escape = '"')
+      AS (${sql})
+      `,
+      params,
+    );
+  }
+
+  private async createExternalTableFromTable(table: string, columns: string) {
     await this.query(
       `
       CREATE TABLE ${table}_csv_export
@@ -440,4 +462,8 @@ export class DatabricksDriver extends JDBCDriver {
       [],
     );
   }
+
+  public capabilities(): DriverCapabilities {
+    return { unloadWithoutTempTable: true };
+  }
 }