feat: Added support for partitioned dml

googleapis · Jan 4, 2024 · c0a2dab · c0a2dab
1 parent 66e7a5f
commit c0a2dab
Show file tree

Hide file tree

Showing 9 changed files with 45 additions and 46 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,6 +2,8 @@
 **/node_modules
 /.coverage
 /coverage
+/.idea
+/.vscode
 /.nyc_output
 /docs/
 /out/

diff --git a/samples/directed-reads.js b/samples/directed-reads.js
@@ -1,4 +1,4 @@
-// Copyright 2023 Google LLC
+// Copyright 2024 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -73,7 +73,7 @@ function main(
               .Type.READ_ONLY,
           },
         ],
-        autoFailover: true,
+        autoFailoverDisabled: true,
       },
     };
 

diff --git a/src/index.ts b/src/index.ts
@@ -112,7 +112,7 @@ export type GetInstanceConfigOperationsCallback = PagedCallback<
  * Disabling leader aware routing would route all requests in RW/PDML transactions to any region.
  * @property {google.spanner.v1.IDirectedReadOptions} [directedReadOptions] Sets the DirectedReadOptions for all ReadRequests and ExecuteSqlRequests for the Client.
  * Indicates which replicas or regions should be used for non-transactional reads or queries.
- * DirectedReadOptions won't be set for readWrite transactions or partitioned dml requests"
+ * DirectedReadOptions won't be set for readWrite transactions"
  */
 export interface SpannerOptions extends GrpcClientOptions {
   apiEndpoint?: string;

diff --git a/src/transaction.ts b/src/transaction.ts
@@ -1336,7 +1336,7 @@ export class Snapshot extends EventEmitter {
     if (
       !directedReadOptions &&
       this._getSpanner().directedReadOptions &&
-      !this._options.readWrite
+      this._options.readOnly
     ) {
       return this._getSpanner().directedReadOptions;
     }

diff --git a/system-test/spanner.ts b/system-test/spanner.ts
@@ -7113,7 +7113,7 @@ describe('Spanner', () => {
                   .ReplicaSelection.Type.READ_ONLY,
               },
             ],
-            autoFailover: true,
+            autoFailoverDisabled: true,
           },
         };
 
@@ -8717,7 +8717,7 @@ describe('Spanner', () => {
                   .ReplicaSelection.Type.READ_WRITE,
               },
             ],
-            autoFailover: true,
+            autoFailoverDisabled: true,
           },
         };
 

diff --git a/test/batch-transaction.ts b/test/batch-transaction.ts
@@ -25,6 +25,7 @@ import * as proxyquire from 'proxyquire';
 import * as sinon from 'sinon';
 
 import {Session, Database, Spanner} from '../src';
