feat(Spanner): Enable Leader Aware Routing (LAR) (#6981)

Changes required for LAR in PHP Spanner Library.

Spanner/src/Connection/Grpc.php

@@ -45,6 +45,7 @@
 use Google\Cloud\Spanner\Admin\Instance\V1\UpdateInstanceMetadata;
 use Google\Cloud\Spanner\Operation;
 use Google\Cloud\Spanner\SpannerClient as ManualSpannerClient;
+use Google\Cloud\Spanner\RequestHeaderTrait;
 use Google\Cloud\Spanner\V1\CreateSessionRequest;
 use Google\Cloud\Spanner\V1\DeleteSessionRequest;
 use Google\Cloud\Spanner\V1\DirectedReadOptions;
@@ -83,6 +84,7 @@ class Grpc implements ConnectionInterface
     use EmulatorTrait;
     use GrpcTrait;
     use OperationResponseTrait;
+    use RequestHeaderTrait;
 
     /**
      * @var InstanceAdminClient|null
@@ -209,6 +211,11 @@ class Grpc implements ConnectionInterface
      */
     private $credentialsWrapper;
 
+    /**
+     * @var bool
+     */
+    private $larEnabled;
+
     /**
      * @param array $config [optional]
      */
@@ -270,6 +277,7 @@ public function __construct(array $config = [])
         //@codeCoverageIgnoreEnd
 
         $this->grpcConfig = $grpcConfig;
+        $this->larEnabled = $this->pluck('routeToLeader', $config, false) ?? true;
     }
 
     /**
@@ -806,6 +814,7 @@ function ($value) {
             );
         }
 
+        $args = $this->addLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'createSession'], [
             $databaseName,
             $this->addResourcePrefixHeader($args, $databaseName)
@@ -825,6 +834,7 @@ public function createSessionAsync(array $args)
     {
         $databaseName = $this->pluck('database', $args);
         $opts = $this->addResourcePrefixHeader([], $databaseName);
+        $opts = $this->addLarHeader($opts, $this->larEnabled);
         $opts['credentialsWrapper'] = $this->credentialsWrapper;
         $transport = $this->spannerClient->getTransport();
 
@@ -860,6 +870,7 @@ public function batchCreateSessions(array $args)
         );
 
         $databaseName = $this->pluck('database', $args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'batchCreateSessions'], [
             $databaseName,
             $this->pluck('sessionCount', $args),
@@ -873,6 +884,7 @@ public function batchCreateSessions(array $args)
     public function getSession(array $args)
     {
         $databaseName = $this->pluck('database', $args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'getSession'], [
             $this->pluck('name', $args),
             $this->addResourcePrefixHeader($args, $databaseName)
@@ -959,6 +971,7 @@ public function executeStreamingSql(array $args)
             );
         }
 
+        $args = $this->conditionallyUnsetLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'executeStreamingSql'], [
             $this->pluck('session', $args),
             $this->pluck('sql', $args),
@@ -986,6 +999,7 @@ public function streamingRead(array $args)
         $args['transaction'] = $this->createTransactionSelector($args);
 
         $databaseName = $this->pluck('database', $args);
+        $args = $this->conditionallyUnsetLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'streamingRead'], [
             $this->pluck('session', $args),
             $this->pluck('table', $args),
@@ -1002,6 +1016,7 @@ public function executeBatchDml(array $args)
     {
         $databaseName = $this->pluck('database', $args);
         $args['transaction'] = $this->createTransactionSelector($args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
 
         $statements = [];
         foreach ($this->pluck('statements', $args) as $statement) {
@@ -1045,9 +1060,11 @@ public function beginTransaction(array $args)
         } elseif (isset($transactionOptions['readWrite'])) {
             $readWrite = new ReadWrite();
             $options->setReadWrite($readWrite);
+            $args = $this->addLarHeader($args, $this->larEnabled);
         } elseif (isset($transactionOptions['partitionedDml'])) {
             $pdml = new PartitionedDml();
             $options->setPartitionedDml($pdml);
+            $args = $this->addLarHeader($args, $this->larEnabled);
         }
 
         $requestOptions = $this->pluck('requestOptions', $args, false) ?: [];
@@ -1135,6 +1152,7 @@ public function commit(array $args)
         }
 
         $databaseName = $this->pluck('database', $args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'commit'], [
             $this->pluck('session', $args),
             $mutations,
@@ -1148,6 +1166,7 @@ public function commit(array $args)
     public function rollback(array $args)
     {
         $databaseName = $this->pluck('database', $args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
         return $this->send([$this->spannerClient, 'rollback'], [
             $this->pluck('session', $args),
             $this->pluck('transactionId', $args),
@@ -1162,6 +1181,7 @@ public function partitionQuery(array $args)
     {
         $args = $this->formatSqlParams($args);
         $args['transaction'] = $this->createTransactionSelector($args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
 
         $args['partitionOptions'] = $this->serializer->decodeMessage(
             new PartitionOptions,
@@ -1185,6 +1205,7 @@ public function partitionRead(array $args)
         $keySet = $this->serializer->decodeMessage(new KeySet, $this->formatKeySet($keySet));
 
         $args['transaction'] = $this->createTransactionSelector($args);
+        $args = $this->addLarHeader($args, $this->larEnabled);
 
         $args['partitionOptions'] = $this->serializer->decodeMessage(
             new PartitionOptions,
@@ -1511,22 +1532,6 @@ private function formatTransactionOptions(array $transactionOptions)
         return $transactionOptions;
     }
 
-    /**
-     * Add the `google-cloud-resource-prefix` header value to the request.
-     *
-     * @param array $args
-     * @param string $value
-     * @return array
-     */
-    private function addResourcePrefixHeader(array $args, $value)
-    {
-        $args['headers'] = [
-            'google-cloud-resource-prefix' => [$value]
-        ];
-
-        return $args;
-    }
-
     /**
      * Allow lazy instantiation of the instance admin client.
      *

Spanner/src/Database.php

@@ -101,6 +101,7 @@ class Database
 {
     use LROTrait;
     use TransactionConfigurationTrait;
+    use RequestHeaderTrait;
 
     const STATE_CREATING = State::CREATING;
     const STATE_READY = State::READY;
@@ -1670,6 +1671,7 @@ public function execute($sql, array $options = [])
             $options['transaction'],
             $options['transactionContext']
         ) = $this->transactionSelector($options);
+        $options = $this->addLarHeader($options, true, $options['transactionContext']);
 
         $options['directedReadOptions'] = $this->configureDirectedReadOptions(
             $options,
@@ -1811,6 +1813,8 @@ public function executePartitionedUpdate($statement, array $options = [])
             ]
         ]);
 
+        $options = $this->addLarHeader($options);
+
         try {
             return $this->operation->executeUpdate($session, $transaction, $statement, [
                 'statsItem' => 'rowCountLowerBound'
@@ -1954,6 +1958,8 @@ public function read($table, KeySet $keySet, array $columns, array $options = []
             $this->directedReadOptions ?? []
         );
 
+        $options = $this->addLarHeader($options, true, $context);
+
         try {
             return $this->operation->read($session, $table, $keySet, $columns, $options);
         } finally {

Spanner/src/RequestHeaderTrait.php

@@ -0,0 +1,84 @@
+<?php
+/**
+ * Copyright 2024 Google Inc.
+ *
+ * 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 Google\Cloud\Spanner;
+
+use Google\Cloud\Spanner\Session\SessionPoolInterface;
+
+/**
+ * Shared functionality for request headers.
+ *
+ * @internal
+ */
+trait RequestHeaderTrait
+{
+    private static $larHeader = 'x-goog-spanner-route-to-leader';
+    private static $resourcePrefixHeader = 'google-cloud-resource-prefix';
+
+    /**
+     * Add the `x-goog-spanner-route-to-leader` header value to the request.
+     *
+     * @param array $args Request arguments.
+     * @param bool $value LAR header value.
+     * @param string $context Transaction context.
+     * @return array
+     */
+    private function addLarHeader(
+        array $args,
+        bool $value = true,
+        string $context = SessionPoolInterface::CONTEXT_READWRITE
+    ) {
+        if (!$value) {
+            return $args;
+        }
+        // If value is true and context is READWRITE, set LAR header.
+        if ($context === SessionPoolInterface::CONTEXT_READWRITE) {
+            $args['headers'][self::$larHeader] = ['true'];
+        }
+        return $args;
+    }
+
+    /**
+     * Conditionally unset the LAR header.
+     *
+     * @param array $args Request arguments.
+     * @param bool $value Whether to set or unset the LAR header.
+     * @return array
+     */
+    private function conditionallyUnsetLarHeader(
+        array $args,
+        bool $value = true
+    ) {
+        if (!$value) {
+            unset($args['headers'][self::$larHeader]);
+        }
+        return $args;
+    }
+
+    /**
+     * Add the `google-cloud-resource-prefix` header value to the request.
+     *
+     * @param array $args Request arguments.
+     * @param string $value Resource prefix header value.
+     * @return array
+     */
+    private function addResourcePrefixHeader(array $args, string $value)
+    {
+        $args['headers'][self::$resourcePrefixHeader] = [$value];
+        return $args;
+    }
+}

