From 4c29341a671141efd3727bad5579abfedbc1ad0f Mon Sep 17 00:00:00 2001
From: yash30201 <54198301+yash30201@users.noreply.github.com>
Date: Tue, 7 May 2024 10:03:21 +0000
Subject: [PATCH] chore(Datastore): Enhance readability of RunQuery and
 RunAggregationQuery APIs

---
 Datastore/src/Operation.php | 115 +++++++++++++++++++-----------------
 1 file changed, 60 insertions(+), 55 deletions(-)

diff --git a/Datastore/src/Operation.php b/Datastore/src/Operation.php
index 8eba08fc5341..05282ff3802c 100644
--- a/Datastore/src/Operation.php
+++ b/Datastore/src/Operation.php
@@ -567,26 +567,9 @@ public function runQuery(QueryInterface $query, array $options = [])
         }
         $runQueryObj = clone $query;
         $runQueryFn = function (array $args = []) use (&$runQueryObj, $options, &$remainingLimit) {
-            $args += [
-                'query' => [],
-            ];
+            $parsedArgs = $this->parseRunQueryArguments($args, $runQueryObj, $options, $remainingLimit);
 
-            // The iterator provides the startCursor for subsequent pages as an argument.
-            $requestQueryArr = $args['query'] + $runQueryObj->queryObject();
-            if (isset($remainingLimit)) {
-                $requestQueryArr['limit'] = $remainingLimit;
-            }
-            $req = [
-                'projectId' => $this->projectId,
-                'partitionId' => $this->partitionId(
-                    $this->projectId,
-                    $options['namespaceId'],
-                    $options['databaseId']
-                ),
-                $runQueryObj->queryKey() => $requestQueryArr,
-            ] + $this->readOptions($options) + $options;
-
-            list($data, $optionalArgs) = $this->splitOptionalArgs($req);
+            list($data, $optionalArgs) = $this->splitOptionalArgs($parsedArgs);
 
             // Remove redundant keys for request.
             $this->pluckArray(
@@ -594,12 +577,6 @@ public function runQuery(QueryInterface $query, array $options = [])
                 $data
             );
 
-            if (isset($data['query'])) {
-                $data['query'] = $this->parseQuery($data['query']);
-            }
-            if (isset($data['gqlQuery'])) {
-                $data['gqlQuery'] = $this->parseGqlQuery($data['gqlQuery']);
-            }
             $request = $this->serializer->decodeMessage(new RunQueryRequest(), $data);
             $res = $this->requestHandler->sendRequest(
                 DatastoreClient::class,
@@ -668,20 +645,9 @@ public function runAggregationQuery(AggregationQuery $runQueryObj, array $option
             'databaseId' => $this->databaseId,
         ];
 
-        $args = [
-            'query' => [],
-        ];
-        $requestQueryArr = $args['query'] + $runQueryObj->queryObject();
-        $req = [
-            'projectId' => $this->projectId,
-            'partitionId' => $this->partitionId(
-                $this->projectId,
-                $options['namespaceId'],
-                $options['databaseId']
-            ),
-        ] + $requestQueryArr + $this->readOptions($options) + $options;
+        $parsedArgs = $this->parseRunAggregationQueryArgs($runQueryObj, $options);
 
-        list($data, $optionalArgs) = $this->splitOptionalArgs($req);
+        list($data, $optionalArgs) = $this->splitOptionalArgs($parsedArgs);
 
         // Remove redundant keys for request.
         $this->pluckArray(
@@ -689,23 +655,6 @@ public function runAggregationQuery(AggregationQuery $runQueryObj, array $option
             $data
         );
 
-        if (isset($data['aggregationQuery'])) {
-            if (isset($data['aggregationQuery']['nestedQuery'])) {
-                $data['aggregationQuery']['nestedQuery'] = $this->parseQuery(
-                    $data['aggregationQuery']['nestedQuery']
-                );
-            }
-
-            $data['aggregationQuery'] = $this->serializer->decodeMessage(
-                new V1AggregationQuery(),
-                $data['aggregationQuery']
-            );
-        }
-
-        if (isset($data['gqlQuery'])) {
-            $data['gqlQuery'] = $this->parseGqlQuery($data['gqlQuery']);
-        }
-
         $request = $this->serializer->decodeMessage(new RunAggregationQueryRequest(), $data);
         $res = $this->requestHandler->sendRequest(
             DatastoreClient::class,
@@ -1231,4 +1180,60 @@ private function toGrpcValue(array $property)
 
         return [$type, $val];
     }
+
+    private function parseRunQueryArguments($args, &$runQueryObj, $options, &$remainingLimit)
+    {
+        $args += ['query' => []];
+
+        // The iterator provides the startCursor for subsequent pages as an argument.
+        $requestQueryArr = $args['query'] + $runQueryObj->queryObject();
+        if (isset($remainingLimit)) {
+            $requestQueryArr['limit'] = $remainingLimit;
+        }
+        $result = [
+            'projectId' => $this->projectId,
+            'partitionId' => $this->partitionId(
+                $this->projectId,
+                $options['namespaceId'],
+                $options['databaseId']
+            ),
+            $runQueryObj->queryKey() => $requestQueryArr,
+        ] + $this->readOptions($options) + $options;
+
+        if (isset($result['query'])) {
+            $result['query'] = $this->parseQuery($result['query']);
+        }
+        if (isset($result['gqlQuery'])) {
+            $result['gqlQuery'] = $this->parseGqlQuery($result['gqlQuery']);
+        }
+
+        return $result;
+    }
+
+    private function parseRunAggregationQueryArgs($runQueryObj, $options)
+    {
+        $requestQueryArr = $runQueryObj->queryObject();
+        $result = [
+            'projectId' => $this->projectId,
+            'partitionId' => $this->partitionId(
+                $this->projectId,
+                $options['namespaceId'],
+                $options['databaseId']
+            ),
+        ] + $requestQueryArr + $this->readOptions($options) + $options;
+
+        if (isset($result['aggregationQuery'])) {
+            if (isset($result['aggregationQuery']['nestedQuery'])) {
+                $result['aggregationQuery']['nestedQuery'] = $this->parseQuery(
+                    $result['aggregationQuery']['nestedQuery']
+                );
+            }
+        }
+
+        if (isset($result['gqlQuery'])) {
+            $result['gqlQuery'] = $this->parseGqlQuery($result['gqlQuery']);
+        }
+
+        return $result;
+    }
 }