From 2292abcc171a73f5b069f90fe659a8dcc90e8adc Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Mon, 28 Sep 2026 10:24:28 +0000 Subject: [PATCH 1/2] feat(dbapi): wire timeout parameter through Connection to execute_sql - Wire timeout parameter to Connection.__init__, connect(), run_statement(), and validate(). - Pass timeout from connection to execute_sql and batch_update calls when specified. - Add unit tests for DBAPI timeout propagation. --- .../cloud/spanner_dbapi/batch_dml_executor.py | 17 +- .../google/cloud/spanner_dbapi/connection.py | 33 +++- .../google/cloud/spanner_dbapi/cursor.py | 22 ++- .../tests/unit/spanner_dbapi/test_timeout.py | 182 ++++++++++++++++++ 4 files changed, 241 insertions(+), 13 deletions(-) create mode 100644 packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py index eede542974ae..6a7364061e42 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py @@ -92,7 +92,9 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]): statements_tuple.append(statement.get_tuple()) if not connection._client_transaction_started: res = connection.database.run_in_transaction( - _do_batch_update_autocommit, statements_tuple + _do_batch_update_autocommit, + statements_tuple, + timeout=connection.timeout, ) many_result_set.add_iter(res) cursor._row_count = sum([max(val, 0) for val in res]) @@ -101,7 +103,10 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]): while True: try: transaction = connection.transaction_checkout() - status, res = transaction.batch_update(statements_tuple) + kwargs = {} + if connection.timeout is not None: + kwargs["timeout"] = connection.timeout + status, res = transaction.batch_update(statements_tuple, **kwargs) if status.code == ABORTED: connection._transaction = None raise Aborted(status.message) @@ -133,10 +138,14 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]): raise ex -def _do_batch_update_autocommit(transaction, statements): +def _do_batch_update_autocommit(transaction, statements, timeout=None): from google.cloud.spanner_dbapi import OperationalError - status, res = transaction.batch_update(statements, last_statement=True) + kwargs = {"last_statement": True} + if timeout is not None: + kwargs["timeout"] = timeout + + status, res = transaction.batch_update(statements, **kwargs) if status.code == ABORTED: raise Aborted(status.message) elif status.code != OK: diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py index d73a036b01a7..a3e78bcc512d 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py @@ -112,6 +112,7 @@ def __init__( read_only=False, data_boost_enabled=False, auto_partition_mode=False, + timeout=None, **kwargs, ): self._instance = instance @@ -133,6 +134,7 @@ def __init__( self._auto_partition_mode = bool(auto_partition_mode) self._staleness = None self.request_priority = None + self.timeout = timeout self._transaction_begin_marked = False self._transaction_isolation_level = None # whether transaction started at Spanner. This means that we had @@ -598,7 +600,10 @@ def run_prior_DDL_statements(self): return self.database.update_ddl(ddl_statements).result() def run_statement( - self, statement: Statement, request_options: RequestOptions = None + self, + statement: Statement, + request_options: RequestOptions = None, + timeout: float = None, ): """Run single SQL statement in begun transaction. @@ -616,17 +621,28 @@ def run_statement( :type request_options: :class:`RequestOptions` :param request_options: Request options to use for this statement. + :type timeout: float + :param timeout: (Optional) The amount of time, in seconds, to wait + for the request to complete. + :rtype: :class:`google.cloud.spanner_v1.streamed.StreamedResultSet`, :class:`google.cloud.spanner_dbapi.checksum.ResultsChecksum` :returns: Streamed result set of the statement and a checksum of this statement results. """ transaction = self.transaction_checkout() + kwargs = { + "param_types": statement.param_types, + "request_options": request_options or self.request_options, + } + effective_timeout = timeout if timeout is not None else self.timeout + if effective_timeout is not None: + kwargs["timeout"] = effective_timeout + return transaction.execute_sql( statement.sql, statement.params, - param_types=statement.param_types, - request_options=request_options or self.request_options, + **kwargs, ) @check_not_closed @@ -646,8 +662,11 @@ def validate(self): """ if self.database is None: raise ValueError("Database needs to be passed for this operation") + kwargs = {} + if self.timeout is not None: + kwargs["timeout"] = self.timeout with self.database.snapshot() as snapshot: - result = list(snapshot.execute_sql("SELECT 1")) + result = list(snapshot.execute_sql("SELECT 1", **kwargs)) if result != [[1]]: raise OperationalError( "The checking query (SELECT 1) returned an unexpected result: %s. " @@ -822,10 +841,15 @@ def connect( auto_partition_mode=False, username=None, password=None, + timeout=None, **kwargs, ): """Creates a connection to a Google Cloud Spanner database. + :type timeout: float + :param timeout: (Optional) The amount of time, in seconds, to wait + for requests to complete. + :type instance_id: str :param instance_id: The ID of the instance to connect to. @@ -1024,6 +1048,7 @@ def connect( database, data_boost_enabled=data_boost_enabled, auto_partition_mode=auto_partition_mode, + timeout=timeout, **kwargs, ) if pool is not None: diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py index c721ff16d282..27035ec1e6f4 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py @@ -232,11 +232,16 @@ def _do_execute_update_in_autocommit(self, transaction, sql, params): """This function should only be used in autocommit mode.""" self.connection._transaction = transaction self.connection._snapshot = None + kwargs = { + "params": params, + "param_types": get_param_types(params), + "last_statement": True, + } + if self.connection.timeout is not None: + kwargs["timeout"] = self.connection.timeout self._result_set = transaction.execute_sql( sql, - params=params, - param_types=get_param_types(params), - last_statement=True, + **kwargs, ) self._itr = BufferedIterator(self._result_set) self._row_count = None @@ -558,11 +563,14 @@ def _fetch(self, cursor_statement_type, size=None): return rows def _handle_DQL_with_snapshot(self, snapshot, sql, params): + kwargs = {"request_options": self.request_options} + if self.connection.timeout is not None: + kwargs["timeout"] = self.connection.timeout self._result_set = snapshot.execute_sql( sql, params, get_param_types(params), - request_options=self.request_options, + **kwargs, ) # Read the first element so that the StreamedResultSet can # return the metadata after a DQL statement. @@ -651,8 +659,12 @@ def run_sql_in_snapshot(self, sql, params=None, param_types=None): raise ValueError("Database needs to be passed for this operation") self.connection.run_prior_DDL_statements() + kwargs = {} + if self.connection.timeout is not None: + kwargs["timeout"] = self.connection.timeout + with self.connection.database.snapshot() as snapshot: - return list(snapshot.execute_sql(sql, params, param_types)) + return list(snapshot.execute_sql(sql, params, param_types, **kwargs)) def get_table_column_schema(self, table_name, schema_name=""): rows = self.run_sql_in_snapshot( diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py new file mode 100644 index 000000000000..9ca1a96bcb6a --- /dev/null +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py @@ -0,0 +1,182 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for timeout handling in Spanner DB-API.""" + +import unittest +from unittest import mock + +from google.auth.credentials import AnonymousCredentials +from google.rpc.code_pb2 import OK + +from google.cloud.spanner_dbapi import Connection, connect +from google.cloud.spanner_dbapi.parsed_statement import ( + ParsedStatement, + Statement, + StatementType, +) + + +class TestDBAPITimeout(unittest.TestCase): + INSTANCE = "test-instance" + DATABASE = "test-database" + + def _make_connection(self, timeout=None, **kwargs): + from google.cloud.spanner_v1.client import Client + from google.cloud.spanner_v1.instance import Instance + + client = Client( + project="test", + credentials=AnonymousCredentials(), + client_options={"api_endpoint": "none"}, + ) + instance = Instance(self.INSTANCE, client=client) + database = instance.database(self.DATABASE) + return Connection(instance, database, timeout=timeout, **kwargs) + + def test_connection_timeout_default_is_none(self): + connection = self._make_connection() + self.assertIsNone(connection.timeout) + + def test_connection_timeout_init(self): + connection = self._make_connection(timeout=30.0) + self.assertEqual(connection.timeout, 30.0) + + def test_connect_accepts_timeout(self): + connection = connect( + self.INSTANCE, + self.DATABASE, + project="test-project", + credentials=AnonymousCredentials(), + client_options={"api_endpoint": "none"}, + timeout=45.0, + ) + self.assertEqual(connection.timeout, 45.0) + + def test_run_statement_passes_timeout(self): + connection = self._make_connection(timeout=25.0) + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + connection._transaction = mock_transaction + + sql = "SELECT 1" + stmt = Statement(sql, [], {}) + connection.run_statement(stmt) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 25.0) + + def test_run_statement_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + connection._transaction = mock_transaction + + sql = "SELECT 1" + stmt = Statement(sql, [], {}) + connection.run_statement(stmt) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertNotIn("timeout", kwargs) + + def test_run_statement_explicit_timeout_override(self): + connection = self._make_connection(timeout=25.0) + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + connection._transaction = mock_transaction + + sql = "SELECT 1" + stmt = Statement(sql, [], {}) + connection.run_statement(stmt, timeout=5.0) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 5.0) + + def test_cursor_dql_snapshot_passes_timeout(self): + connection = self._make_connection(timeout=12.0) + connection.autocommit = True + cursor = connection.cursor() + + mock_snapshot = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_result_set.metadata.transaction.read_timestamp = None + mock_snapshot.execute_sql.return_value = mock_result_set + + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.QUERY, Statement("SELECT 1")), + ): + cursor.execute("SELECT 1") + + mock_snapshot.execute_sql.assert_called_once() + _, kwargs = mock_snapshot.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 12.0) + + def test_cursor_autocommit_dml_passes_timeout(self): + connection = self._make_connection(timeout=18.0) + connection.autocommit = True + cursor = connection.cursor() + + mock_transaction = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_transaction.execute_sql.return_value = mock_result_set + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 18.0) + self.assertTrue(kwargs.get("last_statement")) + + def test_batch_dml_transaction_passes_timeout(self): + connection = self._make_connection(timeout=7.0) + connection.autocommit = True + cursor = connection.cursor() + + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertEqual(kwargs.get("timeout"), 7.0) + self.assertTrue(kwargs.get("last_statement")) From 4196ae159c02b31e327159fb24be3ce00fb38fb3 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Sun, 4 Oct 2026 18:59:53 +0000 Subject: [PATCH 2/2] Addressed comments --- .../cloud/spanner_dbapi/batch_dml_executor.py | 6 +- .../google/cloud/spanner_dbapi/connection.py | 14 +- .../google/cloud/spanner_dbapi/cursor.py | 39 +- .../cloud/spanner_dbapi/transaction_helper.py | 3 + .../tests/unit/spanner_dbapi/test_timeout.py | 358 ++++++++++++++++++ 5 files changed, 407 insertions(+), 13 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py index 6a7364061e42..c13497f3fc76 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/batch_dml_executor.py @@ -94,7 +94,7 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]): res = connection.database.run_in_transaction( _do_batch_update_autocommit, statements_tuple, - timeout=connection.timeout, + timeout=cursor.timeout, ) many_result_set.add_iter(res) cursor._row_count = sum([max(val, 0) for val in res]) @@ -104,8 +104,8 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]): try: transaction = connection.transaction_checkout() kwargs = {} - if connection.timeout is not None: - kwargs["timeout"] = connection.timeout + if cursor.timeout is not None: + kwargs["timeout"] = cursor.timeout status, res = transaction.batch_update(statements_tuple, **kwargs) if status.code == ABORTED: connection._transaction = None diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py index a3e78bcc512d..aa3f02f2d823 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py @@ -15,6 +15,7 @@ """DB-API Connection for the Google Cloud Spanner.""" import warnings +from typing import Optional from google.api_core.client_options import ClientOptions from google.api_core.exceptions import Aborted @@ -102,6 +103,10 @@ class Connection: or autocommit connections are automatically partitioned and executed in parallel. Defaults to False. + :type timeout: Optional[float] + :param timeout: (Optional) The amount of time, in seconds, to wait + for requests to complete. + **kwargs: Initial value for connection variables. """ @@ -603,7 +608,7 @@ def run_statement( self, statement: Statement, request_options: RequestOptions = None, - timeout: float = None, + timeout: Optional[float] = None, ): """Run single SQL statement in begun transaction. @@ -621,9 +626,12 @@ def run_statement( :type request_options: :class:`RequestOptions` :param request_options: Request options to use for this statement. - :type timeout: float + :type timeout: Optional[float] :param timeout: (Optional) The amount of time, in seconds, to wait - for the request to complete. + for the request to complete. Defaults to the + connection's timeout if None; when the connection's + timeout is set, passing None will use the connection's + timeout rather than the underlying RPC default timeout. :rtype: :class:`google.cloud.spanner_v1.streamed.StreamedResultSet`, :class:`google.cloud.spanner_dbapi.checksum.ResultsChecksum` diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py index 27035ec1e6f4..dc12d1b84609 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py @@ -103,6 +103,31 @@ def __init__(self, connection): self._in_retry_mode = False self._batch_dml_rows_count = None self._request_tag = None + self._timeout = None + + @property + def timeout(self): + """The timeout (in seconds) that will be applied to statements on this + cursor. Defaults to the connection's timeout if not set on the cursor. + + Returns: + Optional[float]: The timeout (in seconds) for statements on this + cursor. + """ + return ( + self._timeout + if self._timeout is not None + else getattr(self.connection, "timeout", None) + ) + + @timeout.setter + def timeout(self, value): + """Sets the timeout (in seconds) for statements on this cursor. + + Args: + value (Optional[float]): The timeout (in seconds) for statements. + """ + self._timeout = value @property def request_tag(self): @@ -237,8 +262,8 @@ def _do_execute_update_in_autocommit(self, transaction, sql, params): "param_types": get_param_types(params), "last_statement": True, } - if self.connection.timeout is not None: - kwargs["timeout"] = self.connection.timeout + if self.timeout is not None: + kwargs["timeout"] = self.timeout self._result_set = transaction.execute_sql( sql, **kwargs, @@ -381,7 +406,7 @@ def _execute_in_rw_transaction(self): while True: try: self._result_set = self.connection.run_statement( - statement, self.request_options + statement, self.request_options, self.timeout ) self._itr = PeekIterator(self._result_set) return @@ -564,8 +589,8 @@ def _fetch(self, cursor_statement_type, size=None): def _handle_DQL_with_snapshot(self, snapshot, sql, params): kwargs = {"request_options": self.request_options} - if self.connection.timeout is not None: - kwargs["timeout"] = self.connection.timeout + if self.timeout is not None: + kwargs["timeout"] = self.timeout self._result_set = snapshot.execute_sql( sql, params, @@ -660,8 +685,8 @@ def run_sql_in_snapshot(self, sql, params=None, param_types=None): self.connection.run_prior_DDL_statements() kwargs = {} - if self.connection.timeout is not None: - kwargs["timeout"] = self.connection.timeout + if self.timeout is not None: + kwargs["timeout"] = self.timeout with self.connection.database.snapshot() as snapshot: return list(snapshot.execute_sql(sql, params, param_types, **kwargs)) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/transaction_helper.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/transaction_helper.py index aadaacc61ac1..b9e9f816ca72 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/transaction_helper.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/transaction_helper.py @@ -187,6 +187,9 @@ def retry_transaction(self, default_retry_delay=None): else: cursor = self._connection.cursor() cursor._in_retry_mode = True + cursor._timeout = getattr( + statement_result_details.cursor, "_timeout", None + ) self._cursor_map[statement_result_details.cursor] = cursor try: _handle_statement(statement_result_details, cursor) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py index 9ca1a96bcb6a..54d2dd7d742c 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_timeout.py @@ -180,3 +180,361 @@ def fake_run_in_transaction(func, *args, **kwargs): _, kwargs = mock_transaction.batch_update.call_args self.assertEqual(kwargs.get("timeout"), 7.0) self.assertTrue(kwargs.get("last_statement")) + + def test_cursor_timeout_property(self): + connection = self._make_connection(timeout=30.0) + cursor = connection.cursor() + self.assertEqual(cursor.timeout, 30.0) + + cursor.timeout = 10.0 + self.assertEqual(cursor.timeout, 10.0) + self.assertEqual(connection.timeout, 30.0) + + cursor.timeout = None + self.assertEqual(cursor.timeout, 30.0) + + def test_validate_passes_timeout(self): + connection = self._make_connection(timeout=15.0) + mock_snapshot = mock.MagicMock() + mock_snapshot.execute_sql.return_value = [[1]] + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + connection.validate() + + mock_snapshot.execute_sql.assert_called_once_with("SELECT 1", timeout=15.0) + + def test_validate_without_timeout_omits_kwarg(self): + connection = self._make_connection() + mock_snapshot = mock.MagicMock() + mock_snapshot.execute_sql.return_value = [[1]] + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + connection.validate() + + mock_snapshot.execute_sql.assert_called_once_with("SELECT 1") + + def test_cursor_dql_snapshot_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection.autocommit = True + cursor = connection.cursor() + + mock_snapshot = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_result_set.metadata.transaction.read_timestamp = None + mock_snapshot.execute_sql.return_value = mock_result_set + + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.QUERY, Statement("SELECT 1")), + ): + cursor.execute("SELECT 1") + + mock_snapshot.execute_sql.assert_called_once() + _, kwargs = mock_snapshot.execute_sql.call_args + self.assertNotIn("timeout", kwargs) + + def test_cursor_dql_snapshot_uses_cursor_timeout(self): + connection = self._make_connection(timeout=12.0) + connection.autocommit = True + cursor = connection.cursor() + cursor.timeout = 4.0 + + mock_snapshot = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_result_set.metadata.transaction.read_timestamp = None + mock_snapshot.execute_sql.return_value = mock_result_set + + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.QUERY, Statement("SELECT 1")), + ): + cursor.execute("SELECT 1") + + mock_snapshot.execute_sql.assert_called_once() + _, kwargs = mock_snapshot.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 4.0) + + def test_cursor_autocommit_dml_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection.autocommit = True + cursor = connection.cursor() + + mock_transaction = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_transaction.execute_sql.return_value = mock_result_set + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertNotIn("timeout", kwargs) + self.assertTrue(kwargs.get("last_statement")) + + def test_cursor_autocommit_dml_uses_cursor_timeout(self): + connection = self._make_connection(timeout=18.0) + connection.autocommit = True + cursor = connection.cursor() + cursor.timeout = 6.0 + + mock_transaction = mock.MagicMock() + mock_result_set = mock.MagicMock() + mock_transaction.execute_sql.return_value = mock_result_set + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 6.0) + self.assertTrue(kwargs.get("last_statement")) + + def test_cursor_rw_transaction_execute_passes_timeout(self): + connection = self._make_connection(timeout=14.0) + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.execute_sql.return_value = iter([]) + connection._transaction = mock_transaction + + cursor = connection.cursor() + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 14.0) + + def test_cursor_rw_transaction_execute_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.execute_sql.return_value = iter([]) + connection._transaction = mock_transaction + + cursor = connection.cursor() + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertNotIn("timeout", kwargs) + + def test_cursor_rw_transaction_execute_uses_cursor_timeout(self): + connection = self._make_connection(timeout=14.0) + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.execute_sql.return_value = iter([]) + connection._transaction = mock_transaction + + cursor = connection.cursor() + cursor.timeout = 3.0 + sql = "UPDATE t SET col = 1" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.execute(sql) + + mock_transaction.execute_sql.assert_called_once() + _, kwargs = mock_transaction.execute_sql.call_args + self.assertEqual(kwargs.get("timeout"), 3.0) + + def test_run_sql_in_snapshot_passes_timeout(self): + connection = self._make_connection(timeout=20.0) + cursor = connection.cursor() + + mock_snapshot = mock.MagicMock() + mock_snapshot.execute_sql.return_value = [["row1"]] + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + res = cursor.run_sql_in_snapshot("SELECT 1") + self.assertEqual(res, [["row1"]]) + mock_snapshot.execute_sql.assert_called_once_with( + "SELECT 1", None, None, timeout=20.0 + ) + + def test_run_sql_in_snapshot_without_timeout_omits_kwarg(self): + connection = self._make_connection() + cursor = connection.cursor() + + mock_snapshot = mock.MagicMock() + mock_snapshot.execute_sql.return_value = [["row1"]] + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + res = cursor.run_sql_in_snapshot("SELECT 1") + self.assertEqual(res, [["row1"]]) + mock_snapshot.execute_sql.assert_called_once_with("SELECT 1", None, None) + + def test_run_sql_in_snapshot_uses_cursor_timeout(self): + connection = self._make_connection(timeout=20.0) + cursor = connection.cursor() + cursor.timeout = 8.0 + + mock_snapshot = mock.MagicMock() + mock_snapshot.execute_sql.return_value = [["row1"]] + mock_ctx = mock.MagicMock() + mock_ctx.__enter__.return_value = mock_snapshot + connection.database.snapshot = mock.MagicMock(return_value=mock_ctx) + + res = cursor.run_sql_in_snapshot("SELECT 1") + self.assertEqual(res, [["row1"]]) + mock_snapshot.execute_sql.assert_called_once_with( + "SELECT 1", None, None, timeout=8.0 + ) + + def test_batch_dml_autocommit_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection.autocommit = True + cursor = connection.cursor() + + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertNotIn("timeout", kwargs) + self.assertTrue(kwargs.get("last_statement")) + + def test_batch_dml_autocommit_uses_cursor_timeout(self): + connection = self._make_connection(timeout=7.0) + connection.autocommit = True + cursor = connection.cursor() + cursor.timeout = 2.0 + + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + + def fake_run_in_transaction(func, *args, **kwargs): + return func(mock_transaction, *args, **kwargs) + + connection.database.run_in_transaction = fake_run_in_transaction + + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertEqual(kwargs.get("timeout"), 2.0) + self.assertTrue(kwargs.get("last_statement")) + + def test_batch_dml_in_rw_transaction_passes_timeout(self): + connection = self._make_connection(timeout=11.0) + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + connection._transaction = mock_transaction + + cursor = connection.cursor() + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertEqual(kwargs.get("timeout"), 11.0) + + def test_batch_dml_in_rw_transaction_without_timeout_omits_kwarg(self): + connection = self._make_connection() + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + connection._transaction = mock_transaction + + cursor = connection.cursor() + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertNotIn("timeout", kwargs) + + def test_batch_dml_in_rw_transaction_uses_cursor_timeout(self): + connection = self._make_connection(timeout=11.0) + connection.autocommit = False + connection._spanner_transaction_started = True + mock_transaction = mock.MagicMock() + mock_transaction.batch_update.return_value = [mock.MagicMock(code=OK), [1, 1]] + connection._transaction = mock_transaction + + cursor = connection.cursor() + cursor.timeout = 3.5 + sql = "UPDATE t SET col = %s" + with mock.patch( + "google.cloud.spanner_dbapi.parse_utils.classify_statement", + return_value=ParsedStatement(StatementType.UPDATE, Statement(sql)), + ): + cursor.executemany(sql, [(1,), (2,)]) + + mock_transaction.batch_update.assert_called_once() + _, kwargs = mock_transaction.batch_update.call_args + self.assertEqual(kwargs.get("timeout"), 3.5)