From a406f1253c282ea1ff15102585bb58e6741aa4ca Mon Sep 17 00:00:00 2001
From: Takayasu Oyama <toyama4649@hotmail.com>
Date: Mon, 13 May 2024 16:41:13 +0900
Subject: [PATCH] feat: allow request timeout at query level (#208)

feature/set-request-timeout
---
 CHANGELOG.md                               |  1 +
 README.md                                  |  1 +
 src/Query/Builder.php                      |  9 ++++-
 src/Query/Concerns/SetsRequestTimeouts.php | 44 ++++++++++++++++++++++
 tests/Query/BuilderTest.php                | 13 +++++++
 5 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 src/Query/Concerns/SetsRequestTimeouts.php

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fc6a9df..3d79121 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 Added
 - Support for `INSERT OR IGNORE` (#207)
+- Support adding request timeout at query level (#208)
 
 # v7.3.0 (2024-04-17)
 
diff --git a/README.md b/README.md
index a7ea3be..ddfbb12 100644
--- a/README.md
+++ b/README.md
@@ -193,6 +193,7 @@ $connection->selectWithOptions('SELECT ...', $bindings, ['dataBoostEnabled' => t
 // Using Query Builder
 $queryBuilder
     ->useDataBoost()
+    ->setRequestTimeoutSeconds(60)
     ->get();
 ```
 
diff --git a/src/Query/Builder.php b/src/Query/Builder.php
index 7f13f86..7fb846f 100644
--- a/src/Query/Builder.php
+++ b/src/Query/Builder.php
@@ -17,7 +17,6 @@
 
 namespace Colopl\Spanner\Query;
 
-use Closure;
 use Colopl\Spanner\Connection;
 use Illuminate\Contracts\Support\Arrayable;
 use Illuminate\Database\Query\Builder as BaseBuilder;
@@ -26,6 +25,7 @@
 
 class Builder extends BaseBuilder
 {
+    use Concerns\SetsRequestTimeouts;
     use Concerns\UsesDataBoost;
     use Concerns\UsesMutations;
     use Concerns\UsesPartitionedDml;
@@ -166,8 +166,13 @@ protected function runSelect()
         $bindings = $this->getBindings();
         $options = [];
 
+        $requestTimeoutSeconds = $this->getRequestTimeoutSeconds();
+        if ($requestTimeoutSeconds !== null) {
+            $options['requestTimeout'] = $requestTimeoutSeconds;
+        }
+
         if ($this->dataBoostEnabled()) {
-            $options += ['dataBoostEnabled' => true];
+            $options['dataBoostEnabled'] = true;
         }
 
         if ($this->timestampBound !== null) {
diff --git a/src/Query/Concerns/SetsRequestTimeouts.php b/src/Query/Concerns/SetsRequestTimeouts.php
new file mode 100644
index 0000000..2590b57
--- /dev/null
+++ b/src/Query/Concerns/SetsRequestTimeouts.php
@@ -0,0 +1,44 @@
+<?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;
+
+trait SetsRequestTimeouts
+{
+    /**
+     * @var float|null
+     */
+    protected ?float $requestTimeoutSeconds = null;
+
+    /**
+     * @param float $seconds
+     * @return $this
+     */
+    public function setRequestTimeoutSeconds(float $seconds): static
+    {
+        $this->requestTimeoutSeconds = $seconds;
+        return $this;
+    }
+
+    /**
+     * @return float|null
+     */
+    public function getRequestTimeoutSeconds(): ?float
+    {
+        return $this->requestTimeoutSeconds;
+    }
+}
diff --git a/tests/Query/BuilderTest.php b/tests/Query/BuilderTest.php
index f57f066..99093c7 100644
--- a/tests/Query/BuilderTest.php
+++ b/tests/Query/BuilderTest.php
@@ -23,6 +23,7 @@
 use Colopl\Spanner\Tests\TestCase;
 use Colopl\Spanner\TimestampBound\ExactStaleness;
 use Google\Cloud\Core\Exception\ConflictException;
+use Google\Cloud\Core\Exception\DeadlineExceededException;
 use Google\Cloud\Spanner\Bytes;
 use Google\Cloud\Spanner\Duration;
 use Illuminate\Database\Events\QueryExecuted;
@@ -1073,4 +1074,16 @@ public function test_dataBoost_disabled(): void
         $query = $this->getDefaultConnection()->table('t')->useDataBoost(false);
         $this->assertFalse($query->dataBoostEnabled());
     }
+
+    public function test_setRequestTimeoutSeconds(): void
+    {
+        $query = $this->getDefaultConnection()->table(self::TABLE_NAME_USER);
+        $this->assertNull($query->getRequestTimeoutSeconds());
+        $query->setRequestTimeoutSeconds(0.0001);
+        $this->assertSame(0.0001, $query->getRequestTimeoutSeconds());
+
+        $this->expectException(QueryException::class);
+        $this->expectExceptionMessageMatches('/DEADLINE_EXCEEDED/');
+        $query->get();
+    }
 }