Upgrade spanner-orm to work with google-cloud-spanner v2

supersam654 · supersam654 · commit 71ec2db83dc5 · 2021-12-17T20:15:13.000Z
diff --git a/setup.py b/setup.py
@@ -17,7 +17,7 @@
 
 setup(
     name='spanner-orm',
-    version='0.1.10',
+    version='0.2.0',
     description='Basic ORM for Spanner',
     maintainer='Python Spanner ORM developers',
     maintainer_email='python-spanner-orm@google.com',
@@ -26,7 +26,7 @@
     include_package_data=True,
     python_requires='~=3.7',
     install_requires=[
-        'google-cloud-spanner >= 1.6, <4',
+        'google-cloud-spanner >= 2, <4',
         'immutabledict',
     ],
     tests_require=['absl-py', 'google-api-core', 'portpicker'],
diff --git a/spanner_orm/condition.py b/spanner_orm/condition.py
@@ -30,7 +30,8 @@
 from spanner_orm import relationship
 
 from google.api_core import datetime_helpers
-from google.cloud.spanner_v1.proto import type_pb2
+from google.cloud import spanner
+from google.cloud import spanner_v1
 import immutabledict
 
 T = TypeVar('T')
@@ -105,7 +106,7 @@ def sql(self) -> str:
   def _sql(self) -> str:
     pass
 
-  def types(self) -> Dict[str, type_pb2.Type]:
+  def types(self) -> Dict[str, spanner_v1.Type]:
     """Returns parameter types to be used in the SQL query.
 
     Returns:
@@ -117,7 +118,7 @@ def types(self) -> Dict[str, type_pb2.Type]:
     return self._types()
 
   @abc.abstractmethod
-  def _types(self) -> Dict[str, type_pb2.Type]:
+  def _types(self) -> Dict[str, spanner_v1.Type]:
     raise NotImplementedError
 
   @abc.abstractmethod
@@ -158,7 +159,8 @@ def _validate(self, model_class: Type[Any]) -> None:
 ]
 
 
-def _spanner_type_of_python_object(value: GuessableParamType) -> type_pb2.Type:
+def _spanner_type_of_python_object(
+    value: GuessableParamType) -> spanner_v1.Type:
   """Returns the Cloud Spanner type of the given object.
 
   Args:
@@ -173,31 +175,30 @@ def _spanner_type_of_python_object(value: GuessableParamType) -> type_pb2.Type:
     raise TypeError(
         'Cannot infer type of None, because any SQL type can be NULL.')
   simple_type_code = {
-      bool: type_pb2.BOOL,
-      int: type_pb2.INT64,
-      float: type_pb2.FLOAT64,
-      datetime_helpers.DatetimeWithNanoseconds: type_pb2.TIMESTAMP,
-      datetime.datetime: type_pb2.TIMESTAMP,
-      datetime.date: type_pb2.DATE,
-      bytes: type_pb2.BYTES,
-      str: type_pb2.STRING,
-      decimal.Decimal: type_pb2.NUMERIC,
+      bool: spanner_v1.TypeCode.BOOL,
+      int: spanner_v1.TypeCode.INT64,
+      float: spanner_v1.TypeCode.FLOAT64,
+      datetime_helpers.DatetimeWithNanoseconds: spanner_v1.TypeCode.TIMESTAMP,
+      datetime.datetime: spanner_v1.TypeCode.TIMESTAMP,
+      datetime.date: spanner_v1.TypeCode.DATE,
+      bytes: spanner_v1.TypeCode.BYTES,
+      str: spanner_v1.TypeCode.STRING,
+      decimal.Decimal: spanner_v1.TypeCode.NUMERIC,
   }.get(type(value))
   if simple_type_code is not None:
-    return type_pb2.Type(code=simple_type_code)
+    return spanner_v1.Type(code=simple_type_code)
   elif isinstance(value, (list, tuple)):
     element_types = tuple(
         _spanner_type_of_python_object(item)
         for item in value
         if item is not None)
     unique_element_type_count = len({
-        # Protos aren't hashable, so use their serializations.
-        element_type.SerializeToString(deterministic=True)
-        for element_type in element_types
+        # Protos aren't hashable, so serialize them.
+        str(element_type) for element_type in element_types
     })
     if unique_element_type_count == 1:
-      return type_pb2.Type(
-          code=type_pb2.ARRAY,
+      return spanner_v1.Type(
+          code=spanner_v1.TypeCode.ARRAY,
           array_element_type=element_types[0],
       )
     else:
@@ -211,7 +212,7 @@ def _spanner_type_of_python_object(value: GuessableParamType) -> type_pb2.Type:
 class Param:
   """Parameter for substitution into a SQL query."""
   value: Any
-  type: type_pb2.Type
+  type: spanner_v1.Type
 
   @classmethod
   def from_value(cls: Type[T], value: GuessableParamType) -> T:
@@ -220,14 +221,13 @@ def from_value(cls: Type[T], value: GuessableParamType) -> T:
 
     # BYTES must be base64-encoded, see
     # https://github.com/googleapis/python-spanner/blob/87789c939990794bfd91f5300bedc449fd74bd7e/google/cloud/spanner_v1/proto/type.proto#L108-L110
-    if (isinstance(value, bytes) and
-        guessed_type == type_pb2.Type(code=type_pb2.BYTES)):
+    if (isinstance(value, bytes) and guessed_type == spanner.param_types.BYTES):
       encoded_value = base64.b64encode(value).decode()
     elif (isinstance(value, (list, tuple)) and
           all(isinstance(x, bytes) for x in value if x is not None) and
-          guessed_type == type_pb2.Type(
-              code=type_pb2.ARRAY,
-              array_element_type=type_pb2.Type(code=type_pb2.BYTES),
+          guessed_type == spanner_v1.Type(
+              code=spanner_v1.TypeCode.ARRAY,
+              array_element_type=spanner.param_types.BYTES,
           )):
       encoded_value = tuple(
           None if item is None else base64.b64encode(item).decode()
@@ -299,7 +299,7 @@ def _params(self) -> Dict[str, Any]:
         if isinstance(v, Param)
     }
 
-  def _types(self) -> Dict[str, type_pb2.Type]:
+  def _types(self) -> Dict[str, spanner_v1.Type]:
     """See base class."""
     return {
         self.key(k): v.type
@@ -345,7 +345,7 @@ def _sql(self) -> str:
         other_table=self.destination_model_class.table,
         other_column=self.destination_column)
 
-  def _types(self) -> Dict[str, type_pb2.Type]:
+  def _types(self) -> Dict[str, spanner_v1.Type]:
     return {}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -389,7 +389,7 @@ def segment(self) -> Segment:
   def _sql(self) -> str:
     return '@{{FORCE_INDEX={}}}'.format(self.name)
 
-  def _types(self) -> Dict[str, type_pb2.Type]:
+  def _types(self) -> Dict[str, spanner_v1.Type]:
     return {}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -509,7 +509,7 @@ def segment(self) -> Segment:
   def _sql(self) -> str:
     return ''
 
-  def _types(self) -> Dict[str, type_pb2.Type]:
+  def _types(self) -> Dict[str, spanner_v1.Type]:
     return {}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -568,10 +568,10 @@ def _sql(self) -> str:
           limit_key=self._limit_key, offset_key=self._offset_key)
     return 'LIMIT @{limit_key}'.format(limit_key=self._limit_key)
 
-  def _types(self) -> Dict[str, type_pb2.Type]:
-    types = {self._limit_key: type_pb2.Type(code=type_pb2.INT64)}
+  def _types(self) -> Dict[str, spanner_v1.Type]:
+    types = {self._limit_key: spanner.param_types.INT64}
     if self.offset:
-      types[self._offset_key] = type_pb2.Type(code=type_pb2.INT64)
+      types[self._offset_key] = spanner.param_types.INT64
     return types
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -624,7 +624,7 @@ def _sql(self) -> str:
   def segment(self) -> Segment:
     return Segment.WHERE
 
-  def _types(self) -> type_pb2.Type:
+  def _types(self) -> spanner_v1.Type:
     result = {}
     for condition in self.all_conditions:
       condition.suffix = str(int(self.suffix or 0) + len(result))
@@ -669,7 +669,7 @@ def _sql(self) -> str:
   def segment(self) -> Segment:
     return Segment.ORDER_BY
 
-  def _types(self) -> type_pb2.Type:
+  def _types(self) -> spanner_v1.Type:
     return {}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -714,7 +714,7 @@ def _sql(self) -> str:
         operator=self.operator,
         column_key=self._column_key)
 
-  def _types(self) -> type_pb2.Type:
+  def _types(self) -> spanner_v1.Type:
     return {self._column_key: self.model_class.fields[self.column].grpc_type()}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -740,9 +740,10 @@ def _sql(self) -> str:
         operator=self.operator,
         column_key=self._column_key)
 
