From ddb54cd3cab51994a89cefe9ba083fb20622d559 Mon Sep 17 00:00:00 2001
From: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Date: Sat, 9 Aug 2025 03:33:47 +0530
Subject: [PATCH] support for session params

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
---
 examples/session_params.js         | 29 +++++++++++++++++++++++++++++
 lib/DBSQLClient.ts                 |  1 +
 lib/contracts/IDBSQLClient.ts      |  1 +
 tests/e2e/query_parameters.test.ts |  3 +++
 tests/unit/DBSQLClient.test.ts     | 10 ++++++++++
 5 files changed, 44 insertions(+)
 create mode 100644 examples/session_params.js

diff --git a/examples/session_params.js b/examples/session_params.js
new file mode 100644
index 00000000..d5d31208
--- /dev/null
+++ b/examples/session_params.js
@@ -0,0 +1,29 @@
+const { DBSQLClient } = require('..');
+
+const client = new DBSQLClient();
+
+const host = process.env.DATABRICKS_HOST;
+const path = process.env.DATABRICKS_HTTP_PATH;
+const token = process.env.DATABRICKS_TOKEN;
+
+client
+  .connect({ host, path, token })
+  .then(async (client) => {
+    const session = await client.openSession({
+      configuration: {
+        QUERY_TAGS: 'team:engineering,test:session-params,driver:node',
+        ansi_mode: 'false',
+      },
+    });
+
+    const op = await session.executeStatement('SELECT 1');
+    const rows = await op.fetchAll();
+    console.log(rows);
+    await op.close();
+
+    await session.close();
+    await client.close();
+  })
+  .catch((error) => {
+    console.log(error);
+  });
diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts
index 5dcf8a48..79cd936d 100644
--- a/lib/DBSQLClient.ts
+++ b/lib/DBSQLClient.ts
@@ -221,6 +221,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
     const response = await this.driver.openSession({
       client_protocol_i64: new Int64(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8),
       ...getInitialNamespaceOptions(request.initialCatalog, request.initialSchema),
+      configuration: request.configuration,
       canUseMultipleCatalogs: true,
     });
 
diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts
index 73fedde1..e18b45fb 100644
--- a/lib/contracts/IDBSQLClient.ts
+++ b/lib/contracts/IDBSQLClient.ts
@@ -38,6 +38,7 @@ export type ConnectionOptions = {
 export interface OpenSessionRequest {
   initialCatalog?: string;
   initialSchema?: string;
+  configuration?: { [key: string]: string };
 }
 
 export default interface IDBSQLClient {
diff --git a/tests/e2e/query_parameters.test.ts b/tests/e2e/query_parameters.test.ts
index d9b1a470..61022758 100644
--- a/tests/e2e/query_parameters.test.ts
+++ b/tests/e2e/query_parameters.test.ts
@@ -17,6 +17,9 @@ const openSession = async () => {
   return connection.openSession({
     initialCatalog: config.catalog,
     initialSchema: config.schema,
+    configuration: {
+      QUERY_TAGS: 'test:e2e,driver:node',
+    },
   });
 };
 
diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts
index f942c6b8..b81dbae1 100644
--- a/tests/unit/DBSQLClient.test.ts
+++ b/tests/unit/DBSQLClient.test.ts
@@ -185,6 +185,16 @@ describe('DBSQLClient.openSession', () => {
     }
   });
 
+  it('should pass session configuration to OpenSessionReq', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    const configuration = { QUERY_TAGS: 'team:engineering', ansi_mode: 'true' };
+    await client.openSession({ configuration });
+    expect(thriftClient.openSessionReq?.configuration).to.deep.equal(configuration);
+  });
+
   it('should affect session behavior based on protocol version', async () => {
     const client = new DBSQLClient();
     const thriftClient = new ThriftClientStub();