fix: read time should be used for transaction reads (#1171)

danieljbruce · web-flow · commit 73a0a39b4c04 · 2024-04-04T15:16:34.000-04:00
* Allow datastore projectId to be fetched from clien

Latency is caused by the call to getProjectId from Google auth. This change allows the project id to be retrieved if it is set in the client at creation time thereby reducing call latency.

* Create a file for mocking out commits

A test file is created where we mock out commit in the Gapic layer. The mock allows us to get the results passed to the commit endpoint in the Gapic layer.

* Create a test to measure latency of call.

To prove that the change works to reduce latency, a test is written. The test checks to see that the amount of time that passes between the time when the initial call is made in the user’s code and the time when the call reaches the gapic layer is sufficiently small. It will be a very small amount of time if the program does not need to do an auth lookup.

* Run the linting fixes

Run the linter so that spacing in the PR gets fixed for some of the lines of code.

* Add license header to top of test file

The license header needs to be added to the top of the new test file that is used for mocking out commit.

* Add comment for test for now

This is going to be a test for investigating the latency of the client.

* Add a test for the mock server

Measure the latency between the original call and the mock server.

* Set current retry attempt to 0

* Add a log for call time

Do check external to function after async call. Add log for call time.

* Add another mock

Other mock doesn’t require lazy client initialization.

* Eliminate code from the mock file

Eliminate the fake datastore client because we need to do assertion checks that are specific to each test. This means there is no point in defining runQuery once in a mock because each test will mock it out differently.

* Start off by adding read time to read options

Add the code change that will add read time to read options for transactions.

# Conflicts:
#	test/transaction.ts

* Update the test to use transactions

The idea is to test that read time got passed along for transactions specifically. This will be necessary for snapshot reads to work.

* Remove only

Need the entire test suite to run

* Remove the before hook

The before hook is not necessary. Just mock out the data client at the start.

* Remove unnecessary cherry picked files

Files were cherry-picked that weren’t helpful for solving the problem. Remove them.

* Clean up PR diff

* clean up PR diff

* Update the test so that it is run as a transaction

Right now, providing a transaction id is necessary to run the request as a transaction.

* Add an integration test

The integration test looks at the data from the snapshot read time for transactions and ensures that the read has no data thereby exercising the read time parameter.

* Linting fixing indents

Fix the indents in the system test folder

* Update the header

* Fix unit test

beginTransaction needs to be mocked out now that a transaction will begin if runQuery is called.

* System test changes.

Add a sleep. Instead of changing the current test, add a new test because it means the reader of the PR can be sure that test coverage wasn’t reduced which is better.

* Modify test

Modify the test so that sleeps are long enough to create predictable results and tests actually check for the right values.

* Replace with less precise assert

The test setup sometimes prepares before data with 0 entries and sometimes prepares before data with 1 entry so a less restrictive test is required in order for it to consistently pass.
diff --git a/src/request.ts b/src/request.ts
@@ -1025,9 +1025,13 @@ class DatastoreRequest {
         );
       }
 
