chore: move ClientContext validation to a helper function

Author: olavloite

google/cloud/spanner_v1/_helpers.py

@@ -230,6 +230,25 @@ def _merge_client_context(base, merge):
     return combined
 
 
+def _validate_client_context(client_context):
+    """Validate and convert client_context.
+
+    :type client_context: :class:`~google.cloud.spanner_v1.types.ClientContext`
+        or :class:`dict`
+    :param client_context: (Optional) Client context to use.
+
+    :rtype: :class:`~google.cloud.spanner_v1.types.ClientContext`
+    :returns: Validated ClientContext object or None.
+    :raises TypeError: if client_context is not a ClientContext or a dict.
+    """
+    if client_context is not None:
+        if isinstance(client_context, dict):
+            client_context = ClientContext(client_context)
+        elif not isinstance(client_context, ClientContext):
+            raise TypeError("client_context must be a ClientContext or a dict")
+    return client_context
+
+
 def _merge_request_options(request_options, client_context):
     """Merge RequestOptions and ClientContext.
 

google/cloud/spanner_v1/batch.py

@@ -30,6 +30,7 @@
     _merge_Transaction_Options,
     _merge_client_context,
     _merge_request_options,
+    _validate_client_context,
     AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
@@ -66,13 +67,7 @@ def __init__(self, session, client_context=None):
         self.committed = None
         """Timestamp at which the batch was successfully committed."""
         self.commit_stats: Optional[CommitResponse.CommitStats] = None
-
-        if client_context is not None:
-            if isinstance(client_context, dict):
-                client_context = ClientContext(client_context)
-            elif not isinstance(client_context, ClientContext):
-                raise TypeError("client_context must be a ClientContext or a dict")
-        self._client_context = client_context
+        self._client_context = _validate_client_context(client_context)
 
     def insert(self, table, columns, values):
         """Insert one or more new table rows.

google/cloud/spanner_v1/client.py

@@ -53,6 +53,7 @@
 from google.cloud.spanner_v1.types import ClientContext
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
+from google.cloud.spanner_v1._helpers import _validate_client_context
 from google.cloud.spanner_v1.instance import Instance
 from google.cloud.spanner_v1.metrics.constants import (
     METRIC_EXPORT_INTERVAL_MS,
@@ -293,13 +294,7 @@ def __init__(
 
         # Environment flag config has higher precedence than application config.
         self._query_options = _merge_query_options(query_options, env_query_options)
-
-        if client_context is not None:
-            if isinstance(client_context, dict):
-                client_context = ClientContext(client_context)
-            elif not isinstance(client_context, ClientContext):
-                raise TypeError("client_context must be a ClientContext or a dict")
-        self._client_context = client_context
+        self._client_context = _validate_client_context(client_context)
 
         if self._emulator_host is not None and (
             "http://" in self._emulator_host or "https://" in self._emulator_host

google/cloud/spanner_v1/snapshot.py

@@ -50,6 +50,7 @@
     _SessionWrapper,
     AtomicCounter,
     _augment_error_with_request_id,
+    _validate_client_context,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call, add_span_event
 from google.cloud.spanner_v1.streamed import StreamedResultSet
@@ -212,13 +213,7 @@ class _SnapshotBase(_SessionWrapper):
     def __init__(self, session, client_context=None):
         super().__init__(session)
 
-        if client_context is not None:
-            if isinstance(client_context, dict):
-                client_context = ClientContext(client_context)
-            elif not isinstance(client_context, ClientContext):
-                raise TypeError("client_context must be a ClientContext or a dict")
-        self._client_context = client_context
-
+        self._client_context = _validate_client_context(client_context)
         # Counts for execute SQL requests and total read requests (including
         # execute SQL requests). Used to provide sequence numbers for
         # :class:`google.cloud.spanner_v1.types.ExecuteSqlRequest` and to