Visitar URL original
feat(spanner): add DataBoost and auto_partition_mode support to DBAPI… · googleapis/google-cloud-python@63b2f31 · GitHub
Skip to content

Commit 63b2f31

Browse files
sakthivelmaniisakthivelmanii
committed
feat(spanner): add DataBoost and auto_partition_mode support to DBAPI driver
Add support for Cloud Spanner DataBoost and `auto_partition_mode` in the Python DBAPI (PEP 249) driver (`google.cloud.spanner_dbapi`). Key updates: * Add `data_boost_enabled: bool = False` and `auto_partition_mode: bool = False` arguments to `spanner_dbapi.connect()` and `Connection.__init__()`. * Add `@property def data_boost_enabled(self)` and `@property def auto_partition_mode(self)` getters and setters 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) - `SET AUTO_PARTITION_MODE = TRUE|FALSE` - `SHOW VARIABLE AUTO_PARTITION_MODE` (returns column `AUTO_PARTITION_MODE` with `BOOL` type) * Add automatic query partitioning in `Cursor._execute()` when `auto_partition_mode=True`, routing queries to `run_partitioned_query` and forwarding `data_boost_enabled`. * Forward `data_boost_enabled` in `Connection.partition_query()` and `Connection.run_partitioned_query()`. * Add comprehensive unit and gRPC mock server test coverage across DBAPI connection, cursor, parser, statement executor, and mock server test suites.
1 parent 7f8a552 commit 63b2f31

11 files changed

Lines changed: 740 additions & 4 deletions

File tree

‎packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@
1515

1616
from google.cloud.spanner_v1 import TransactionOptions
1717

18+
from google.cloud.spanner_dbapi.exceptions import ProgrammingError
19+
1820
if TYPE_CHECKING:
19-
from google.cloud.spanner_dbapi import ProgrammingError
2021
from google.cloud.spanner_dbapi.cursor import Cursor
2122

