8000 refactor: update classify_stmt() helper and cover it with unit tests … · FirePing32/python-spanner-django@e6f6179 · GitHub
[go: up one dir, main page]
Skip to content

Commit e6f6179

Browse files
author
Ilya Gurov
authored
refactor: update classify_stmt() helper and cover it with unit tests (googleapis#488)
1 parent 6c2f292 commit e6f6179

File tree

2 files changed

+51
-37
lines changed

2 files changed

+51
-37
lines changed

‎google/cloud/spanner_dbapi/parse_utils.py

Lines changed: 43 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
# license that can be found in the LICENSE file or at
55
# https://developers.google.com/open-source/licenses/bsd
66

7+
"SQL parsing and classification utils."
8+
79
import datetime
810
import decimal
911
import re
@@ -23,51 +25,65 @@
2325
STMT_INSERT = "INSERT"
2426

2527
# Heuristic for identifying statements that don't need to be run as updates.
26-
re_NON_UPDATE = re.compile(r"^\s*(SELECT)", re.IGNORECASE)
28+
RE_NON_UPDATE = re.compile(r"^\s*(SELECT)", re.IGNORECASE)
2729

28-
re_WITH = re.compile(r"^\s*(WITH)", re.IGNORECASE)
30+
RE_WITH = re.compile(r"^\s*(WITH)", re.IGNORECASE)
2931

3032
# DDL statements follow
3133
# https://cloud.google.com/spanner/docs/data-definition-language
32-
re_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP)", re.IGNORECASE | re.DOTALL)
33-
34-
re_IS_INSERT = re.compile(r"^\s*(INSERT)", re.IGNORECASE | re.DOTALL)
34+
RE_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP)", re.IGNORECASE | re.DOTALL)
3535

36+
RE_IS_INSERT = re.compile(r"^\s*(INSERT)", re.IGNORECASE | re.DOTALL)
3637

37-
def classify_stmt(sql):
38-
if re_DDL.match(sql):
39-
return STMT_DDL
40-
elif re_IS_INSERT.match(sql):
41-
return STMT_INSERT
42-
elif re_NON_UPDATE.match(sql):
43-
return STMT_NON_UPDATING
44-
elif re_WITH.match(sql):
45-
# As of Fri-13th-March-2020, Cloud Spanner only supports WITH for DQL
46-
# statements and doesn't yet support WITH for DML statements.
47-
# When WITH for DML is added, we'll need to update this classifier
48-
# accordingly.
49-
return STMT_NON_UPDATING
50-
else:
51-
return STMT_UPDATING
52-
53-
54-
re_INSERT = re.compile(
38+
RE_INSERT = re.compile(
5539
# Only match the `INSERT INTO <table_name> (columns...)
5640
# otherwise the rest of the statement could be a complex
5741
# operation.
5842
r"^\s*INSERT INTO (?P<table_name>[^\s\(\)]+)\s*\((?P<columns>[^\(\)]+)\)",
5943
re.IGNORECASE | re.DOTALL,
6044
)
6145

62-
re_VALUES_TILL_END = re.compile(r"VALUES\s*\(.+$", re.IGNORECASE | re.DOTALL)
46+
RE_VALUES_TILL_END = re.compile(r"VALUES\s*\(.+$", re.IGNORECASE | re.DOTALL)
6347

64-
re_VALUES_PYFORMAT = re.compile(
48+
RE_VALUES_PYFORMAT = re.compile(
6549
# To match: (%s, %s,....%s)
6650
r"(\(\s*%s[^\(\)]+\))",
6751
re.DOTALL,
6852
)
6953

7054

55+
def classify_stmt(query):
56+
"""Determine SQL query type.
57+
58+
:type query: :class:`str`
59+
:param query: SQL query.
60+
61+
:rtype: :class:`str`
62+
:returns: Query type name.
63+
"""
64+
if RE_DDL.match(query):
65+
return STMT_DDL
66+
67+
if RE_IS_INSERT.match(query):
68+
return STMT_INSERT
69+
70+
if RE_NON_UPDATE.match(query) or RE_WITH.match(query):
71+
# As of 13-March-2020, Cloud Spanner only supports WITH for DQL
72+
# statements and doesn't yet support WITH for DML statements.
73+
return STMT_NON_UPDATING
74+
75+
return STMT_UPDATING
76+
77+
78+
def strip_backticks(name):
79+
"""
80+
Strip backticks off of quoted names For example, '`no`' (a Spanner reserved
81+
word) becomes 'no'.
82+
"""
83+
has_quotes = name.startswith("`") and name.endswith("`")
84+
return name[1:-1] if has_quotes else name
85+
86+
7187
def parse_insert(insert_sql, params):
7288
"""
7389
Parse an INSERT statement an generate a list of tuples of the form:
@@ -126,14 +142,14 @@ def parse_insert(insert_sql, params):
126142
],
127143
}
128144
""" # noqa
129-
match = re_INSERT.search(insert_sql)
145+
match = RE_INSERT.search(insert_sql)
130146

131147
if not match:
132148
raise ProgrammingError(
133149
"Could not parse an INSERT statement from %s" % insert_sql
134150
)
135151

136-
after_values_sql = re_VALUES_TILL_END.findall(insert_sql)
152+
after_values_sql = RE_VALUES_TILL_END.findall(insert_sql)
137153
if not after_values_sql:
138154
# Case b)
139155
insert_sql = sanitize_literals_for_upload(insert_sql)

‎tests/spanner_dbapi/test_parse_utils.py

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
from google.cloud.spanner_dbapi.exceptions import Error, ProgrammingError
1313
from google.cloud.spanner_dbapi.parse_utils import (
1414
STMT_DDL,
15+
STMT_INSERT,
1516
STMT_NON_UPDATING,
17+
STMT_UPDATING,
1618
DateStr,
1719
TimestampStr,
1820
classify_stmt,
@@ -28,7 +30,7 @@
2830

2931
class ParseUtilsTests(TestCase):
3032
def test_classify_stmt(self):
31-
cases = [
33+
cases = (
3234
("SELECT 1", STMT_NON_UPDATING),
3335
("SELECT s.SongName FROM Songs AS s", STMT_NON_UPDATING),
3436
(
@@ -50,16 +52,12 @@ def test_classify_stmt(self):
5052
"CREATE INDEX AlbumsByAlbumTitle2 ON Albums(AlbumTitle) STORING (MarketingBudget)",
5153
STMT_DDL,
5254
),
53-
]
55+
("INSERT INTO table (col1) VALUES (1)", STMT_INSERT),
56+
("UPDATE table SET col1 = 1 WHERE col1 = NULL", STMT_UPDATING),
57+
)
5458

55-
for tt in cases:
56-
sql, want_classification = tt
57-
got_classification = classify_stmt(sql)
58-
self.assertEqual(
59-
got_classification,
60-
want_classification,
61-
"Classification mismatch",
62-
)
59+
for query, want_class in cases:
60+
self.assertEqual(classify_stmt(query), want_class)
6361

6462
def test_parse_insert(self):
6563
cases = [

0 commit comments

Comments
 (0)
0