fix(spanner/dbapi): bypass sqlparse for RUN PARTITION commands

Author: sinhasubham

google/cloud/spanner_dbapi/parse_utils.py

@@ -233,6 +233,11 @@ def classify_statement(query, args=None):
     :rtype: ParsedStatement
     :returns: parsed statement attributes.
     """
+    # Check for RUN PARTITION command to avoid sqlparse processing it.
+    # sqlparse fails with "Maximum grouping depth exceeded" on long partition IDs.
+    if re.match(r"^\s*RUN\s+PARTITION\s+.+", query, re.IGNORECASE):
+        return client_side_statement_parser.parse_stmt(query.strip())
+
     # sqlparse will strip Cloud Spanner comments,
     # still, special commenting styles, like
     # PostgreSQL dollar quoted comments are not

tests/unit/spanner_dbapi/test_parse_utils.py

@@ -200,6 +200,29 @@ def test_run_partition_classify_stmt(self):
             ),
         )
 
+    def test_run_partition_classify_stmt_long_id(self):
+        # Regression test for "Maximum grouping depth exceeded" with sqlparse
+        long_id = "a" * 5000
+        query = f"RUN PARTITION {long_id}"
+        parsed_statement = classify_statement(query)
+        self.assertEqual(
+            parsed_statement,
+            ParsedStatement(
+                StatementType.CLIENT_SIDE,
+                Statement(query),
+                ClientSideStatementType.RUN_PARTITION,
+                [long_id],
+            ),
+        )
+
+    def test_run_partition_classify_stmt_incomplete(self):
+        # "RUN PARTITION" without ID should be classified as UNKNOWN (not None)
+        # because it falls through the specific check and sqlparse handles it.
+        query = "RUN PARTITION"
+        parsed_statement = classify_statement(query)
+        self.assertEqual(parsed_statement.statement_type, StatementType.UNKNOWN)
+        self.assertEqual(parsed_statement.statement.sql, query)
+
     def test_run_partitioned_query_classify_stmt(self):
         parsed_statement = classify_statement(
             " RUN PARTITIONED  QUERY  SELECT s.SongName FROM Songs AS s  "