feat: Support snapshot queries (#215)

Co-authored-by: halnique <shunsuke4dev@gmail.com>
colopl · Sep 18, 2024 · 72c61a5 · matthewjumpsoffbuildings · Sep 21, 2024 · 72c61a5
1 parent 6ba4a72
commit 72c61a5
Show file tree

Hide file tree

Showing 7 changed files with 235 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,6 @@
 # v8.3.0 (2024-09-02)
 
+- add support for snapshot queries (#215)
 - deprecate Connection::getDatabaseContext() and move logic to UseMutations::getMutationExecutor() (#227)
 - add support for `Query\Builder::whereNotInUnnest(...)` (#225)
 - `Query\Builder::whereIn` will now wrap values in `UNNEST` if the number of values exceeds the limit (950). (#)

diff --git a/compose.yaml b/compose.yaml
@@ -8,4 +8,4 @@ services:
     depends_on:
       - emulator
   emulator:
-    image: "gcr.io/cloud-spanner-emulator/emulator:1.5.17"
+    image: "gcr.io/cloud-spanner-emulator/emulator:1.5.23"
diff --git a/src/Concerns/ManagesSnapshots.php b/src/Concerns/ManagesSnapshots.php
@@ -0,0 +1,60 @@
+<?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\Concerns;
+
+use Closure;
+use Colopl\Spanner\TimestampBound\TimestampBoundInterface;
+use Google\Cloud\Spanner\Snapshot;
+use LogicException;
+
+trait ManagesSnapshots
+{
+    /**
+     * @var Snapshot|null
+     */
+    protected ?Snapshot $currentSnapshot = null;
+
+    /**
+     * @template TReturn
+     * @param TimestampBoundInterface $timestampBound
+     * @param Closure(): TReturn $callback
+     * @return TReturn
+     */
+    public function snapshot(TimestampBoundInterface $timestampBound, Closure $callback): mixed
+    {
+        if ($this->currentSnapshot !== null) {
+            throw new LogicException('Nested snapshots are not supported.');
+        }
+
+        $options = $timestampBound->transactionOptions();
+        try {
+            $this->currentSnapshot = $this->getSpannerDatabase()->snapshot($options);
+            return $callback();
+        } finally {
+            $this->currentSnapshot = null;
+        }
+    }
+
+    /**
+     * @return bool
+     */
+    public function inSnapshot(): bool
+    {
+        return $this->currentSnapshot !== null;
+    }
+}
diff --git a/src/Concerns/ManagesTransactions.php b/src/Concerns/ManagesTransactions.php
@@ -22,6 +22,7 @@
 use Google\Cloud\Core\Exception\AbortedException;
 use Google\Cloud\Spanner\Database;
 use Google\Cloud\Spanner\Transaction;
+use LogicException;
 use Throwable;
 
 /**
@@ -69,6 +70,10 @@ public function transaction(Closure $callback, $attempts = -1)
         return $this->withSessionNotFoundHandling(function () use ($callback, $options) {
             $return = $this->getSpannerDatabase()->runTransaction(function (Transaction $tx) use ($callback) {
                 try {
+                    if ($this->inSnapshot()) {
+                        throw new LogicException('Calling transaction() inside a snapshot is not supported.');
+                    }
+
                     $this->currentTransaction = $tx;
 
                     $this->transactions++;

diff --git a/src/Connection.php b/src/Connection.php
@@ -40,6 +40,7 @@
 use Illuminate\Database\Connection as BaseConnection;
 use Illuminate\Database\Query\Grammars\Grammar as BaseQueryGrammar;
 use Illuminate\Database\QueryException;
+use Illuminate\Support\Arr;
 use InvalidArgumentException;
 use LogicException;
 use Psr\Cache\CacheItemPoolInterface;
@@ -52,6 +53,7 @@ class Connection extends BaseConnection
         Concerns\ManagesMutations,
         Concerns\ManagesPartitionedDml,
         Concerns\ManagesSessionPool,
+        Concerns\ManagesSnapshots,
         Concerns\ManagesTagging,
         Concerns\ManagesTransactions,
         Concerns\MarksAsNotSupported;
@@ -579,14 +581,11 @@ protected function executeQuery(string $query, array $bindings, array $options):
             $options['requestOptions']['requestTag'] = $tag;
         }
 
-        $forceReadOnlyTransaction =
-            ($options['exactStaleness'] ?? false) ||
-            ($options['maxStaleness'] ?? false) ||
-            ($options['minReadTimestamp'] ?? false) ||
-            ($options['readTimestamp'] ?? false) ||
-            ($options['strong'] ?? false);
+        if ($this->inSnapshot()) {
+            return $this->executeSnapshotQuery($query, $options);
+        }
 
-        if (!$forceReadOnlyTransaction && $transaction = $this->getCurrentTransaction()) {
+        if ($this->canExecuteAsReadWriteTransaction($options) && $transaction = $this->getCurrentTransaction()) {
             return $transaction->execute($query, $options)->rows();
         }
 
@@ -611,6 +610,18 @@ protected function executePartitionedQuery(string $query, array $options): Gener
         }
     }
 
+    /**
+     * @param string $query
+     * @param array<string, mixed> $options
+     * @return Generator<int, array<array-key, mixed>>
+     */
+    protected function executeSnapshotQuery(string $query, array $options): Generator
+    {
+        $executeOptions = Arr::only($options, ['parameters', 'types', 'queryOptions', 'requestOptions']);
+        assert($this->currentSnapshot !== null);
+        return $this->currentSnapshot->execute($query, $executeOptions)->rows();
+    }
+
     /**
      * @param Transaction $transaction
      * @param string $query
@@ -651,6 +662,29 @@ protected function executeBatchDml(Transaction $transaction, string $query, arra
         return $rowCount;
     }
 
+    /**
+     * @param array<string, mixed> $options
+     * @return bool
+     */
+    protected function canExecuteAsReadWriteTransaction(array $options): bool
+    {
+        $readOnlyTriggers = [
+            'singleUse',
+            'exactStaleness',
+            'maxStaleness',
+            'minReadTimestamp',
+            'readTimestamp',
+            'strong',
+        ];
+
+        foreach ($readOnlyTriggers as $option) {
+            if ($options[$option] ?? false) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     /**
      * @param string $query
      * @return bool

diff --git a/tests/Query/BuilderTest.php b/tests/Query/BuilderTest.php
@@ -22,7 +22,6 @@
 use Colopl\Spanner\Schema\Blueprint;
 use Colopl\Spanner\Tests\TestCase;
 use Colopl\Spanner\TimestampBound\ExactStaleness;
-use Google\Cloud\Core\Exception\BadRequestException;
 use Google\Cloud\Spanner\Bytes;
 use Google\Cloud\Spanner\Duration;
 use Illuminate\Database\QueryException;
@@ -1093,7 +1092,6 @@ public function test_whereIn_with_unnest_overflow_flag_turned_on(): void
         $this->assertSame([], $query->get()->all());
     }
 
-
     public function test_whereIn_with_unnest_overflow_flag_turned_off(): void
     {
         $this->expectExceptionMessage('Number of parameters in query exceeds the maximum allowed limit of 950.');

diff --git a/tests/SnapshotTest.php b/tests/SnapshotTest.php
@@ -0,0 +1,127 @@
+<?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\Tests;
+
+use Colopl\Spanner\TimestampBound\ExactStaleness;
+use Colopl\Spanner\TimestampBound\StrongRead;
+use LogicException;
+use RuntimeException;
+
+class SnapshotTest extends TestCase
+{
+    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());
+        $result = $conn->snapshot(new StrongRead(), function () use ($conn) {
+            $this->assertTrue($conn->inSnapshot());
+            // call it multiple times
+            $this->assertSame('t', $conn->table(self::TABLE_NAME_USER)->value('name'));
+            $this->assertSame('t', $conn->table(self::TABLE_NAME_USER)->value('name'));
+
+            return 'ok';
+        });
+
+        $this->assertSame('ok', $result);
+    }
+
+    public function test_snapshot_with_staleness(): void
+    {
+        $conn = $this->getDefaultConnection();
+
+        $conn->transaction(function () use ($conn) {
+            $conn->table(self::TABLE_NAME_USER)->insert(['userId' => $this->generateUuid(), 'name' => 't']);
+        });
+
+        $conn->snapshot(new ExactStaleness(10), function () use ($conn) {
+            $this->assertNull($conn->table(self::TABLE_NAME_USER)->first());
+            $this->assertSame(0, $conn->table(self::TABLE_NAME_USER)->count());
+        });
+
+        $conn->snapshot(new StrongRead(), function () use ($conn) {
+            $this->assertNotNull($conn->table(self::TABLE_NAME_USER)->first());
+            $this->assertSame(1, $conn->table(self::TABLE_NAME_USER)->count());
+        });
+    }
+
+    public function test_snapshot_can_call_after_error(): void
+    {
+        $conn = $this->getDefaultConnection();
+
+        try {
+            $conn->snapshot(new ExactStaleness(10), function () use ($conn) {
+                $this->assertSame(0, $conn->table(self::TABLE_NAME_USER)->count());
+                throw new RuntimeException('error');
+            });
+        } catch (RuntimeException $e) {
+            // ignore
+        }
+
+        $conn->transaction(function () use ($conn) {
+            $conn->table(self::TABLE_NAME_USER)->insert(['userId' => $this->generateUuid(), 'name' => 't']);
+        });
+
+        $conn->snapshot(new ExactStaleness(0), function () use ($conn) {
+            $this->assertSame(1, $conn->table(self::TABLE_NAME_USER)->count());
+        });
+    }
+
+    public function test_snapshot_fails_on_nested(): void
+    {
+        $this->expectException(LogicException::class);
+        $this->expectExceptionMessage('Nested snapshots are not supported.');
+
+        $conn = $this->getDefaultConnection();
+        $conn->snapshot(new ExactStaleness(10), function () use ($conn) {
+            $conn->snapshot(new StrongRead(), function () {
+            });
+        });
+    }
+
+    public function test_snapshot_fails_in_transaction(): void
+    {
+        $this->expectException(LogicException::class);
+        $this->expectExceptionMessage('Nested transactions are not supported by this client.');
+
+        $conn = $this->getDefaultConnection();
+        $conn->transaction(function () use ($conn) {
+            $conn->snapshot(new StrongRead(), function () use ($conn) {
+                $conn->select('SELECT 1');
+            });
+        });
+    }
+
+    public function test_snapshot_fails_when_transaction_called_inside(): void
+    {
+        $this->expectException(LogicException::class);
+        $this->expectExceptionMessage('Calling transaction() inside a snapshot is not supported.');
+
+        $conn = $this->getDefaultConnection();
+        $conn->snapshot(new StrongRead(), function () use ($conn) {
+            $conn->transaction(function () use ($conn) {
+                $conn->select('SELECT 1');
+            });
+        });
+    }
+}