Skip to content
This repository was archived by the owner on Mar 31, 2026. It is now read-only.

Commit bedb240

Browse files
ankiagaankiaga
committed
Implementing client side statement in dbapi starting with commit
1 parent 07fbc45 commit bedb240

File tree

8 files changed

+241
-83
lines changed

8 files changed

+241
-83
lines changed

‎google/cloud/spanner_dbapi/client_side_statement_executor.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright 20203 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
from google.cloud.spanner_dbapi.parsed_statement import (
15+
ParsedStatement,
16+
ClientSideStatementType,
17+
)
18+
19+
20+
class StatementExecutor(object):
21+
def __init__(self, connection):
22+
self.connection = connection
23+
24+
def execute(self, parsed_statement: ParsedStatement):
25+
"""Executes the client side statements by calling the relevant method
26+
27+
:type parsed_statement: ParsedStatement
28+
:param parsed_statement: parsed_statement based on the sql query
29+
"""
30+
if (
31+
parsed_statement.client_side_statement_type
32+
== ClientSideStatementType.COMMIT
33+
):
34+
self.connection.commit()

‎google/cloud/spanner_dbapi/client_side_statement_parser.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# Copyright 20203 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import re
16+
17+
from google.cloud.spanner_dbapi.parsed_statement import (
18+
ParsedStatement,
19+
StatementType,
20+
ClientSideStatementType,
21+
)
22+
23+
RE_COMMIT = re.compile(r"^\s*(COMMIT)(TRANSACTION)?", re.IGNORECASE)
24+
25+
RE_BEGIN = re.compile(r"^\s*(BEGIN|START)(TRANSACTION)?", re.IGNORECASE)
26+
27+
28+
def parse_stmt(query):
29+
"""Parses the sql query to check if it matches with any of the client side
30+
statement regex.
31+
32+
:type query: str
33+
:param query: sql query
34+
35+
:rtype: ParsedStatement
36+
:returns: ParsedStatement object.
37+
"""
38+
if RE_COMMIT.match(query):
39+
return ParsedStatement(
40+
StatementType.CLIENT_SIDE, query, ClientSideStatementType.COMMIT
41+
)
42+
if RE_BEGIN.match(query):
43+
return ParsedStatement(
44+
StatementType.CLIENT_SIDE, query, ClientSideStatementType.BEGIN
45+
)
46+
return None

‎google/cloud/spanner_dbapi/cursor.py‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727

2828
from google.cloud import spanner_v1 as spanner
2929
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
30+
from google.cloud.spanner_dbapi.client_side_statement_executor import StatementExecutor
3031
from google.cloud.spanner_dbapi.exceptions import IntegrityError
3132
from google.cloud.spanner_dbapi.exceptions import InterfaceError
3233
from google.cloud.spanner_dbapi.exceptions import OperationalError
@@ -39,6 +40,7 @@
3940
from google.cloud.spanner_dbapi import parse_utils
4041
from google.cloud.spanner_dbapi.parse_utils import get_param_types
4142
from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner
43+
from google.cloud.spanner_dbapi.parsed_statement import StatementType
4244
from google.cloud.spanner_dbapi.utils import PeekIterator
4345
from google.cloud.spanner_dbapi.utils import StreamedManyResultSets
4446

@@ -85,6 +87,7 @@ def __init__(self, connection):
8587
self._row_count = _UNSET_COUNT
8688
self.lastrowid = None
8789
self.connection = connection
90+
self.client_side_statement_executor = StatementExecutor(connection)
8891
self._is_closed = False
8992
# the currently running SQL statement results checksum
9093
self._checksum = None
@@ -210,7 +213,7 @@ def _batch_DDLs(self, sql):
210213
for ddl in sqlparse.split(sql):
211214
if ddl:
212215
ddl = ddl.rstrip(";")
213-
if parse_utils.classify_stmt(ddl) != parse_utils.STMT_DDL:
216+
if parse_utils.classify_stmt(ddl).statement_type != StatementType.DDL:
214217
raise ValueError("Only DDL statements may be batched.")
215218

216219
statements.append(ddl)
@@ -239,8 +242,11 @@ def execute(self, sql, args=None):
239242
self._handle_DQL(sql, args or None)
240243
return
241244

242-
class_ = parse_utils.classify_stmt(sql)
243-
if class_ == parse_utils.STMT_DDL:
245+
parsed_statement = parse_utils.classify_stmt(sql)
246+
if parsed_statement.statement_type == StatementType.CLIENT_SIDE:
247+
self.client_side_statement_executor.execute(parsed_statement)
248+
return
249+
if parsed_statement.statement_type == StatementType.DDL:
244250
self._batch_DDLs(sql)
245251
if self.connection.autocommit:
246252
self.connection.run_prior_DDL_statements()
@@ -251,7 +257,7 @@ def execute(self, sql, args=None):
251257
# self._run_prior_DDL_statements()
252258
self.connection.run_prior_DDL_statements()
253259

