From 40edb728f93bc3531a45a14b01c71859eaf2796a Mon Sep 17 00:00:00 2001
From: Shivam Raj <shivam.raj@databricks.com>
Date: Wed, 2 Apr 2025 04:15:09 +0530
Subject: [PATCH 1/3] make lz4 compression configurable

---
 lib/DBSQLSession.ts            | 2 +-
 lib/contracts/IDBSQLSession.ts | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts
index 768b6b5e..009a8c8b 100644
--- a/lib/DBSQLSession.ts
+++ b/lib/DBSQLSession.ts
@@ -202,7 +202,7 @@ export default class DBSQLSession implements IDBSQLSession {
       ...getArrowOptions(clientConfig),
       canDownloadResult: options.useCloudFetch ?? clientConfig.useCloudFetch,
       parameters: getQueryParameters(options.namedParameters, options.ordinalParameters),
-      canDecompressLZ4Result: clientConfig.useLZ4Compression && Boolean(LZ4),
+      canDecompressLZ4Result: (options.useLZ4Compression ?? clientConfig.useLZ4Compression) && Boolean(LZ4),
     });
     const response = await this.handleResponse(operationPromise);
     const operation = this.createOperation(response);
diff --git a/lib/contracts/IDBSQLSession.ts b/lib/contracts/IDBSQLSession.ts
index 8a0d8bf0..0f751714 100644
--- a/lib/contracts/IDBSQLSession.ts
+++ b/lib/contracts/IDBSQLSession.ts
@@ -17,6 +17,7 @@ export type ExecuteStatementOptions = {
   runAsync?: boolean;
   maxRows?: number | bigint | Int64 | null;
   useCloudFetch?: boolean;
+  useLZ4Compression?: boolean;
   stagingAllowedLocalPath?: string | string[];
   namedParameters?: Record<string, DBSQLParameter | DBSQLParameterValue>;
   ordinalParameters?: Array<DBSQLParameter | DBSQLParameterValue>;

From d34b96baa3f02589e9623ba0452f6dd58e96cc47 Mon Sep 17 00:00:00 2001
From: Shivam Raj <shivam.raj@databricks.com>
Date: Wed, 2 Apr 2025 04:24:58 +0530
Subject: [PATCH 2/3] added unit tests with lz4 disabled

---
 tests/unit/result/ArrowResultHandler.test.ts  | 17 ++++++++++++++++
 .../result/CloudFetchResultHandler.test.ts    | 20 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/tests/unit/result/ArrowResultHandler.test.ts b/tests/unit/result/ArrowResultHandler.test.ts
index c657b16b..715f9fe0 100644
--- a/tests/unit/result/ArrowResultHandler.test.ts
+++ b/tests/unit/result/ArrowResultHandler.test.ts
@@ -66,6 +66,10 @@ const sampleRowSet4: TRowSet = {
   ],
 };
 
+declare global {
+  var LZ4: typeof import('lz4') | undefined; // Declare LZ4 on globalThis with proper type
+}
+
 describe('ArrowResultHandler', () => {
   it('should return data', async () => {
     const rowSetProvider = new ResultsProviderStub([sampleRowSet1], undefined);
@@ -229,4 +233,17 @@ describe('ArrowResultHandler', () => {
     });
     expect(await result.hasMore()).to.be.false;
   });
+
+  it('should handle data without LZ4 compression', async () => {
+    const rowSetProvider = new ResultsProviderStub([sampleRowSet1], undefined);
+    const result = new ArrowResultHandler(new ClientContextStub(), rowSetProvider, {
+      status: { statusCode: TStatusCode.SUCCESS_STATUS },
+      arrowSchema: sampleArrowSchema,
+      lz4Compressed: false,
+    });
+
+    const { batches } = await result.fetchNext({ limit: 10000 });
+    const expectedBatches = sampleRowSet1.arrowBatches?.map(({ batch }) => batch) ?? []; // Ensure iterable
+    expect(batches).to.deep.eq([sampleArrowSchema, ...expectedBatches]);
+  });
 });
diff --git a/tests/unit/result/CloudFetchResultHandler.test.ts b/tests/unit/result/CloudFetchResultHandler.test.ts
index 7927ee41..77950d6a 100644
--- a/tests/unit/result/CloudFetchResultHandler.test.ts
+++ b/tests/unit/result/CloudFetchResultHandler.test.ts
@@ -379,4 +379,24 @@ describe('CloudFetchResultHandler', () => {
       expect(context.invokeWithRetryStub.callCount).to.be.equal(1);
     }
   });
+
+  it('should handle data without LZ4 compression', async () => {
+    const context = new ClientContextStub();
+    const rowSetProvider = new ResultsProviderStub([sampleRowSet1], undefined);
+
+    const result = new CloudFetchResultHandler(context, rowSetProvider, {
+      lz4Compressed: false,
+      status: { statusCode: TStatusCode.SUCCESS_STATUS },
+    });
+
+    context.invokeWithRetryStub.callsFake(async () => ({
+      request: new Request('localhost'),
+      response: new Response(sampleArrowBatch, { status: 200 }), // Return only the batch
+    }));
+
+    const { batches } = await result.fetchNext({ limit: 10000 });
+
+    // Ensure the batches array matches the expected structure
+    expect(batches).to.deep.eq([sampleArrowBatch]);
+  });
 });

From 4736816e1475490ffedd96eb383a1bfc5724c146 Mon Sep 17 00:00:00 2001
From: Shivam Raj <shivam.raj@databricks.com>
Date: Wed, 2 Apr 2025 04:26:22 +0530
Subject: [PATCH 3/3] removed unnecessary declaration

---
 tests/unit/result/ArrowResultHandler.test.ts | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/unit/result/ArrowResultHandler.test.ts b/tests/unit/result/ArrowResultHandler.test.ts
index 715f9fe0..3003083f 100644
--- a/tests/unit/result/ArrowResultHandler.test.ts
+++ b/tests/unit/result/ArrowResultHandler.test.ts
@@ -66,10 +66,6 @@ const sampleRowSet4: TRowSet = {
   ],
 };
 
-declare global {
-  var LZ4: typeof import('lz4') | undefined; // Declare LZ4 on globalThis with proper type
-}
-
 describe('ArrowResultHandler', () => {
   it('should return data', async () => {
     const rowSetProvider = new ResultsProviderStub([sampleRowSet1], undefined);