-      reqOpts.readOptions = {
-        transaction: this.id,
-      };
+      if (reqOpts.readOptions) {
+        Object.assign(reqOpts.readOptions, {transaction: this.id});
+      } else {
+        reqOpts.readOptions = {
+          transaction: this.id,
+        };
+      }
     }
 
     datastore.auth.getProjectId((err, projectId) => {
diff --git a/system-test/datastore.ts b/system-test/datastore.ts
@@ -1931,6 +1931,29 @@ async.each(
         });
       });
       describe('transactions', () => {
+        before(async () => {
+          // This 'sleep' function is used to ensure that when data is saved to datastore,
+          // the time on the server is far enough ahead to be sure to be later than timeBeforeDataCreation
+          // so that when we read at timeBeforeDataCreation we get a snapshot of data before the save.
+          const key = datastore.key(['Company', 'Google']);
+          function sleep(ms: number) {
+            return new Promise(resolve => setTimeout(resolve, ms));
+          }
+          // Save for a key so that a read time can be accessed for snapshot reads.
+          const emptyData = {
+            key,
+            data: {},
+          };
+          await datastore.save(emptyData);
+          // Sleep for 10 seconds to ensure timeBeforeDataCreation includes the empty data
+          await sleep(10000);
+          timeBeforeDataCreation = await getReadTime([
+            {kind: 'Company', name: 'Google'},
+          ]);
+          // Sleep for 10 seconds so that any future reads will be later than timeBeforeDataCreation.
+          await sleep(10000);
+        });
+
         it('should run in a transaction', async () => {
           const key = datastore.key(['Company', 'Google']);
           const obj = {
@@ -2031,6 +2054,25 @@ async.each(
           await transaction.commit();
         });
 
+        it('should query within a transaction at a previous read time', async () => {
+          const transaction = datastore.transaction();
+          await transaction.run();
+          const query = transaction.createQuery('Company');
+          let entitiesBefore;
+          let entitiesNow;
+          try {
+            [entitiesBefore] = await query.run({
+              readTime: timeBeforeDataCreation,
+            });
+            [entitiesNow] = await query.run({});
+          } catch (e) {
+            await transaction.rollback();
+            return;
+          }
+          assert(entitiesBefore!.length < entitiesNow!.length);
+          await transaction.commit();
+        });
+
         describe('aggregate query within a transaction', async () => {
           it('should run a query and return the results', async () => {
             // Add a test here to verify what the data is at this time.
diff --git a/test/gapic-mocks/runQuery.ts b/test/gapic-mocks/runQuery.ts
@@ -0,0 +1,118 @@
+// 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.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import * as assert from 'assert';
+import {describe} from 'mocha';
+import {DatastoreClient, Datastore} from '../../src';
+import * as protos from '../../protos/protos';
+import {Callback} from 'google-gax';
+
+describe('Run Query', () => {
+  const PROJECT_ID = 'project-id';
+  const NAMESPACE = 'namespace';
+  const clientName = 'DatastoreClient';
+  const options = {
+    projectId: PROJECT_ID,
+    namespace: NAMESPACE,
+  };
+  const datastore = new Datastore(options);
+  // By default, datastore.clients_ is an empty map.
+  // To mock out commit we need the map to contain the Gapic data client.
+  // Normally a call to the data client through the datastore object would initialize it.
+  // We don't want to make this call because it would make a grpc request.
+  // So we just add the data client to the map.
+  const gapic = Object.freeze({
+    v1: require('../../src/v1'),
+  });
+  datastore.clients_.set(clientName, new gapic.v1[clientName](options));
+
+  // This function is used for doing assertion checks.
+  // The idea is to check that the right request gets passed to the commit function in the Gapic layer.
+  function setRunQueryComparison(
+    compareFn: (request: protos.google.datastore.v1.IRunQueryRequest) => void
+  ) {
+    const dataClient = datastore.clients_.get(clientName);
+    if (dataClient) {
+      dataClient.runQuery = (
+        request: any,
+        options: any,
+        callback: (
+          err?: unknown,
+          res?: protos.google.datastore.v1.IRunQueryResponse
+        ) => void
+      ) => {
+        try {
+          compareFn(request);
+        } catch (e) {
+          callback(e);
+        }
+        callback(null, {
+          batch: {
+            moreResults:
+              protos.google.datastore.v1.QueryResultBatch.MoreResultsType
+                .NO_MORE_RESULTS,
+          },
+        });
+      };
+    }
+  }
+
+  it('should pass read time into runQuery for transactions', async () => {
+    // First mock out beginTransaction
+    const dataClient = datastore.clients_.get(clientName);
+    const testId = Buffer.from(Array.from(Array(100).keys()));
+    if (dataClient) {
+      dataClient.beginTransaction = (
+        request: protos.google.datastore.v1.IBeginTransactionRequest,
+        options: protos.google.datastore.v1.IBeginTransactionResponse,
+        callback: Callback<
+          protos.google.datastore.v1.IBeginTransactionResponse,
+          | protos.google.datastore.v1.IBeginTransactionRequest
+          | null
+          | undefined,
+          {} | null | undefined
+        >
+      ) => {
+        callback(null, {
+          transaction: testId,
+        });
+      };
+    }
+    setRunQueryComparison(
+      (request: protos.google.datastore.v1.IRunQueryRequest) => {
+        assert.deepStrictEqual(request, {
+          readOptions: {
+            transaction: testId,
+            readTime: {
+              seconds: 77,
+            },
+          },
+          partitionId: {
+            namespaceId: 'namespace',
+          },
+          query: {
+            distinctOn: [],
+            kind: [{name: 'Task'}],
+            order: [],
+            projection: [],
+          },
+          projectId: 'project-id',
+        });
+      }
+    );
+    const transaction = datastore.transaction();
+    const query = datastore.createQuery('Task');
+    await transaction.runQuery(query, {readTime: 77000});
+  });
+});
