Expose session and query ID (#250)

Signed-off-by: Levko Kravets <levko.ne@gmail.com>
databricks · Apr 19, 2024 · 38db1b0 · 38db1b0
1 parent 1255fad
commit 38db1b0
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 15 deletions.
diff --git a/lib/DBSQLOperation.ts b/lib/DBSQLOperation.ts
@@ -1,4 +1,4 @@
-import { stringify, NIL, parse } from 'uuid';
+import { stringify, NIL } from 'uuid';
 import IOperation, {
   FetchOptions,
   FinishedOptions,
@@ -88,11 +88,12 @@ export default class DBSQLOperation implements IOperation {
       useOnlyPrefetchedResults,
     );
     this.closeOperation = directResults?.closeOperation;
-    this.context.getLogger().log(LogLevel.debug, `Operation created with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Operation created with id: ${this.id}`);
   }
 
-  public getId() {
-    return stringify(this.operationHandle?.operationId?.guid || parse(NIL));
+  public get id() {
+    const operationId = this.operationHandle?.operationId?.guid;
+    return operationId ? stringify(operationId) : NIL;
   }
 
   /**
@@ -119,7 +120,7 @@ export default class DBSQLOperation implements IOperation {
       const chunk = await this.fetchChunk(fetchChunkOptions);
       data.push(chunk);
     } while (await this.hasMoreRows()); // eslint-disable-line no-await-in-loop
-    this.context.getLogger().log(LogLevel.debug, `Fetched all data from operation with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Fetched all data from operation with id: ${this.id}`);
 
     return data.flat();
   }
@@ -173,7 +174,7 @@ export default class DBSQLOperation implements IOperation {
       .getLogger()
       .log(
         LogLevel.debug,
-        `Fetched chunk of size: ${options?.maxRows || defaultMaxRows} from operation with id: ${this.getId()}`,
+        `Fetched chunk of size: ${options?.maxRows || defaultMaxRows} from operation with id: ${this.id}`,
       );
     return result;
   }
@@ -185,7 +186,7 @@ export default class DBSQLOperation implements IOperation {
    */
   public async status(progress: boolean = false): Promise<TGetOperationStatusResp> {
     await this.failIfClosed();
-    this.context.getLogger().log(LogLevel.debug, `Fetching status for operation with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Fetching status for operation with id: ${this.id}`);
 
     if (this.operationStatus) {
       return this.operationStatus;
@@ -209,7 +210,7 @@ export default class DBSQLOperation implements IOperation {
       return Status.success();
     }
 
-    this.context.getLogger().log(LogLevel.debug, `Cancelling operation with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Cancelling operation with id: ${this.id}`);
 
     const driver = await this.context.getDriver();
     const response = await driver.cancelOperation({
@@ -233,7 +234,7 @@ export default class DBSQLOperation implements IOperation {
       return Status.success();
     }
 
-    this.context.getLogger().log(LogLevel.debug, `Closing operation with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Closing operation with id: ${this.id}`);
 
     const driver = await this.context.getDriver();
     const response =
@@ -274,7 +275,7 @@ export default class DBSQLOperation implements IOperation {
 
     await this.waitUntilReady(options);
 
-    this.context.getLogger().log(LogLevel.debug, `Fetching schema for operation with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Fetching schema for operation with id: ${this.id}`);
     const metadata = await this.fetchMetadata();
     return metadata.schema ?? null;
   }

diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts
@@ -2,7 +2,7 @@ import * as fs from 'fs';
 import * as path from 'path';
 import stream from 'node:stream';
 import util from 'node:util';
-import { stringify, NIL, parse } from 'uuid';
+import { stringify, NIL } from 'uuid';
 import fetch, { HeadersInit } from 'node-fetch';
 import {
   TSessionHandle,
@@ -140,11 +140,12 @@ export default class DBSQLSession implements IDBSQLSession {
   constructor({ handle, context }: DBSQLSessionConstructorOptions) {
     this.sessionHandle = handle;
     this.context = context;
-    this.context.getLogger().log(LogLevel.debug, `Session created with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Session created with id: ${this.id}`);
   }
 
-  public getId() {
-    return stringify(this.sessionHandle?.sessionId?.guid || parse(NIL));
+  public get id() {
+    const sessionId = this.sessionHandle?.sessionId?.guid;
+    return sessionId ? stringify(sessionId) : NIL;
   }
 
   /**
@@ -531,7 +532,7 @@ export default class DBSQLSession implements IDBSQLSession {
     this.onClose?.();
     this.isOpen = false;
 
-    this.context.getLogger().log(LogLevel.debug, `Session closed with id: ${this.getId()}`);
+    this.context.getLogger().log(LogLevel.debug, `Session closed with id: ${this.id}`);
     return new Status(response.status);
   }
 

diff --git a/lib/contracts/IDBSQLSession.ts b/lib/contracts/IDBSQLSession.ts
@@ -112,6 +112,11 @@ export type CrossReferenceRequest = {
 };
 
 export default interface IDBSQLSession {
+  /**
+   * Session identifier
+   */
+  readonly id: string;
+
   /**
    * Returns general information about the data source
    *

diff --git a/lib/contracts/IOperation.ts b/lib/contracts/IOperation.ts
@@ -24,6 +24,11 @@ export interface GetSchemaOptions extends WaitUntilReadyOptions {
 }
 
 export default interface IOperation {
+  /**
+   * Operation identifier
+   */
+  readonly id: string;
+
   /**
    * Fetch a portion of data
    */