From fe20d41e6ee333290fa461f40811c791ca1655fd Mon Sep 17 00:00:00 2001
From: Sanjeev Bhatt <bhatt4982@gmail.com>
Date: Mon, 22 Jul 2024 10:25:04 +0530
Subject: [PATCH] chore(spanner): Issue#1163 Remove dependency of spanner dbapi
 from spanner_v1 (#1164)

- Moved BatchTransactionId dataclass from spanner_dbapi to spanner_v1.transactions

Co-authored-by: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com>
---
 google/cloud/spanner_dbapi/partition_helper.py | 9 ++-------
 google/cloud/spanner_v1/__init__.py            | 2 ++
 google/cloud/spanner_v1/database.py            | 2 +-
 google/cloud/spanner_v1/transaction.py         | 9 +++++++++
 4 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/google/cloud/spanner_dbapi/partition_helper.py b/google/cloud/spanner_dbapi/partition_helper.py
index 94b396c801..a130e29721 100644
--- a/google/cloud/spanner_dbapi/partition_helper.py
+++ b/google/cloud/spanner_dbapi/partition_helper.py
@@ -19,6 +19,8 @@
 import pickle
 import base64
 
+from google.cloud.spanner_v1 import BatchTransactionId
+
 
 def decode_from_string(encoded_partition_id):
     gzip_bytes = base64.b64decode(bytes(encoded_partition_id, "utf-8"))
@@ -33,13 +35,6 @@ def encode_to_string(batch_transaction_id, partition_result):
     return str(base64.b64encode(gzip_bytes), "utf-8")
 
 
-@dataclass
-class BatchTransactionId:
-    transaction_id: str
-    session_id: str
-    read_timestamp: Any
-
-
 @dataclass
 class PartitionId:
     batch_transaction_id: BatchTransactionId
diff --git a/google/cloud/spanner_v1/__init__.py b/google/cloud/spanner_v1/__init__.py
index deba096163..d2e7a23938 100644
--- a/google/cloud/spanner_v1/__init__.py
+++ b/google/cloud/spanner_v1/__init__.py
@@ -64,6 +64,7 @@
 from .types.type import TypeAnnotationCode
 from .types.type import TypeCode
 from .data_types import JsonObject
+from .transaction import BatchTransactionId
 
 from google.cloud.spanner_v1 import param_types
 from google.cloud.spanner_v1.client import Client
@@ -147,4 +148,5 @@
     # google.cloud.spanner_v1.services
     "SpannerClient",
     "SpannerAsyncClient",
+    "BatchTransactionId",
 )
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
index 5b7c27b236..6bd4f3703e 100644
--- a/google/cloud/spanner_v1/database.py
+++ b/google/cloud/spanner_v1/database.py
@@ -40,7 +40,7 @@
 from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
 from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
-from google.cloud.spanner_dbapi.partition_helper import BatchTransactionId
+from google.cloud.spanner_v1.transaction import BatchTransactionId
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import Type
 from google.cloud.spanner_v1 import TypeCode
diff --git a/google/cloud/spanner_v1/transaction.py b/google/cloud/spanner_v1/transaction.py
index ee1dd8ef3b..c872cc380d 100644
--- a/google/cloud/spanner_v1/transaction.py
+++ b/google/cloud/spanner_v1/transaction.py
@@ -36,6 +36,8 @@
 from google.cloud.spanner_v1 import RequestOptions
 from google.api_core import gapic_v1
 from google.api_core.exceptions import InternalServerError
+from dataclasses import dataclass
+from typing import Any
 
 
 class Transaction(_SnapshotBase, _BatchBase):
@@ -554,3 +556,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self.commit()
         else:
             self.rollback()
+
+
+@dataclass
+class BatchTransactionId:
+    transaction_id: str
+    session_id: str
+    read_timestamp: Any