2223
from google.cloud.spanner_dbapi.parsed_statement import (
@@ -108,6 +109,47 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
108109
return connection.run_partitioned_query(parsed_statement)
109110
if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE:
110111
return connection._set_autocommit_dml_mode(parsed_statement)
112+
if statement_type == ClientSideStatementType.SET_DATA_BOOST_ENABLED:
113+
val_str = (
114+
parsed_statement.client_side_statement_params[0]
115+
.strip()
116+
.strip("'\"")
117+
.lower()
118+
)
119+
if val_str not in ("true", "false"):
120+
raise ProgrammingError(
121+
f"Invalid value for DATA_BOOST_ENABLED: '{parsed_statement.client_side_statement_params[0]}'. Expected TRUE or FALSE."
122+
)
123+
connection.data_boost_enabled = val_str == "true"
124+
return None
125+
if statement_type == ClientSideStatementType.SHOW_DATA_BOOST_ENABLED:
126+
column_values.append(connection.data_boost_enabled)
127+
return _get_streamed_result_set(
128+
"DATA_BOOST_ENABLED",
129+
TypeCode.BOOL,
130+
column_values,
131+
)
132+
if statement_type == ClientSideStatementType.SET_AUTO_PARTITION_MODE:
133+
val_str = (
134+
parsed_statement.client_side_statement_params[0]
135+
.strip()
136+
.strip("'\"")
137+
.lower()
138+
)
139+
if val_str not in ("true", "false"):
140+
raise ProgrammingError(
141+
f"Invalid value for AUTO_PARTITION_MODE: '{parsed_statement.client_side_statement_params[0]}'. Expected TRUE or FALSE."
142+
)
143+
connection.auto_partition_mode = val_str == "true"
144+
return None
145+
if statement_type == ClientSideStatementType.SHOW_AUTO_PARTITION_MODE:
146+
column_values.append(connection.auto_partition_mode)
147+
return _get_streamed_result_set(
148+
"AUTO_PARTITION_MODE",
149+
TypeCode.BOOL,
150+
column_values,
151+
)
152+
return None
111153

112154

113155
def _get_streamed_result_set(column_name, type_code, column_values):

‎packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,18 @@
4444
RE_SET_AUTOCOMMIT_DML_MODE = re.compile(
4545
r"^\s*(SET)\s+(AUTOCOMMIT_DML_MODE)\s+(=)\s+(.+)", re.IGNORECASE
4646
)
47+
RE_SET_DATA_BOOST_ENABLED = re.compile(
48+
r"^\s*(SET)\s+(DATA_BOOST_ENABLED)\s+(=)\s+(.+)", re.IGNORECASE
49+
)
50+
RE_SHOW_DATA_BOOST_ENABLED = re.compile(
51+
r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*$", re.IGNORECASE
52+
)
53+
RE_SET_AUTO_PARTITION_MODE = re.compile(
54+
r"^\s*(SET)\s+(AUTO_PARTITION_MODE)\s+(=)\s+(.+)", re.IGNORECASE
55+
)
56+
RE_SHOW_AUTO_PARTITION_MODE = re.compile(
57+
r"^\s*(SHOW)\s+(VARIABLE)\s+(AUTO_PARTITION_MODE)\s*$", re.IGNORECASE
58+
)
4759

4860

4961
def parse_stmt(query):
@@ -68,6 +80,10 @@ def parse_stmt(query):
6880
client_side_statement_type = ClientSideStatementType.SHOW_COMMIT_TIMESTAMP
6981
elif RE_SHOW_READ_TIMESTAMP.match(query):
7082
client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP
83+
elif RE_SHOW_DATA_BOOST_ENABLED.match(query):
84+
client_side_statement_type = ClientSideStatementType.SHOW_DATA_BOOST_ENABLED
85+
elif RE_SHOW_AUTO_PARTITION_MODE.match(query):
86+
client_side_statement_type = ClientSideStatementType.SHOW_AUTO_PARTITION_MODE
7187
elif RE_START_BATCH_DML.match(query):
7288
client_side_statement_type = ClientSideStatementType.START_BATCH_DML
7389
elif RE_BEGIN.match(query):
@@ -96,6 +112,14 @@ def parse_stmt(query):
96112
match = re.search(RE_SET_AUTOCOMMIT_DML_MODE, query)
97113
client_side_statement_params.append(match.group(4))
98114
client_side_statement_type = ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE
115+
elif RE_SET_DATA_BOOST_ENABLED.match(query):
116+
match = re.search(RE_SET_DATA_BOOST_ENABLED, query)
117+
client_side_statement_params.append(match.group(4))
118+
client_side_statement_type = ClientSideStatementType.SET_DATA_BOOST_ENABLED
119+
elif RE_SET_AUTO_PARTITION_MODE.match(query):
120+
match = re.search(RE_SET_AUTO_PARTITION_MODE, query)
121+
client_side_statement_params.append(match.group(4))
122+
client_side_statement_type = ClientSideStatementType.SET_AUTO_PARTITION_MODE
99123
if client_side_statement_type is not None:
100124
return ParsedStatement(
101125
StatementType.CLIENT_SIDE,

‎packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py‎

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,15 @@ class Connection:
9494
**kwargs: Initial value for connection variables.
9595
"""
9696

97-
def __init__(self, instance, database=None, read_only=False, **kwargs):
97+
def __init__(
98+
self,
99+
instance,
100+
database=None,
101+
read_only=False,
102+
data_boost_enabled=False,
103+
auto_partition_mode=False,
104+
**kwargs,
105+
):
98106
self._instance = instance
99107
self._database = database
100108
self._ddl_statements = []
@@ -110,6 +118,8 @@ def __init__(self, instance, database=None, read_only=False, **kwargs):
110118
# connection close
111119
self._own_pool = True
112120
self._read_only = read_only
121+
self._data_boost_enabled = bool(data_boost_enabled)
122+
self._auto_partition_mode = bool(auto_partition_mode)
113123
self._staleness = None
114124
self.request_priority = None
115125
self._transaction_begin_marked = False
@@ -123,6 +133,47 @@ def __init__(self, instance, database=None, read_only=False, **kwargs):
123133
self._autocommit_dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
124134
self._connection_variables = kwargs
125135

136+
@property
137+
def data_boost_enabled(self):
138+
"""Flag: whether DataBoost is enabled for partitioned queries on this connection.
139+
140+
Note that DataBoost is only supported for partitioned query execution.
141+
142+
Returns:
143+
bool: True if DataBoost is enabled, False otherwise.
144+
"""
145+
return self._data_boost_enabled
146+
147+
@data_boost_enabled.setter
148+
def data_boost_enabled(self, value):
149+
"""Change the DataBoost enablement state for partitioned queries on this connection.
150+
151+
:type value: bool
152+
:param value: New data_boost_enabled state.
153+
"""
154+
self._data_boost_enabled = bool(value)
155+
156+
@property
157+
def auto_partition_mode(self):
158+
"""Flag: whether auto partition mode is enabled for queries on this connection.
159+
160+
When enabled, standard queries executed on read-only or autocommit connections
161+
are automatically partitioned and executed in parallel via run_partitioned_query.
162+
163+
Returns:
164+
bool: True if auto partition mode is enabled, False otherwise.
165+
"""
166+
return self._auto_partition_mode
167+
168+
@auto_partition_mode.setter
169+
def auto_partition_mode(self, value):
170+
"""Change the auto partition mode enablement state for this connection.
171+
172+
:type value: bool
173+
:param value: New auto_partition_mode state.
174+
"""
175+
self._auto_partition_mode = bool(value)
176+
126177
@property
127178
def spanner_client(self):
128179
"""Client for interacting with Cloud Spanner API. This property exposes
@@ -638,11 +689,15 @@ def partition_query(
638689
self,
639690
parsed_statement: ParsedStatement,
640691
query_options=None,
692+
data_boost_enabled=None,
641693
):
642694
statement = parsed_statement.statement
643695
partitioned_query = parsed_statement.client_side_statement_params[0]
644696
self._partitioned_query_validation(partitioned_query, statement)
645697

698+
if data_boost_enabled is None:
699+
data_boost_enabled = self.data_boost_enabled
700+
646701
batch_snapshot = self._database.batch_snapshot()
647702
partition_ids = []
648703
partitions = list(
@@ -651,6 +706,7 @@ def partition_query(
651706
statement.params,
652707
statement.param_types,
653708
query_options=query_options,
709+
data_boost_enabled=data_boost_enabled,
654710
)
655711
)
656712

@@ -684,7 +740,10 @@ def run_partitioned_query(
684740
self._partitioned_query_validation(partitioned_query, statement)
685741
batch_snapshot = self._database.batch_snapshot()
686742
return batch_snapshot.run_partitioned_query(
687-
partitioned_query, statement.params, statement.param_types
743+
partitioned_query,
744+
statement.params,
745+
statement.param_types,
746+
data_boost_enabled=self.data_boost_enabled,
688747
)
689748

690749
@check_not_closed
@@ -748,6 +807,8 @@ def connect(
748807
client_certificate=None,
749808
client_key=None,
750809
instance_type=None,
810+
data_boost_enabled=False,
811+
auto_partition_mode=False,
751812
**kwargs,
752813
):
753814
"""Creates a connection to a Google Cloud Spanner database.
@@ -795,6 +856,17 @@ def connect(
795856
:param database_role: (Optional) The database role to connect as when using
796857
fine-grained access controls.
797858
859+
:type data_boost_enabled: bool
860+
:param data_boost_enabled: (Optional) Whether to enable DataBoost for
861+
partitioned queries executed via this connection. Defaults to False.
862+
Note that DataBoost is only supported for partitioned query execution.
863+
864+
:type auto_partition_mode: bool
865+
:param auto_partition_mode: (Optional) Whether to enable auto partition mode
866+
for queries executed via this connection. When True, queries on read-only
867+
or autocommit connections are automatically partitioned and executed in parallel.
868+
Defaults to False.
869+
798870
**kwargs: Initial value for connection variables.
799871
800872
@@ -909,7 +981,13 @@ def connect(
909981
database = instance.database(
910982
database_id, pool=pool, database_role=database_role, logger=logger
911983
)
912-
conn = Connection(instance, database, **kwargs)
984+
conn = Connection(
985+
instance,
986+
database,
987+
data_boost_enabled=data_boost_enabled,
988+
auto_partition_mode=auto_partition_mode,
989+
**kwargs,
990+
)
913991
if pool is not None:
914992
conn._own_pool = False
915993

‎packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,11 @@ def _execute(self, sql, args=None, call_from_execute_many=False):
308308
self._itr = PeekIterator(self._result_set)
309309
elif self.connection._batch_mode == BatchMode.DML:
310310
self.connection.execute_batch_dml_statement(self._parsed_statement)
311+
elif (
312+
self._parsed_statement.statement_type == StatementType.QUERY
313+
and self.connection.auto_partition_mode
314+
):
315+
self._handle_auto_partition_query(sql, args or None)
311316
elif self.connection.read_only or (
312317
not self.connection._client_transaction_started
313318
and self._parsed_statement.statement_type == StatementType.QUERY
@@ -581,6 +586,24 @@ def _handle_DQL(self, sql, params):
581586
self.connection._transaction = None
582587
self._handle_DQL_with_snapshot(snapshot, sql, params)
583588

589+
def _handle_auto_partition_query(self, sql, params):
590+
if self.connection.database is None:
591+
raise ValueError("Database needs to be passed for this operation")
592+
if not self.connection.read_only and self.connection._client_transaction_started:
593+
raise ProgrammingError(
594+
"Partitioned query is not supported, because the connection is in a read/write transaction."
595+
)
596+
sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params)
597+
batch_snapshot = self.connection.database.batch_snapshot()
598+
self._result_set = batch_snapshot.run_partitioned_query(
599+
sql,
600+
params=params,
601+
param_types=get_param_types(params),
602+
data_boost_enabled=self.connection.data_boost_enabled,
603+
)
604+
self._itr = self._result_set
605+
self._row_count = None
606+
584607
def __enter__(self):
585608
return self
586609

‎packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ class ClientSideStatementType(Enum):
3838
RUN_PARTITION = 10
3939
RUN_PARTITIONED_QUERY = 11
4040
SET_AUTOCOMMIT_DML_MODE = 12
41+
SET_DATA_BOOST_ENABLED = 13
42+
SHOW_DATA_BOOST_ENABLED = 14
43+
SET_AUTO_PARTITION_MODE = 15
44+
SHOW_AUTO_PARTITION_MODE = 16
4145

4246

4347
class AutocommitDmlMode(Enum):

0 commit comments

Comments
 (0)