-  def _types(self) -> type_pb2.Type:
+  def _types(self) -> spanner_v1.Type:
     grpc_type = self.model_class.fields[self.column].grpc_type()
-    list_type = type_pb2.Type(code=type_pb2.ARRAY, array_element_type=grpc_type)
+    list_type = spanner_v1.Type(
+        code=spanner_v1.TypeCode.ARRAY, array_element_type=grpc_type)
     return {self._column_key: list_type}
 
   def _validate(self, model_class: Type[Any]) -> None:
@@ -782,7 +783,7 @@ def _sql(self) -> str:
           operator=self.nullable_operator)
     return super()._sql()
 
-  def _types(self) -> type_pb2.Type:
+  def _types(self) -> spanner_v1.Type:
     if self.is_null():
       return {}
     return super()._types()
diff --git a/spanner_orm/field.py b/spanner_orm/field.py
@@ -20,7 +20,8 @@
 import datetime
 from typing import Any, Type
 
-from google.cloud.spanner_v1.proto import type_pb2
+from google.cloud import spanner
+from google.cloud import spanner_v1
 from spanner_orm import error
 
 
@@ -34,7 +35,7 @@ def ddl() -> str:
 
   @staticmethod
   @abc.abstractmethod
-  def grpc_type() -> type_pb2.Type:
+  def grpc_type() -> spanner_v1.Type:
     raise NotImplementedError
 
   @staticmethod
