Visitar URL original
feat(spanner): add auto_partition_mode support to DBAPI driver and cl… · googleapis/google-cloud-python@8b24d43 · GitHub
Skip to content

Commit 8b24d43

Browse files
sakthivelmaniisakthivelmanii
committed
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.
1 parent cff1650 commit 8b24d43

11 files changed

Lines changed: 324 additions & 11 deletions

File tree

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

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -110,17 +110,10 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
110110
if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE:
111111
return connection._set_autocommit_dml_mode(parsed_statement)
112112
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()
113+
connection.data_boost_enabled = _parse_bool(
114+
parsed_statement.client_side_statement_params[0],
115+
"DATA_BOOST_ENABLED",
118116
)
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"
124117
return None
125118
if statement_type == ClientSideStatementType.SHOW_DATA_BOOST_ENABLED:
126119
column_values.append(connection.data_boost_enabled)
@@ -129,9 +122,34 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
129122
TypeCode.BOOL,
130123
column_values,
131124
)
125+
if statement_type == ClientSideStatementType.SET_AUTO_PARTITION_MODE:
126+
connection.auto_partition_mode = _parse_bool(
127+
parsed_statement.client_side_statement_params[0],
128+
"AUTO_PARTITION_MODE",
129+
)
130+
return None
131+
if statement_type == ClientSideStatementType.SHOW_AUTO_PARTITION_MODE:
132+
column_values.append(connection.auto_partition_mode)
133+
return _get_streamed_result_set(
134+
"AUTO_PARTITION_MODE",
135+
TypeCode.BOOL,
136+
column_values,
137+
)
132138
return None
133139

134140

141+
_BOOL_MAP = {"true": True, "false": False}
142+
143+
144+
def _parse_bool(raw_val: str, var_name: str) -> bool:
145+
cleaned = raw_val.strip().strip("'\"").lower()
146+
if cleaned not in _BOOL_MAP:
147+
raise ProgrammingError(
148+
f"Invalid value for {var_name}: '{raw_val}'. Expected TRUE or FALSE."
149+
)
150+
return _BOOL_MAP[cleaned]
151+
152+
135153
def _get_streamed_result_set(column_name, type_code, column_values):
136154
struct_type_pb = StructType(
137155
fields=[StructType.Field(name=column_name, type_=Type(code=type_code))]

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,12 @@
5050
RE_SHOW_DATA_BOOST_ENABLED = re.compile(
5151
r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*$", re.IGNORECASE
5252
)
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+
)
5359

5460