+import {protos} from '../src';
 import * as bt from '../src/batch-transaction';
 import {PartialResultStream} from '../src/partial-result-stream';
 import {
@@ -144,19 +145,35 @@ describe('BatchTransaction', () => {
 
   describe('createQueryPartitions', () => {
     const GAX_OPTS = {a: 'b'};
+
+    const fakeDirectedReadOptionsForRequest = {
+      includeReplicas: {
+        replicaSelections: [
+          {
+            location: 'us-west1',
+            type: protos.google.spanner.v1.DirectedReadOptions.ReplicaSelection
+              .Type.READ_WRITE,
+          },
+        ],
+        autoFailoverDisabled: true,
+      },
+    };
+
     const QUERY = {
       sql: 'SELECT * FROM Singers',
       gaxOptions: GAX_OPTS,
       params: {},
       types: {},
       dataBoostEnabled: true,
+      directedReadOptions: fakeDirectedReadOptionsForRequest,
     };
 
     it('should make the correct request', () => {
       const fakeParams = {
         params: {a: 'b'},
         paramTypes: {a: 'string'},
         dataBoostEnabled: true,
+        directedReadOptions: fakeDirectedReadOptionsForRequest,
       };
 
       const expectedQuery = Object.assign({sql: QUERY.sql}, fakeParams);
@@ -301,12 +318,27 @@ describe('BatchTransaction', () => {
 
   describe('createReadPartitions', () => {
     const GAX_OPTS = {};
+
+    const fakeDirectedReadOptionsForRequest = {
+      includeReplicas: {
+        replicaSelections: [
+          {
+            location: 'us-west1',
+            type: protos.google.spanner.v1.DirectedReadOptions.ReplicaSelection
+              .Type.READ_WRITE,
+          },
+        ],
+        autoFailoverDisabled: true,
+      },
+    };
+
     const QUERY = {
       table: 'abc',
       keys: ['a', 'b'],
       ranges: [{}, {}],
       gaxOptions: GAX_OPTS,
       dataBoostEnabled: true,
+      directedReadOptions: fakeDirectedReadOptionsForRequest,
     };
 
     it('should make the correct request', () => {
@@ -315,6 +347,7 @@ describe('BatchTransaction', () => {
         table: QUERY.table,
         keySet: fakeKeySet,
         dataBoostEnabled: true,
+        directedReadOptions: fakeDirectedReadOptionsForRequest,
       };
 
       const stub = sandbox.stub(batchTransaction, 'createPartitions_');

diff --git a/test/database.ts b/test/database.ts
@@ -2463,7 +2463,7 @@ describe('Database', () => {
               .Type.READ_WRITE,
           },
         ],
-        autoFailover: true,
+        autoFailoverDisabled: true,
       },
     };
 
@@ -2582,19 +2582,9 @@ describe('Database', () => {
       assert.ok(fakeCallback.calledOnce);
     });
 
-    it('should override directedReadOptions set for client when passed', () => {
+    it('should ignore directedReadOptions set for client', () => {
       const fakeCallback = sandbox.spy();
 
-      const fakeDirectedReadOptionsForRequest = {
-        includeReplicas: {
-          replicaSelections: [
-            {
-              location: 'us-east1',
-            },
-          ],
-        },
-      };
-
       database.parent.parent = {
         routeToLeaderEnabled: true,
         directedReadOptions: fakeDirectedReadOptions,
@@ -2605,31 +2595,6 @@ describe('Database', () => {
           sql: QUERY.sql,
           params: QUERY.params,
           requestOptions: {priority: RequestOptions.Priority.PRIORITY_LOW},
-          directedReadOptions: fakeDirectedReadOptionsForRequest,
-        },
-        fakeCallback
-      );
-
-      const [query] = runUpdateStub.lastCall.args;
-
-      assert.deepStrictEqual(query, {
-        sql: QUERY.sql,
-        params: QUERY.params,
-        requestOptions: {priority: RequestOptions.Priority.PRIORITY_LOW},
-        directedReadOptions: fakeDirectedReadOptionsForRequest,
-      });
-      assert.ok(fakeCallback.calledOnce);
-    });
-
-    it('should accept requestOptions for client', () => {
-      const fakeCallback = sandbox.spy();
-
-      database.runPartitionedUpdate(
-        {
-          sql: QUERY.sql,
-          params: QUERY.params,
-          requestOptions: {priority: RequestOptions.Priority.PRIORITY_LOW},
-          directedReadOptions: fakeDirectedReadOptions,
         },
         fakeCallback
       );
@@ -2640,7 +2605,6 @@ describe('Database', () => {
         sql: QUERY.sql,
         params: QUERY.params,
         requestOptions: {priority: RequestOptions.Priority.PRIORITY_LOW},
-        directedReadOptions: fakeDirectedReadOptions,
       });
       assert.ok(fakeCallback.calledOnce);
     });

diff --git a/test/index.ts b/test/index.ts
@@ -292,7 +292,7 @@ describe('Spanner', () => {
                 .ReplicaSelection.Type.READ_ONLY,
             },
           ],
-          autoFailover: true,
+          autoFailoverDisabled: true,
         },
       };
 

diff --git a/test/transaction.ts b/test/transaction.ts
@@ -77,7 +77,7 @@ describe('Transaction', () => {
             .Type.READ_ONLY,
         },
       ],
-      autoFailover: true,
+      autoFailoverDisabled: true,
     },
   };