[PECOBLR-306] Make LZ4 compression configurable (#288)

shivam2680 · web-flow · commit b754b77fccd0 · 2025-04-04T14:18:14.000+05:30
diff --git 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
@@ -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>;
diff --git a/tests/unit/result/ArrowResultHandler.test.ts b/tests/unit/result/ArrowResultHandler.test.ts
@@ -229,4 +229,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
@@ -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]);
+  });
 });
