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..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 @@ -13,10 +13,10 @@ # limitations under the License. from typing import TYPE_CHECKING, Union +from google.cloud.spanner_dbapi.exceptions import ProgrammingError from google.cloud.spanner_v1 import TransactionOptions 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 +108,45 @@ 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: + connection.data_boost_enabled = _parse_bool( + parsed_statement.client_side_statement_params[0], + "DATA_BOOST_ENABLED", + ) + 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, + ) + 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().rstrip(";").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): 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..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 @@ -44,6 +44,18 @@ 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*;?\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*;?\s*$", re.IGNORECASE +) def parse_stmt(query): @@ -68,6 +80,10 @@ 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_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): @@ -96,6 +112,14 @@ 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 + 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 eb0ef4217411..ba4ff6613333 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 @@ -91,10 +90,29 @@ 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. """ - def __init__(self, instance, database=None, read_only=False, **kwargs): + def __init__( + self, + instance, + database=None, + read_only=False, + data_boost_enabled=False, + auto_partition_mode=False, + **kwargs, + ): self._instance = instance self._database = database self._ddl_statements = [] @@ -110,6 +128,8 @@ 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._auto_partition_mode = bool(auto_partition_mode) self._staleness = None self.request_priority = None self._transaction_begin_marked = False @@ -123,6 +143,47 @@ 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 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 @@ -638,11 +699,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 +716,7 @@ def partition_query( statement.params, statement.param_types, query_options=query_options, + data_boost_enabled=data_boost_enabled, ) ) @@ -684,7 +750,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 +817,8 @@ def connect( client_certificate=None, client_key=None, instance_type=None, + data_boost_enabled=False, + auto_partition_mode=False, **kwargs, ): """Creates a connection to a Google Cloud Spanner database. @@ -795,6 +866,17 @@ 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. + + :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. @@ -909,7 +991,13 @@ 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, + 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 a8d03f6fa410..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 @@ -38,6 +38,10 @@ 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 + 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 new file mode 100644 index 000000000000..4c234d47dc95 --- /dev/null +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py @@ -0,0 +1,267 @@ +# 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) + + 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 888f81e830f7..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 @@ -52,3 +52,127 @@ 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 = 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) + + 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) + 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 = 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 f80c3c3e52b5..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 @@ -157,3 +157,29 @@ 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) + + 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 3d47f57fb4fb..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" @@ -145,6 +145,26 @@ 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_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 @@ -924,6 +944,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_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 29a08edb4264..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 @@ -89,6 +89,10 @@ 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), + ("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), @@ -250,6 +254,108 @@ 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, + [], + ), + ) + + 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( + 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"], + ), + ) + + 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( + parsed_statement, + ParsedStatement( + StatementType.CLIENT_SIDE, + Statement("show variable auto_partition_mode"), + ClientSideStatementType.SHOW_AUTO_PARTITION_MODE, + [], + ), + ) + + 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