@@ -88,8 +89,8 @@ def ddl() -> str:
     return 'BOOL'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.BOOL)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.BOOL
 
   @staticmethod
   def validate_type(value: Any) -> None:
@@ -105,8 +106,8 @@ def ddl() -> str:
     return 'INT64'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.INT64)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.INT64
 
   @staticmethod
   def validate_type(value: Any) -> None:
@@ -122,8 +123,8 @@ def ddl() -> str:
     return 'FLOAT64'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.FLOAT64)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.FLOAT64
 
   @staticmethod
   def validate_type(value: Any) -> None:
@@ -139,8 +140,8 @@ def ddl() -> str:
     return 'STRING(MAX)'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.STRING)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.STRING
 
   @staticmethod
   def validate_type(value) -> None:
@@ -156,8 +157,8 @@ def ddl() -> str:
     return 'ARRAY<STRING(MAX)>'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.ARRAY)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.Array(spanner.param_types.STRING)
 
   @staticmethod
   def validate_type(value: Any) -> None:
@@ -176,8 +177,8 @@ def ddl() -> str:
     return 'TIMESTAMP'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.TIMESTAMP)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.TIMESTAMP
 
   @staticmethod
   def validate_type(value: Any) -> None:
@@ -193,8 +194,8 @@ def ddl() -> str:
     return 'BYTES(MAX)'
 
   @staticmethod
-  def grpc_type() -> type_pb2.Type:
-    return type_pb2.Type(code=type_pb2.BYTES)
+  def grpc_type() -> spanner_v1.Type:
+    return spanner.param_types.BYTES
 
   @staticmethod
   def validate_type(value) -> None:
diff --git a/spanner_orm/table_apis.py b/spanner_orm/table_apis.py
@@ -19,8 +19,8 @@
 
 # TODO(https://github.com/google/pytype/issues/1081): Remove pytype disable.
 from google.cloud import spanner  # pytype: disable=import-error
+from google.cloud import spanner_v1
 from google.cloud.spanner_v1 import transaction as spanner_transaction
-from google.cloud.spanner_v1.proto import type_pb2
 
 _logger = logging.getLogger(__name__)
 
@@ -50,9 +50,10 @@ def find(transaction: spanner_transaction.Transaction, table_name: str,
   return list(stream_results)
 
 
-def sql_query(transaction: spanner_transaction.Transaction, query: str,
-              parameters: Dict[str, Any],
-              parameter_types: Dict[str, type_pb2.Type]) -> List[Sequence[Any]]:
+def sql_query(
+    transaction: spanner_transaction.Transaction, query: str,
+    parameters: Dict[str, Any],
+    parameter_types: Dict[str, spanner_v1.Type]) -> List[Sequence[Any]]:
   """Executes a given SQL query against the Spanner database.
 
   This isn't technically read-only, but it's necessary to implement the read-
diff --git a/spanner_orm/testlib/spanner_emulator/testlib.py b/spanner_orm/testlib/spanner_emulator/testlib.py
@@ -55,9 +55,9 @@ def _get_instance(spanner_client: client.Client) -> instance.Instance:
   Args:
     spanner_client: An initialized spanner client.
   """
-  existing_instances = list(spanner_client.list_instances())
-  if existing_instances:
-    return existing_instances[0]
+  existing_instances_pb = list(spanner_client.list_instances())
+  if existing_instances_pb:
+    return instance.Instance.from_pb(existing_instances_pb[0], spanner_client)
 
   # The emulator has one default config.
   config = list(spanner_client.list_instance_configs())[0]
diff --git a/spanner_orm/tests/condition_test.py b/spanner_orm/tests/condition_test.py
diff --git a/spanner_orm/tests/query_test.py b/spanner_orm/tests/query_test.py
