Allow any number type (number, bigint, Int64) for maxRows and `quer…

…yTimeout` options

Signed-off-by: Levko Kravets <[email protected]>

lib/DBSQLSession.ts

@@ -49,14 +49,29 @@ interface OperationResponseShape {
   directResults?: TSparkDirectResults;
 }
 
-function getDirectResultsOptions(maxRows: number | null = defaultMaxRows) {
+export function numberToInt64(value: number | bigint | Int64): Int64 {
+  if (value instanceof Int64) {
+    return value;
+  }
+
+  if (typeof value === 'bigint') {
+    const buffer = new ArrayBuffer(BigInt64Array.BYTES_PER_ELEMENT);
+    const view = new DataView(buffer);
+    view.setBigInt64(0, value, false); // `false` to use big-endian order
+    return new Int64(Buffer.from(buffer));
+  }
+
+  return new Int64(value);
+}
+
+function getDirectResultsOptions(maxRows: number | bigint | Int64 | null = defaultMaxRows) {
   if (maxRows === null) {
     return {};
   }
 
   return {
     getDirectResults: {
-      maxRows: new Int64(maxRows),
+      maxRows: numberToInt64(maxRows),
     },
   };
 }
@@ -184,7 +199,7 @@ export default class DBSQLSession implements IDBSQLSession {
     const operationPromise = driver.executeStatement({
       sessionHandle: this.sessionHandle,
       statement,
-      queryTimeout: options.queryTimeout,
+      queryTimeout: options.queryTimeout ? numberToInt64(options.queryTimeout) : undefined,
       runAsync: true,
       ...getDirectResultsOptions(options.maxRows),
       ...getArrowOptions(clientConfig),

lib/contracts/IDBSQLSession.ts

@@ -5,12 +5,12 @@ import InfoValue from '../dto/InfoValue';
 import { DBSQLParameter, DBSQLParameterValue } from '../DBSQLParameter';
 
 export type ExecuteStatementOptions = {
-  queryTimeout?: Int64;
+  queryTimeout?: number | bigint | Int64;
   /**
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
   useCloudFetch?: boolean;
   stagingAllowedLocalPath?: string | string[];
   namedParameters?: Record<string, DBSQLParameter | DBSQLParameterValue>;
@@ -22,15 +22,15 @@ export type TypeInfoRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type CatalogsRequest = {
   /**
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type SchemasRequest = {
@@ -40,7 +40,7 @@ export type SchemasRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type TablesRequest = {
@@ -52,15 +52,15 @@ export type TablesRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type TableTypesRequest = {
   /**
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type ColumnsRequest = {
@@ -72,7 +72,7 @@ export type ColumnsRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type FunctionsRequest = {
@@ -83,7 +83,7 @@ export type FunctionsRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type PrimaryKeysRequest = {
@@ -94,7 +94,7 @@ export type PrimaryKeysRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export type CrossReferenceRequest = {
@@ -108,7 +108,7 @@ export type CrossReferenceRequest = {
    * @deprecated This option is no longer supported and will be removed in future releases
    */
   runAsync?: boolean;
-  maxRows?: number | null;
+  maxRows?: number | bigint | Int64 | null;
 };
 
 export default interface IDBSQLSession {

tests/unit/DBSQLSession.test.js

@@ -1,7 +1,8 @@
 const { expect, AssertionError } = require('chai');
 const { DBSQLLogger, LogLevel } = require('../../dist');
 const sinon = require('sinon');
-const DBSQLSession = require('../../dist/DBSQLSession').default;
+const Int64 = require('node-int64');
+const { default: DBSQLSession, numberToInt64 } = require('../../dist/DBSQLSession');
 const InfoValue = require('../../dist/dto/InfoValue').default;
 const Status = require('../../dist/dto/Status').default;
 const DBSQLOperation = require('../../dist/DBSQLOperation').default;
@@ -62,6 +63,30 @@ async function expectFailure(fn) {
 }
 
 describe('DBSQLSession', () => {
+  describe('numberToInt64', () => {
+    it('should convert regular number to Int64', () => {
+      const num = Math.random() * 1000000;
+      const value = numberToInt64(num);
+      expect(value.equals(new Int64(num))).to.be.true;
+    });
+
+    it('should return Int64 values as is', () => {
+      const num = new Int64(Math.random() * 1000000);
+      const value = numberToInt64(num);
+      expect(value).to.equal(num);
+    });
+
+    it('should convert BigInt to Int64', () => {
+      // This case is especially important, because Int64 has no native methods to convert
+      // between Int64 and BigInt. This conversion involves some byte operations, and it's
+      // important to make sure we don't mess up with things like byte order
+
+      const num = BigInt(Math.round(Math.random() * 10000)) * BigInt(Math.round(Math.random() * 10000));
+      const value = numberToInt64(num);
+      expect(value.toString()).equal(num.toString());
+    });
+  });
+
   describe('getInfo', () => {
     it('should run operation', async () => {
       const session = createSession();