Merge branch 'main' into mock-server-tests

Author: olavloite

google/cloud/spanner_dbapi/connection.py

@@ -89,9 +89,11 @@ class Connection:
         committed by other transactions since the start of the read-only transaction. Commit or rolling back
         the read-only transaction is semantically the same, and only indicates that the read-only transaction
         should end a that a new one should be started when the next statement is executed.
+
+    **kwargs: Initial value for connection variables.
     """
 
-    def __init__(self, instance, database=None, read_only=False):
+    def __init__(self, instance, database=None, read_only=False, **kwargs):
         self._instance = instance
         self._database = database
         self._ddl_statements = []
@@ -117,6 +119,7 @@ def __init__(self, instance, database=None, read_only=False):
         self._batch_dml_executor: BatchDmlExecutor = None
         self._transaction_helper = TransactionRetryHelper(self)
         self._autocommit_dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
+        self._connection_variables = kwargs
 
     @property
     def spanner_client(self):
@@ -206,6 +209,10 @@ def _client_transaction_started(self):
         """
         return (not self._autocommit) or self._transaction_begin_marked
 
+    @property
+    def _ignore_transaction_warnings(self):
+        return self._connection_variables.get("ignore_transaction_warnings", False)
+
     @property
     def instance(self):
         """Instance to which this connection relates.
@@ -232,7 +239,7 @@ def read_only(self, value):
         Args:
             value (bool): True for ReadOnly mode, False for ReadWrite.
         """
-        if self._spanner_transaction_started:
+        if self._read_only != value and self._spanner_transaction_started:
             raise ValueError(
                 "Connection read/write mode can't be changed while a transaction is in progress. "
                 "Commit or rollback the current transaction and try again."
@@ -398,9 +405,10 @@ def commit(self):
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
         if not self._client_transaction_started:
-            warnings.warn(
-                CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
-            )
+            if not self._ignore_transaction_warnings:
+                warnings.warn(
+                    CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
+                )
             return
 
         self.run_prior_DDL_statements()
@@ -418,9 +426,10 @@ def rollback(self):
         This is a no-op if there is no active client transaction.
         """
         if not self._client_transaction_started:
-            warnings.warn(
-                CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
-            )
+            if not self._ignore_transaction_warnings:
+                warnings.warn(
+                    CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
+                )
             return
         try:
             if self._spanner_transaction_started and not self._read_only:
@@ -654,6 +663,7 @@ def connect(
     user_agent=None,
     client=None,
     route_to_leader_enabled=True,
+    **kwargs,
 ):
     """Creates a connection to a Google Cloud Spanner database.
 
@@ -696,6 +706,8 @@ def connect(
         disable leader aware routing. Disabling leader aware routing would
         route all requests in RW/PDML transactions to the closest region.
 
+    **kwargs: Initial value for connection variables.
+
 
     :rtype: :class:`google.cloud.spanner_dbapi.connection.Connection`
     :returns: Connection object associated with the given Google Cloud Spanner

google/cloud/spanner_v1/_helpers.py

@@ -266,66 +266,69 @@ def _parse_value_pb(value_pb, field_type, field_name, column_info=None):
     :returns: value extracted from value_pb
     :raises ValueError: if unknown type is passed
     """
+    decoder = _get_type_decoder(field_type, field_name, column_info)
+    return _parse_nullable(value_pb, decoder)
+
+
+def _get_type_decoder(field_type, field_name, column_info=None):
+    """Returns a function that converts a Value protobuf to cell data.
+
+    :type field_type: :class:`~google.cloud.spanner_v1.types.Type`
+    :param field_type: type code for the value
+
+    :type field_name: str
+    :param field_name: column name
+
+    :type column_info: dict
+    :param column_info: (Optional) dict of column name and column information.
+            An object where column names as keys and custom objects as corresponding
+            values for deserialization. It's specifically useful for data types like
+            protobuf where deserialization logic is on user-specific code. When provided,
+            the custom object enables deserialization of backend-received column data.
+            If not provided, data remains serialized as bytes for Proto Messages and
+            integer for Proto Enums.
+
+    :rtype: a function that takes a single protobuf value as an input argument
+    :returns: a function that can be used to extract a value from a protobuf value
+    :raises ValueError: if unknown type is passed
+    """
+
     type_code = field_type.code
-    if value_pb.HasField("null_value"):
-        return None
     if type_code == TypeCode.STRING:
-        return value_pb.string_value
+        return _parse_string
     elif type_code == TypeCode.BYTES:
-        return value_pb.string_value.encode("utf8")
+        return _parse_bytes
     elif type_code == TypeCode.BOOL:
-        return value_pb.bool_value
+        return _parse_bool
     elif type_code == TypeCode.INT64:
-        return int(value_pb.string_value)
+        return _parse_int64
     elif type_code == TypeCode.FLOAT64:
-        if value_pb.HasField("string_value"):
-            return float(value_pb.string_value)
-        else:
-            return value_pb.number_value
+        return _parse_float
     elif type_code == TypeCode.FLOAT32:
-        if value_pb.HasField("string_value"):
-            return float(value_pb.string_value)
-        else:
-            return value_pb.number_value
+        return _parse_float
     elif type_code == TypeCode.DATE:
-        return _date_from_iso8601_date(value_pb.string_value)
+        return _parse_date
     elif type_code == TypeCode.TIMESTAMP:
-        DatetimeWithNanoseconds = datetime_helpers.DatetimeWithNanoseconds
-        return DatetimeWithNanoseconds.from_rfc3339(value_pb.string_value)
-    elif type_code == TypeCode.ARRAY:
-        return [
-            _parse_value_pb(
-                item_pb, field_type.array_element_type, field_name, column_info
-            )
-            for item_pb in value_pb.list_value.values
-        ]
-    elif type_code == TypeCode.STRUCT:
-        return [
-            _parse_value_pb(
-                item_pb, field_type.struct_type.fields[i].type_, field_name, column_info
-            )
-            for (i, item_pb) in enumerate(value_pb.list_value.values)
-        ]
+        return _parse_timestamp
     elif type_code == TypeCode.NUMERIC:
-        return decimal.Decimal(value_pb.string_value)
+        return _parse_numeric
     elif type_code == TypeCode.JSON:
-        return JsonObject.from_str(value_pb.string_value)
+        return _parse_json
     elif type_code == TypeCode.PROTO:
-        bytes_value = base64.b64decode(value_pb.string_value)
-        if column_info is not None and column_info.get(field_name) is not None:
-            default_proto_message = column_info.get(field_name)
-            if isinstance(default_proto_message, Message):
-                proto_message = type(default_proto_message)()
-                proto_message.ParseFromString(bytes_value)
-                return proto_message
-        return bytes_value
+        return lambda value_pb: _parse_proto(value_pb, column_info, field_name)
     elif type_code == TypeCode.ENUM:
-        int_value = int(value_pb.string_value)
-        if column_info is not None and column_info.get(field_name) is not None:
-            proto_enum = column_info.get(field_name)
-            if isinstance(proto_enum, EnumTypeWrapper):
-                return proto_enum.Name(int_value)
-        return int_value
+        return lambda value_pb: _parse_proto_enum(value_pb, column_info, field_name)
+    elif type_code == TypeCode.ARRAY:
+        element_decoder = _get_type_decoder(
+            field_type.array_element_type, field_name, column_info
+        )
+        return lambda value_pb: _parse_array(value_pb, element_decoder)
+    elif type_code == TypeCode.STRUCT:
+        element_decoders = [
+            _get_type_decoder(item_field.type_, field_name, column_info)
+            for item_field in field_type.struct_type.fields
+        ]
+        return lambda value_pb: _parse_struct(value_pb, element_decoders)
     else:
         raise ValueError("Unknown type: %s" % (field_type,))
 
@@ -351,6 +354,87 @@ def _parse_list_value_pbs(rows, row_type):
     return result
 
 
+def _parse_string(value_pb) -> str:
+    return value_pb.string_value
+
+
+def _parse_bytes(value_pb):
+    return value_pb.string_value.encode("utf8")
+
+
+def _parse_bool(value_pb) -> bool:
+    return value_pb.bool_value
+
+
+def _parse_int64(value_pb) -> int:
+    return int(value_pb.string_value)
+
+
+def _parse_float(value_pb) -> float:
+    if value_pb.HasField("string_value"):
+        return float(value_pb.string_value)
+    else:
+        return value_pb.number_value
+
+
+def _parse_date(value_pb):
+    return _date_from_iso8601_date(value_pb.string_value)
+
+
+def _parse_timestamp(value_pb):
+    DatetimeWithNanoseconds = datetime_helpers.DatetimeWithNanoseconds
+    return DatetimeWithNanoseconds.from_rfc3339(value_pb.string_value)
+
+
+def _parse_numeric(value_pb):
+    return decimal.Decimal(value_pb.string_value)
+
+
+def _parse_json(value_pb):
+    return JsonObject.from_str(value_pb.string_value)
+
+
+def _parse_proto(value_pb, column_info, field_name):
+    bytes_value = base64.b64decode(value_pb.string_value)
+    if column_info is not None and column_info.get(field_name) is not None:
+        default_proto_message = column_info.get(field_name)
+        if isinstance(default_proto_message, Message):
+            proto_message = type(default_proto_message)()
+            proto_message.ParseFromString(bytes_value)
+            return proto_message
+    return bytes_value
+
+
+def _parse_proto_enum(value_pb, column_info, field_name):
+    int_value = int(value_pb.string_value)
+    if column_info is not None and column_info.get(field_name) is not None:
+        proto_enum = column_info.get(field_name)
+        if isinstance(proto_enum, EnumTypeWrapper):
+            return proto_enum.Name(int_value)
+    return int_value
+
+
+def _parse_array(value_pb, element_decoder) -> []:
+    return [
+        _parse_nullable(item_pb, element_decoder)
+        for item_pb in value_pb.list_value.values
+    ]
+
+
+def _parse_struct(value_pb, element_decoders):
+    return [
+        _parse_nullable(item_pb, element_decoders[i])
+        for (i, item_pb) in enumerate(value_pb.list_value.values)
+    ]
+
+
+def _parse_nullable(value_pb, decoder):
+    if value_pb.HasField("null_value"):
+        return None
+    else:
+        return decoder(value_pb)
+
+
 class _SessionWrapper(object):
     """Base class for objects wrapping a session.