254-
if class_ == parse_utils.STMT_UPDATING:
260+
if parsed_statement.statement_type == StatementType.UPDATE:
255261
sql = parse_utils.ensure_where_clause(sql)
256262

257263
sql, args = sql_pyformat_args_to_spanner(sql, args or None)
@@ -276,7 +282,7 @@ def execute(self, sql, args=None):
276282
self.connection.retry_transaction()
277283
return
278284

279-
if class_ == parse_utils.STMT_NON_UPDATING:
285+
if parsed_statement.statement_type == StatementType.QUERY:
280286
self._handle_DQL(sql, args or None)
281287
else:
282288
self.connection.database.run_in_transaction(
@@ -309,19 +315,26 @@ def executemany(self, operation, seq_of_params):
309315
self._result_set = None
310316
self._row_count = _UNSET_COUNT
311317

312-
class_ = parse_utils.classify_stmt(operation)
313-
if class_ == parse_utils.STMT_DDL:
318+
parsed_statement = parse_utils.classify_stmt(operation)
319+
if parsed_statement.statement_type == StatementType.DDL:
314320
raise ProgrammingError(
315321
"Executing DDL statements with executemany() method is not allowed."
316322
)
317323

318324
# For every operation, we've got to ensure that any prior DDL
319325
# statements were run.
320326
self.connection.run_prior_DDL_statements()
327+
if parsed_statement.statement_type == StatementType.CLIENT_SIDE:
328+
raise ProgrammingError(
329+
"Executing ClientSide statements with executemany() method is not allowed."
330+
)
321331

322332
many_result_set = StreamedManyResultSets()
323333

324-
if class_ in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING):
334+
if parsed_statement.statement_type in (
335+
StatementType.INSERT,
336+
StatementType.UPDATE,
337+
):
325338
statements = []
326339

327340
for params in seq_of_params:

‎google/cloud/spanner_dbapi/parse_utils.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,10 @@
2121
import sqlparse
2222
from google.cloud import spanner_v1 as spanner
2323
from google.cloud.spanner_v1 import JsonObject
24+
from . import client_side_statement_parser
2425

2526
from .exceptions import Error
27+
from .parsed_statement import ParsedStatement, StatementType
2628
from .types import DateStr, TimestampStr
2729
from .utils import sanitize_literals_for_upload
2830

@@ -139,11 +141,6 @@
139141
"WITHIN",
140142
}
141143

142-
STMT_DDL = "DDL"
143-
STMT_NON_UPDATING = "NON_UPDATING"
144-
STMT_UPDATING = "UPDATING"
145-
STMT_INSERT = "INSERT"
146-
147144
# Heuristic for identifying statements that don't need to be run as updates.
148145
RE_NON_UPDATE = re.compile(r"^\W*(SELECT)", re.IGNORECASE)
149146

