feat: add query statistics package support (#129)

* feat: add query statistics package support

* style: fix lint

* test: reorder env mocks

Co-authored-by: larkee <larkee@users.noreply.github.com>

Author: larkee

google/cloud/spanner_v1/_helpers.py

@@ -84,15 +84,19 @@ def _merge_query_options(base, merge):
     combined = base or ExecuteSqlRequest.QueryOptions()
     if type(combined) == dict:
         combined = ExecuteSqlRequest.QueryOptions(
-            optimizer_version=combined.get("optimizer_version", "")
+            optimizer_version=combined.get("optimizer_version", ""),
+            optimizer_statistics_package=combined.get(
+                "optimizer_statistics_package", ""
+            ),
         )
     merge = merge or ExecuteSqlRequest.QueryOptions()
     if type(merge) == dict:
         merge = ExecuteSqlRequest.QueryOptions(
-            optimizer_version=merge.get("optimizer_version", "")
+            optimizer_version=merge.get("optimizer_version", ""),
+            optimizer_statistics_package=merge.get("optimizer_statistics_package", ""),
         )
     type(combined).pb(combined).MergeFrom(type(merge).pb(merge))
-    if not combined.optimizer_version:
+    if not combined.optimizer_version and not combined.optimizer_statistics_package:
         return None
     return combined
 

google/cloud/spanner_v1/client.py

@@ -64,6 +64,7 @@
 ) % ((EMULATOR_ENV_VAR,) * 3)
 SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin"
 OPTIMIZER_VERSION_ENV_VAR = "SPANNER_OPTIMIZER_VERSION"
+OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR = "SPANNER_OPTIMIZER_STATISTICS_PACKAGE"
 
 
 def _get_spanner_emulator_host():
@@ -74,6 +75,10 @@ def _get_spanner_optimizer_version():
     return os.getenv(OPTIMIZER_VERSION_ENV_VAR, "")
 
 
+def _get_spanner_optimizer_statistics_package():
+    return os.getenv(OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR, "")
+
+
 class Client(ClientWithProject):
     """Client for interacting with Cloud Spanner API.
 
@@ -160,7 +165,8 @@ def __init__(
         self._client_info = client_info
 
         env_query_options = ExecuteSqlRequest.QueryOptions(
-            optimizer_version=_get_spanner_optimizer_version()
+            optimizer_version=_get_spanner_optimizer_version(),
+            optimizer_statistics_package=_get_spanner_optimizer_statistics_package(),
         )
 
         # Environment flag config has higher precedence than application config.

tests/unit/test__helpers.py

@@ -30,9 +30,15 @@ def test_base_none_and_merge_none(self):
     def test_base_dict_and_merge_none(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
-        base = {"optimizer_version": "2"}
+        base = {
+            "optimizer_version": "2",
+            "optimizer_statistics_package": "auto_20191128_14_47_22UTC",
+        }
         merge = None
-        expected = ExecuteSqlRequest.QueryOptions(optimizer_version="2")
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="2",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         result = self._callFUT(base, merge)
         self.assertEqual(result, expected)
 
@@ -48,7 +54,10 @@ def test_base_none_merge_object(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
         base = None
-        merge = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
+        merge = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="3",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         result = self._callFUT(base, merge)
         self.assertEqual(result, merge)
 
@@ -64,9 +73,15 @@ def test_base_none_merge_dict(self):
     def test_base_object_merge_dict(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
-        base = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        base = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         merge = {"optimizer_version": "3"}
-        expected = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="3",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         result = self._callFUT(base, merge)
         self.assertEqual(result, expected)
 

tests/unit/test_client.py

@@ -179,30 +179,42 @@ def test_constructor_custom_query_options_client_config(self):
 
         expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
         creds = _make_credentials()
+        query_options = expected_query_options = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         self._constructor_test_helper(
             expected_scopes,
             creds,
-            query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"),
-            expected_query_options=ExecuteSqlRequest.QueryOptions(
-                optimizer_version="1"
-            ),
+            query_options=query_options,
+            expected_query_options=expected_query_options,
         )
 
+    @mock.patch(
+        "google.cloud.spanner_v1.client._get_spanner_optimizer_statistics_package"
+    )
     @mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version")
-    def test_constructor_custom_query_options_env_config(self, mock_ver):
+    def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
         from google.cloud.spanner_v1 import client as MUT
 
         expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
         creds = _make_credentials()
         mock_ver.return_value = "2"
+        mock_stats.return_value = "auto_20191128_14_47_22UTC"
+        query_options = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="auto_20191128_10_47_22UTC",
+        )
+        expected_query_options = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="2",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
         self._constructor_test_helper(
             expected_scopes,
             creds,
-            query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"),
-            expected_query_options=ExecuteSqlRequest.QueryOptions(
-                optimizer_version="2"
-            ),
+            query_options=query_options,
+            expected_query_options=expected_query_options,
         )
 
     @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host")