feat(druid-driver): Retrieve types for columns with new Druid versions (#7414)

Author: ovr

.github/actions/integration/druid.sh

@@ -6,7 +6,7 @@ export DEBUG=testcontainers
 
 export TEST_POSTGRES_VERSION=13
 export TEST_ZOOKEEPER_VERSION=3.5
-export TEST_DRUID_VERSION=0.19.0
+export TEST_DRUID_VERSION=27.0.0
 
 echo "::group::Druid ${TEST_DRUID_VERSION}";
 

packages/cubejs-druid-driver/docker-compose.yml

@@ -33,7 +33,7 @@ services:
     restart: always
 
   coordinator:
-    image: apache/druid:${TEST_DRUID_VERSION:-0.19.0}
+    image: apache/druid:${TEST_DRUID_VERSION:-27.0.0}
     container_name: coordinator
     volumes:
       - ./storage:/opt/data
@@ -56,7 +56,7 @@ services:
       start_period: 60s
 
   broker:
-    image: apache/druid:${TEST_DRUID_VERSION:-0.19.0}
+    image: apache/druid:${TEST_DRUID_VERSION:-27.0.0}
     container_name: broker
     volumes:
       - broker_var:/opt/druid/var
@@ -79,7 +79,7 @@ services:
       start_period: 60s
 
   historical:
-    image: apache/druid:${TEST_DRUID_VERSION:-0.19.0}
+    image: apache/druid:${TEST_DRUID_VERSION:-27.0.0}
     container_name: historical
     volumes:
       - ./storage:/opt/data
@@ -103,7 +103,7 @@ services:
       start_period: 60s
 
   middlemanager:
-    image: apache/druid:${TEST_DRUID_VERSION:-0.19.0}
+    image: apache/druid:${TEST_DRUID_VERSION:-27.0.0}
     container_name: middlemanager
     volumes:
       - ./storage:/opt/data
@@ -127,7 +127,7 @@ services:
       start_period: 60s
 
   router:
-    image: apache/druid:${TEST_DRUID_VERSION:-0.19.0}
+    image: apache/druid:${TEST_DRUID_VERSION:-27.0.0}
     container_name: router
     volumes:
       - router_var:/opt/druid/var

packages/cubejs-druid-driver/src/DruidClient.ts

@@ -41,11 +41,11 @@ export class DruidClient {
     });
   }
 
-  public async query<R = unknown>(query: string, parameters: { type: string, value: unknown }[]): Promise<R[]> {
+  public async query<R = unknown>(query: string, parameters: { type: string, value: unknown }[]): Promise<{ rows: R[], columns: Record<string, { sqlType: string }> | null }> {
     let cancelled = false;
     const cancelObj: any = {};
 
-    const promise: Promise<R[]> & { cancel?: () => void } = (async () => {
+    const promise: Promise<{ rows: R[], columns: any }> & { cancel?: () => void } = (async () => {
       cancelObj.cancel = async () => {
         cancelled = true;
       };
@@ -57,6 +57,8 @@ export class DruidClient {
           data: {
             query,
             parameters,
+            header: true,
+            sqlTypesHeader: true,
             resultFormat: 'object',
           },
         });
@@ -67,7 +69,19 @@ export class DruidClient {
           throw new Error('Query cancelled');
         }
 
-        return response.data;
+        if (response.headers['x-druid-sql-header-included']) {
+          const [columns, ...rows] = response.data;
+
+          return {
+            columns,
+            rows
+          };
+        } else {
+          return {
+            columns: null,
+            rows: response.data,
+          };
+        }
       } catch (e: any) {
         if (cancelled) {
           throw new Error('Query cancelled');

packages/cubejs-druid-driver/src/DruidDriver.ts

@@ -8,7 +8,12 @@ import {
   getEnv,
   assertDataSource,
 } from '@cubejs-backend/shared';
-import { BaseDriver, TableQueryResult } from '@cubejs-backend/base-driver';
+import {
+  BaseDriver,
+  DownloadQueryResultsOptions,
+  DownloadQueryResultsResult,
+  TableQueryResult, TableStructure,
+} from '@cubejs-backend/base-driver';
 import { DruidClient, DruidClientBaseConfiguration, DruidClientConfiguration } from './DruidClient';
 import { DruidQuery } from './DruidQuery';
 
@@ -105,7 +110,8 @@ export class DruidDriver extends BaseDriver {
   }
 
   public async query<R = unknown>(query: string, values: unknown[] = []): Promise<Array<R>> {
-    return this.client.query(query, this.normalizeQueryValues(values));
+    const result = await this.client.query<R>(query, this.normalizeQueryValues(values));
+    return result.rows;
   }
 
   public informationSchemaQuery() {
@@ -130,6 +136,29 @@ export class DruidDriver extends BaseDriver {
     ]);
   }
 
+  public async downloadQueryResults(query: string, values: unknown[], _options: DownloadQueryResultsOptions): Promise<DownloadQueryResultsResult> {
+    const { rows, columns } = await this.client.query<any>(query, this.normalizeQueryValues(values));
+    if (!columns) {
+      throw new Error(
+        'You are using an old version of Druid. Unable to detect column types in readOnly mode.'
+      );
+    }
+
+    const types: TableStructure = [];
+
+    for (const [name, meta] of Object.entries(columns)) {
+      types.push({
+        name,
+        type: this.toGenericType(meta.sqlType.toLowerCase()),
+      });
+    }
+
+    return {
+      rows,
+      types,
+    };
+  }
+
   protected normalizeQueryValues(values: unknown[]) {
     return values.map((value) => ({
       value,

packages/cubejs-druid-driver/test/druid-driver.test.ts

@@ -6,7 +6,7 @@ import path from 'path';
 import { DruidDriver, DruidDriverConfiguration } from '../src/DruidDriver';
 
 describe('DruidDriver', () => {
-  let env: StartedDockerComposeEnvironment|null = null;
+  let env: StartedDockerComposeEnvironment | null = null;
   let config: DruidDriverConfiguration;
 
   const doWithDriver = async (callback: (driver: DruidDriver) => Promise<any>) => {
@@ -23,6 +23,8 @@ describe('DruidDriver', () => {
 
       config = {
         url: `http://${host}:${port}`,
+        user: 'admin',
+        password: 'password1',
       };
 
       return;
@@ -85,4 +87,28 @@ describe('DruidDriver', () => {
       }]);
     });
   });
+
+  it('downloadQueryResults', async () => {
+    jest.setTimeout(10 * 1000);
+
+    return doWithDriver(async (driver) => {
+      const result = await driver.downloadQueryResults(
+        'SELECT 1 as id, true as finished, \'netherlands\' as country, CAST(\'2020-01-01T01:01:01.111Z\' as timestamp) as created UNION ALL SELECT 2 as id, false as finished, \'spain\' as country, CAST(\'2020-01-01T01:01:01.111Z\' as timestamp) as created',
+        [],
+        { highWaterMark: 1 }
+      );
+      expect(result).toEqual({
+        rows: [
+          { country: 'netherlands', created: '2020-01-01T01:01:01.111Z', finished: true, id: 1 },
+          { country: 'spain', created: '2020-01-01T01:01:01.111Z', finished: false, id: 2 }
+        ],
+        types: [
+          { name: 'id', type: 'int' },
+          { name: 'finished', type: 'boolean' },
+          { name: 'country', type: 'text' },
+          { name: 'created', type: 'timestamp' }
+        ]
+      });
+    });
+  });
 });