@@ -180,27 +177,29 @@ def classify_stmt(query):
180177
:type query: str
181178
:param query: A SQL query.
182179
183-
:rtype: str
184-
:returns: The query type name.
180+
:rtype: ParsedStatement
181+
:returns: parsed statement attributes.
185182
"""
186183
# sqlparse will strip Cloud Spanner comments,
187184
# still, special commenting styles, like
188185
# PostgreSQL dollar quoted comments are not
189186
# supported and will not be stripped.
190187
query = sqlparse.format(query, strip_comments=True).strip()
191-
188+
parsed_statement = client_side_statement_parser.parse_stmt(query)
189+
if parsed_statement is not None:
190+
return parsed_statement
192191
if RE_DDL.match(query):
193-
return STMT_DDL
192+
return ParsedStatement(StatementType.DDL, query)
194193

195194
if RE_IS_INSERT.match(query):
196-
return STMT_INSERT
195+
return ParsedStatement(StatementType.INSERT, query)
197196

198197
if RE_NON_UPDATE.match(query) or RE_WITH.match(query):
199198
# As of 13-March-2020, Cloud Spanner only supports WITH for DQL
200199
# statements and doesn't yet support WITH for DML statements.
201-
return STMT_NON_UPDATING
200+
return ParsedStatement(StatementType.QUERY, query)
202201

203-
return STMT_UPDATING
202+
return ParsedStatement(StatementType.UPDATE, query)
204203

205204

206205
def sql_pyformat_args_to_spanner(sql, params):

‎google/cloud/spanner_dbapi/parsed_statement.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Copyright 20203 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from dataclasses import dataclass
16+
from enum import Enum
17+
18+
19+
class StatementType(Enum):
20+
CLIENT_SIDE = 1
21+
DDL = 2
22+
QUERY = 3
23+
UPDATE = 4
24+
INSERT = 5
25+
26+
27+
class ClientSideStatementType(Enum):
28+
COMMIT = 1
29+
BEGIN = 2
30+
31+
32+
@dataclass
33+
class ParsedStatement:
34+
statement_type: StatementType
35+
query: str
36+
client_side_statement_type: ClientSideStatementType = None

‎tests/system/test_dbapi.py‎

Lines changed: 52 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020

2121
from google.cloud import spanner_v1
2222
from google.cloud._helpers import UTC
23+
24+
from google.cloud.spanner_dbapi import Cursor
2325
from google.cloud.spanner_dbapi.connection import connect
2426
from google.cloud.spanner_dbapi.connection import Connection
2527
from google.cloud.spanner_dbapi.exceptions import ProgrammingError
@@ -72,37 +74,11 @@ def dbapi_database(raw_database):
7274

7375
def test_commit(shared_instance, dbapi_database):
7476
"""Test committing a transaction with several statements."""
75-
want_row = (
76-
1,
77-
"updated-first-name",
78-
"last-name",
79-
"test.email_updated@domen.ru",
80-
)
8177
# connect to the test database
8278
conn = Connection(shared_instance, dbapi_database)
8379
cursor = conn.cursor()
8480

85-
# execute several DML statements within one transaction
86-
cursor.execute(
87-
"""
88-
INSERT INTO contacts (contact_id, first_name, last_name, email)
89-
VALUES (1, 'first-name', 'last-name', 'test.email@domen.ru')
90-
"""
91-
)
92-
cursor.execute(
93-
"""
94-
UPDATE contacts
95-
SET first_name = 'updated-first-name'
96-
WHERE first_name = 'first-name'
97-
"""
98-
)
99-
cursor.execute(
100-
"""
101-
UPDATE contacts
102-
SET email = 'test.email_updated@domen.ru'
103-
WHERE email = 'test.email@domen.ru'
104-
"""
105-
)
81+
want_row = _execute_common_precommit_statements(cursor)
10682
conn.commit()
10783

10884
# read the resulting data from the database
@@ -116,6 +92,25 @@ def test_commit(shared_instance, dbapi_database):
11692
conn.close()
11793

11894

95+
def test_commit_client_side(shared_instance, dbapi_database):
96+
"""Test committing a transaction with several statements."""
97+
# connect to the test database
98+
conn = Connection(shared_instance, dbapi_database)
99+
cursor = conn.cursor()
100+
101+
want_row = _execute_common_precommit_statements(cursor)
102+
cursor.execute("""COMMIT""")
103+
104+
# read the resulting data from the database
105+
cursor.execute("SELECT * FROM contacts")
106+
got_rows = cursor.fetchall()
107+
conn.commit()
108+
cursor.close()
109+
conn.close()
110+
111+
assert got_rows == [want_row]
112+
113+
119114
def test_rollback(shared_instance, dbapi_database):
120115
"""Test rollbacking a transaction with several statements."""
121116
want_row = (2, "first-name", "last-name", "test.email@domen.ru")
@@ -810,3 +805,33 @@ def test_dml_returning_delete(shared_instance, dbapi_database, autocommit):
810805
assert cur.fetchone() == (1, "first-name")
811806
assert cur.rowcount == 1
812807
conn.commit()
808+
809+
810+
def _execute_common_precommit_statements(cursor: Cursor):
811+
# execute several DML statements within one transaction
812+
cursor.execute(
813+
"""
814+
INSERT INTO contacts (contact_id, first_name, last_name, email)
815+
VALUES (1, 'first-name', 'last-name', 'test.email@domen.ru')
816+
"""
817+
)
818+
cursor.execute(
819+
"""
820+
UPDATE contacts
821+
SET first_name = 'updated-first-name'
822+
WHERE first_name = 'first-name'
823+
"""
824+
)
825+
cursor.execute(
826+
"""
827+
UPDATE contacts
828+
SET email = 'test.email_updated@domen.ru'
829+
WHERE email = 'test.email@domen.ru'
830+
"""
831+
)
832+
return (
833+
1,
834+
"updated-first-name",
835+
"last-name",
836+
"test.email_updated@domen.ru",
837+
)

0 commit comments

Comments
 (0)