feat: Add snapshot methods to query Builder/Model (#267)

Author: matthewjumpsoffbuildings

README.md

@@ -183,6 +183,32 @@ $queryBuilder
 
 Stale reads always runs as read-only transaction with `singleUse` option. So you can not run as read-write transaction.
 
+### Snapshot reads
+
+You can use explicit Snapshot reads, either on `Connection`, or on `Model` or `Builder` instances. When running `snapshot()` on `Connection`, you pass a `Closure` that you can use to run multiple reads from within the same Snapshot.
+
+```php
+$timestampBound = new ExactStaleness(10);
+
+// by Connection
+$connection->snapshot($timestampBound, function() use ($connection) {
+    $result1 = $connection->table('foo')->get();
+    $result2 = $connection->table('bar')->get();
+
+    return [$result1, $result2];
+);
+
+// by Model
+User::where('foo', 'bar')
+    ->snapshot($timestampBound)
+    ->get();
+
+// by Query Builder
+$queryBuilder
+    ->snapshot($timestampBound)
+    ->get();
+```
+
 ### Data Boost
 
 Data boost creates snapshot and runs the query in parallel without affecting existing workloads.

src/Connection.php

@@ -25,6 +25,7 @@
 use Colopl\Spanner\Query\Processor as QueryProcessor;
 use Colopl\Spanner\Schema\Builder as SchemaBuilder;
 use Colopl\Spanner\Schema\Grammar as SchemaGrammar;
+use Colopl\Spanner\TimestampBound\TimestampBoundInterface;
 use DateTimeImmutable;
 use DateTimeInterface;
 use Exception;
@@ -625,6 +626,12 @@ protected function executeQuery(string $query, array $bindings, array $options):
             $options['requestOptions']['requestTag'] = $tag;
         }
 
+        if (isset($options['snapshotEnabled'])) {
+            $timestamp = $options['snapshotTimestampBound'];
+            assert($timestamp instanceof TimestampBoundInterface);
+            return $this->snapshot($timestamp, fn() => $this->executeSnapshotQuery($query, $options));
+        }
+
         if ($this->inSnapshot()) {
             return $this->executeSnapshotQuery($query, $options);
         }

src/Query/Builder.php

@@ -32,6 +32,7 @@ class Builder extends BaseBuilder
     use Concerns\UsesFullTextSearch;
     use Concerns\UsesMutations;
     use Concerns\UsesPartitionedDml;
+    use Concerns\UsesSnapshots;
     use Concerns\UsesStaleReads;
 
     public const PARAMETER_LIMIT = 950;
@@ -242,6 +243,11 @@ protected function runSelect()
             $options['dataBoostEnabled'] = true;
         }
 
+        if ($this->snapshotEnabled()) {
+            $options['snapshotEnabled'] = $this->snapshotEnabled();
+            $options['snapshotTimestampBound'] = $this->snapshotTimestampBound();
+        }
+
         if ($this->timestampBound !== null) {
             $options += $this->timestampBound->transactionOptions();
         }

src/Query/Concerns/UsesSnapshots.php

@@ -0,0 +1,54 @@
+<?php
+/**
+ * Copyright 2019 Colopl Inc. All Rights Reserved.
+ *
+ * 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.
+ */
+
+namespace Colopl\Spanner\Query\Concerns;
+
+use Colopl\Spanner\TimestampBound\TimestampBoundInterface;
+
+trait UsesSnapshots
+{
+    /**
+     * @var TimestampBoundInterface
+     */
+    protected ?TimestampBoundInterface $snapshotTimestampBound;
+
+    /**
+     * @param TimestampBoundInterface $timestamp
+     * @return $this
+     */
+    public function snapshot(TimestampBoundInterface $timestamp): static
+    {
+        $this->snapshotTimestampBound = $timestamp;
+        return $this;
+    }
+
+    /**
+     * @return bool
+     */
+    public function snapshotEnabled(): bool
+    {
+        return isset($this->snapshotTimestampBound);
+    }
+
+    /**
+     * @return TimestampBoundInterface|null
+     */
+    public function snapshotTimestampBound(): TimestampBoundInterface|null
+    {
+        return $this->snapshotTimestampBound;
+    }
+}

tests/Query/BuilderTest.php

@@ -24,6 +24,7 @@
 use Colopl\Spanner\Schema\TokenizerFunction;
 use Colopl\Spanner\Tests\TestCase;
 use Colopl\Spanner\TimestampBound\ExactStaleness;
+use Colopl\Spanner\TimestampBound\StrongRead;
 use Google\Cloud\Spanner\Bytes;
 use Google\Cloud\Spanner\Duration;
 use Illuminate\Database\QueryException;
@@ -1074,6 +1075,34 @@ public function test_dataBoost_disabled(): void
         $this->assertFalse($query->dataBoostEnabled());
     }
 
+    public function test_snapshot(): void
+    {
+        $conn = $this->getDefaultConnection();
+
+        $conn->transaction(function () use ($conn) {
+            $this->assertFalse($conn->inSnapshot());
+            $conn->table(self::TABLE_NAME_USER)->insert(['userId' => $this->generateUuid(), 'name' => 't']);
+        });
+
+        $this->assertFalse($conn->inSnapshot());
+        $query = $conn->table(self::TABLE_NAME_USER)->snapshot(new ExactStaleness(5));
+        $result = $query->first();
+
+        $this->assertTrue($query->snapshotEnabled());
+        $this->assertInstanceOf(ExactStaleness::class, $query->snapshotTimestampBound());
+        $this->assertFalse($conn->inSnapshot());
+        $this->assertNull($result);
+
+        $query = $conn->table(self::TABLE_NAME_USER)->snapshot(new StrongRead());
+        $result = $query->first();
+
+        $this->assertTrue($query->snapshotEnabled());
+        $this->assertInstanceOf(StrongRead::class, $query->snapshotTimestampBound());
+        $this->assertFalse($conn->inSnapshot());
+        $this->assertNotNull($result);
+        $this->assertSame('t', $result['name']);
+    }
+
     public function test_setRequestTimeoutSeconds(): void
     {
         $query = $this->getDefaultConnection()->table(self::TABLE_NAME_USER);