fix(doris): guard _fetch_native_df disconnect retry and reject CTE-prefixed write retries

Author: brucearctor

sqlmesh/core/engine_adapter/doris.py

@@ -595,7 +595,16 @@ def _fetch_native_df(
                 isinstance(getattr(e, "__context__", None), TimeoutError)
                 or "timed out" in str(e).lower()
             )
-            if is_disconnect and not is_timeout and not self._connection_pool.is_transaction_active:
+            if (
+                is_disconnect
+                and not is_timeout
+                and not self._connection_pool.is_transaction_active
+                and self._is_safe_to_retry_sql(
+                    self._to_sql(query, quote=quote_identifiers)
+                    if isinstance(query, exp.Expr)
+                    else query
+                )
+            ):
                 logger.warning(
                     "Doris connection lost (%s: %s) during fetchdf, reconnecting and retrying once...",
                     err_name,
@@ -616,14 +625,23 @@ def _fetch_native_df(
             raise
 
     _SAFE_RETRY_SQL_PREFIX_RE = re.compile(
-        r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(?:SELECT|WITH|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH)\b",
+        r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(SELECT|WITH|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH)\b",
         re.IGNORECASE | re.DOTALL,
     )
 
     @classmethod
     def _is_safe_to_retry_sql(cls, sql: str) -> bool:
         """Returns True only for read-only or idempotent catalog/database switch statements that cannot duplicate writes."""
-        return bool(cls._SAFE_RETRY_SQL_PREFIX_RE.match(sql))
+        match = cls._SAFE_RETRY_SQL_PREFIX_RE.match(sql)
+        if not match:
+            return False
+        if match.group(1).upper() == "WITH":
+            try:
+                parsed = parse_one(sql, read=cls.DIALECT)
+            except Exception:
+                return False
+            return isinstance(parsed, (exp.Select, exp.SetOperation))
+        return True
 
     def ping(self) -> None:
         self._connection_pool.get().ping(reconnect=False)

tests/core/engine_adapter/test_doris.py

@@ -951,23 +951,27 @@ def timeout_side_effect(*args, **kwargs):
     conn_mock.ping.assert_not_called()
     close_spy.assert_called_once()
 
-    # Write statements (e.g. INSERT) must NOT be retried after OperationalError(2013) to avoid duplicating rows
-    conn_mock.reset_mock()
-    close_spy.reset_mock()
-    write_calls = 0
+    # Write statements (e.g. INSERT or CTE-prefixed INSERT) must NOT be retried after OperationalError(2013) to avoid duplicating rows
+    for unsafe_write_sql in (
+        "INSERT INTO `tbl` VALUES (1)",
+        "WITH `cte` AS (SELECT 1) INSERT INTO `tbl` SELECT * FROM `cte`",
+    ):
+        conn_mock.reset_mock()
+        close_spy.reset_mock()
+        write_calls = 0
 
-    def write_side_effect(*args, **kwargs):
-        nonlocal write_calls
-        write_calls += 1
-        raise OperationalError(2013, "Lost connection to MySQL server during query")
+        def write_side_effect(*args, **kwargs):
+            nonlocal write_calls
+            write_calls += 1
+            raise OperationalError(2013, "Lost connection to MySQL server during query")
 
-    adapter.cursor.execute.side_effect = write_side_effect
-    with pytest.raises(OperationalError):
-        adapter._execute("INSERT INTO `tbl` VALUES (1)")
+        adapter.cursor.execute.side_effect = write_side_effect
+        with pytest.raises(OperationalError):
+            adapter._execute(unsafe_write_sql)
 
-    assert write_calls == 1
-    conn_mock.ping.assert_not_called()
-    close_spy.assert_called_once()
+        assert write_calls == 1
+        conn_mock.ping.assert_not_called()
+        close_spy.assert_called_once()
 
 
 def test_external_catalog_view_error_and_drop_noop(
@@ -2203,6 +2207,14 @@ def fake_read_sql_query(sql, con, *args, **kwargs):
     assert fetch_calls == 2
     assert close_spy.call_count == 1
 
+    # 4. _fetch_native_df does NOT retry non-read-only queries on disconnect
+    fetch_calls = 0
+    close_spy.reset_mock()
+    with pytest.raises(OperationalError):
+        adapter.fetchdf("WITH `cte` AS (SELECT 1) INSERT INTO `tbl` SELECT * FROM `cte`")
+    assert fetch_calls == 1
+    assert close_spy.call_count == 1
+
 
 def test_strip_for_update_and_multi_col_index_and_partition_by_alias_and_deferred_mv_audit_refresh(
     make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],