From eeb7836b6350aa9626dfb733208e6827d38bb9c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Wed, 4 Dec 2024 12:30:04 +0100
Subject: [PATCH] feat: add connection variable for ignoring transaction
 warnings (#1249)

Adds a connection variable for ignoring transaction warnings. Also adds
a **kwargs argument to the connect function. This will be used for
further connection variables in the future.

Fixes https://github.com/googleapis/python-spanner-sqlalchemy/issues/494
---
 google/cloud/spanner_dbapi/connection.py    | 26 +++++++++++++++------
 tests/unit/spanner_dbapi/test_connection.py | 13 +++++++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/google/cloud/spanner_dbapi/connection.py b/google/cloud/spanner_dbapi/connection.py
index b02d62ea27..65afcd4a2a 100644
--- a/google/cloud/spanner_dbapi/connection.py
+++ b/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.
@@ -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
diff --git a/tests/unit/spanner_dbapi/test_connection.py b/tests/unit/spanner_dbapi/test_connection.py
index d0fa521f8f..62867bbd2e 100644
--- a/tests/unit/spanner_dbapi/test_connection.py
+++ b/tests/unit/spanner_dbapi/test_connection.py
@@ -300,6 +300,19 @@ def test_commit_in_autocommit_mode(self, mock_warn):
             CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
         )
 
+    @mock.patch.object(warnings, "warn")
+    def test_commit_in_autocommit_mode_with_ignore_warnings(self, mock_warn):
+        conn = self._make_connection(
+            DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
+            ignore_transaction_warnings=True,
+        )
+        assert conn._ignore_transaction_warnings
+        conn._autocommit = True
+
+        conn.commit()
+
+        assert not mock_warn.warn.called
+
     def test_commit_database_error(self):
         from google.cloud.spanner_dbapi import Connection