Merge pull request #234 from Canner/feature/support-data-lake-api

onlyjackfrost · web-flow · commit e6c563a1c393 · 2023-07-17T15:06:31.000+08:00
feature: support reading data from data lake
diff --git a/packages/extension-driver-duckdb/README.md b/packages/extension-driver-duckdb/README.md
@@ -29,4 +29,15 @@
        log-queries: false
        # Optional: Whether log query requests' parameters, please be aware that query parameters might contain sensitive data (default: false)
        log-parameters: false
+       # Optional: Contains the configuration parameters DuckDB need (ex: for duckdb extension "httpfs", it will needs region, accessKeyId, ...)
+       # You can read more in the [duckdb extension page](https://duckdb.org/docs/extensions/overview) 
+       configuration-parameters: 
+         region?: string
+         accessKeyId?: string
+         secretAccessKey?: string
+         # alternative option for accessKeyId and secretAccessKey
+         sessionToken?: string
+         endpoint?: string
+         url_style?: string
+         use_ssl?: boolean
    ```
diff --git a/packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts b/packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts
@@ -15,6 +15,7 @@ import {
 } from '@vulcan-sql/core';
 import * as path from 'path';
 import { buildSQL } from './sqlBuilder';
+import { DuckDBExtensionLoader } from './duckdbExtensionLoader';
 
 const getType = (value: any) => {
   const jsType = typeof value;
@@ -28,18 +29,36 @@ const getType = (value: any) => {
   }
 };
 
+// read more configuration in DuckDB document: https://duckdb.org/docs/extensions/httpfs
+export interface ConfigurationParameters {
+  region?: string;
+  accessKeyId?: string;
+  secretAccessKey?: string;
+  // sessionToken: alternative option for accessKeyId and secretAccessKey
+  sessionToken?: string;
+  endpoint?: string;
+  url_style?: string;
+  use_ssl?: boolean;
+}
+
 export interface DuckDBOptions {
   'persistent-path'?: string;
   'log-queries'?: boolean;
   'log-parameters'?: boolean;
+  'configuration-parameters'?: ConfigurationParameters;
 }
 
 @VulcanExtensionId('duckdb')
 export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
   // Use protected method for unit test
   protected dbMapping = new Map<
     string,
-    { db: duckdb.Database; logQueries: boolean; logParameters: boolean }
+    {
+      db: duckdb.Database;
+      logQueries: boolean;
+      logParameters: boolean;
+      configurationParameters: ConfigurationParameters;
+    }
   >();
   private logger = this.getLogger();
 
@@ -56,13 +75,15 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
       // Only reuse db instance when not using in-memory only db
       let db = dbPath !== ':memory:' ? dbByPath.get(dbPath) : undefined;
       if (!db) {
-        db = new duckdb.Database(dbPath);
+        db = await this.initDatabase(dbPath);
         dbByPath.set(dbPath, db);
       }
       this.dbMapping.set(profile.name, {
         db,
         logQueries: profile.connection?.['log-queries'] || false,
         logParameters: profile.connection?.['log-parameters'] || false,
+        configurationParameters:
+          profile.connection?.['configuration-parameters'] || {},
       });
     }
   }
@@ -76,10 +97,12 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
     if (!this.dbMapping.has(profileName)) {
       throw new InternalError(`Profile instance ${profileName} not found`);
     }
-    const { db, ...options } = this.dbMapping.get(profileName)!;
+    const { db, configurationParameters, ...options } =
+      this.dbMapping.get(profileName)!;
     const builtSQL = buildSQL(sql, operations);
     // create new connection for each query
     const connection = db.connect();
+    await this.loadExtensions(connection, configurationParameters);
     const statement = connection.prepare(builtSQL);
     const parameters = Array.from(bindParams.values());
     this.logRequest(builtSQL, parameters, options);
@@ -153,9 +176,10 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
     if (!fs.existsSync(directory!))
       throw new InternalError(`The directory ${directory} not exists`);
 
-    const { db } = this.dbMapping.get(profileName)!;
+    const { db, configurationParameters } = this.dbMapping.get(profileName)!;
     // create new connection for export
     const connection = db.connect();
+    await this.loadExtensions(connection, configurationParameters);
     const formatTypeMapper = {
       [CacheLayerStoreFormatType.parquet.toString()]: 'parquet',
     };
@@ -221,4 +245,34 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
       );
     });
   }
+
+  private async initDatabase(dbPath: string) {
+    const db = new duckdb.Database(dbPath);
+    const conn = db.connect();
+    await this.installExtensions(conn);
+    return db;
+  }
+
+  private async installExtensions(
+    connection: duckdb.Connection
+  ): Promise<void> {
+    // allows DuckDB to read remote/write remote files
+    return await new Promise((resolve, reject) => {
+      connection.run('INSTALL httpfs', (err: any) => {
+        if (err) reject(err);
+        this.logger.debug('Installed httpfs extension');
+        resolve();
+      });
+    });
+  }
+
+  private async loadExtensions(
+    connection: duckdb.Connection,
+    configurationParameters: ConfigurationParameters
+  ) {
+    const configurationLoader = new DuckDBExtensionLoader(
+      configurationParameters
+    );
+    await configurationLoader.loadExtension(connection, 'httpfs');
+  }
 }
diff --git a/packages/extension-driver-duckdb/src/lib/duckdbExtensionLoader.ts b/packages/extension-driver-duckdb/src/lib/duckdbExtensionLoader.ts
@@ -0,0 +1,72 @@
+import * as duckdb from 'duckdb';
+import { getLogger } from '@vulcan-sql/core';
+import { ConfigurationParameters } from './duckdbDataSource';
+
+// DuckDB parameter name : configuration name
+export const HTTPFS_CONFIGURATIONS = Object.freeze({
+  s3_region: 'region',
+  s3_access_key_id: 'accessKeyId',
+  s3_secret_access_key: 'secretAccessKey',
+  // sessionToken= alternative option for accessKeyId and secretAccessKey
+  s3_session_token: 'sessionToken',
+  s3_endpoint: 'endpoint',
+  s3_url_style: 'url_style',
+  s3_use_ssl: 'use_ssl',
+});
+
+export class DuckDBExtensionLoader {
+  private configurations: ConfigurationParameters;
+  private configurationMap: ReadonlyMap<string, Record<string, string>> =
+    new Map([['httpfs', HTTPFS_CONFIGURATIONS]]);
+
+  private logger = getLogger({ scopeName: 'CORE' });
+
+  constructor(configurations: ConfigurationParameters) {
+    this.configurations = configurations;
+  }
+
+  public async loadExtension(
+    conn: duckdb.Connection,
+    extensionName: string
+  ): Promise<void> {
+    const extensionConfigurations = this.configurationMap.get(extensionName);
+    if (!extensionConfigurations) {
+      this.logger.debug(
+        `Can not find duckdb extension ${extensionName} or no configuration need to be set.`
+      );
+      return;
+    }
+    try {
+      conn.run(`LOAD ${extensionName}`);
+    } catch (error) {
+      this.logger.debug(`Error when loading extension:${extensionName}`);
+      throw error;
+    }
+
+    Object.entries(extensionConfigurations).forEach(
+      ([dbParameterName, configurationKey]) => {
+        const configurationValue =
+          this.configurations[
+            configurationKey as keyof ConfigurationParameters
+          ];
+        // if configuration is not undefined
+        if (configurationValue !== undefined) {
+          conn.run(
+            `SET ${dbParameterName}='${configurationValue}'`,
+            (err: any) => {
+              if (err) throw err;
+              this.logger.debug(
+                `Configuration error "${dbParameterName}": ${err}`
+              );
+            }
+          );
+          this.logger.debug(`Configuration parameter "${dbParameterName}" set`);
+        } else {
+          this.logger.debug(
+            `Configuration "${dbParameterName}" has not been set`
+          );
+        }
+      }
+    );
+  }
+}
diff --git a/packages/extension-driver-duckdb/tests/configurationLoader.spec.ts b/packages/extension-driver-duckdb/tests/configurationLoader.spec.ts
@@ -0,0 +1,117 @@
+import * as duckdb from 'duckdb';
+import {
+  DuckDBExtensionLoader,
+  HTTPFS_CONFIGURATIONS,
+} from '../src/lib/duckdbExtensionLoader';
+
+const getQueryResults = (conn: duckdb.Connection, sql: string) =>
+  new Promise<any[]>((resolve, reject) => {
+    conn.all(sql, (err: any, result: any[]) => {
+      err ? reject(err) : resolve(result);
+    });
+  });
+
+it('Given Configurations should been set into DuckDB session after loadExtension', async () => {
+  // Arrange
+  const db = new duckdb.Database(':memory:');
+  const connection = db.connect();
+  connection.run('INSTALL httpfs');
+  const configurations = {
+    use_ssl: false,
+    region: 'us-east-1',
+    accessKeyId: 'accessKeyId',
+    secretAccessKey: 'secretAccessKey',
+    endpoint: 'http://localhost:4566',
+    url_style: 'path',
+    sessionToken: 'sessionToken',
+  };
+  const configurationLoader = new DuckDBExtensionLoader(configurations);
+
+  // Act
+  await configurationLoader.loadExtension(connection, 'httpfs');
+
+  // Assert
+  let dbConfigurations = await getQueryResults(
+    connection,
+    'SELECT * FROM duckdb_settings();'
+  );
+  dbConfigurations = dbConfigurations.filter(({ name }) =>
+    Object.keys(HTTPFS_CONFIGURATIONS).includes(name)
+  );
+
+  Object.entries(HTTPFS_CONFIGURATIONS).forEach(
+    ([dbParameterName, configurationKey]) => {
+      if (Object.keys(configurations).includes(configurationKey)) {
+        const dbValue = dbConfigurations.find(
+          ({ name }) => name === dbParameterName
+        ).value;
+        const expectedValue =
+          configurations[configurationKey as keyof typeof configurations];
+        if (configurationKey === 'use_ssl') {
+          expect(dbValue).toEqual('false');
+        } else {
+          expect(dbValue).toEqual(expectedValue);
+        }
+      }
+    }
+  );
+}, 500000);
+
+it('Unprovided Configurations should not been set into DuckDB session after loadExtension', async () => {
+  // Arrange
+  const db = new duckdb.Database(':memory:');
+  const connection = db.connect();
+  connection.run('INSTALL httpfs');
+  const configurations = {
+    use_ssl: false,
+    region: 'us-east-1',
+    //   accessKeyId: 'accessKeyId', // not provided, default is ''
+    //   secretAccessKey: 'secretAccessKey', // not provided, default is ''
+    endpoint: 'http://localhost:4566',
+    url_style: 'path',
+    sessionToken: 'sessionToken',
+  };
+  const unprovidedConfigurations = ['accessKeyId', 'secretAccessKey'];
+  const configurationLoader = new DuckDBExtensionLoader(configurations);
+
+  // Act
+  await configurationLoader.loadExtension(connection, 'httpfs');
+
+  // Assert
+  let dbConfigurations = await getQueryResults(
+    connection,
+    'SELECT * FROM duckdb_settings();'
+  );
+  dbConfigurations = dbConfigurations.filter(({ name }) =>
+    unprovidedConfigurations.includes(name)
+  );
+  dbConfigurations.forEach(({ name, value }) => {
+    if (unprovidedConfigurations.includes(name)) {
+      expect(value).toEqual('');
+    }
+  });
+}, 500000);
+
+it('HTTPFS_CONFIGURATIONS key should be a DuckDB parameter', async () => {
+  // Arrange
+  const db = new duckdb.Database(':memory:');
+  const connection = db.connect();
+  //   duckdb will have the configuration after install extension
+  await new Promise((resolve, reject) => {
+    connection.run('INSTALL httpfs');
+    connection.run('LOAD httpfs', (err: any) => {
+      if (err) reject(err);
+      resolve(true);
+    });
+  });
+
+  // Assert
+  const dbConfigurations = await getQueryResults(
+    connection,
+    'SELECT * FROM duckdb_settings();'
+  );
+  const dbConfigurationNames = dbConfigurations.map(({ name }) => name);
+  Object.keys(HTTPFS_CONFIGURATIONS).forEach((configurationKey) => {
+    expect(dbConfigurationNames).toContain(configurationKey);
+  });
+}, 500000);
diff --git a/packages/extension-driver-duckdb/tests/duckdbDataSource.spec.ts b/packages/extension-driver-duckdb/tests/duckdbDataSource.spec.ts
@@ -275,8 +275,7 @@ it('Should print queries without binding when log-queries = true', async () => {
     profileName: 'mocked-profile',
   });
   // Assert
-  expect(logs[0].length).toBe(1);
-  expect(logs[0][0]).toBe(`select $1::INTEGER as test`);
+  expect(logs.slice(-1)[0][0]).toBe(`select $1::INTEGER as test`);
 });
 
 it('Should print queries with binding when log-queries = true and log-parameters = true', async () => {
@@ -317,9 +316,8 @@ it('Should print queries with binding when log-queries = true and log-parameters
     profileName: 'mocked-profile',
   });
   // Assert
-  expect(logs[0].length).toBe(2);
-  expect(logs[0][0]).toBe(`select $1::INTEGER as test`);
-  expect(logs[0][1]).toEqual([1234]);
+  expect(logs.slice(-1)[0][0]).toBe(`select $1::INTEGER as test`);
+  expect(logs.slice(-1)[0][1]).toEqual([1234]);
 });
 
 it('Should share db instances for same path besides in-memory only db', async () => {
