iris-pgwire 1.2.33__py3-none-any.whl → 1.2.34__py3-none-any.whl

iris_pgwire/__init__.py

@@ -6,7 +6,7 @@ Based on the specification from docs/iris_pgwire_plan.md and proven patterns fro
 caretdev/sqlalchemy-iris.
 """
 
-__version__ = "1.2.33"
+__version__ = "1.2.34"
 __author__ = "Thomas Dyar <thomas.dyar@intersystems.com>"
 
 # Don't import server/protocol in __init__ to avoid sys.modules conflicts

iris_pgwire/dbapi_executor.py

@@ -14,6 +14,7 @@ Contract: contracts/dbapi-executor-contract.md
 """
 
 import asyncio
+import datetime as dt
 import re
 import time
 from typing import Any
@@ -25,12 +26,43 @@ from iris_pgwire.dbapi_connection_pool import IRISConnectionPool
 from iris_pgwire.models.backend_config import BackendConfig
 from iris_pgwire.models.connection_pool_state import ConnectionPoolState
 from iris_pgwire.models.vector_query_request import VectorQueryRequest
+from iris_pgwire.schema_mapper import IRIS_SCHEMA
 from iris_pgwire.sql_translator import SQLPipeline
 from iris_pgwire.sql_translator.parser import get_parser
 
 logger = structlog.get_logger(__name__)
 
 
+class MockResult:
+    """Mock result object for RETURNING emulation"""
+
+    def __init__(self, rows, meta=None):
+        self._rows = rows if rows is not None else []
+        self._meta = meta
+        self.description = meta
+        self.rowcount = len(self._rows)
+        self._index = 0
+
+    def __iter__(self):
+        return iter(self._rows)
+
+    def fetchall(self):
+        return self._rows
+
+    def fetchone(self):
+        if self._index < len(self._rows):
+            row = self._rows[self._index]
+            self._index += 1
+            return row
+        return None
+
+    def fetch(self):
+        return self._rows
+
+    def close(self):
+        pass
+
+
 class DBAPIExecutor:
     """
     Execute SQL queries against IRIS via DBAPI backend.
@@ -103,11 +135,9 @@ class DBAPIExecutor:
     def _convert_value_for_iris(self, value: Any) -> Any:
         """Helper to convert a single value."""
         if isinstance(value, str):
-            # FR-004: Normalize ISO 8601 timestamp strings for IRIS
-            # Handles: YYYY-MM-DD[T ]HH:MM:SS[.fff][Z|[+-]HH:MM]
+            # Check for ISO 8601 timestamp: 2026-01-29T21:27:38.111Z
+            # or 2026-01-29T21:27:38.111+00:00
             # IRIS rejects the 'T' and 'Z' or offset in %PosixTime/TIMESTAMP
-            import re
-
             ts_match = re.match(
                 r"^(\d{4}-\d{2}-\d{2})[T ](\d{2}:\d{2}:\d{2}(?:\.\d+)?)(?:Z|[+-]\d{2}:?(\d{2})?)?$",
                 value,
@@ -151,21 +181,55 @@ class DBAPIExecutor:
             # Acquire connection from pool
             conn_wrapper = await self.pool.acquire()
 
+            # Detect RETURNING clause
+            has_returning = self.has_returning_clause(sql)
+
             # Execute query in thread pool (DBAPI is synchronous)
             def execute_in_thread():
                 cursor = conn_wrapper.connection.cursor()  # type: ignore
                 try:
-                    # Strip trailing semicolon for IRIS compatibility
-                    clean_sql = sql.strip().rstrip(";")
-
                     # Feature 034: Apply per-session namespace if set
                     if session_id and session_id in self.session_namespaces:
                         ns = self.session_namespaces[session_id]
-                        # In DBAPI, we switch namespace by executing a command if supported,
-                        # but IRIS DBAPI connection is usually fixed to a namespace.
-                        # For now, we log it.
                         logger.debug(f"Session {session_id} using namespace {ns}")
 
+                    # Handle RETURNING emulation
+                    if has_returning:
+                        op, table, cols, where, stripped_sql = self._parse_returning_clause(sql)
+                        if op and table:
+                            # Strip trailing semicolon
+                            clean_sql = stripped_sql.strip().rstrip(";")
+
+                            # For DELETE, we must fetch BEFORE deleting
+                            delete_rows = []
+                            delete_meta = None
+                            if op == "DELETE":
+                                delete_rows, delete_meta = self._emulate_returning_sync(
+                                    cursor, op, table, cols, where, converted_params, sql
+                                )
+
+                            # Execute the main statement
+                            if converted_params:
+                                cursor.execute(clean_sql, converted_params)
+                            else:
+                                cursor.execute(clean_sql)
+
+                            # Emulate RETURNING result
+                            if op == "DELETE":
+                                rows = delete_rows
+                                columns = delete_meta
+                            else:
+                                rows, columns = self._emulate_returning_sync(
+                                    cursor, op, table, cols, where, converted_params, sql
+                                )
+
+                            row_count = len(rows)
+                            return rows, columns, row_count
+
+                    # Standard execution path
+                    # Strip trailing semicolon for IRIS compatibility
+                    clean_sql = sql.strip().rstrip(";")
+
                     if converted_params:
                         cursor.execute(clean_sql, converted_params)
                     else:
@@ -260,13 +324,8 @@ class DBAPIExecutor:
         """
         Execute SQL with multiple parameter sets for batch operations.
 
