Merge pull request #310 from Canner/feature/cache-impersonate

Feature: extends CacheLayerInfo interface and Profile interface

Author: kokokuo

packages/core/src/lib/cache-layer/cacheLayerLoader.ts

@@ -42,7 +42,14 @@ export class CacheLayerLoader implements ICacheLayerLoader {
     templateName: string,
     cache: CacheLayerInfo
   ): Promise<void> {
-    const { cacheTableName, sql, profile, indexes, folderSubpath } = cache;
+    const {
+      cacheTableName,
+      sql,
+      profile,
+      indexes,
+      folderSubpath,
+      options: cacheOptions,
+    } = cache;
     const type = this.options.type!;
     const dataSource = this.dataSourceFactory(profile);
 
@@ -81,6 +88,7 @@ export class CacheLayerLoader implements ICacheLayerLoader {
         directory,
         profileName: profile,
         type,
+        options: cacheOptions,
       });
     } else {
       this.logger.debug(

packages/core/src/models/artifact.ts

@@ -118,6 +118,8 @@ export class CacheLayerInfo {
   indexes?: Record<string, string>;
   // cache folder subpath
   folderSubpath?: string;
+  // options pass to the data source
+  options?: any;
 }
 
 export class APISchema {

packages/core/src/models/extensions/dataSource.ts

@@ -19,6 +19,8 @@ export interface ExportOptions {
   directory: string;
   // The profile name to select to export data
   profileName: string;
+  // data source options
+  options?: any;
   // export file format type
   type: CacheLayerStoreFormatType | string;
 }

packages/core/src/models/profile.ts

@@ -29,4 +29,6 @@ export interface Profile<C = Record<string, any>> {
   cache?: C;
   /** What users have access to this profile */
   allow: ProfileAllowConstraints;
+  /** Properties that can be used when involking the dataSource method */
+  properties?: Record<string, any>;
 }

packages/extension-driver-canner/src/lib/cannerAdapter.ts

@@ -31,15 +31,23 @@ export class CannerAdapter {
   // When querying Canner enterprise, the Canner enterprise will save the query result as parquet files,
   // and store them in S3. This method will return the S3 urls of the query result.
   // For more Canner API ref: https://docs.cannerdata.com/reference/restful
-  public async createAsyncQueryResultUrls(sql: string): Promise<string[]> {
+  public async createAsyncQueryResultUrls(
+    sql: string,
+    headers?: Record<string, string>
+  ): Promise<string[]> {
     this.logger.debug(`Create async request to Canner.`);
-    let data = await this.getWorkspaceRequestData('post', '/v2/async-queries', {
-      data: {
-        sql,
-        timeout: 600,
-        noLimit: true,
+    let data = await this.getWorkspaceRequestData(
+      'post',
+      '/v2/async-queries',
+      {
+        data: {
+          sql,
+          timeout: 600,
+          noLimit: true,
+        },
       },
-    });
+      headers
+    );
 
     const { id: requestId } = data;
     this.logger.debug(`Wait Async request to finished.`);
@@ -60,14 +68,13 @@ export class CannerAdapter {
   private async getWorkspaceRequestData(
     method: string,
     urlPath: string,
-    options?: Record<string, any>
+    options?: Record<string, any>,
+    headers?: Record<string, string>
   ) {
     await this.prepare();
     try {
       const response = await axios({
-        headers: {
-          Authorization: `Token ${this.PAT}`,
-        },
+        headers: { ...headers, Authorization: `Token ${this.PAT}` },
         params: {
           workspaceSqlName: this.workspaceSqlName,
         },
@@ -78,7 +85,9 @@ export class CannerAdapter {
       return response.data;
     } catch (error: any) {
       const message = error.response
-        ? `response status: ${error.response.status}, response data: ${error.response.data}`
+        ? `response status: ${
+            error.response.status
+          }, response data: ${JSON.stringify(error.response.data)}`
         : `remote server does not response. request ${error.toJSON()}}`;
       throw new InternalError(
         `Failed to get workspace request "${urlPath}" data, ${message}`

packages/extension-driver-canner/src/lib/cannerDataSource.ts

@@ -26,7 +26,7 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
   private logger = this.getLogger();
   protected poolMapping = new Map<
     string,
-    { pool: Pool; options?: PGOptions }
+    { pool: Pool; options?: PGOptions; properties?: Record<string, any> }
   >();
   protected UserPool = new Map<string, Pool>();
 
@@ -52,6 +52,7 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
       this.poolMapping.set(profile.name, {
         pool,
         options: profile.connection,
+        properties: profile.properties,
       });
       this.logger.debug(`Profile ${profile.name} initialized`);
     }
@@ -61,6 +62,7 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
     sql,
     directory,
     profileName,
+    options: cannerOptions,
   }: ExportOptions): Promise<void> {
     if (!this.poolMapping.has(profileName)) {
       throw new InternalError(`Profile instance ${profileName} not found`);
@@ -69,12 +71,16 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
     if (!fs.existsSync(directory)) {
       throw new InternalError(`Directory ${directory} not found`);
     }
-    const { options: connection } = this.poolMapping.get(profileName)!;
-
+    const { options: connection, properties } =
+      this.poolMapping.get(profileName)!;
     const cannerAdapter = new CannerAdapter(connection);
     try {
       this.logger.debug('Send the async query to the Canner Enterprise');
-      const presignedUrls = await cannerAdapter.createAsyncQueryResultUrls(sql);
+      const header = this.getCannerRequestHeader(properties, cannerOptions);
+      const presignedUrls = await cannerAdapter.createAsyncQueryResultUrls(
+        sql,
+        header
+      );
       this.logger.debug(
         'Start fetching the query result parquet files from URLs'
       );
@@ -85,6 +91,21 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
       throw error;
     }
   }
+  private getCannerRequestHeader(
+    properties?: Record<string, any>,
+    cannerOptions?: any
+  ) {
+    const header: Record<string, string> = {};
+    const userId = cannerOptions?.userId;
+    const rootUserId = properties?.['rootUserId'];
+    if (userId && rootUserId) {
+      header[
+        'x-trino-session'
+      ] = `root_user_id=${rootUserId}, canner_user_id=${userId}`;
+      this.logger.debug(`Impersonate used: ${userId}`);
+    }
+    return header;
+  }
 
   private async downloadFiles(urls: string[], directory: string) {
     await Promise.all(

packages/extension-driver-canner/test/cannerDataSource.spec.ts

@@ -55,6 +55,7 @@ it('Data source should export successfully', async () => {
       sql: 'select 1',
       directory,
       profileName: 'profile1',
+      options: {},
     } as ExportOptions)
   ).resolves.not.toThrow();
   expect(fs.readdirSync(directory).length).toBe(1);
@@ -86,6 +87,7 @@ it('Data source should throw error when fail to export data', async () => {
       sql: 'select 1',
       directory,
       profileName: 'profile1',
+      options: {},
     } as ExportOptions)
   ).rejects.toThrow();
   expect(fs.readdirSync(directory).length).toBe(0);
@@ -105,6 +107,7 @@ it('Data source should throw error when given directory is not exist', async ()
       sql: 'select 1',
       directory: directory,
       profileName: 'profile1',
+      options: {},
     } as ExportOptions)
   ).rejects.toThrow();
 }, 100000);
@@ -121,6 +124,7 @@ it('Data source should throw error when given profile name is not exist', async
       sql: 'select 1',
       directory,
       profileName: 'profile not exist',
+      options: {},
     } as ExportOptions)
   ).rejects.toThrow();
 }, 100000);

packages/extension-driver-canner/test/cannerServer.ts

@@ -20,6 +20,7 @@ export class CannerServer {
         database: process.env['CANNER_WORKSPACE_SQL_NAME'],
       } as PGOptions,
       allow: '*',
+      properties: {},
     };
   }
 }

packages/extension-store-canner/README.md

@@ -63,6 +63,8 @@ export PROFILE_CANNER_DRIVER_PASSWORD=<password>
 export PROFILE_CANNER_DRIVER_HOST=<host>
 # Canner enterprise driver port, the default is 7432
 export PROFILE_CANNER_DRIVER_PORT=<port>
+# Canner enterprise root user id
+export PROFILE_CANNER_DRIVER_ROOT_USER_ID=<userId>
 ```
 
 ### Connect Canner Enterprise used storage.

packages/extension-store-canner/src/lib/canner/profileReader.ts

@@ -44,6 +44,7 @@ export class CannerProfileReader extends ProfileReader {
 
     // generate profiles from the indicator files of each workspaces
     const { user, password, host, port, max } = this.envConfig.profile;
+    const { rootUserId } = this.envConfig.properties;
     if (!user || !password || !host)
       throw new ConfigurationError(
         'Canner profile reader needs username, password, host properties.'
@@ -67,6 +68,9 @@ export class CannerProfileReader extends ProfileReader {
             max,
           },
           allow: '*',
+          properties: {
+            rootUserId,
+          },
         } as Profile<Record<string, any>>;
         this.logger.debug(`created "${profile.name}".`);
         return profile;