From cff1650f7191ba30f68d9c945b6582afcd8394e3 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 19 Aug 2026 11:50:24 +0000 Subject: [PATCH 1/3] feat(spanner): add DataBoost support to DBAPI driver and client-side statements Add support for Cloud Spanner DataBoost in the Python DBAPI (PEP 249) driver (`google.cloud.spanner_dbapi`) for partitioned queries. Key updates: * Add `data_boost_enabled: bool = False` argument to `spanner_dbapi.connect()` and `Connection.__init__()`. * Add `@property def data_boost_enabled(self)` getter and setter on `Connection`. * Add client-side statement parsing and execution for: - `SET DATA_BOOST_ENABLED = TRUE|FALSE` - `SHOW VARIABLE DATA_BOOST_ENABLED` (returns column `DATA_BOOST_ENABLED` with `BOOL` type) * Forward `data_boost_enabled` in `Connection.partition_query()` and `Connection.run_partitioned_query()`. * Add unit and gRPC mock server test coverage across DBAPI connection, cursor, parser, statement executor, and mock server test suites. --- .../client_side_statement_executor.py | 24 ++- .../client_side_statement_parser.py | 12 ++ .../google/cloud/spanner_dbapi/connection.py | 50 ++++- .../cloud/spanner_dbapi/parsed_statement.py | 2 + .../mockserver_tests/test_dbapi_databoost.py | 194 ++++++++++++++++++ .../test_client_side_statement_executor.py | 48 +++++ .../tests/unit/spanner_dbapi/test_connect.py | 13 ++ .../unit/spanner_dbapi/test_connection.py | 50 +++++ .../unit/spanner_dbapi/test_parse_utils.py | 36 ++++ 9 files changed, 425 insertions(+), 4 deletions(-) create mode 100644 packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py index 45302cb7a01a..3334bbfaac83 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py @@ -15,8 +15,9 @@ from google.cloud.spanner_v1 import TransactionOptions +from google.cloud.spanner_dbapi.exceptions import ProgrammingError + if TYPE_CHECKING: - from google.cloud.spanner_dbapi import ProgrammingError from google.cloud.spanner_dbapi.cursor import Cursor from google.cloud.spanner_dbapi.parsed_statement import ( @@ -108,6 +109,27 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement): return connection.run_partitioned_query(parsed_statement) if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE: return connection._set_autocommit_dml_mode(parsed_statement) + if statement_type == ClientSideStatementType.SET_DATA_BOOST_ENABLED: + val_str = ( + parsed_statement.client_side_statement_params[0] + .strip() + .strip("'\"") + .lower() + ) + if val_str not in ("true", "false"): + raise ProgrammingError( + f"Invalid value for DATA_BOOST_ENABLED: '{parsed_statement.client_side_statement_params[0]}'. Expected TRUE or FALSE." + ) + connection.data_boost_enabled = val_str == "true" + return None + if statement_type == ClientSideStatementType.SHOW_DATA_BOOST_ENABLED: + column_values.append(connection.data_boost_enabled) + return _get_streamed_result_set( + "DATA_BOOST_ENABLED", + TypeCode.BOOL, + column_values, + ) + return None def _get_streamed_result_set(column_name, type_code, column_values): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py index 51dfdb63ad62..a1f09b69e823 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py @@ -44,6 +44,12 @@ RE_SET_AUTOCOMMIT_DML_MODE = re.compile( r"^\s*(SET)\s+(AUTOCOMMIT_DML_MODE)\s+(=)\s+(.+)", re.IGNORECASE ) +RE_SET_DATA_BOOST_ENABLED = re.compile( + r"^\s*(SET)\s+(DATA_BOOST_ENABLED)\s+(=)\s+(.+)", re.IGNORECASE +) +RE_SHOW_DATA_BOOST_ENABLED = re.compile( + r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*$", re.IGNORECASE +) def parse_stmt(query): @@ -68,6 +74,8 @@ def parse_stmt(query): client_side_statement_type = ClientSideStatementType.SHOW_COMMIT_TIMESTAMP elif RE_SHOW_READ_TIMESTAMP.match(query): client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP + elif RE_SHOW_DATA_BOOST_ENABLED.match(query): + client_side_statement_type = ClientSideStatementType.SHOW_DATA_BOOST_ENABLED elif RE_START_BATCH_DML.match(query): client_side_statement_type = ClientSideStatementType.START_BATCH_DML elif RE_BEGIN.match(query): @@ -96,6 +104,10 @@ def parse_stmt(query): match = re.search(RE_SET_AUTOCOMMIT_DML_MODE, query) client_side_statement_params.append(match.group(4)) client_side_statement_type = ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE + elif RE_SET_DATA_BOOST_ENABLED.match(query): + match = re.search(RE_SET_DATA_BOOST_ENABLED, query) + client_side_statement_params.append(match.group(4)) + client_side_statement_type = ClientSideStatementType.SET_DATA_BOOST_ENABLED if client_side_statement_type is not None: return ParsedStatement( StatementType.CLIENT_SIDE, 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 eb0ef4217411..7fc0b7d575af 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py @@ -94,7 +94,14 @@ class Connection: **kwargs: Initial value for connection variables. """ - def __init__(self, instance, database=None, read_only=False, **kwargs): + def __init__( + self, + instance, + database=None, + read_only=False, + data_boost_enabled=False, + **kwargs, + ): self._instance = instance self._database = database self._ddl_statements = [] @@ -110,6 +117,7 @@ def __init__(self, instance, database=None, read_only=False, **kwargs): # connection close self._own_pool = True self._read_only = read_only + self._data_boost_enabled = bool(data_boost_enabled) self._staleness = None self.request_priority = None self._transaction_begin_marked = False @@ -123,6 +131,26 @@ def __init__(self, instance, database=None, read_only=False, **kwargs): self._autocommit_dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL self._connection_variables = kwargs + @property + def data_boost_enabled(self): + """Flag: whether DataBoost is enabled for partitioned queries on this connection. + + Note that DataBoost is only supported for partitioned query execution. + + Returns: + bool: True if DataBoost is enabled, False otherwise. + """ + return self._data_boost_enabled + + @data_boost_enabled.setter + def data_boost_enabled(self, value): + """Change the DataBoost enablement state for partitioned queries on this connection. + + :type value: bool + :param value: New data_boost_enabled state. + """ + self._data_boost_enabled = bool(value) + @property def spanner_client(self): """Client for interacting with Cloud Spanner API. This property exposes @@ -638,11 +666,15 @@ def partition_query( self, parsed_statement: ParsedStatement, query_options=None, + data_boost_enabled=None, ): statement = parsed_statement.statement partitioned_query = parsed_statement.client_side_statement_params[0] self._partitioned_query_validation(partitioned_query, statement) + if data_boost_enabled is None: + data_boost_enabled = self.data_boost_enabled + batch_snapshot = self._database.batch_snapshot() partition_ids = [] partitions = list( @@ -651,6 +683,7 @@ def partition_query( statement.params, statement.param_types, query_options=query_options, + data_boost_enabled=data_boost_enabled, ) ) @@ -684,7 +717,10 @@ def run_partitioned_query( self._partitioned_query_validation(partitioned_query, statement) batch_snapshot = self._database.batch_snapshot() return batch_snapshot.run_partitioned_query( - partitioned_query, statement.params, statement.param_types + partitioned_query, + statement.params, + statement.param_types, + data_boost_enabled=self.data_boost_enabled, ) @check_not_closed @@ -748,6 +784,7 @@ def connect( client_certificate=None, client_key=None, instance_type=None, + data_boost_enabled=False, **kwargs, ): """Creates a connection to a Google Cloud Spanner database. @@ -795,6 +832,11 @@ def connect( :param database_role: (Optional) The database role to connect as when using fine-grained access controls. + :type data_boost_enabled: bool + :param data_boost_enabled: (Optional) Whether to enable DataBoost for + partitioned queries executed via this connection. Defaults to False. + Note that DataBoost is only supported for partitioned query execution. + **kwargs: Initial value for connection variables. @@ -909,7 +951,9 @@ def connect( database = instance.database( database_id, pool=pool, database_role=database_role, logger=logger ) - conn = Connection(instance, database, **kwargs) + conn = Connection( + instance, database, data_boost_enabled=data_boost_enabled, **kwargs + ) if pool is not None: conn._own_pool = False diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py index a8d03f6fa410..6a5eee1ba068 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py @@ -38,6 +38,8 @@ class ClientSideStatementType(Enum): RUN_PARTITION = 10 RUN_PARTITIONED_QUERY = 11 SET_AUTOCOMMIT_DML_MODE = 12 + SET_DATA_BOOST_ENABLED = 13 + SHOW_DATA_BOOST_ENABLED = 14 class AutocommitDmlMode(Enum): diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py new file mode 100644 index 000000000000..8ca5a6e29f77 --- /dev/null +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py @@ -0,0 +1,194 @@ +# 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. + +from google.cloud.spanner_dbapi import Connection +from google.cloud.spanner_dbapi.parsed_statement import ( + ClientSideStatementType, + ParsedStatement, + Statement, + StatementType, +) +from google.cloud.spanner_v1 import ( + ExecuteSqlRequest, + TypeCode, +) +from google.cloud.spanner_v1.types import spanner as spanner_types +from tests.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_single_result, +) + + +class TestDbapiDataBoost(MockServerTestBase): + def setUp(self): + super().setUp() + add_single_result( + "select name from singers", "name", TypeCode.STRING, [("Some Singer",)] + ) + + def test_select_with_data_boost_enabled_autocommit(self): + # Non-partitioned queries should NOT have data_boost_enabled=True on ExecuteSqlRequest, + # because Spanner rejects data_boost_enabled without a partition_token. + connection = Connection(self.instance, self.database, data_boost_enabled=True) + connection.autocommit = True + with connection.cursor() as cursor: + cursor.execute("select name from singers") + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + self.assertEqual("Some Singer", rows[0][0]) + + requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(requests)) + self.assertFalse(requests[0].data_boost_enabled) + + def test_select_with_data_boost_enabled_read_only(self): + connection = Connection( + self.instance, self.database, read_only=True, data_boost_enabled=True + ) + with connection.cursor() as cursor: + cursor.execute("select name from singers") + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + self.assertEqual("Some Singer", rows[0][0]) + + requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(requests)) + self.assertFalse(requests[0].data_boost_enabled) + + def test_select_with_data_boost_disabled_by_default(self): + connection = Connection(self.instance, self.database) + connection.autocommit = True + with connection.cursor() as cursor: + cursor.execute("select name from singers") + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + + requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(requests)) + self.assertFalse(requests[0].data_boost_enabled) + + def test_select_with_set_data_boost_statement(self): + connection = Connection(self.instance, self.database, data_boost_enabled=False) + connection.autocommit = True + with connection.cursor() as cursor: + # Enable DataBoost via SQL client-side statement + cursor.execute("SET DATA_BOOST_ENABLED = TRUE") + self.assertTrue(cursor.connection.data_boost_enabled) + cursor.execute("select name from singers") + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + + requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(requests)) + # Non-partitioned query should still be False + self.assertFalse(requests[0].data_boost_enabled) + + def test_partition_query_and_run_partition_with_data_boost_enabled(self): + sql = "SELECT name FROM users WHERE active = true" + + partition_response = spanner_types.PartitionResponse() + partition_response.partitions.extend( + [ + spanner_types.Partition(partition_token=b"mock-token-1"), + ] + ) + self.spanner_service.mock_spanner.add_partition_result(sql, partition_response) + add_single_result(sql, "name", TypeCode.STRING, [("Alice",)]) + + connection = Connection( + self.instance, self.database, read_only=True, data_boost_enabled=True + ) + + parsed = ParsedStatement( + statement_type=StatementType.CLIENT_SIDE, + statement=Statement(sql), + client_side_statement_type=ClientSideStatementType.PARTITION_QUERY, + client_side_statement_params=[sql], + ) + + partition_ids = connection.partition_query(parsed) + self.assertEqual(1, len(partition_ids)) + + # Execute the partition and verify the ExecuteSqlRequest has data_boost_enabled=True + result_stream = connection.run_partition(partition_ids[0]) + rows = list(result_stream) + self.assertEqual(1, len(rows)) + self.assertEqual("Alice", rows[0][0]) + + execute_requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(execute_requests)) + self.assertTrue(execute_requests[0].data_boost_enabled) + self.assertEqual(b"mock-token-1", execute_requests[0].partition_token) + + def test_run_partitioned_query_with_data_boost_enabled(self): + sql = "SELECT name FROM users WHERE active = true" + + partition_response = spanner_types.PartitionResponse() + partition_response.partitions.extend( + [ + spanner_types.Partition(partition_token=b"mock-token-1"), + ] + ) + self.spanner_service.mock_spanner.add_partition_result(sql, partition_response) + add_single_result(sql, "name", TypeCode.STRING, [("Alice",)]) + + connection = Connection( + self.instance, self.database, read_only=True, data_boost_enabled=True + ) + + parsed = ParsedStatement( + statement_type=StatementType.CLIENT_SIDE, + statement=Statement(sql), + client_side_statement_type=ClientSideStatementType.RUN_PARTITIONED_QUERY, + client_side_statement_params=[sql], + ) + + result_set = connection.run_partitioned_query(parsed) + rows = list(result_set) + self.assertEqual(1, len(rows)) + self.assertEqual("Alice", rows[0][0]) + + execute_requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(execute_requests)) + self.assertTrue(execute_requests[0].data_boost_enabled) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py index 888f81e830f7..57e8ac1b39b3 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py @@ -52,3 +52,51 @@ def test_get_isolation_level(self): ) ), ) + + +class TestClientSideStatementExecutor(unittest.TestCase): + def test_execute_set_data_boost_enabled(self): + from unittest import mock + from google.cloud.spanner_dbapi.client_side_statement_executor import execute + from google.cloud.spanner_dbapi.exceptions import ProgrammingError + + cursor = mock.MagicMock() + cursor.connection.is_closed = False + cursor.connection.data_boost_enabled = False + + stmt = classify_statement("SET DATA_BOOST_ENABLED = TRUE") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertTrue(cursor.connection.data_boost_enabled) + + stmt = classify_statement("SET DATA_BOOST_ENABLED = FALSE") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertFalse(cursor.connection.data_boost_enabled) + + stmt = classify_statement("SET DATA_BOOST_ENABLED = INVALID") + with self.assertRaises(ProgrammingError): + execute(cursor, stmt) + + def test_execute_show_data_boost_enabled(self): + from unittest import mock + from google.cloud.spanner_dbapi.client_side_statement_executor import execute + from google.cloud.spanner_v1 import TypeCode + + cursor = mock.MagicMock() + cursor.connection.is_closed = False + cursor.connection.data_boost_enabled = True + + stmt = classify_statement("SHOW VARIABLE DATA_BOOST_ENABLED") + res = execute(cursor, stmt) + rows = list(res) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0][0], True) + self.assertEqual(res.fields[0].name, "DATA_BOOST_ENABLED") + self.assertEqual(res.fields[0].type_.code, TypeCode.BOOL) + + cursor.connection.data_boost_enabled = False + res = execute(cursor, stmt) + rows = list(res) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0][0], False) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py index f80c3c3e52b5..c6b51c322b51 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py @@ -157,3 +157,16 @@ def test_with_kwargs(self, mock_client): self.assertIsInstance(connection, Connection) self.assertTrue(connection._ignore_transaction_warnings) + + def test_w_data_boost_enabled(self, mock_client): + from google.cloud.spanner_dbapi import Connection, connect + + client = mock_client.return_value + instance = client.instance.return_value + database = instance.database.return_value + self.assertIsNotNone(database) + + connection = connect(INSTANCE, DATABASE, data_boost_enabled=True) + + self.assertIsInstance(connection, Connection) + self.assertTrue(connection.data_boost_enabled) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py index 3d47f57fb4fb..abae1803a578 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py @@ -145,6 +145,16 @@ def test_read_only_connection(self): connection.read_only = False self.assertFalse(connection.read_only) + def test_property_data_boost_enabled(self): + connection = self._make_connection() + self.assertFalse(connection.data_boost_enabled) + + connection.data_boost_enabled = True + self.assertTrue(connection.data_boost_enabled) + + connection.data_boost_enabled = False + self.assertFalse(connection.data_boost_enabled) + def test__session_checkout_read_only(self): connection = build_connection(read_only=True) database = connection._database @@ -924,6 +934,46 @@ def test_connect_w_invalid_instance_type_raises_error(self): "instance_type must be one of 'cloud' or 'omni'", str(ctx.exception) ) + def test_partition_query_data_boost_enabled(self): + from google.cloud.spanner_dbapi.parse_utils import classify_statement + + connection = self._make_connection(read_only=True, data_boost_enabled=True) + batch_snapshot = mock.MagicMock() + batch_snapshot.generate_query_batches.return_value = [] + batch_snapshot.get_batch_transaction_id.return_value = mock.MagicMock( + transaction_id=b"tx", session_id="sess", read_timestamp=None + ) + connection.database.batch_snapshot = mock.MagicMock(return_value=batch_snapshot) + + parsed_stmt = classify_statement("PARTITION SELECT 1") + res = connection.partition_query(parsed_stmt) + self.assertEqual(res, []) + batch_snapshot.generate_query_batches.assert_called_once_with( + "SELECT 1", + None, + None, + query_options=None, + data_boost_enabled=True, + ) + + def test_run_partitioned_query_data_boost_enabled(self): + from google.cloud.spanner_dbapi.parse_utils import classify_statement + + connection = self._make_connection(read_only=True, data_boost_enabled=True) + batch_snapshot = mock.MagicMock() + batch_snapshot.run_partitioned_query.return_value = "merged_result_set" + connection.database.batch_snapshot = mock.MagicMock(return_value=batch_snapshot) + + parsed_stmt = classify_statement("RUN PARTITIONED QUERY SELECT 1") + res = connection.run_partitioned_query(parsed_stmt) + self.assertEqual(res, "merged_result_set") + batch_snapshot.run_partitioned_query.assert_called_once_with( + "SELECT 1", + None, + None, + data_boost_enabled=True, + ) + def exit_ctx_func(self, exc_type, exc_value, traceback): """Context __exit__ method mock.""" diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py index 29a08edb4264..46a847c9f47d 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py @@ -89,6 +89,8 @@ def test_classify_stmt(self): (" rollback TRANSACTION ", StatementType.CLIENT_SIDE), (" SHOW VARIABLE COMMIT_TIMESTAMP ", StatementType.CLIENT_SIDE), ("SHOW VARIABLE READ_TIMESTAMP", StatementType.CLIENT_SIDE), + ("SET DATA_BOOST_ENABLED = TRUE", StatementType.CLIENT_SIDE), + ("SHOW VARIABLE DATA_BOOST_ENABLED", StatementType.CLIENT_SIDE), ("GRANT SELECT ON TABLE Singers TO ROLE parent", StatementType.DDL), ("REVOKE SELECT ON TABLE Singers TO ROLE parent", StatementType.DDL), ("GRANT ROLE parent TO ROLE child", StatementType.DDL), @@ -250,6 +252,40 @@ def test_set_autocommit_dml_mode_stmt(self): ), ) + def test_set_data_boost_enabled_stmt(self): + parsed_statement = classify_statement(" set data_boost_enabled = true ") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("set data_boost_enabled = true"), + ClientSideStatementType.SET_DATA_BOOST_ENABLED, + ["true"], + ), + ) + parsed_statement = classify_statement("SET DATA_BOOST_ENABLED = FALSE") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("SET DATA_BOOST_ENABLED = FALSE"), + ClientSideStatementType.SET_DATA_BOOST_ENABLED, + ["FALSE"], + ), + ) + + def test_show_data_boost_enabled_stmt(self): + parsed_statement = classify_statement(" show variable data_boost_enabled ") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("show variable data_boost_enabled"), + ClientSideStatementType.SHOW_DATA_BOOST_ENABLED, + [], + ), + ) + @unittest.skipIf(skip_condition, skip_message) def test_sql_pyformat_args_to_spanner(self): from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner From 9c53e3f9a8196872383d58709168bbd662a79a07 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 19 Aug 2026 15:42:16 +0000 Subject: [PATCH 2/3] feat(spanner): add auto_partition_mode support to DBAPI driver and client-side statements Add support for `auto_partition_mode` in the Python DBAPI driver (`google.cloud.spanner_dbapi`). Key updates: * Add `auto_partition_mode: bool = False` parameter to `spanner_dbapi.connect()` and `Connection.__init__()`. * Add `@property def auto_partition_mode(self)` getter and setter on `Connection`. * Add client-side statement parsing and execution for: - `SET AUTO_PARTITION_MODE = TRUE|FALSE` - `SHOW VARIABLE AUTO_PARTITION_MODE` (returns column `AUTO_PARTITION_MODE` with `BOOL` type) * Automatically route queries in `Cursor._execute()` to `run_partitioned_query` when `auto_partition_mode=True`. * Add unit and gRPC mock server test coverage for automatic query partitioning. --- .../client_side_statement_executor.py | 41 ++++++++--- .../client_side_statement_parser.py | 12 +++ .../google/cloud/spanner_dbapi/connection.py | 37 +++++++++- .../google/cloud/spanner_dbapi/cursor.py | 27 ++++++- .../cloud/spanner_dbapi/parsed_statement.py | 2 + .../mockserver_tests/test_dbapi_databoost.py | 73 +++++++++++++++++++ .../test_client_side_statement_executor.py | 50 +++++++++++++ .../tests/unit/spanner_dbapi/test_connect.py | 13 ++++ .../unit/spanner_dbapi/test_connection.py | 12 ++- .../tests/unit/spanner_dbapi/test_cursor.py | 54 ++++++++++++-- .../unit/spanner_dbapi/test_parse_utils.py | 37 ++++++++++ 11 files changed, 334 insertions(+), 24 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py index 3334bbfaac83..6d4258745a5b 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py @@ -13,9 +13,8 @@ # limitations under the License. from typing import TYPE_CHECKING, Union -from google.cloud.spanner_v1 import TransactionOptions - from google.cloud.spanner_dbapi.exceptions import ProgrammingError +from google.cloud.spanner_v1 import TransactionOptions if TYPE_CHECKING: from google.cloud.spanner_dbapi.cursor import Cursor @@ -110,17 +109,10 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement): if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE: return connection._set_autocommit_dml_mode(parsed_statement) if statement_type == ClientSideStatementType.SET_DATA_BOOST_ENABLED: - val_str = ( - parsed_statement.client_side_statement_params[0] - .strip() - .strip("'\"") - .lower() + connection.data_boost_enabled = _parse_bool( + parsed_statement.client_side_statement_params[0], + "DATA_BOOST_ENABLED", ) - if val_str not in ("true", "false"): - raise ProgrammingError( - f"Invalid value for DATA_BOOST_ENABLED: '{parsed_statement.client_side_statement_params[0]}'. Expected TRUE or FALSE." - ) - connection.data_boost_enabled = val_str == "true" return None if statement_type == ClientSideStatementType.SHOW_DATA_BOOST_ENABLED: column_values.append(connection.data_boost_enabled) @@ -129,9 +121,34 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement): TypeCode.BOOL, column_values, ) + if statement_type == ClientSideStatementType.SET_AUTO_PARTITION_MODE: + connection.auto_partition_mode = _parse_bool( + parsed_statement.client_side_statement_params[0], + "AUTO_PARTITION_MODE", + ) + return None + if statement_type == ClientSideStatementType.SHOW_AUTO_PARTITION_MODE: + column_values.append(connection.auto_partition_mode) + return _get_streamed_result_set( + "AUTO_PARTITION_MODE", + TypeCode.BOOL, + column_values, + ) return None +_BOOL_MAP = {"true": True, "false": False} + + +def _parse_bool(raw_val: str, var_name: str) -> bool: + cleaned = raw_val.strip().strip("'\"").lower() + if cleaned not in _BOOL_MAP: + raise ProgrammingError( + f"Invalid value for {var_name}: '{raw_val}'. Expected TRUE or FALSE." + ) + return _BOOL_MAP[cleaned] + + def _get_streamed_result_set(column_name, type_code, column_values): struct_type_pb = StructType( fields=[StructType.Field(name=column_name, type_=Type(code=type_code))] diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py index a1f09b69e823..963263eac2b4 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py @@ -50,6 +50,12 @@ RE_SHOW_DATA_BOOST_ENABLED = re.compile( r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*$", re.IGNORECASE ) +RE_SET_AUTO_PARTITION_MODE = re.compile( + r"^\s*(SET)\s+(AUTO_PARTITION_MODE)\s+(=)\s+(.+)", re.IGNORECASE +) +RE_SHOW_AUTO_PARTITION_MODE = re.compile( + r"^\s*(SHOW)\s+(VARIABLE)\s+(AUTO_PARTITION_MODE)\s*$", re.IGNORECASE +) def parse_stmt(query): @@ -76,6 +82,8 @@ def parse_stmt(query): client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP elif RE_SHOW_DATA_BOOST_ENABLED.match(query): client_side_statement_type = ClientSideStatementType.SHOW_DATA_BOOST_ENABLED + elif RE_SHOW_AUTO_PARTITION_MODE.match(query): + client_side_statement_type = ClientSideStatementType.SHOW_AUTO_PARTITION_MODE elif RE_START_BATCH_DML.match(query): client_side_statement_type = ClientSideStatementType.START_BATCH_DML elif RE_BEGIN.match(query): @@ -108,6 +116,10 @@ def parse_stmt(query): match = re.search(RE_SET_DATA_BOOST_ENABLED, query) client_side_statement_params.append(match.group(4)) client_side_statement_type = ClientSideStatementType.SET_DATA_BOOST_ENABLED + elif RE_SET_AUTO_PARTITION_MODE.match(query): + match = re.search(RE_SET_AUTO_PARTITION_MODE, query) + client_side_statement_params.append(match.group(4)) + client_side_statement_type = ClientSideStatementType.SET_AUTO_PARTITION_MODE if client_side_statement_type is not None: return ParsedStatement( StatementType.CLIENT_SIDE, 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 7fc0b7d575af..19d975a13453 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py @@ -20,7 +20,6 @@ from google.api_core.exceptions import Aborted from google.api_core.gapic_v1.client_info import ClientInfo from google.auth.credentials import AnonymousCredentials - from google.cloud import spanner_v1 as spanner from google.cloud.spanner_dbapi import partition_helper from google.cloud.spanner_dbapi.batch_dml_executor import BatchDmlExecutor, BatchMode @@ -100,6 +99,7 @@ def __init__( database=None, read_only=False, data_boost_enabled=False, + auto_partition_mode=False, **kwargs, ): self._instance = instance @@ -118,6 +118,7 @@ def __init__( self._own_pool = True self._read_only = read_only self._data_boost_enabled = bool(data_boost_enabled) + self._auto_partition_mode = bool(auto_partition_mode) self._staleness = None self.request_priority = None self._transaction_begin_marked = False @@ -151,6 +152,27 @@ def data_boost_enabled(self, value): """ self._data_boost_enabled = bool(value) + @property + def auto_partition_mode(self): + """Flag: whether auto partition mode is enabled for queries on this connection. + + When enabled, standard queries executed on read-only or autocommit connections + are automatically partitioned and executed in parallel via run_partitioned_query. + + Returns: + bool: True if auto partition mode is enabled, False otherwise. + """ + return self._auto_partition_mode + + @auto_partition_mode.setter + def auto_partition_mode(self, value): + """Change the auto partition mode enablement state for this connection. + + :type value: bool + :param value: New auto_partition_mode state. + """ + self._auto_partition_mode = bool(value) + @property def spanner_client(self): """Client for interacting with Cloud Spanner API. This property exposes @@ -785,6 +807,7 @@ def connect( client_key=None, instance_type=None, data_boost_enabled=False, + auto_partition_mode=False, **kwargs, ): """Creates a connection to a Google Cloud Spanner database. @@ -837,6 +860,12 @@ def connect( partitioned queries executed via this connection. Defaults to False. Note that DataBoost is only supported for partitioned query execution. + :type auto_partition_mode: bool + :param auto_partition_mode: (Optional) Whether to enable auto partition mode + for queries executed via this connection. When True, queries on read-only + or autocommit connections are automatically partitioned and executed in parallel. + Defaults to False. + **kwargs: Initial value for connection variables. @@ -952,7 +981,11 @@ def connect( database_id, pool=pool, database_role=database_role, logger=logger ) conn = Connection( - instance, database, data_boost_enabled=data_boost_enabled, **kwargs + instance, + database, + data_boost_enabled=data_boost_enabled, + auto_partition_mode=auto_partition_mode, + **kwargs, ) if pool is not None: conn._own_pool = False 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 f090a8a7e051..ceb26b823f1f 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py @@ -25,7 +25,6 @@ InvalidArgument, OutOfRange, ) - from google.cloud import spanner_v1 as spanner from google.cloud.spanner_dbapi import ( _helpers, @@ -308,6 +307,11 @@ def _execute(self, sql, args=None, call_from_execute_many=False): self._itr = PeekIterator(self._result_set) elif self.connection._batch_mode == BatchMode.DML: self.connection.execute_batch_dml_statement(self._parsed_statement) + elif ( + self._parsed_statement.statement_type == StatementType.QUERY + and self.connection.auto_partition_mode + ): + self._handle_auto_partition_query(sql, args or None) elif self.connection.read_only or ( not self.connection._client_transaction_started and self._parsed_statement.statement_type == StatementType.QUERY @@ -581,6 +585,27 @@ def _handle_DQL(self, sql, params): self.connection._transaction = None self._handle_DQL_with_snapshot(snapshot, sql, params) + def _handle_auto_partition_query(self, sql, params): + if self.connection.database is None: + raise ValueError("Database needs to be passed for this operation") + if ( + not self.connection.read_only + and self.connection._client_transaction_started + ): + raise ProgrammingError( + "Partitioned query is not supported, because the connection is in a read/write transaction." + ) + sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params) + batch_snapshot = self.connection.database.batch_snapshot() + self._result_set = batch_snapshot.run_partitioned_query( + sql, + params=params, + param_types=get_param_types(params), + data_boost_enabled=self.connection.data_boost_enabled, + ) + self._itr = self._result_set + self._row_count = None + def __enter__(self): return self diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py index 6a5eee1ba068..3a2defb76bd3 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py @@ -40,6 +40,8 @@ class ClientSideStatementType(Enum): SET_AUTOCOMMIT_DML_MODE = 12 SET_DATA_BOOST_ENABLED = 13 SHOW_DATA_BOOST_ENABLED = 14 + SET_AUTO_PARTITION_MODE = 15 + SHOW_AUTO_PARTITION_MODE = 16 class AutocommitDmlMode(Enum): diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py index 8ca5a6e29f77..4c234d47dc95 100644 --- a/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py @@ -24,6 +24,7 @@ TypeCode, ) from google.cloud.spanner_v1.types import spanner as spanner_types + from tests.mockserver_tests.mock_server_test_base import ( MockServerTestBase, add_single_result, @@ -192,3 +193,75 @@ def test_run_partitioned_query_with_data_boost_enabled(self): ) self.assertEqual(1, len(execute_requests)) self.assertTrue(execute_requests[0].data_boost_enabled) + + def test_auto_partition_mode_with_data_boost_enabled(self): + sql = "SELECT name FROM users WHERE active = true" + + partition_response = spanner_types.PartitionResponse() + partition_response.partitions.extend( + [ + spanner_types.Partition(partition_token=b"mock-token-auto-1"), + ] + ) + self.spanner_service.mock_spanner.add_partition_result(sql, partition_response) + add_single_result(sql, "name", TypeCode.STRING, [("Alice",)]) + + connection = Connection( + self.instance, + self.database, + read_only=True, + auto_partition_mode=True, + data_boost_enabled=True, + ) + + with connection.cursor() as cursor: + # Plain cursor.execute automatically runs as partitioned query with DataBoost + cursor.execute(sql) + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + self.assertEqual("Alice", rows[0][0]) + + execute_requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(execute_requests)) + self.assertTrue(execute_requests[0].data_boost_enabled) + self.assertEqual(b"mock-token-auto-1", execute_requests[0].partition_token) + + def test_auto_partition_mode_via_statement(self): + sql = "SELECT name FROM users WHERE active = true" + + partition_response = spanner_types.PartitionResponse() + partition_response.partitions.extend( + [ + spanner_types.Partition(partition_token=b"mock-token-auto-2"), + ] + ) + self.spanner_service.mock_spanner.add_partition_result(sql, partition_response) + add_single_result(sql, "name", TypeCode.STRING, [("Bob",)]) + + connection = Connection(self.instance, self.database, read_only=True) + + with connection.cursor() as cursor: + cursor.execute("SET AUTO_PARTITION_MODE = TRUE") + cursor.execute("SET DATA_BOOST_ENABLED = TRUE") + self.assertTrue(cursor.connection.auto_partition_mode) + self.assertTrue(cursor.connection.data_boost_enabled) + + cursor.execute(sql) + rows = cursor.fetchall() + self.assertEqual(1, len(rows)) + self.assertEqual("Bob", rows[0][0]) + + execute_requests = list( + filter( + lambda msg: isinstance(msg, ExecuteSqlRequest), + self.spanner_service.requests, + ) + ) + self.assertEqual(1, len(execute_requests)) + self.assertTrue(execute_requests[0].data_boost_enabled) + self.assertEqual(b"mock-token-auto-2", execute_requests[0].partition_token) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py index 57e8ac1b39b3..084e1fb90ebb 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py @@ -57,6 +57,7 @@ def test_get_isolation_level(self): class TestClientSideStatementExecutor(unittest.TestCase): def test_execute_set_data_boost_enabled(self): from unittest import mock + from google.cloud.spanner_dbapi.client_side_statement_executor import execute from google.cloud.spanner_dbapi.exceptions import ProgrammingError @@ -80,6 +81,7 @@ def test_execute_set_data_boost_enabled(self): def test_execute_show_data_boost_enabled(self): from unittest import mock + from google.cloud.spanner_dbapi.client_side_statement_executor import execute from google.cloud.spanner_v1 import TypeCode @@ -100,3 +102,51 @@ def test_execute_show_data_boost_enabled(self): rows = list(res) self.assertEqual(len(rows), 1) self.assertEqual(rows[0][0], False) + + def test_execute_set_auto_partition_mode(self): + from unittest import mock + + from google.cloud.spanner_dbapi.client_side_statement_executor import execute + from google.cloud.spanner_dbapi.exceptions import ProgrammingError + + cursor = mock.MagicMock() + cursor.connection.is_closed = False + cursor.connection.auto_partition_mode = False + + stmt = classify_statement("SET AUTO_PARTITION_MODE = TRUE") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertTrue(cursor.connection.auto_partition_mode) + + stmt = classify_statement("SET AUTO_PARTITION_MODE = FALSE") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertFalse(cursor.connection.auto_partition_mode) + + stmt = classify_statement("SET AUTO_PARTITION_MODE = INVALID") + with self.assertRaises(ProgrammingError): + execute(cursor, stmt) + + def test_execute_show_auto_partition_mode(self): + from unittest import mock + + from google.cloud.spanner_dbapi.client_side_statement_executor import execute + from google.cloud.spanner_v1 import TypeCode + + cursor = mock.MagicMock() + cursor.connection.is_closed = False + cursor.connection.auto_partition_mode = True + + stmt = classify_statement("SHOW VARIABLE AUTO_PARTITION_MODE") + res = execute(cursor, stmt) + rows = list(res) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0][0], True) + self.assertEqual(res.fields[0].name, "AUTO_PARTITION_MODE") + self.assertEqual(res.fields[0].type_.code, TypeCode.BOOL) + + cursor.connection.auto_partition_mode = False + res = execute(cursor, stmt) + rows = list(res) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0][0], False) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py index c6b51c322b51..4496884027f7 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py @@ -170,3 +170,16 @@ def test_w_data_boost_enabled(self, mock_client): self.assertIsInstance(connection, Connection) self.assertTrue(connection.data_boost_enabled) + + def test_w_auto_partition_mode(self, mock_client): + from google.cloud.spanner_dbapi import Connection, connect + + client = mock_client.return_value + instance = client.instance.return_value + database = instance.database.return_value + self.assertIsNotNone(database) + + connection = connect(INSTANCE, DATABASE, auto_partition_mode=True) + + self.assertIsInstance(connection, Connection) + self.assertTrue(connection.auto_partition_mode) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py index abae1803a578..d51d20e3779b 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py @@ -21,7 +21,6 @@ import mock import pytest from google.auth.credentials import AnonymousCredentials - from google.cloud.spanner_admin_database_v1 import DatabaseDialect from google.cloud.spanner_dbapi import Connection from google.cloud.spanner_dbapi.batch_dml_executor import BatchMode @@ -39,6 +38,7 @@ StatementType, ) from google.cloud.spanner_v1.database_sessions_manager import TransactionType + from tests._builders import build_connection, build_session PROJECT = "test-project" @@ -155,6 +155,16 @@ def test_property_data_boost_enabled(self): connection.data_boost_enabled = False self.assertFalse(connection.data_boost_enabled) + def test_property_auto_partition_mode(self): + connection = self._make_connection() + self.assertFalse(connection.auto_partition_mode) + + connection.auto_partition_mode = True + self.assertTrue(connection.auto_partition_mode) + + connection.auto_partition_mode = False + self.assertFalse(connection.auto_partition_mode) + def test__session_checkout_read_only(self): connection = build_connection(read_only=True) database = connection._database diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py index 60acc37e29d2..6ad4cf5572ca 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py @@ -20,14 +20,13 @@ from google.api_core.exceptions import Aborted from google.auth.credentials import AnonymousCredentials -from google.rpc.code_pb2 import ABORTED - from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.parsed_statement import ( ParsedStatement, Statement, StatementType, ) +from google.rpc.code_pb2 import ABORTED class TestCursor(unittest.TestCase): @@ -455,7 +454,6 @@ def test_execute_statement_exception_with_cursor_not_in_retry_mode(self): def test_execute_integrity_error(self): from google.api_core import exceptions - from google.cloud.spanner_dbapi.exceptions import IntegrityError connection = self._make_connection(self.INSTANCE, mock.MagicMock()) @@ -488,7 +486,6 @@ def test_execute_integrity_error(self): def test_execute_invalid_argument(self): from google.api_core import exceptions - from google.cloud.spanner_dbapi.exceptions import ProgrammingError connection = self._make_connection(self.INSTANCE, mock.MagicMock()) @@ -503,7 +500,6 @@ def test_execute_invalid_argument(self): def test_execute_internal_server_error(self): from google.api_core import exceptions - from google.cloud.spanner_dbapi.exceptions import OperationalError connection = self._make_connection(self.INSTANCE, mock.MagicMock()) @@ -767,11 +763,10 @@ def test_executemany_insert_batch_autocommit(self): transaction.commit.assert_called_once() def test_executemany_insert_batch_failed(self): - from google.rpc.code_pb2 import UNKNOWN - from google.cloud.spanner_dbapi import connect from google.cloud.spanner_dbapi.exceptions import OperationalError from google.cloud.spanner_v1.types.spanner import Session + from google.rpc.code_pb2 import UNKNOWN sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)""" err_details = "Details here" @@ -1008,6 +1003,50 @@ def test_handle_dql_priority(self, MockedPeekIterator): sql, None, None, request_options=RequestOptions(priority=1) ) + def test_execute_query_with_auto_partition_mode(self): + from google.cloud import spanner + + connection = self._make_connection( + self.INSTANCE, + mock.MagicMock(), + read_only=True, + data_boost_enabled=True, + auto_partition_mode=True, + ) + batch_snapshot = connection.database.batch_snapshot.return_value = ( + mock.MagicMock() + ) + mock_result_set = mock.MagicMock() + batch_snapshot.run_partitioned_query.return_value = mock_result_set + + cursor = self._make_one(connection) + cursor.execute("SELECT * FROM table WHERE col = %s", (10,)) + + self.assertEqual(cursor._result_set, mock_result_set) + self.assertEqual(cursor._itr, mock_result_set) + batch_snapshot.run_partitioned_query.assert_called_once_with( + "SELECT * FROM table WHERE col = @a0", + params={"a0": 10}, + param_types={"a0": spanner.param_types.INT64}, + data_boost_enabled=True, + ) + + def test_execute_query_with_auto_partition_mode_rw_transaction_error(self): + from google.cloud.spanner_dbapi.exceptions import ProgrammingError + + connection = self._make_connection( + self.INSTANCE, + mock.MagicMock(), + read_only=False, + auto_partition_mode=True, + ) + connection._autocommit = False + connection._transaction_begin_marked = True + + cursor = self._make_one(connection) + with self.assertRaises(ProgrammingError): + cursor.execute("SELECT * FROM table") + def test_handle_dql_database_error(self): connection = self._make_connection(self.INSTANCE) cursor = self._make_one(connection) @@ -1123,7 +1162,6 @@ def test_peek_iterator_aborted(self, mock_client): while streaming the first element with a PeekIterator. """ from google.api_core.exceptions import Aborted - from google.cloud.spanner_dbapi.connection import connect connection = connect("test-instance", "test-database") diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py index 46a847c9f47d..7aae52970294 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py @@ -91,6 +91,8 @@ def test_classify_stmt(self): ("SHOW VARIABLE READ_TIMESTAMP", StatementType.CLIENT_SIDE), ("SET DATA_BOOST_ENABLED = TRUE", StatementType.CLIENT_SIDE), ("SHOW VARIABLE DATA_BOOST_ENABLED", StatementType.CLIENT_SIDE), + ("SET AUTO_PARTITION_MODE = TRUE", StatementType.CLIENT_SIDE), + ("SHOW VARIABLE AUTO_PARTITION_MODE", StatementType.CLIENT_SIDE), ("GRANT SELECT ON TABLE Singers TO ROLE parent", StatementType.DDL), ("REVOKE SELECT ON TABLE Singers TO ROLE parent", StatementType.DDL), ("GRANT ROLE parent TO ROLE child", StatementType.DDL), @@ -286,6 +288,41 @@ def test_show_data_boost_enabled_stmt(self): ), ) + def test_set_auto_partition_mode_stmt(self): + parsed_statement = classify_statement("SET AUTO_PARTITION_MODE = TRUE") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("SET AUTO_PARTITION_MODE = TRUE"), + ClientSideStatementType.SET_AUTO_PARTITION_MODE, + ["TRUE"], + ), + ) + + parsed_statement = classify_statement("set auto_partition_mode = false") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("set auto_partition_mode = false"), + ClientSideStatementType.SET_AUTO_PARTITION_MODE, + ["false"], + ), + ) + + def test_show_auto_partition_mode_stmt(self): + parsed_statement = classify_statement(" show variable auto_partition_mode ") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("show variable auto_partition_mode"), + ClientSideStatementType.SHOW_AUTO_PARTITION_MODE, + [], + ), + ) + @unittest.skipIf(skip_condition, skip_message) def test_sql_pyformat_args_to_spanner(self): from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner From 85de4a70ebeb80957c92ba608c972ad0db43662b Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Thu, 20 Aug 2026 12:49:42 +0000 Subject: [PATCH 3/3] fix(spanner): document dbapi connection arguments and support trailing semicolons in client-side statements Address code review comments: - Document `data_boost_enabled` and `auto_partition_mode` arguments in the `Connection` class docstring. - Support trailing semicolons and quotes in `_parse_bool` when executing `SET DATA_BOOST_ENABLED` and `SET AUTO_PARTITION_MODE`. - Allow optional trailing semicolons in `SHOW VARIABLE DATA_BOOST_ENABLED` and `SHOW VARIABLE AUTO_PARTITION_MODE` regex parser patterns. - Add test cases for statements ending with semicolons. --- .../client_side_statement_executor.py | 2 +- .../client_side_statement_parser.py | 4 +-- .../google/cloud/spanner_dbapi/connection.py | 11 +++++++ .../test_client_side_statement_executor.py | 26 +++++++++++++++ .../unit/spanner_dbapi/test_parse_utils.py | 33 +++++++++++++++++++ 5 files changed, 73 insertions(+), 3 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py index 6d4258745a5b..4c0f591dc7e3 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py @@ -141,7 +141,7 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement): def _parse_bool(raw_val: str, var_name: str) -> bool: - cleaned = raw_val.strip().strip("'\"").lower() + cleaned = raw_val.strip().rstrip(";").strip().strip("'\"").lower() if cleaned not in _BOOL_MAP: raise ProgrammingError( f"Invalid value for {var_name}: '{raw_val}'. Expected TRUE or FALSE." diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py index 963263eac2b4..bce3842927e5 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py @@ -48,13 +48,13 @@ r"^\s*(SET)\s+(DATA_BOOST_ENABLED)\s+(=)\s+(.+)", re.IGNORECASE ) RE_SHOW_DATA_BOOST_ENABLED = re.compile( - r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*$", re.IGNORECASE + r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*;?\s*$", re.IGNORECASE ) RE_SET_AUTO_PARTITION_MODE = re.compile( r"^\s*(SET)\s+(AUTO_PARTITION_MODE)\s+(=)\s+(.+)", re.IGNORECASE ) RE_SHOW_AUTO_PARTITION_MODE = re.compile( - r"^\s*(SHOW)\s+(VARIABLE)\s+(AUTO_PARTITION_MODE)\s*$", re.IGNORECASE + r"^\s*(SHOW)\s+(VARIABLE)\s+(AUTO_PARTITION_MODE)\s*;?\s*$", re.IGNORECASE ) 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 19d975a13453..ba4ff6613333 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py @@ -90,6 +90,17 @@ class Connection: 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. + :type data_boost_enabled: bool + :param data_boost_enabled: (Optional) Whether to enable DataBoost for + partitioned queries executed via this connection. Defaults to False. + Note that DataBoost is only supported for partitioned query execution. + + :type auto_partition_mode: bool + :param auto_partition_mode: (Optional) Whether to enable auto partition mode + for queries executed via this connection. When True, queries on read-only + or autocommit connections are automatically partitioned and executed in parallel. + Defaults to False. + **kwargs: Initial value for connection variables. """ diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py index 084e1fb90ebb..2a5d2c30ca24 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py @@ -75,6 +75,16 @@ def test_execute_set_data_boost_enabled(self): self.assertIsNone(res) self.assertFalse(cursor.connection.data_boost_enabled) + stmt = classify_statement("SET DATA_BOOST_ENABLED = TRUE;") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertTrue(cursor.connection.data_boost_enabled) + + stmt = classify_statement("SET DATA_BOOST_ENABLED = 'FALSE';") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertFalse(cursor.connection.data_boost_enabled) + stmt = classify_statement("SET DATA_BOOST_ENABLED = INVALID") with self.assertRaises(ProgrammingError): execute(cursor, stmt) @@ -97,6 +107,12 @@ def test_execute_show_data_boost_enabled(self): self.assertEqual(res.fields[0].name, "DATA_BOOST_ENABLED") self.assertEqual(res.fields[0].type_.code, TypeCode.BOOL) + stmt_semicolon = classify_statement("SHOW VARIABLE DATA_BOOST_ENABLED;") + res_semicolon = execute(cursor, stmt_semicolon) + rows_semicolon = list(res_semicolon) + self.assertEqual(len(rows_semicolon), 1) + self.assertEqual(rows_semicolon[0][0], True) + cursor.connection.data_boost_enabled = False res = execute(cursor, stmt) rows = list(res) @@ -123,6 +139,16 @@ def test_execute_set_auto_partition_mode(self): self.assertIsNone(res) self.assertFalse(cursor.connection.auto_partition_mode) + stmt = classify_statement("SET AUTO_PARTITION_MODE = TRUE;") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertTrue(cursor.connection.auto_partition_mode) + + stmt = classify_statement("SET AUTO_PARTITION_MODE = 'FALSE';") + res = execute(cursor, stmt) + self.assertIsNone(res) + self.assertFalse(cursor.connection.auto_partition_mode) + stmt = classify_statement("SET AUTO_PARTITION_MODE = INVALID") with self.assertRaises(ProgrammingError): execute(cursor, stmt) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py index 7aae52970294..71c4bb4d7fda 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py @@ -288,6 +288,17 @@ def test_show_data_boost_enabled_stmt(self): ), ) + parsed_statement = classify_statement("SHOW VARIABLE DATA_BOOST_ENABLED;") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("SHOW VARIABLE DATA_BOOST_ENABLED;"), + ClientSideStatementType.SHOW_DATA_BOOST_ENABLED, + [], + ), + ) + def test_set_auto_partition_mode_stmt(self): parsed_statement = classify_statement("SET AUTO_PARTITION_MODE = TRUE") self.assertEqual( @@ -311,6 +322,17 @@ def test_set_auto_partition_mode_stmt(self): ), ) + parsed_statement = classify_statement("SET AUTO_PARTITION_MODE = TRUE;") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("SET AUTO_PARTITION_MODE = TRUE;"), + ClientSideStatementType.SET_AUTO_PARTITION_MODE, + ["TRUE;"], + ), + ) + def test_show_auto_partition_mode_stmt(self): parsed_statement = classify_statement(" show variable auto_partition_mode ") self.assertEqual( @@ -323,6 +345,17 @@ def test_show_auto_partition_mode_stmt(self): ), ) + parsed_statement = classify_statement("SHOW VARIABLE AUTO_PARTITION_MODE;") + self.assertEqual( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("SHOW VARIABLE AUTO_PARTITION_MODE;"), + ClientSideStatementType.SHOW_AUTO_PARTITION_MODE, + [], + ), + ) + @unittest.skipIf(skip_condition, skip_message) def test_sql_pyformat_args_to_spanner(self): from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner