Enhance DML detection and command tag consistency

- Centralize SQL statement categorization in IRISSQLParser.
- Add support for MERGE and TRUNCATE command tags.
- Use centralized parser in protocol handler to ensure consistent batch flushing.
- Refactor embedded execution path to use shared command tag logic.

Author: tom-dyar

src/iris_pgwire/__init__.py

@@ -6,7 +6,7 @@
 caretdev/sqlalchemy-iris.
 """
 
-__version__ = "1.2.5"
+__version__ = "1.2.6"
 __author__ = "Thomas Dyar <thomas.dyar@intersystems.com>"
 
 # Don't import server/protocol in __init__ to avoid sys.modules conflicts

src/iris_pgwire/iris_executor.py

@@ -30,6 +30,7 @@
     SQLTranslator,  # Feature 021: PostgreSQL→IRIS normalization
     TransactionTranslator,
 )  # Feature 022: PostgreSQL transaction verb translation
+from .sql_translator.parser import get_parser
 from .sql_translator.alias_extractor import AliasExtractor  # Column alias preservation
 from .sql_translator.performance_monitor import MetricType, PerformanceTracker, get_monitor
 from .type_mapping import (
@@ -104,6 +105,7 @@ def __init__(self, iris_config: dict[str, Any], server=None):
         self.sql_pipeline = SQLPipeline()
         self.sql_interceptor = SQLInterceptor(self)
         self.sql_translator = self.sql_pipeline.translator
+        self.sql_parser = get_parser()
         self.transaction_translator = TransactionTranslator()
 
         # Connection pool management
@@ -121,9 +123,9 @@ def __init__(self, iris_config: dict[str, Any], server=None):
 
         logger.info(
             "IRIS executor initialized",
-            host=iris_config.get("host"),
-            port=iris_config.get("port"),
-            namespace=iris_config.get("namespace"),
+            host=self.iris_config.get("host"),
+            port=self.iris_config.get("port"),
+            namespace=self.iris_config.get("namespace"),
             embedded_mode=self.embedded_mode,
         )
 
@@ -3241,15 +3243,7 @@ def _sync_execute():
                         session_id=session_id,
                     )
                     # Determine command tag from SQL
-                    sql_upper = sql.strip().upper()
-                    if sql_upper.startswith("DELETE"):
-                        command_tag = "DELETE"
-                    elif sql_upper.startswith("UPDATE"):
-                        command_tag = "UPDATE"
-                    elif sql_upper.startswith("INSERT"):
-                        command_tag = "INSERT"
-                    else:
-                        command_tag = "UNKNOWN"
+                    command_tag = self._determine_command_tag(sql, 0)
 
                     return {
                         "success": True,  # SQLCODE 100 is success!
@@ -4380,30 +4374,27 @@ def _map_iris_type_to_oid(self, iris_type: str) -> int:
 
     def _determine_command_tag(self, sql: str, row_count: int) -> str:
         """Determine PostgreSQL command tag from SQL"""
-        sql_upper = sql.upper().strip()
+        # Normalize: strip and get first word
+        sql_clean = sql.strip().upper()
+        if not sql_clean:
+            return "UNKNOWN"
 
-        if sql_upper.startswith("SELECT"):
+        first_word = sql_clean.split()[0] if sql_clean.split() else ""
+
+        if first_word == "SELECT":
             return "SELECT"
-        elif sql_upper.startswith("INSERT"):
+        elif first_word == "INSERT":
             return f"INSERT 0 {row_count}"
-        elif sql_upper.startswith("UPDATE"):
+        elif first_word == "UPDATE":
             return f"UPDATE {row_count}"
-        elif sql_upper.startswith("DELETE"):
+        elif first_word == "DELETE":
             return f"DELETE {row_count}"
-        elif sql_upper.startswith("CREATE"):
-            return "CREATE"
-        elif sql_upper.startswith("DROP"):
-            return "DROP"
-        elif sql_upper.startswith("ALTER"):
-            return "ALTER"
-        elif sql_upper.startswith("BEGIN"):
-            return "BEGIN"
-        elif sql_upper.startswith("COMMIT"):
-            return "COMMIT"
-        elif sql_upper.startswith("ROLLBACK"):
-            return "ROLLBACK"
-        elif sql_upper.startswith("SHOW"):
-            return "SHOW"
+        elif first_word == "MERGE":
+            return f"MERGE {row_count}"
+        elif first_word == "TRUNCATE":
+            return "TRUNCATE"
+        elif first_word in ("CREATE", "DROP", "ALTER", "BEGIN", "COMMIT", "ROLLBACK", "SHOW"):
+            return first_word
         else:
             return "UNKNOWN"
 

src/iris_pgwire/protocol.py

@@ -2906,24 +2906,26 @@ async def handle_describe_message(self, body: bytes):
                     connection_id=self.connection_id,
                     statement_name=name,
                     query_preview=query[:100],
-                    is_select=query_upper.startswith("SELECT"),
-                    is_show=query_upper.startswith("SHOW"),
+                    is_select=self.iris_executor.sql_parser.is_select_statement(query),
+                    is_show=self.iris_executor.sql_parser.is_show_statement(query),
                 )
 
                 # Check if query has RETURNING clause (INSERT/UPDATE/DELETE with RETURNING)
-                has_returning = "RETURNING" in query_upper
+                has_returning = self.iris_executor.sql_parser.has_returning_clause(query)
 
                 if (
-                    query_upper.startswith("SELECT")
-                    or query_upper.startswith("SHOW")
+                    self.iris_executor.sql_parser.is_select_statement(query)
+                    or self.iris_executor.sql_parser.is_show_statement(query)
                     or has_returning
                 ):
                     # Execute metadata discovery to get column information
                     # Use LIMIT 0 pattern to avoid fetching actual data
                     # For RETURNING queries, we'll send synthetic column metadata based on RETURNING columns
                     try:
                         # Special handling for RETURNING queries - extract columns from RETURNING clause
-                        if has_returning and not query_upper.startswith("SELECT"):
+                        if has_returning and not self.iris_executor.sql_parser.is_select_statement(
+                            query
+                        ):
                             import re
 
                             # Extract RETURNING columns from query
@@ -3085,14 +3087,19 @@ async def handle_describe_message(self, body: bytes):
 
                     # Batch Execution Interception:
                     # Short-circuit Describe portal for DML statements.
-                    is_dml = any(query_upper.startswith(k) for k in ["INSERT", "UPDATE", "DELETE"])
-                    has_returning = "RETURNING" in query_upper
-
-                    if is_dml and not has_returning:
+                    # DML portals (without RETURNING) don't have RowDescription.
+                    is_dml = self.iris_executor.sql_parser.is_dml_statement(query)
+                    if is_dml:
+                        logger.info(
+                            "Describe portal: DML statement (no row metadata)", portal_name=name
+                        )
                         await self.send_no_data()
                         return
 
-                    if query_upper.startswith("SELECT") or query_upper.startswith("SHOW"):
+                    # Metadata discovery for SELECT/SHOW
+                    if self.iris_executor.sql_parser.is_select_statement(
+                        query
+                    ) or self.iris_executor.sql_parser.is_show_statement(query):
                         try:
                             result = await self.iris_executor.execute_query(
                                 query, params=portal.get("params", [])
@@ -3277,9 +3284,8 @@ async def handle_execute_message(self, body: bytes):
             # Intercept INSERT/UPDATE/DELETE (DML) without RETURNING for protocol-level batching.
             # Standard PostgreSQL clients (psycopg3) send Sync every 5 rows, which is slow.
             # We buffer parameters and send synthetic CommandComplete to keep client pipe full.
-            query_upper = query.strip().upper()
-            is_dml = any(query_upper.startswith(k) for k in ["INSERT", "UPDATE", "DELETE"])
-            has_returning = "RETURNING" in query_upper
+            is_dml = self.iris_executor.sql_parser.is_dml_statement(query)
+            has_returning = self.iris_executor.sql_parser.has_returning_clause(query)
 
             if is_dml and not has_returning:
                 # Store SQL if first row in batch
@@ -3291,6 +3297,7 @@ async def handle_execute_message(self, body: bytes):
 
                 # Send synthetic CommandComplete immediately to client
                 # This tricks the client into sending the next row immediately.
+                query_upper = query.strip().upper()
                 tag = f"{query_upper.split()[0]} 0 1\x00".encode()
                 msg_len = 4 + len(tag)
                 self.writer.write(struct.pack("!cI", MSG_COMMAND_COMPLETE, msg_len) + tag)

src/iris_pgwire/sql_translator/parser.py

@@ -593,10 +593,18 @@ def is_select_statement(self, sql: str) -> bool:
         sql_clean = sql.strip().upper()
         return sql_clean.startswith("SELECT") or sql_clean.startswith("WITH")
 
+    def is_show_statement(self, sql: str) -> bool:
+        """Check if SQL is a SHOW statement"""
+        return sql.strip().upper().startswith("SHOW")
+
     def is_dml_statement(self, sql: str) -> bool:
-        """Check if SQL is a DML statement (INSERT, UPDATE, DELETE)"""
+        """Check if SQL is a DML statement (INSERT, UPDATE, DELETE, MERGE)"""
         sql_clean = sql.strip().upper()
-        return any(sql_clean.startswith(stmt) for stmt in ["INSERT", "UPDATE", "DELETE"])
+        return any(sql_clean.startswith(stmt) for stmt in ["INSERT", "UPDATE", "DELETE", "MERGE"])
+
+    def has_returning_clause(self, sql: str) -> bool:
+        """Check if SQL has a RETURNING clause"""
+        return "RETURNING" in sql.upper()
 
     def is_ddl_statement(self, sql: str) -> bool:
         """Check if SQL is a DDL statement"""