Spanner/src/SpannerClient.php

@@ -93,6 +93,13 @@
  * $spanner = new SpannerClient($directedOptions);
  * ```
  *
+ * ```
+ * use Google\Cloud\Spanner\SpannerClient;
+ *
+ * $config = ['routeToLeader' => false];
+ * $spanner = new SpannerClient($config);
+ * ```
+ *
  * @method resumeOperation() {
  *     Resume a Long Running Operation
  *
@@ -193,6 +200,8 @@ class SpannerClient
      *           {@see \Google\Cloud\Spanner\V1\DirectedReadOptions}
      *           If using the `replicaSelection::type` setting, utilize the constants available in
      *           {@see \Google\Cloud\Spanner\V1\DirectedReadOptions\ReplicaSelection\Type} to set a value.
+     *     @type bool $routeToLeader Enable/disable Leader Aware Routing.
+     *           **Defaults to** `true` (enabled).
      * }
      * @throws GoogleException If the gRPC extension is not enabled.
      */

Spanner/src/Transaction.php

@@ -742,6 +742,6 @@ private function buildUpdateOptions(array $options): array
         $selector = $this->transactionSelector($options);
         $options['transaction'] = $selector[0];
 
-        return $options;
+        return $this->addLarHeader($options);
     }
 }

Spanner/src/TransactionalReadTrait.php

@@ -26,6 +26,7 @@
 trait TransactionalReadTrait
 {
     use TransactionConfigurationTrait;
+    use RequestHeaderTrait;
 
     /**
      * @var Operation
@@ -302,6 +303,8 @@ public function execute($sql, array $options = [])
             $this->directedReadOptions ?? []
         );
 
+        $options = $this->addLarHeader($options, true, $this->context);
+
         $result = $this->operation->execute($this->session, $sql, $options);
         if (empty($this->id()) && $result->transaction()) {
             $this->setId($result->transaction()->id());
@@ -380,6 +383,8 @@ public function read($table, KeySet $keySet, array $columns, array $options = []
             $this->directedReadOptions ?? []
         );
 
+        $options = $this->addLarHeader($options, true, $this->context);
+
         $result = $this->operation->read($this->session, $table, $keySet, $columns, $options);
         if (empty($this->id()) && $result->transaction()) {
             $this->setId($result->transaction()->id());

Spanner/tests/Snippet/SpannerClientTest.php

@@ -305,4 +305,11 @@ public function testClassWithDirectedRead()
         $res = $snippet->invoke('spanner');
         $this->assertInstanceOf(SpannerClient::class, $res->returnVal());
     }
+
+    public function testClassWithLar()
+    {
+        $snippet = $this->snippetFromClass(SpannerClient::class, 3);
+        $res = $snippet->invoke('spanner');
+        $this->assertInstanceOf(SpannerClient::class, $res->returnVal());
+    }
 }