feat: Improved latency for executeStreaminSql calls (#7254)

Spanner/composer.json

@@ -7,7 +7,7 @@
         "php": "^8.0",
         "ext-grpc": "*",
         "google/cloud-core": "^1.52.7",
-        "google/gax": "^1.30"
+        "google/gax": "^1.32"
     },
     "require-dev": {
         "phpunit/phpunit": "^9.0",

Spanner/src/Connection/Grpc.php

@@ -55,6 +55,7 @@
 use Google\Cloud\Spanner\V1\Mutation;
 use Google\Cloud\Spanner\V1\Mutation\Delete;
 use Google\Cloud\Spanner\V1\Mutation\Write;
+use Google\Cloud\Spanner\V1\PartialResultSet;
 use Google\Cloud\Spanner\V1\PartitionOptions;
 use Google\Cloud\Spanner\V1\RequestOptions;
 use Google\Cloud\Spanner\V1\Session;
@@ -68,6 +69,7 @@
 use Google\Protobuf;
 use Google\Protobuf\FieldMask;
 use Google\Protobuf\GPBEmpty;
+use Google\Protobuf\Internal\RepeatedField;
 use Google\Protobuf\ListValue;
 use Google\Protobuf\Struct;
 use Google\Protobuf\Value;
@@ -235,6 +237,24 @@ public function __construct(array $config = [])
             'google.protobuf.Timestamp' => function ($v) {
                 return $this->formatTimestampFromApi($v);
             }
+        ], [], [], [
+            // A custom encoder that short-circuits the encodeMessage in Serializer class,
+            // but only if the argument is of the type PartialResultSet.
+            PartialResultSet::class => function ($msg) {
+                $data = json_decode($msg->serializeToJsonString(), true);
+
+                // The transaction id is serialized as a base64 encoded string in $data. So, we
+                // add a step to get the transaction id using a getter instead of the serialized value.
+                // The null-safe operator is used to handle edge cases where the relevant fields are not present.
+                $data['metadata']['transaction']['id'] = (string) $msg?->getMetadata()?->getTransaction()?->getId();
+
+                // Helps convert metadata enum values from string types to their respective code/annotation
+                // pairs. Ex: INT64 is converted to {code: 2, typeAnnotation: 0}.
+                $fields = $msg->getMetadata()?->getRowType()?->getFields();
+                $data['metadata']['rowType']['fields'] = $this->getFieldDataFromRepeatedFields($fields);
+
+                return $data;
+            }
         ]);
         //@codeCoverageIgnoreEnd
 
@@ -1620,4 +1640,71 @@ private function setDirectedReadOptions(array &$args)
             );
         }
     }
+
+    /**
+     * Utiltiy method to take in a Google\Cloud\Spanner\V1\Type value and return
+     * the data as an array. The method takes care of array and struct elements.
+     *
+     * @param Type $type The "type" object
+     *
+     * @return array The formatted data.
+     */
+    private function getTypeData(Type $type): array
+    {
+        $data = [
+            'code' => $type->getCode(),
+            'typeAnnotation' => $type->getTypeAnnotation(),
+            'protoTypeFqn' => $type->getProtoTypeFqn()
+        ];
+
+        // If this is a struct field, then recursisevly call getTypeData
+        if ($type->hasStructType()) {
+            $nestedType = $type->getStructType();
+            $fields = $nestedType->getFields();
+            $data['structType'] = [
+                'fields' => $this->getFieldDataFromRepeatedFields($fields)
+            ];
+        }
+        // If this is an array field, then recursisevly call getTypeData
+        if ($type->hasArrayElementType()) {
+            $nestedType = $type->getArrayElementType();
+            $data['arrayElementType'] = $this->getTypeData($nestedType);
+        }
+
+        return $data;
+    }
+
+    /**
+     * Utility method to return "fields data" in the format:
+     * [
+     *   "name" => ""
+     *   "type" => []
+     * ].
+     *
+     * The type is converted from a string like INT64 to ["code" => 2, "typeAnnotation" => 0]
+     * conforming with the Google\Cloud\Spanner\V1\TypeCode class.
+     *
+     * @param ?RepeatedField $fields The array contain list of fields.
+     *
+     * @return array The formatted fields data.
+     */
+    private function getFieldDataFromRepeatedFields(?RepeatedField $fields): array
+    {
+        if (is_null($fields)) {
+            return [];
+        }
+
+        $fieldsData = [];
+        foreach ($fields as $key => $field) {
+            $type = $field->getType();
+            $typeData = $this->getTypeData($type);
+
+            $fieldsData[$key] = [
+                'name' => $field->getName(),
+                'type' => $typeData
+            ];
+        }
+
+        return $fieldsData;
+    }
 }

Spanner/tests/Unit/Connection/GrpcTest.php

@@ -21,6 +21,7 @@
 use Google\ApiCore\CredentialsWrapper;
 use Google\ApiCore\OperationResponse;
 use Google\ApiCore\Serializer;
+use Google\ApiCore\Testing\MockResponse;
 use Google\ApiCore\Transport\TransportInterface;
 use Google\Cloud\Core\GrpcRequestWrapper;
 use Google\Cloud\Core\GrpcTrait;
@@ -40,6 +41,7 @@
 use Google\Cloud\Spanner\V1\Mutation;
 use Google\Cloud\Spanner\V1\Mutation\Delete;
 use Google\Cloud\Spanner\V1\Mutation\Write;
+use Google\Cloud\Spanner\V1\PartialResultSet;
 use Google\Cloud\Spanner\V1\PartitionOptions;
 use Google\Cloud\Spanner\V1\RequestOptions;
 use Google\Cloud\Spanner\V1\Session;
@@ -1418,6 +1420,37 @@ public function larOptions()
         ];
     }
 
+    public function testPartialResultSetCustomEncoder()
+    {
+        $partialResultSet = new PartialResultSet();
+        $partialResultSet->mergeFromJsonString(json_encode([
+            'metadata' => [
+                'transaction' => [
+                    'id' => base64_encode(0b00010100) // bytedata is represented as a base64-encoded string in JSON
+                ],
+                'rowType' => [
+                    'fields' => [
+                        ['type' => ['code' => 'INT64']] // enums are represented as their string equivalents in JSON
+                    ]
+                ],
+            ],
+        ]));
+
+        $this->assertEquals(0b00010100, $partialResultSet->getMetadata()->getTransaction()->getId());
+        $this->assertEquals(2, $partialResultSet->getMetadata()->getRowType()->getFields()[0]->getType()->getCode());
+
+        // decode the message and ensure it's decoded as expected
+        $grpc = new Grpc();
+        $serializerProp = new \ReflectionProperty($grpc, 'serializer');
+        $serializerProp->setAccessible(true);
+        $serializer = $serializerProp->getValue($grpc);
+        $arr = $serializer->encodeMessage($partialResultSet);
+
+        // We expect this to be the binary string
+        $this->assertEquals(0b00010100, $arr['metadata']['transaction']['id']);
+        // We expect this to be the integer
+        $this->assertEquals(2, $arr['metadata']['rowType']['fields'][0]['type']['code']);
+    }
     private function assertCallCorrect(
         $method,
         array $args,

composer.json

@@ -50,7 +50,7 @@
         "monolog/monolog": "^2.9||^3.0",
         "psr/http-message": "^1.0|^2.0",
         "ramsey/uuid": "^4.0",
-        "google/gax": "^1.30",
+        "google/gax": "^1.32",
         "google/common-protos": "^4.4",
         "google/auth": "^1.34"
     },