googleapis · alkatrivedi · May 6, 2025 · Apr 20, 2025 · Apr 20, 2025 · Apr 20, 2025
@@ -18,11 +18,12 @@ import {PreciseDate} from '@google-cloud/precise-date';
 import {promisifyAll} from '@google-cloud/promisify';
 import * as extend from 'extend';
 import * as is from 'is';
-import {Snapshot} from './transaction';
+import {ExecuteSqlRequest, Snapshot} from './transaction';
 import {google} from '../protos/protos';
 import {Session, Database} from '.';
 import {
   CLOUD_RESOURCE_HEADER,
+  ResourceCallback,
   addLeaderAwareRoutingHeader,
 } from '../src/common';
 import {startTrace, setSpanError, traceConfig} from './instrument';
@@ -34,6 +35,16 @@ export interface TransactionIdentifier {
   timestamp?: google.protobuf.ITimestamp;
 }
 
+export type CreateQueryPartitionsResponse = [
+  google.spanner.v1.IPartitionQueryRequest,
+  google.spanner.v1.IPartitionResponse,
+];
+
+export type CreateQueryPartitionsCallback = ResourceCallback<
+  google.spanner.v1.IPartitionQueryRequest,
+  google.spanner.v1.IPartitionResponse
+>;
+
 /**
  * Use a BatchTransaction object to create partitions and read/query against
  * your Cloud Spanner database.
@@ -126,20 +137,27 @@ class BatchTransaction extends Snapshot {
    * @example <caption>include:samples/batch.js</caption>
    * region_tag:spanner_batch_client
    */
-  createQueryPartitions(query, callback) {
-    if (is.string(query)) {
-      query = {
-        sql: query,
-      };
-    }
+  createQueryPartitions(
+    query: string | ExecuteSqlRequest,
+  ): Promise<CreateQueryPartitionsResponse>;
+  createQueryPartitions(
+    query: string | ExecuteSqlRequest,
+    callback: CreateQueryPartitionsCallback,
+  ): void;
+  createQueryPartitions(
+    query: string | ExecuteSqlRequest,
+    cb?: CreateQueryPartitionsCallback,
+  ): void | Promise<CreateQueryPartitionsResponse> {
+    const request: ExecuteSqlRequest =
+      typeof query === 'string' ? {sql: query} : query;
 
-    const reqOpts = Object.assign({}, query, Snapshot.encodeParams(query));
+    const reqOpts = Object.assign({}, request, Snapshot.encodeParams(request));
 
-    delete reqOpts.gaxOptions;
-    delete reqOpts.types;
+    delete (reqOpts as any).gaxOptions;
+    delete (reqOpts as any).types;
 
     const traceConfig: traceConfig = {
-      sql: query,
+      sql: request.sql,
       opts: this._observabilityOptions,
       dbName: this.getDBName(),
     };
@@ -157,7 +175,7 @@ class BatchTransaction extends Snapshot {
             client: 'SpannerClient',
             method: 'partitionQuery',
             reqOpts,
-            gaxOpts: query.gaxOptions,
+            gaxOpts: request.gaxOptions,
             headers: injectRequestIDIntoHeaders(headers, this.session),
           },
           (err, partitions, resp) => {
@@ -166,7 +184,7 @@ class BatchTransaction extends Snapshot {
             }
 
             span.end();
-            callback(err, partitions, resp);
+            cb!(err, partitions, resp);
           },
         );
       },

@@ -32,6 +32,8 @@ import {
   CLOUD_RESOURCE_HEADER,
   LEADER_AWARE_ROUTING_HEADER,
 } from '../src/common';
