From 0ecca862fd93fc668a3dd3f5c93c54724e43f055 Mon Sep 17 00:00:00 2001
From: danieljbruce <danieljbruce@users.noreply.github.com>
Date: Tue, 26 Jul 2022 09:52:01 -0400
Subject: [PATCH] feat: Snapshot reads (#963)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Initial implementation of read time parameter

* linting fix

* system test works

* Almost there

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Finished unit test for snapshot reads

* PR follow-up

* PR changes

* system test fix

* system test fix again

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
---
 src/query.ts             |  1 +
 src/request.ts           | 17 +++++++++++++++-
 system-test/datastore.ts | 36 ++++++++++++++++++++++++++++++++++
 test/request.ts          | 42 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/src/query.ts b/src/query.ts
index c062a4c64..d4891f58f 100644
--- a/src/query.ts
+++ b/src/query.ts
@@ -561,6 +561,7 @@ export interface IntegerTypeCastOptions {
 
 export interface RunQueryOptions {
   consistency?: 'strong' | 'eventual';
+  readTime?: number;
   gaxOptions?: CallOptions;
   wrapNumbers?: boolean | IntegerTypeCastOptions;
 }
diff --git a/src/request.ts b/src/request.ts
index 261abb95b..e93dd1154 100644
--- a/src/request.ts
+++ b/src/request.ts
@@ -54,6 +54,7 @@ import {
   RunQueryCallback,
 } from './query';
 import {Datastore} from '.';
+import ITimestamp = google.protobuf.ITimestamp;
 
 /**
  * A map of read consistency values to proto codes.
@@ -284,6 +285,16 @@ class DatastoreRequest {
           readConsistency: code,
         };
       }
+      if (options.readTime) {
+        if (reqOpts.readOptions === undefined) {
+          reqOpts.readOptions = {};
+        }
+        const readTime = options.readTime;
+        const seconds = readTime / 1000;
+        reqOpts.readOptions.readTime = {
+          seconds: Math.floor(seconds),
+        };
+      }
 
       this.request_(
         {
@@ -1058,7 +1069,11 @@ export interface RequestConfig {
 export interface RequestOptions {
   mutations?: google.datastore.v1.IMutation[];
   keys?: Entity;
-  readOptions?: {readConsistency?: number; transaction?: string};
+  readOptions?: {
+    readConsistency?: number;
+    transaction?: string;
+    readTime?: ITimestamp;
+  };
   partitionId?: google.datastore.v1.IPartitionId | null;
   transactionOptions?: {
     readOnly?: {};
diff --git a/system-test/datastore.ts b/system-test/datastore.ts
index 3615ce206..1a778c02d 100644
--- a/system-test/datastore.ts
+++ b/system-test/datastore.ts
@@ -308,6 +308,42 @@ describe('Datastore', () => {
       await datastore.delete(postKey);
     });
 
+    it('should save/get/delete from a snapshot', async () => {
+      function sleep(ms: number) {
+        return new Promise(resolve => setTimeout(resolve, ms));
+      }
+      const post2 = {
+        title: 'Another way to make pizza',
+        tags: ['pizza', 'grill'],
+        publishedAt: new Date(),
+        author: 'Silvano',
+        isDraft: false,
+        wordCount: 400,
+        rating: 5.0,
+        likes: null,
+        metadata: {
+          views: 100,
+        },
+      };
+      const path = ['Post', 'post1'];
+      const postKey = datastore.key(path);
+      await datastore.save({key: postKey, data: post});
+      await sleep(1000);
+      const savedTime = Date.now();
+      await sleep(1000);
+      // Save new post2 data, but then verify the timestamp read has post1 data
+      await datastore.save({key: postKey, data: post2});
+      const [entity] = await datastore.get(postKey, {readTime: savedTime});
+      assert.deepStrictEqual(entity[datastore.KEY], postKey);
+      const [entityNoOptions] = await datastore.get(postKey);
+      assert.deepStrictEqual(entityNoOptions[datastore.KEY], postKey);
+      delete entity[datastore.KEY];
+      assert.deepStrictEqual(entity, post);
+      delete entityNoOptions[datastore.KEY];
+      assert.deepStrictEqual(entityNoOptions, post2);
+      await datastore.delete(postKey);
+    });
+
     it('should save/get/delete with a numeric key id', async () => {
       const postKey = datastore.key(['Post', 123456789]);
       await datastore.save({key: postKey, data: post});
diff --git a/test/request.ts b/test/request.ts
index 158929001..532966420 100644
--- a/test/request.ts
+++ b/test/request.ts
@@ -34,6 +34,7 @@ import {
   PrepareEntityObjectResponse,
   CommitResponse,
   GetResponse,
+  RequestCallback,
 } from '../src/request';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -672,6 +673,47 @@ describe('Request', () => {
   });
 
   describe('get', () => {
+    it('should pass along readTime for reading snapshots', done => {
+      const savedTime = Date.now();
+      request.request_ = (config: RequestConfig, callback: RequestCallback) => {
+        assert.deepStrictEqual(config, {
+          client: 'DatastoreClient',
+          method: 'lookup',
+          gaxOpts: undefined,
+          reqOpts: {
+            keys: [
+              {
+                path: [
+                  {
+                    kind: 'Company',
+                    id: 123,
+                  },
+                ],
+                partitionId: {namespaceId: 'namespace'},
+              },
+            ],
+            readOptions: {
+              readTime: {
+                seconds: Math.floor(savedTime / 1000),
+              },
+            },
+          },
+        });
+        callback(null, {
+          deferred: [],
+          found: [],
+          missing: [],
+          readTime: {seconds: Math.floor(savedTime / 1000), nanos: 0},
+        });
+      };
+      request.get(key, {readTime: savedTime}, (err: any) => {
+        if (err) {
+          throw err;
+        }
+        done();
+      });
+    });
+
     describe('success', () => {
       const keys = [key];
       const fakeEntities = [{a: 'a'}, {b: 'b'}];