fix: improve client-side regex statement parser (#1328)

* fix: improve client-side regex statement parser

The client-side regex-based statement parser contained multiple minor errors, like:
- BEGIN would match any string as BEGIN TRANSACTION (including stuff like `BEGIN foo`)
- COMMIT and ROLLBACK had the same problem as BEGIN.
- Mismatches were reported as UPDATE. They are now returned as UNKNOWN.
- DLL missed the ANALYZE keyword

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: olavloite

google/cloud/spanner_dbapi/batch_dml_executor.py

@@ -54,9 +54,12 @@ def execute_statement(self, parsed_statement: ParsedStatement):
         """
         from google.cloud.spanner_dbapi import ProgrammingError
 
+        # Note: Let the server handle it if the client-side parser did not
+        # recognize the type of statement.
         if (
             parsed_statement.statement_type != StatementType.UPDATE
             and parsed_statement.statement_type != StatementType.INSERT
+            and parsed_statement.statement_type != StatementType.UNKNOWN
         ):
             raise ProgrammingError("Only DML statements are allowed in batch DML mode.")
         self._statements.append(parsed_statement.statement)

google/cloud/spanner_dbapi/client_side_statement_parser.py

@@ -21,18 +21,18 @@
     Statement,
 )
 
-RE_BEGIN = re.compile(r"^\s*(BEGIN|START)(TRANSACTION)?", re.IGNORECASE)
-RE_COMMIT = re.compile(r"^\s*(COMMIT)(TRANSACTION)?", re.IGNORECASE)
-RE_ROLLBACK = re.compile(r"^\s*(ROLLBACK)(TRANSACTION)?", re.IGNORECASE)
+RE_BEGIN = re.compile(r"^\s*(BEGIN|START)(\s+TRANSACTION)?\s*$", re.IGNORECASE)
+RE_COMMIT = re.compile(r"^\s*(COMMIT)(\s+TRANSACTION)?\s*$", re.IGNORECASE)
+RE_ROLLBACK = re.compile(r"^\s*(ROLLBACK)(\s+TRANSACTION)?\s*$", re.IGNORECASE)
 RE_SHOW_COMMIT_TIMESTAMP = re.compile(
-    r"^\s*(SHOW)\s+(VARIABLE)\s+(COMMIT_TIMESTAMP)", re.IGNORECASE
+    r"^\s*(SHOW)\s+(VARIABLE)\s+(COMMIT_TIMESTAMP)\s*$", re.IGNORECASE
 )
 RE_SHOW_READ_TIMESTAMP = re.compile(
-    r"^\s*(SHOW)\s+(VARIABLE)\s+(READ_TIMESTAMP)", re.IGNORECASE
+    r"^\s*(SHOW)\s+(VARIABLE)\s+(READ_TIMESTAMP)\s*$", re.IGNORECASE
 )
-RE_START_BATCH_DML = re.compile(r"^\s*(START)\s+(BATCH)\s+(DML)", re.IGNORECASE)
-RE_RUN_BATCH = re.compile(r"^\s*(RUN)\s+(BATCH)", re.IGNORECASE)
-RE_ABORT_BATCH = re.compile(r"^\s*(ABORT)\s+(BATCH)", re.IGNORECASE)
+RE_START_BATCH_DML = re.compile(r"^\s*(START)\s+(BATCH)\s+(DML)\s*$", re.IGNORECASE)
+RE_RUN_BATCH = re.compile(r"^\s*(RUN)\s+(BATCH)\s*$", re.IGNORECASE)
+RE_ABORT_BATCH = re.compile(r"^\s*(ABORT)\s+(BATCH)\s*$", re.IGNORECASE)
 RE_PARTITION_QUERY = re.compile(r"^\s*(PARTITION)\s+(.+)", re.IGNORECASE)
 RE_RUN_PARTITION = re.compile(r"^\s*(RUN)\s+(PARTITION)\s+(.+)", re.IGNORECASE)
 RE_RUN_PARTITIONED_QUERY = re.compile(

google/cloud/spanner_dbapi/connection.py

@@ -20,11 +20,7 @@
 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 BatchMode, BatchDmlExecutor
-from google.cloud.spanner_dbapi.parse_utils import _get_statement_type
-from google.cloud.spanner_dbapi.parsed_statement import (
-    StatementType,
-    AutocommitDmlMode,
-)
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from google.cloud.spanner_dbapi.partition_helper import PartitionId
 from google.cloud.spanner_dbapi.parsed_statement import ParsedStatement, Statement
 from google.cloud.spanner_dbapi.transaction_helper import TransactionRetryHelper
@@ -702,10 +698,6 @@ def set_autocommit_dml_mode(
         self._autocommit_dml_mode = autocommit_dml_mode
 
     def _partitioned_query_validation(self, partitioned_query, statement):
-        if _get_statement_type(Statement(partitioned_query)) is not StatementType.QUERY:
-            raise ProgrammingError(
-                "Only queries can be partitioned. Invalid statement: " + statement.sql
-            )
         if self.read_only is not True and self._client_transaction_started is True:
             raise ProgrammingError(
                 "Partitioned query is not supported, because the connection is in a read/write transaction."

google/cloud/spanner_dbapi/cursor.py

@@ -404,9 +404,12 @@ def executemany(self, operation, seq_of_params):
             # For every operation, we've got to ensure that any prior DDL
             # statements were run.
             self.connection.run_prior_DDL_statements()
+            # Treat UNKNOWN statements as if they are DML and let the server
+            # determine what is wrong with it.
             if self._parsed_statement.statement_type in (
                 StatementType.INSERT,
                 StatementType.UPDATE,
+                StatementType.UNKNOWN,
             ):
                 statements = []
                 for params in seq_of_params:

google/cloud/spanner_dbapi/parse_utils.py

@@ -155,25 +155,30 @@
 STMT_INSERT = "INSERT"
 
 # Heuristic for identifying statements that don't need to be run as updates.
+# TODO: This and the other regexes do not match statements that start with a hint.
 RE_NON_UPDATE = re.compile(r"^\W*(SELECT|GRAPH|FROM)", re.IGNORECASE)
 
 RE_WITH = re.compile(r"^\s*(WITH)", re.IGNORECASE)
 
 # DDL statements follow
 # https://cloud.google.com/spanner/docs/data-definition-language
 RE_DDL = re.compile(
-    r"^\s*(CREATE|ALTER|DROP|GRANT|REVOKE|RENAME)", re.IGNORECASE | re.DOTALL
+    r"^\s*(CREATE|ALTER|DROP|GRANT|REVOKE|RENAME|ANALYZE)", re.IGNORECASE | re.DOTALL
 )
 
-RE_IS_INSERT = re.compile(r"^\s*(INSERT)", re.IGNORECASE | re.DOTALL)
+# TODO: These do not match statements that start with a hint.
+RE_IS_INSERT = re.compile(r"^\s*(INSERT\s+)", re.IGNORECASE | re.DOTALL)
+RE_IS_UPDATE = re.compile(r"^\s*(UPDATE\s+)", re.IGNORECASE | re.DOTALL)
+RE_IS_DELETE = re.compile(r"^\s*(DELETE\s+)", re.IGNORECASE | re.DOTALL)
 
 RE_INSERT = re.compile(
     # Only match the `INSERT INTO <table_name> (columns...)
     # otherwise the rest of the statement could be a complex
     # operation.
-    r"^\s*INSERT INTO (?P<table_name>[^\s\(\)]+)\s*\((?P<columns>[^\(\)]+)\)",
+    r"^\s*INSERT(?:\s+INTO)?\s+(?P<table_name>[^\s\(\)]+)\s*\((?P<columns>[^\(\)]+)\)",
     re.IGNORECASE | re.DOTALL,
 )
+"""Deprecated: Use the RE_IS_INSERT, RE_IS_UPDATE, and RE_IS_DELETE regexes"""
 
 RE_VALUES_TILL_END = re.compile(r"VALUES\s*\(.+$", re.IGNORECASE | re.DOTALL)
 
@@ -259,8 +264,13 @@ def _get_statement_type(statement):
         # statements and doesn't yet support WITH for DML statements.
         return StatementType.QUERY
 
-    statement.sql = ensure_where_clause(query)
-    return StatementType.UPDATE
+    if RE_IS_UPDATE.match(query) or RE_IS_DELETE.match(query):
+        # TODO: Remove this? It makes more sense to have this in SQLAlchemy and
+        #       Django than here.
+        statement.sql = ensure_where_clause(query)
+        return StatementType.UPDATE
+
+    return StatementType.UNKNOWN
 
 
 def sql_pyformat_args_to_spanner(sql, params):
@@ -355,7 +365,7 @@ def get_param_types(params):
 def ensure_where_clause(sql):
     """
     Cloud Spanner requires a WHERE clause on UPDATE and DELETE statements.
-    Add a dummy WHERE clause if non detected.
+    Add a dummy WHERE clause if not detected.
 
     :type sql: str
     :param sql: SQL code to check.

google/cloud/spanner_dbapi/parsed_statement.py

@@ -17,6 +17,7 @@
 
 
 class StatementType(Enum):
+    UNKNOWN = 0
     CLIENT_SIDE = 1
     DDL = 2
     QUERY = 3

tests/unit/spanner_dbapi/test_parse_utils.py

@@ -74,11 +74,31 @@ def test_classify_stmt(self):
             ("REVOKE SELECT ON TABLE Singers TO ROLE parent", StatementType.DDL),
             ("GRANT ROLE parent TO ROLE child", StatementType.DDL),
             ("INSERT INTO table (col1) VALUES (1)", StatementType.INSERT),
+            ("INSERT table (col1) VALUES (1)", StatementType.INSERT),
+            ("INSERT OR UPDATE table (col1) VALUES (1)", StatementType.INSERT),
+            ("INSERT OR IGNORE table (col1) VALUES (1)", StatementType.INSERT),
             ("UPDATE table SET col1 = 1 WHERE col1 = NULL", StatementType.UPDATE),
+            ("delete from table WHERE col1 = 2", StatementType.UPDATE),
+            ("delete from table WHERE col1 in (select 1)", StatementType.UPDATE),
+            ("dlete from table where col1 = 2", StatementType.UNKNOWN),
+            ("udpate table set col2=1 where col1 = 2", StatementType.UNKNOWN),
+            ("begin foo", StatementType.UNKNOWN),
+            ("begin transaction foo", StatementType.UNKNOWN),
+            ("commit foo", StatementType.UNKNOWN),
+            ("commit transaction foo", StatementType.UNKNOWN),
+            ("rollback foo", StatementType.UNKNOWN),
+            ("rollback transaction foo", StatementType.UNKNOWN),
+            ("show variable", StatementType.UNKNOWN),
+            ("show variable read_timestamp foo", StatementType.UNKNOWN),
+            ("INSERTs INTO table (col1) VALUES (1)", StatementType.UNKNOWN),
+            ("UPDATEs table SET col1 = 1 WHERE col1 = NULL", StatementType.UNKNOWN),
+            ("DELETEs from table WHERE col1 = 2", StatementType.UNKNOWN),
         )
 
         for query, want_class in cases:
-            self.assertEqual(classify_statement(query).statement_type, want_class)
+            self.assertEqual(
+                classify_statement(query).statement_type, want_class, query
+            )
 
     def test_partition_query_classify_stmt(self):
         parsed_statement = classify_statement(