5561
def parse_stmt(query):
@@ -76,6 +82,8 @@ def parse_stmt(query):
7682
client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP
7783
elif RE_SHOW_DATA_BOOST_ENABLED.match(query):
7884
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
7987
elif RE_START_BATCH_DML.match(query):
8088
client_side_statement_type = ClientSideStatementType.START_BATCH_DML
8189
elif RE_BEGIN.match(query):
@@ -108,6 +116,10 @@ def parse_stmt(query):
108116
match = re.search(RE_SET_DATA_BOOST_ENABLED, query)
109117
client_side_statement_params.append(match.group(4))
110118
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
111123
if client_side_statement_type is not None:
112124
return ParsedStatement(
113125
StatementType.CLIENT_SIDE,

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

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ def __init__(
100100
database=None,
101101
read_only=False,
102102
data_boost_enabled=False,
103+
auto_partition_mode=False,
103104
**kwargs,
104105
):
105106
self._instance = instance
@@ -118,6 +119,7 @@ def __init__(
118119
self._own_pool = True
119120
self._read_only = read_only
120121
self._data_boost_enabled = bool(data_boost_enabled)
122+
self._auto_partition_mode = bool(auto_partition_mode)
121123
self._staleness = None
122124
self.request_priority = None
123125
self._transaction_begin_marked = False
@@ -151,6 +153,27 @@ def data_boost_enabled(self, value):
151153
"""
152154
self._data_boost_enabled = bool(value)
153155

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+
154177
@property
155178
def spanner_client(self):
156179
"""Client for interacting with Cloud Spanner API. This property exposes
@@ -785,6 +808,7 @@ def connect(
785808
client_key=None,
786809
instance_type=None,
787810
data_boost_enabled=False,
811+
auto_partition_mode=False,
788812
**kwargs,
789813
):
790814
"""Creates a connection to a Google Cloud Spanner database.
@@ -837,6 +861,12 @@ def connect(
837861
partitioned queries executed via this connection. Defaults to False.
838862
Note that DataBoost is only supported for partitioned query execution.
839863
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+
840870
**kwargs: Initial value for connection variables.
841871
842872
@@ -952,7 +982,11 @@ def connect(
952982
database_id, pool=pool, database_role=database_role, logger=logger
953983
)
954984
conn = Connection(
955-
instance, database, data_boost_enabled=data_boost_enabled, **kwargs
985+
instance,
986+
database,
987+
data_boost_enabled=data_boost_enabled,
988+
auto_partition_mode=auto_partition_mode,
989+
**kwargs,
956990
)
957991
if pool is not None:
958992
conn._own_pool = False

‎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: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ class ClientSideStatementType(Enum):
4040
SET_AUTOCOMMIT_DML_MODE = 12
4141
SET_DATA_BOOST_ENABLED = 13
4242
SHOW_DATA_BOOST_ENABLED = 14
43+
SET_AUTO_PARTITION_MODE = 15
44+
SHOW_AUTO_PARTITION_MODE = 16
4345

4446

4547
class AutocommitDmlMode(Enum):

‎packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_databoost.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,3 +192,75 @@ def test_run_partitioned_query_with_data_boost_enabled(self):
192192
)
193193
self.assertEqual(1, len(execute_requests))
194194
self.assertTrue(execute_requests[0].data_boost_enabled)
195+
196+
def test_auto_partition_mode_with_data_boost_enabled(self):
197+
sql = "SELECT name FROM users WHERE active = true"
198+
199+
partition_response = spanner_types.PartitionResponse()
200+
partition_response.partitions.extend(
201+
[
202+
spanner_types.Partition(partition_token=b"mock-token-auto-1"),
203+
]
204+
)
205+
self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
206+
add_single_result(sql, "name", TypeCode.STRING, [("Alice",)])
207+
208+
connection = Connection(
209+
self.instance,
210+
self.database,
211+
read_only=True,
212+
auto_partition_mode=True,
213+
data_boost_enabled=True,
214+
)
215+
216+
with connection.cursor() as cursor:
217+
# Plain cursor.execute automatically runs as partitioned query with DataBoost
218+
cursor.execute(sql)
219+
rows = cursor.fetchall()
220+
self.assertEqual(1, len(rows))
221+
self.assertEqual("Alice", rows[0][0])
222+
223+
execute_requests = list(
224+
filter(
225+
lambda msg: isinstance(msg, ExecuteSqlRequest),
226+
self.spanner_service.requests,
227+
)
228+
)
229+
self.assertEqual(1, len(execute_requests))
230+
self.assertTrue(execute_requests[0].data_boost_enabled)
231+
self.assertEqual(b"mock-token-auto-1", execute_requests[0].partition_token)
232+
233+
def test_auto_partition_mode_via_statement(self):
234+
sql = "SELECT name FROM users WHERE active = true"
235+
236+
partition_response = spanner_types.PartitionResponse()
237+
partition_response.partitions.extend(
238+
[
239+
spanner_types.Partition(partition_token=b"mock-token-auto-2"),
240+
]
241+
)
242+
self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
243+
add_single_result(sql, "name", TypeCode.STRING, [("Bob",)])
244+
245+
connection = Connection(self.instance, self.database, read_only=True)
246+
247+
with connection.cursor() as cursor:
248+
cursor.execute("SET AUTO_PARTITION_MODE = TRUE")
249+
cursor.execute("SET DATA_BOOST_ENABLED = TRUE")
250+
self.assertTrue(cursor.connection.auto_partition_mode)
251+
self.assertTrue(cursor.connection.data_boost_enabled)
252+
253+
cursor.execute(sql)
254+
rows = cursor.fetchall()
255+
self.assertEqual(1, len(rows))
256+
self.assertEqual("Bob", rows[0][0])
257+
258+
execute_requests = list(
259+
filter(
260+
lambda msg: isinstance(msg, ExecuteSqlRequest),
261+
self.spanner_service.requests,
262+
)
263+
)
264+
self.assertEqual(1, len(execute_requests))
265+
self.assertTrue(execute_requests[0].data_boost_enabled)
266+
self.assertEqual(b"mock-token-auto-2", execute_requests[0].partition_token)

‎packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_client_side_statement_executor.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,3 +100,49 @@ def test_execute_show_data_boost_enabled(self):
100100
rows = list(res)
101101
self.assertEqual(len(rows), 1)
102102
self.assertEqual(rows[0][0], False)
103+
104+
def test_execute_set_auto_partition_mode(self):
105+
from unittest import mock
106+
from google.cloud.spanner_dbapi.client_side_statement_executor import execute
107+
from google.cloud.spanner_dbapi.exceptions import ProgrammingError
108+
109+
cursor = mock.MagicMock()
110+
cursor.connection.is_closed = False
111+
cursor.connection.auto_partition_mode = False
112+
113+
stmt = classify_statement("SET AUTO_PARTITION_MODE = TRUE")
114+
res = execute(cursor, stmt)
115+
self.assertIsNone(res)
116+
self.assertTrue(cursor.connection.auto_partition_mode)
117+
118+
stmt = classify_statement("SET AUTO_PARTITION_MODE = FALSE")
119+
res = execute(cursor, stmt)
120+
self.assertIsNone(res)
121+
self.assertFalse(cursor.connection.auto_partition_mode)
122+
123+
stmt = classify_statement("SET AUTO_PARTITION_MODE = INVALID")
124+
with self.assertRaises(ProgrammingError):
125+
execute(cursor, stmt)
126+
127+
def test_execute_show_auto_partition_mode(self):
128+
from unittest import mock
129+
from google.cloud.spanner_dbapi.client_side_statement_executor import execute
130+
from google.cloud.spanner_v1 import TypeCode
131+
132+
cursor = mock.MagicMock()
133+
cursor.connection.is_closed = False
134+
cursor.connection.auto_partition_mode = True
135+
136+
stmt = classify_statement("SHOW VARIABLE AUTO_PARTITION_MODE")
137+
res = execute(cursor, stmt)
138+
rows = list(res)
139+
self.assertEqual(len(rows), 1)
140+
self.assertEqual(rows[0][0], True)
141+
self.assertEqual(res.fields[0].name, "AUTO_PARTITION_MODE")
142+
self.assertEqual(res.fields[0].type_.code, TypeCode.BOOL)
143+
144+
cursor.connection.auto_partition_mode = False
145+
res = execute(cursor, stmt)
146+
rows = list(res)
147+
self.assertEqual(len(rows), 1)
148+
self.assertEqual(rows[0][0], False)

‎packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connect.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,3 +170,16 @@ def test_w_data_boost_enabled(self, mock_client):
170170

171171
self.assertIsInstance(connection, Connection)
172172
self.assertTrue(connection.data_boost_enabled)
173+
174+
def test_w_auto_partition_mode(self, mock_client):
175+
from google.cloud.spanner_dbapi import Connection, connect
176+
177+
client = mock_client.return_value
178+
instance = client.instance.return_value
179+
database = instance.database.return_value
180+
self.assertIsNotNone(database)
181+
182+
connection = connect(INSTANCE, DATABASE, auto_partition_mode=True)
183+
184+
self.assertIsInstance(connection, Connection)
185+
self.assertTrue(connection.auto_partition_mode)

‎packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_connection.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,16 @@ def test_property_data_boost_enabled(self):
155155
connection.data_boost_enabled = False
156156
self.assertFalse(connection.data_boost_enabled)
157157

158+
def test_property_auto_partition_mode(self):
159+
connection = self._make_connection()
160+
self.assertFalse(connection.auto_partition_mode)
161+
162+
connection.auto_partition_mode = True
163+
self.assertTrue(connection.auto_partition_mode)
164+
165+
connection.auto_partition_mode = False
166+
self.assertFalse(connection.auto_partition_mode)
167+
158168
def test__session_checkout_read_only(self):
159169
connection = build_connection(read_only=True)
160170
database = connection._database

0 commit comments

Comments
 (0)