-        Args:
-            sql: SQL query string (usually INSERT)
-            params_list: List of parameter tuples/lists
-            session_id: Optional session identifier
-
-        Returns:
-            Dict with execution results (rows_affected, execution_time_ms, etc.)
+        RETURNING SUPPORT: When SQL contains RETURNING clause, executes each statement
+        individually and aggregates the returned rows.
         """
         start_time = time.perf_counter()
         conn_wrapper = None
@@ -275,6 +334,9 @@ class DBAPIExecutor:
             # Translate placeholders ($1 -> ?)
             sql = self._translate_placeholders(sql)
 
+            # Detect RETURNING clause
+            has_returning = self.has_returning_clause(sql)
+
             # Acquire connection from pool
             conn_wrapper = await self.pool.acquire()
 
@@ -285,6 +347,38 @@ class DBAPIExecutor:
                     # Strip trailing semicolon for IRIS compatibility
                     clean_sql = sql.strip().rstrip(";")
 
+                    if has_returning:
+                        op, table, cols, where, stripped_sql = self._parse_returning_clause(sql)
+                        if op and table:
+                            all_rows = []
+                            all_meta = None
+
+                            for params in params_list:
+                                converted_params = self._convert_params_for_iris(params)
+                                # For DELETE, capture before
+                                if op == "DELETE":
+                                    rows, meta = self._emulate_returning_sync(
+                                        cursor, op, table, cols, where, converted_params, sql
+                                    )
+                                    all_rows.extend(rows)
+                                    if not all_meta:
+                                        all_meta = meta
+
+                                # Execute statement
+                                cursor.execute(stripped_sql.strip().rstrip(";"), converted_params)
+
+                                # For INSERT/UPDATE, capture after
+                                if op != "DELETE":
+                                    rows, meta = self._emulate_returning_sync(
+                                        cursor, op, table, cols, where, converted_params, sql
+                                    )
+                                    all_rows.extend(rows)
+                                    if not all_meta:
+                                        all_meta = meta
+
+                            return all_rows, all_meta or [], len(params_list)
+
+                    # Standard batch execution
                     # Pre-process parameters (e.g. convert lists to IRIS vector strings)
                     final_params_list = []
                     for p_set in params_list:
@@ -307,11 +401,11 @@ class DBAPIExecutor:
                     rows_affected = (
                         cursor.rowcount if hasattr(cursor, "rowcount") else len(params_list)
                     )
-                    return rows_affected
+                    return [], [], rows_affected
                 finally:
                     cursor.close()
 
-            rows_affected = await asyncio.to_thread(execute_batch_in_thread)
+            rows, columns, rows_affected = await asyncio.to_thread(execute_batch_in_thread)
 
             # Record metrics
             elapsed_ms = (time.perf_counter() - start_time) * 1000
@@ -332,9 +426,11 @@ class DBAPIExecutor:
                 "rows_affected": rows_affected,
                 "execution_time_ms": elapsed_ms,
                 "batch_size": len(params_list),
-                "rows": [],
-                "columns": [],
-                "_execution_path": "dbapi_executemany",
+                "rows": rows,
+                "columns": columns,
+                "_execution_path": (
+                    "dbapi_executemany_returning" if has_returning else "dbapi_executemany"
+                ),
             }
 
         except Exception as e:
@@ -433,14 +529,18 @@ class DBAPIExecutor:
         }
 
     def has_returning_clause(self, query: str) -> bool:
-        """Check if query has a RETURNING clause."""
+        """
+        Check if query has a RETURNING clause.
+        """
         if not query:
             return False
         return bool(re.search(r"\bRETURNING\b", query, re.IGNORECASE | re.DOTALL))
 
     def get_returning_columns(self, query: str) -> list[str]:
-        """Extract column names from RETURNING clause."""
-        match = re.search(r"RETURNING\s+(.+)$", query, re.IGNORECASE | re.DOTALL)
+        """
+        Extract column names from RETURNING clause.
+        """
+        match = re.search(r"RETURNING\s+(.+?)(?=$|;)", query, re.IGNORECASE | re.DOTALL)
         if not match:
             return []
         cols_str = match.group(1).strip()
@@ -448,6 +548,356 @@ class DBAPIExecutor:
             return ["*"]
         return [c.strip() for c in cols_str.split(",")]
 
+    def _get_table_columns_from_schema(self, table: str, cursor=None) -> list[str]:
+        """
+        Query INFORMATION_SCHEMA.COLUMNS for the given table.
+        Returns the list of column names in order.
+        """
+        if self.strict_single_connection or cursor is None:
+            return []
+        try:
+            table_clean = table.strip('"').strip("'")
+            metadata_sql = f"""
+                SELECT COLUMN_NAME
+                FROM INFORMATION_SCHEMA.COLUMNS
+                WHERE LOWER(TABLE_NAME) = LOWER('{table_clean}')
+                AND LOWER(TABLE_SCHEMA) = LOWER('{IRIS_SCHEMA}')
+                ORDER BY ORDINAL_POSITION
+            """
+            cursor.execute(metadata_sql)
+            rows = cursor.fetchall()
+            return [row[0] for row in rows]
+        except Exception as e:
+            logger.debug(f"Failed to get columns from schema for {table}: {e}")
+        return []
+
+    def _get_column_type_from_schema(self, table: str, column: str, cursor=None) -> int | None:
+        """
+        Query INFORMATION_SCHEMA.COLUMNS for the given table and column.
+        Returns the PostgreSQL type OID.
+        """
+        if self.strict_single_connection or cursor is None:
+            return None
+        try:
+            table_clean = table.strip('"').strip("'")
+            column_clean = column.strip('"').strip("'")
+            metadata_sql = f"""
+                SELECT DATA_TYPE
+                FROM INFORMATION_SCHEMA.COLUMNS
+                WHERE LOWER(TABLE_NAME) = LOWER('{table_clean}')
+                AND LOWER(COLUMN_NAME) = LOWER('{column_clean}')
+                AND LOWER(TABLE_SCHEMA) = LOWER('{IRIS_SCHEMA}')
+            """
+            cursor.execute(metadata_sql)
+            row = cursor.fetchone()
+            if row:
+                iris_type = row[0]
+                return self._map_iris_type_to_oid(iris_type)
+        except Exception as e:
+            logger.debug(f"Failed to get type from schema for {table}.{column}: {e}")
+        return None
+
+    def _infer_type_from_value(self, value, column_name: str | None = None) -> int:
+        """
+        Infer PostgreSQL type OID from Python value
+        """
+        from decimal import Decimal
+
+        if value is None:
+            return 1043  # VARCHAR
+        elif isinstance(value, bool):
+            return 16  # BOOL
+        elif isinstance(value, int):
+            if column_name and any(k in column_name.lower() for k in ("id", "key")):
+                return 20  # BIGINT
+            return 23  # INTEGER
+        elif isinstance(value, float):
+            return 701  # FLOAT8
+        elif isinstance(value, Decimal):
+            return 1700  # NUMERIC
+        elif isinstance(value, dt.datetime):
+            return 1114
+        elif isinstance(value, dt.date):
+            return 1082
+        elif isinstance(value, str):
+            return 1043  # VARCHAR
+        else:
+            return 1043
+
+    def _serialize_value(self, value: Any, type_oid: int) -> Any:
+        """
+        Robust value serialization for PostgreSQL wire protocol compatibility.
+        """
+        if value is None:
+            return None
+
+        if type_oid == 1114:  # TIMESTAMP
+            if isinstance(value, dt.datetime):
+                return value.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+            elif isinstance(value, str):
+                return value  # Already a string
+
+        return value
+
+    def _parse_returning_clause(
+        self, sql: str
+    ) -> tuple[str | None, str | None, Any, str | None, str]:
+        """
+        Parse RETURNING clause from SQL and return metadata.
+        Returns: (operation, table, columns, where_clause, stripped_sql)
+        """
+        returning_operation = None
+        returning_table = None
+        returning_columns = None
+        returning_where_clause = None
+
+        returning_pattern = r"\s+RETURNING\s+(.*?)($|;)"
+        returning_match = re.search(returning_pattern, sql, re.IGNORECASE | re.DOTALL)
+
+        if not returning_match:
+            return None, None, None, None, sql
+
+        returning_clause = returning_match.group(1).strip()
+
+        if returning_clause == "*":
+            returning_columns = "*"
+        else:
+            # Better column parsing that preserves expressions and aliases
+            # Split by commas but respect parentheses
+            returning_columns = []
+            current_col = ""
+            depth = 0
+            for char in returning_clause:
+                if char == "(":
+                    depth += 1
+                elif char == ")":
+                    depth -= 1
+
+                if char == "," and depth == 0:
+                    col = current_col.strip()
+                    # Extract last part of identifier if it's schema-qualified
+                    # e.g. public.users.id -> id, or "public"."users"."id" -> id
+                    col_match = re.search(r'"?(\w+)"?\s*$', col)
+                    if col_match:
+                        returning_columns.append(col_match.group(1).lower())
+                    else:
+                        returning_columns.append(col.lower())
+                    current_col = ""
+                else:
+                    current_col += char
+            if current_col.strip():
+                col = current_col.strip()
+                col_match = re.search(r'"?(\w+)"?\s*$', col)
+                if col_match:
+                    returning_columns.append(col_match.group(1).lower())
+                else:
+                    returning_columns.append(col.lower())
+
+        sql_upper = sql.upper().strip()
+        # Robust table extraction regex for all operations
+        table_regex = r'(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:(?:"?\w+"?)\s*\.\s*)*"?(\w+)"?'
+        table_match = re.search(table_regex, sql, re.IGNORECASE)
+        if table_match:
+            returning_table = table_match.group(1).upper()
+
+        if sql_upper.startswith("INSERT"):
+            returning_operation = "INSERT"
+        elif sql_upper.startswith("UPDATE"):
+            returning_operation = "UPDATE"
+            where_match = re.search(
+                r"\bWHERE\s+(.+?)\s+RETURNING\b",
+                sql,
+                re.IGNORECASE | re.DOTALL,
+            )
+            if where_match:
+                returning_where_clause = where_match.group(1).strip()
+        elif sql_upper.startswith("DELETE"):
+            returning_operation = "DELETE"
+            where_match = re.search(
+                r"\bWHERE\s+(.+?)\s+RETURNING\b",
+                sql,
+                re.IGNORECASE | re.DOTALL,
+            )
+            if where_match:
+                returning_where_clause = where_match.group(1).strip()
+
+        stripped_sql = re.sub(
+            r"\s+RETURNING\s+.*?(?=$|;)",
+            "",
+            sql,
+            flags=re.IGNORECASE | re.DOTALL,
+            count=1,
+        )
+
+        return (
+            returning_operation,
+            returning_table,
+            returning_columns,
+            returning_where_clause,
+            stripped_sql,
+        )
+
+    def _expand_select_star(self, sql: str, expected_columns: int, cursor=None) -> list[str] | None:
+        """
+        Expand SELECT * or RETURNING * into explicit column names using INFORMATION_SCHEMA.
+        """
+        try:
+            table_name = None
+            sql_upper = sql.upper()
+
+            if "RETURNING" in sql_upper:
+                table_regex = (
+                    r'(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:(?:"?\w+"?)\s*\.\s*)*"?(\w+)"?'
+                )
+                table_match = re.search(table_regex, sql, re.IGNORECASE)
+                if table_match:
+                    table_name = table_match.group(1)
+            else:
+                from_match = re.search(r"FROM\s+([^\s,;()]+)", sql, re.IGNORECASE)
+                if from_match:
+                    table_name = from_match.group(1)
+
+            if table_name:
+                if "." in table_name:
+                    table_name = table_name.split(".")[-1]
+                table_name = table_name.strip('"').strip("'")
+
+                schema_columns = self._get_table_columns_from_schema(table_name, cursor)
+                if schema_columns:
+                    if expected_columns == 0 or len(schema_columns) == expected_columns:
+                        return schema_columns
+            return None
+        except Exception as e:
+            logger.debug(f"Failed to expand SELECT *: {e}")
+            return None
+
+    def _extract_insert_id_from_sql(
+        self, sql: str, params: list | None, session_id: str | None = None
+    ) -> tuple[str | None, Any]:
+        """
+        Extract the ID value from an INSERT statement.
+        """
+        col_match = re.search(r"INSERT\s+INTO\s+[^\s(]+\s*\(\s*([^)]+)\s*\)", sql, re.IGNORECASE)
+        if not col_match:
+            return None, None
+
+        columns_str = col_match.group(1)
+        columns = [c.strip().strip('"').strip("'").lower() for c in columns_str.split(",")]
+
+        id_col_names = ["id", "uuid", "_id"]
+        id_col_idx = None
+        id_col_name = None
+        for i, col in enumerate(columns):
+            if col in id_col_names:
+                id_col_idx = i
+                id_col_name = col
+                break
+
+        if id_col_idx is None:
+            return None, None
+
+        if params and len(params) > id_col_idx:
+            return id_col_name, params[id_col_idx]
+
+        return None, None
+
+    def _emulate_returning_sync(
+        self,
+        cursor,
+        operation: str,
+        table: str,
+        columns: list[str] | str,
+        where_clause: str | None,
+        params: list | None,
+        original_sql: str | None = None,
+    ) -> tuple[list[Any], Any]:
+        """
+        Synchronous emulation of RETURNING clause.
+        """
+        table_normalized = table.upper() if table else table
+        if columns == "*":
+            # Expand * using table schema
+            expanded_cols = self._get_table_columns_from_schema(table_normalized, cursor)
+            if expanded_cols:
+                col_list = ", ".join([f'"{col}"' for col in expanded_cols])
+                columns = expanded_cols  # Update columns for metadata generation
+            else:
+                col_list = "*"
+        else:
+            col_list = ", ".join([f'"{col}"' for col in columns])
+
+        rows = []
+        meta = None
+
+        try:
+            if operation == "INSERT":
+                # Method 1: LAST_IDENTITY()
+                cursor.execute("SELECT LAST_IDENTITY()")
+                id_row = cursor.fetchone()
+                last_id = id_row[0] if id_row else None
+
+                if last_id:
+                    cursor.execute(
+                        f'SELECT {col_list} FROM {IRIS_SCHEMA}."{table_normalized}" WHERE %ID = ?',
+                        (last_id,),
+                    )
+                    rows = cursor.fetchall()
+                    meta = cursor.description
+
+                # Method 2: Extract from SQL if still no rows
+                if not rows and original_sql:
+                    id_col_name, id_value = self._extract_insert_id_from_sql(original_sql, params)
+                    if id_col_name and id_value:
+                        cursor.execute(
+                            f'SELECT {col_list} FROM {IRIS_SCHEMA}."{table_normalized}" WHERE "{id_col_name}" = ?',
+                            (id_value,),
+                        )
+                        rows = cursor.fetchall()
+                        meta = cursor.description
+
+            elif operation in ("UPDATE", "DELETE"):
+                if where_clause:
+                    # Translate schema references in WHERE clause
+                    translated_where = re.sub(
+                        r'"public"\s*\.\s*"(\w+)"',
+                        rf'{IRIS_SCHEMA}."\1"',
+                        where_clause,
+                        flags=re.IGNORECASE,
+                    )
+                    translated_where = re.sub(
+                        r'\bpublic\s*\.\s*"(\w+)"',
+                        rf'{IRIS_SCHEMA}."\1"',
+                        translated_where,
+                        flags=re.IGNORECASE,
+                    )
+
+                    # Very basic where clause parameter extraction
+                    where_param_count = translated_where.count("?")
+                    where_params = (
+                        params[-where_param_count:] if params and where_param_count > 0 else None
+                    )
+
+                    cursor.execute(
+                        f'SELECT {col_list} FROM {IRIS_SCHEMA}."{table_normalized}" WHERE {translated_where}',
+                        where_params or (),
+                    )
+                    rows = cursor.fetchall()
+                    meta = cursor.description
+
+            # Build metadata if needed
+            if meta and not any(isinstance(m, dict) and "type_oid" in m for m in meta):
+                new_meta = []
+                for i, desc in enumerate(meta):
+                    col_name = desc[0]
+                    col_oid = self._map_dbapi_type_to_oid(desc[1])
+                    new_meta.append({"name": col_name, "type_oid": col_oid, "format_code": 0})
+                meta = new_meta
+
+        except Exception as e:
+            logger.error(f"RETURNING emulation failed: {e}")
+
+        return rows, meta
+
     def _map_dbapi_type_to_oid(self, dbapi_type: Any) -> int:
         """Map DBAPI type to PostgreSQL OID."""
         # Simple mapping for now, can be expanded
@@ -462,6 +912,41 @@ class DBAPIExecutor:
             return 1114
         return 1043  # Default to VARCHAR
 
+    def _map_iris_type_to_oid(self, iris_type: str) -> int:
+        """
+        Map IRIS data type to PostgreSQL type OID.
+
+        Args:
+            iris_type: IRIS data type (e.g., 'INT', 'VARCHAR', 'DATE')
+
+        Returns:
+            PostgreSQL type OID
+        """
+        type_map = {
+            "INT": 23,  # int4
+            "INTEGER": 23,  # int4
+            "BIGINT": 20,  # int8
+            "SMALLINT": 21,  # int2
+            "VARCHAR": 1043,  # varchar
+            "CHAR": 1042,  # char
+            "TEXT": 25,  # text
+            "DATE": 1082,  # date
+            "TIME": 1083,  # time
+            "TIMESTAMP": 1114,  # timestamp
+            "DOUBLE": 701,  # float8
+            "FLOAT": 701,  # float8
+            "NUMERIC": 1700,  # numeric
+            "DECIMAL": 1700,  # numeric
+            "BIT": 1560,  # bit
+            "BOOLEAN": 16,  # bool
+            "VARBINARY": 17,  # bytea
+        }
+
+        # Normalize type name (remove size, etc.)
+        normalized_type = iris_type.upper().split("(")[0].strip()
+
+        return type_map.get(normalized_type, 1043)  # Default to VARCHAR (OID 1043)
+
     def _determine_command_tag(self, sql: str, row_count: int) -> str:
         """Determine PostgreSQL command tag from SQL"""
         sql_clean = sql.strip().upper()

iris_pgwire/iris_executor.py

@@ -1623,41 +1623,49 @@ class IRISExecutor:
         if returning_clause == "*":
             returning_columns = "*"
         else:
-            # Parse column names from RETURNING clause
-            # Or just: col1, col2, ...
-            raw_cols = [c.strip() for c in returning_clause.split(",")]
+            # Better column parsing that preserves expressions and aliases
+            # Split by commas but respect parentheses
             returning_columns = []
-            for col in raw_cols:
-                # Extract just the column name (last part after dots)
-                # Handle aliased columns like "col AS alias"
-                if " AS " in col.upper():
+            current_col = ""
+            depth = 0
+            for char in returning_clause:
+                if char == "(":
+                    depth += 1
+                elif char == ")":
+                    depth -= 1
+
+                if char == "," and depth == 0:
+                    col = current_col.strip()
+                    # Extract last part of identifier if it's schema-qualified
+                    # e.g. public.users.id -> id, or "public"."users"."id" -> id
                     col_match = re.search(r'"?(\w+)"?\s*$', col)
+                    if col_match:
+                        returning_columns.append(col_match.group(1).lower())
+                    else:
+                        returning_columns.append(col.lower())
+                    current_col = ""
                 else:
-                    col_match = re.search(r'"?(\w+)"?\s*$', col)
-
+                    current_col += char
+            if current_col.strip():
+                col = current_col.strip()
+                col_match = re.search(r'"?(\w+)"?\s*$', col)
                 if col_match:
                     returning_columns.append(col_match.group(1).lower())
+                else:
+                    returning_columns.append(col.lower())
 
         # Determine operation type and extract table/where clause
         sql_upper = sql.upper().strip()
+        # Robust table extraction regex for all operations
+        table_regex = r'(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:(?:"?\w+"?)\s*\.\s*)*"?(\w+)"?'
+        table_match = re.search(table_regex, sql, re.IGNORECASE)
+        if table_match:
+            returning_table = table_match.group(1).upper()
+
         if sql_upper.startswith("INSERT"):
             returning_operation = "INSERT"
-            table_match = re.search(
-                rf'INSERT\s+INTO\s+(?:{re.escape(IRIS_SCHEMA)}\s*\.\s*)?"?(\w+)"?',
-                sql,
-                re.IGNORECASE,
-            )
-            if table_match:
-                returning_table = table_match.group(1)
         elif sql_upper.startswith("UPDATE"):
             returning_operation = "UPDATE"
-            table_match = re.search(
-                rf'UPDATE\s+(?:{re.escape(IRIS_SCHEMA)}\s*\.\s*)?"?(\w+)"?',
-                sql,
-                re.IGNORECASE,
-            )
-            if table_match:
-                returning_table = table_match.group(1)
             # Extract WHERE clause (everything between WHERE and RETURNING)
             where_match = re.search(
                 r"\bWHERE\s+(.+?)\s+RETURNING\b",
@@ -1668,13 +1676,6 @@ class IRISExecutor:
                 returning_where_clause = where_match.group(1).strip()
         elif sql_upper.startswith("DELETE"):
             returning_operation = "DELETE"
-            table_match = re.search(
-                rf'DELETE\s+FROM\s+(?:{re.escape(IRIS_SCHEMA)}\s*\.\s*)?"?(\w+)"?',
-                sql,
-                re.IGNORECASE,
-            )
-            if table_match:
-                returning_table = table_match.group(1)
             # Extract WHERE clause
             where_match = re.search(
                 r"\bWHERE\s+(.+?)\s+RETURNING\b",
@@ -1684,13 +1685,15 @@ class IRISExecutor:
             if where_match:
                 returning_where_clause = where_match.group(1).strip()
 
+        # Strip RETURNING clause from SQL for execution
+        # Use a non-greedy match to avoid stripping important parts if multiple statements exist
         stripped_sql = re.sub(
             r"\s+RETURNING\s+.*?(?=$|;)",
             "",
             sql,
             flags=re.IGNORECASE | re.DOTALL,
             count=1,
-        )
+        ).strip()
 
         return (
             returning_operation,
@@ -1805,14 +1808,31 @@ class IRISExecutor:
         import re
 
         # CRITICAL FIX: Normalize table name to UPPERCASE for IRIS compatibility
-        # IRIS stores table names in uppercase in INFORMATION_SCHEMA
         table_normalized = table.upper() if table else table
 
         # Handle columns as list or '*'
         if columns == "*":
-            col_list = "*"
+            # Expand * early to get real column names
+            expanded_cols = self._expand_select_star(
+                f"SELECT * FROM {IRIS_SCHEMA}.{table_normalized}", 0, session_id=session_id
+            )
+            if expanded_cols:
+                columns = expanded_cols
+                col_list = ", ".join([f'"{col}"' for col in columns])
+            else:
+                col_list = "*"
         else:
-            col_list = ", ".join([f'"{col}"' for col in columns])
+            # columns is a list of expressions/names. Preserve them but quote simple identifiers.
+            processed_cols = []
+            for col in columns:
+                if re.match(r"^\"?\w+\"?$", col):
+                    # Simple identifier - quote it
+                    clean_col = col.strip('"')
+                    processed_cols.append(f'"{clean_col}"')
+                else:
+                    # Expression - leave as is
+                    processed_cols.append(col)
+            col_list = ", ".join(processed_cols)
 
         rows = []
         meta = None
@@ -1855,7 +1875,7 @@ class IRISExecutor:
                         f'SELECT {col_list} FROM {IRIS_SCHEMA}."{table_normalized}" WHERE %ID = ?',
                         [last_id],
                     )
-                    if not rows and columns != "*":
+                    if not rows and isinstance(columns, list):
                         id_cols = [
                             c
                             for c in columns
@@ -1929,42 +1949,31 @@ class IRISExecutor:
                     )
                     rows, meta = _fetch_results(select_sql, where_params)
 
-            # Feature 036: Build proper metadata with actual IRIS types if needed
+            # Build/Fix metadata
             if meta is None or not any("type_oid" in c for c in meta if isinstance(c, dict)):
-                if columns == "*":
-                    # For SELECT *, expand from schema
-                    col_names = self._expand_select_star(
-                        f"RETURNING * FROM {IRIS_SCHEMA}.{table}", 0, session_id=session_id
-                    )
-                    if col_names:
-                        new_meta = []
-                        for col in col_names:
-                            col_oid = self._get_column_type_from_schema(
-                                table, col, session_id=session_id
-                            )
-                            new_meta.append(
-                                {
-                                    "name": col,
-                                    "type_oid": col_oid or 1043,
-                                    "type_size": -1,
-                                    "type_modifier": -1,
-                                    "format_code": 0,
-                                }
-                            )
-                        meta = new_meta
-                else:
+                if isinstance(columns, list):
                     new_meta = []
                     for i, col in enumerate(columns):
+                        # Extract alias or column name
+                        col_name = col
+                        alias_match = re.search(r"\s+AS\s+\"?(\w+)\"?$", col, re.IGNORECASE)
+                        if alias_match:
+                            col_name = alias_match.group(1)
+                        else:
+                            col_name = col_name.strip('"')
+                            if "." in col_name:
+                                col_name = col_name.split(".")[-1]
+
                         col_oid = self._get_column_type_from_schema(
-                            table, col, session_id=session_id
+                            table, col_name, session_id=session_id
                         )
                         if col_oid is None and rows:
                             # Fallback to inference from value
-                            col_oid = self._infer_type_from_value(rows[0][i], col)
+                            col_oid = self._infer_type_from_value(rows[0][i], col_name)
 
                         new_meta.append(
                             {
-                                "name": col,
+                                "name": col_name,
                                 "type_oid": col_oid or 1043,
                                 "type_size": -1,
                                 "type_modifier": -1,
@@ -1972,6 +1981,9 @@ class IRISExecutor:
                             }
                         )
                     meta = new_meta
+                else:
+                    # columns is '*' and expansion failed
+                    pass
 
         except Exception as e:
             logger.error(f"RETURNING emulation failed for {operation}", error=str(e))
@@ -2304,6 +2316,7 @@ class IRISExecutor:
                             returning_where_clause,
                             optimized_params,
                             is_embedded=True,
+                            session_id=session_id,
                             original_sql=sql,  # Pass original SQL for UUID extraction
                         )
                         result = MockResult(rows, meta)
@@ -2324,25 +2337,18 @@ class IRISExecutor:
                         # Get original IRIS column name
                         iris_col_name = col_info.get("name", "")
                         iris_type = col_info.get("type", "VARCHAR")
+                        precomputed_oid = col_info.get("type_oid")
 
                         # CRITICAL: Normalize IRIS column names to PostgreSQL conventions
-                        # IRIS generates HostVar_1, Expression_1, Aggregate_1 for unnamed columns
-                        # PostgreSQL uses ?column?, type names (int4), or function names (count)
                         col_name = self._normalize_iris_column_name(
                             iris_col_name, optimized_sql, iris_type
                         )
 
-                        # DEBUG: Log IRIS type for arithmetic expressions
-                        logger.info(
-                            "🔍 IRIS metadata type discovery",
-                            original_column_name=iris_col_name,
-                            normalized_column_name=col_name,
-                            iris_type=iris_type,
-                            col_info=col_info,
-                        )
-
                         # Get PostgreSQL type OID
-                        type_oid = self._iris_type_to_pg_oid(iris_type)
+                        if precomputed_oid is not None:
+                            type_oid = precomputed_oid
+                        else:
+                            type_oid = self._iris_type_to_pg_oid(iris_type)
 
                         # CRITICAL FIX: IRIS type code 2 means NUMERIC, but for decimal literals
 
@@ -3094,6 +3100,13 @@ class IRISExecutor:
                     connection=conn,
                 )
 
+                # Commit for non-SELECT statements to ensure visibility for emulation and durability
+                if not statements[-1].upper().strip().startswith("SELECT"):
+                    try:
+                        conn.commit()
+                    except Exception as commit_err:
+                        logger.warning(f"Failed to commit {statements[-1][:50]}: {commit_err}")
+
                 # RETURNING emulation
                 if returning_operation and returning_columns:
                     if returning_operation == "DELETE":
@@ -3109,6 +3122,7 @@ class IRISExecutor:
                             optimized_params,
                             is_embedded=False,
                             connection=conn,
+                            session_id=session_id,
                             original_sql=sql,  # Pass original SQL for UUID extraction
                         )
                         cursor = MockResult(rows, meta)
@@ -3564,20 +3578,12 @@ class IRISExecutor:
 
             if "RETURNING" in sql_upper:
                 # For INSERT/UPDATE/DELETE ... RETURNING *, extract table from INTO/UPDATE/FROM
-                # INSERT INTO table_name ...
-                insert_match = re.search(r"INSERT\s+INTO\s+([^\s(]+)", sql, re.IGNORECASE)
-                if insert_match:
-                    table_name = insert_match.group(1)
-                else:
-                    # UPDATE table_name SET ...
-                    update_match = re.search(r"UPDATE\s+([^\s]+)\s+SET", sql, re.IGNORECASE)
-                    if update_match:
-                        table_name = update_match.group(1)
-                    else:
-                        # DELETE FROM table_name ...
-                        delete_match = re.search(r"DELETE\s+FROM\s+([^\s]+)", sql, re.IGNORECASE)
-                        if delete_match:
-                            table_name = delete_match.group(1)
+                table_regex = (
+                    r'(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:(?:"?\w+"?)\s*\.\s*)*"?(\w+)"?'
+                )
+                table_match = re.search(table_regex, sql, re.IGNORECASE)
+                if table_match:
+                    table_name = table_match.group(1)
             else:
                 # SELECT * FROM table_name ...
                 from_match = re.search(r"FROM\s+([^\s,;()]+)", sql, re.IGNORECASE)

{iris_pgwire-1.2.33.dist-info → iris_pgwire-1.2.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iris-pgwire
-Version: 1.2.33
+Version: 1.2.34
 Summary: PostgreSQL Wire Protocol Server for InterSystems IRIS - Connect BI tools, Python frameworks, and PostgreSQL clients to IRIS databases
 Project-URL: Homepage, https://github.com/intersystems-community/iris-pgwire
 Project-URL: Documentation, https://github.com/intersystems-community/iris-pgwire#readme

{iris_pgwire-1.2.33.dist-info → iris_pgwire-1.2.34.dist-info}/RECORD

@@ -1,4 +1,4 @@
-iris_pgwire/__init__.py,sha256=Iad0ZJMoEwmN2gP0I685oKfsCTCxM3FV7k3BEzP9EF8,1097
+iris_pgwire/__init__.py,sha256=XdG-aGnecQsLdu9Y3YON1Nnx8NGljGiYgAo1IC8WUaw,1097
 iris_pgwire/backend_selector.py,sha256=e0DymV-TsGkvp2iItcxkSWSUM8Nf_Agw-LDo78itW04,9315
 iris_pgwire/bulk_executor.py,sha256=nUp5e1ARUzRzFYmctqCS5Zh1gh09IV8H5-Sr2IbKbWg,12742
 iris_pgwire/column_validator.py,sha256=0T2UlLzwH0mEsDkkPXqiUGz4t49o6ySBGtLRyeiGyQc,6806
@@ -7,12 +7,12 @@ iris_pgwire/constitutional.py,sha256=4rrl0DdZRzzcT8mppie93ajiZ_WI-TyBDxZ_Z_ZLGY8
 iris_pgwire/copy_handler.py,sha256=uxvlOcdLiK91qASMcYDlY7rtDc2fmgDmZS9tjTE0WQw,15647
 iris_pgwire/csv_processor.py,sha256=tUv-HOvrQ9KdlBk6MbYtOFvX5Ckz6WIWgJXOhDlky4I,8564
 iris_pgwire/dbapi_connection_pool.py,sha256=84hsbGKscQoj3o2zmNKPDgQMftN7JK0ez_WQpx1P98o,17147
-iris_pgwire/dbapi_executor.py,sha256=XihdSX4VK2lTslQdd2rnOa63EE-gx04ae4qupCpX-3M,22466
+iris_pgwire/dbapi_executor.py,sha256=y1gqHD_UJno8ryrVEv6jUJsT9IsUeiUyEy7V7zlaXEs,40917
 iris_pgwire/debug_tracer.py,sha256=9w6Bve8cnzFPPa69ux3OcR5DMLtjI_NdEFw8TvuwlPM,11112
 iris_pgwire/health_checker.py,sha256=5U7HHKmpzn_G9gcExbRTlJZdD47v4a2FqACVNMpURys,7094
 iris_pgwire/integratedml.py,sha256=TjGX3Z8eXzszqABFrfcVnKp-vyzg5yOBQN8KkTws_o0,16031
 iris_pgwire/iris_constructs.py,sha256=HJ3wNbZyxdlL_ZOganH0830AshPTbctxNTkt6GishfA,31739
-iris_pgwire/iris_executor.py,sha256=7ZHt1C_UybCBX43SZRGmr6_QjBW1mx6NPUNoumALDRY,188614
+iris_pgwire/iris_executor.py,sha256=t91AYl3OOITbJVFgAJV_Wju8LzF3GeBtxFjMfQL84nA,188932
 iris_pgwire/iris_log_handler.py,sha256=FrxCyXDwni7aVsmyAx8zlGGbPTAuRSH3j8FNY5OQdLU,3056
 iris_pgwire/iris_user_management.py,sha256=eGN3CCJML-UVfgR4W0mvNHYQJmgMY5NZXT_SAljONdM,23780
 iris_pgwire/observability.py,sha256=FElftKlwf4JNL3RAR1IiK7aQv2I1w_SCWk-uc6DggIk,3308
@@ -96,8 +96,8 @@ iris_pgwire/sql_translator/mappings/document_filters.py,sha256=gjkwxspizx9ObHUHJ
 iris_pgwire/sql_translator/mappings/functions.py,sha256=9-lHfvGKbHlTMOdKtu0k6wH-mUth3alGZYQ5WXexWQI,19982
 iris_pgwire/testing/__init__.py,sha256=_HBL11torW_QjZh2X3DusLlkmFp9YNb0fYoo_0_Jo3U,46
 iris_pgwire/testing/base_fixture_builder.py,sha256=gVMGEXWjPkMPCpmvFzX9Yc5ZGLKtY8c1yj0YruHcwd8,11593
-iris_pgwire-1.2.33.dist-info/METADATA,sha256=4YRm0r5zWt9Z5-jfc4w4d0n1GtBWFkb1wsYsu8sDM5A,14660
-iris_pgwire-1.2.33.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-iris_pgwire-1.2.33.dist-info/entry_points.txt,sha256=pUOj368R5kSiRFbk4sqKUymPgBRNK3MZ9oYWAX5WZ3k,56
-iris_pgwire-1.2.33.dist-info/licenses/LICENSE,sha256=7C-kc5Ll8ZAWzCD882nr7ptzHdwhc-LXm-wpQ3AxRaU,1099
-iris_pgwire-1.2.33.dist-info/RECORD,,
+iris_pgwire-1.2.34.dist-info/METADATA,sha256=FWOonpOtIoSbwOaLFND6rewkb1AmK45DZ1oZGSDJ0Gc,14660
+iris_pgwire-1.2.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+iris_pgwire-1.2.34.dist-info/entry_points.txt,sha256=pUOj368R5kSiRFbk4sqKUymPgBRNK3MZ9oYWAX5WZ3k,56
+iris_pgwire-1.2.34.dist-info/licenses/LICENSE,sha256=7C-kc5Ll8ZAWzCD882nr7ptzHdwhc-LXm-wpQ3AxRaU,1099
+iris_pgwire-1.2.34.dist-info/RECORD,,