+import {ExecuteSqlRequest} from '../src/transaction';
+import {CallOptions} from 'google-gax';
 
 let promisified = false;
 const fakePfy = extend({}, pfy, {
@@ -159,16 +161,16 @@ describe('BatchTransaction', () => {
       },
     };
 
-    const QUERY = {
+    const QUERY: ExecuteSqlRequest = {
       sql: 'SELECT * FROM Singers',
-      gaxOptions: GAX_OPTS,
+      gaxOptions: GAX_OPTS as CallOptions,
       params: {},
       types: {},
       dataBoostEnabled: true,
       directedReadOptions: fakeDirectedReadOptionsForRequest,
     };
 
-    it('should make the correct request', () => {
+    it('should make the correct request using callback', () => {
       const fakeParams = {
         params: {a: 'b'},
         paramTypes: {a: 'string'},
@@ -196,7 +198,7 @@ describe('BatchTransaction', () => {
       );
     });
 
-    it('should accept query as string', () => {
+    it('should accept query as string in a callback based request to createQueryPartitions', () => {
       const query = 'SELECT * FROM Singers';
 
       const expectedQuery = Object.assign({}, {sql: query});
@@ -214,6 +216,34 @@ describe('BatchTransaction', () => {
       assert.deepStrictEqual(reqOpts, expectedQuery);
       assert.strictEqual(gaxOpts, undefined);
     });
+
+    it('should make the correct request using await', async () => {
+      const fakeParams = {
+        params: {a: 'b'},
+        paramTypes: {a: 'string'},
+        dataBoostEnabled: true,
+        directedReadOptions: fakeDirectedReadOptionsForRequest,
+      };
+
+      const expectedQuery = Object.assign({sql: QUERY.sql}, fakeParams);
+      const stub = sandbox.stub(batchTransaction, 'createPartitions_');
+
+      (sandbox.stub(FakeTransaction, 'encodeParams') as sinon.SinonStub)
+        .withArgs(QUERY)
+        .returns(fakeParams);
+
+      await batchTransaction.createQueryPartitions(QUERY);
+
+      const {client, method, reqOpts, gaxOpts, headers} = stub.lastCall.args[0];
+      assert.strictEqual(client, 'SpannerClient');
+      assert.strictEqual(method, 'partitionQuery');
+      assert.deepStrictEqual(reqOpts, expectedQuery);
+      assert.strictEqual(gaxOpts, GAX_OPTS);
+      assert.deepStrictEqual(
+        headers,
+        Object.assign({[LEADER_AWARE_ROUTING_HEADER]: 'true'}),
+      );
+    });
   });
 
   describe('createPartitions_', () => {

@@ -264,6 +264,7 @@ export class MockSpanner {
 
     this.executeBatchDml = this.executeBatchDml.bind(this);
     this.executeStreamingSql = this.executeStreamingSql.bind(this);
+    this.partitionQuery = this.partitionQuery.bind(this);
 
     this.read = this.read.bind(this);
     this.streamingRead = this.streamingRead.bind(this);
@@ -974,11 +975,21 @@ export class MockSpanner {
   }
 
   partitionQuery(
-    call: grpc.ServerUnaryCall<protobuf.PartitionQueryRequest, {}>,
+    call: grpc.ServerUnaryCall<
+      protobuf.PartitionQueryRequest,
+      protobuf.PartitionResponse
+    >,
     callback: protobuf.Spanner.PartitionQueryCallback,
   ) {
     this.pushRequest(call.request!, call.metadata);
-    callback(createUnimplementedError('PartitionQuery is not yet implemented'));
+    this.simulateExecutionTime(this.partitionQuery.name)
+      .then(() => {
+        const response = protobuf.PartitionResponse.create({
+          partitions: [{partitionToken: Buffer.from('mock-token')}],
+        });
+        callback(null, response);
+      })
+      .catch(err => callback(err));
   }
 
   partitionRead(

@@ -3651,6 +3651,21 @@ describe('Spanner with mock server', () => {
       });
     });
 
+    describe('createQueryPartitions', () => {
+      it('should create set of query partitions', async () => {
+        const database = newTestDatabase({min: 0, incStep: 1});
+        const query = {
+          sql: select1,
+        };
+        const [transaction] = await database.createBatchTransaction();
+        const [partitions] = await transaction.createQueryPartitions(query);
+        assert.strictEqual(Object.keys(partitions).length, 1);
+        assert.strictEqual(partitions[0].sql, select1);
+        transaction.close();
+        await database.close();
+      });
+    });
+
     describe('pdml', () => {
       it('should retry on aborted error', async () => {
         const database = newTestDatabase();