fix: a statement that fails to compile no longer aborts the transaction

fakesnow/conn.py

@@ -6,6 +6,7 @@
 from types import TracebackType
 from typing import Any, Self
 
+import duckdb
 import snowflake.connector
 from duckdb import DuckDBPyConnection
 from snowflake.connector.cursor import DictCursor, SnowflakeCursor
@@ -124,17 +125,19 @@ def __exit__(
         exc_value: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        if not self._autocommit:
-            if exc_tb is None:
-                self.commit()
-            else:
-                self.rollback()
-
-        self.close()
+        try:
+            if not self._autocommit:
+                if exc_tb is None:
+                    self.commit()
+                else:
+                    self.rollback()
+        finally:
+            self.close()
 
     def autocommit(self, mode: bool) -> None:
         # see https://docs.snowflake.com/en/sql-reference/transactions#autocommit
         if not self._autocommit:
+            self._raise_if_transaction_aborted()
             self._duck_conn.commit()
             self._in_transaction = False
 
@@ -172,6 +175,30 @@ def is_closed(self) -> bool:
     def rollback(self) -> None:
         self.cursor().execute("ROLLBACK")
 
+    def _raise_if_transaction_aborted(self) -> None:
+        """Reject a commit if a failed statement aborted the transaction.
+
+        Snowflake rolls back only the failed statement, see
+        https://docs.snowflake.com/en/sql-reference/transactions#failed-statements-within-a-transaction
+        but DuckDB has no savepoints, so a statement that fails at runtime (eg: a conversion or constraint
+        error) aborts the whole transaction, and committing it would silently discard the earlier writes.
+        """
+        try:
+            # DuckDB's Python API exposes no transaction-validity flag. This query raises in
+            # an aborted transaction and does not leave a transaction open outside BEGIN.
+            self._duck_conn.execute("SELECT 1")
+        except duckdb.TransactionException as e:
+            if "Current transaction is aborted" not in str(e):
+                raise
+            # Keep it aborted until an explicit rollback (or close). In particular, the real
+            # connector's autocommit() suppresses errors from ALTER SESSION; a later commit
+            # must still report the lost writes instead of silently succeeding.
+            msg = (
+                "The transaction was aborted by a failed statement. Roll back before starting a new transaction. "
+                "Snowflake rolls back only the failed statement, but fakesnow can't."
+            )
+            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=9999, sqlstate="99999") from e
+
     @property
     def schema(self) -> str | None:
         return "INFORMATION_SCHEMA" if self._schema == "_FS_INFORMATION_SCHEMA" else self._schema

fakesnow/cursor.py

@@ -463,6 +463,7 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
             if isinstance(transformed, (exp.Alter, exp.Comment, exp.Create, exp.Drop)) or transformed.args.get(
                 "create_stage_name"
             ):
+                self._conn._raise_if_transaction_aborted()
                 self._duck_conn.commit()
                 self._conn._in_transaction = False
 
@@ -481,11 +482,16 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
             # In transaction book keeping
             elif isinstance(transformed, exp.Transaction):
                 self._conn._in_transaction = True
-            elif isinstance(transformed, (exp.Commit, exp.Rollback)):
-                self._conn._in_transaction = False
 
         affected_count = None
         try:
+            if isinstance(transformed, exp.Commit):
+                self._conn._raise_if_transaction_aborted()
+            elif (set_autocommit := transformed.args.get("set_autocommit")) is not None:
+                # ALTER SESSION is transformed to a success SELECT. Commit before that
+                # query, which would otherwise fail first in an aborted transaction.
+                self._conn.autocommit(set_autocommit)
+
             if isinstance(transformed, exp.Copy):
                 sql = copy_into(
                     self._duck_conn,
@@ -534,6 +540,9 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
             msg = cast(str, e.args[0]).split("\n")[0]
             raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=100035, sqlstate="22007") from e
 
+        if isinstance(transformed, (exp.Commit, exp.Rollback)):
+            self._conn._in_transaction = False
+
         self._conn.table_stages.record_temporary_table(transformed, self._conn.database, self._conn.schema)
 
         if set_database := transformed.args.get("set_database"):
@@ -547,9 +556,6 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
             self._conn.schema_set = True
             result_sql = SQL_SUCCESS
 
-        elif (set_autocommit := transformed.args.get("set_autocommit")) is not None:
-            self._conn.autocommit(set_autocommit)
-
         elif create_db_name := transformed.args.get("create_db_name"):
             # we created a new database, so create the info schema extensions
             self._duck_conn.execute(info_schema.per_db_creation_sql(create_db_name))

fakesnow/instance.py

@@ -28,7 +28,16 @@ def __init__(
         self.results_cache: dict[str, tuple] = {}
         # SQLGlot's DuckDB transpilation still emits legacy arrow lambdas; keep the
         # DuckDB 2.0 compatibility setting enabled.
-        self.duck_conn = duckdb.connect(database=":memory:", config={"lambda_syntax": "ENABLE_SINGLE_ARROW"})
+        # Snowflake keeps a transaction active when a statement in it fails, see
+        # https://docs.snowflake.com/en/sql-reference/transactions#failed-statements-within-a-transaction
+        # so like DuckDB 1.x don't abort the transaction on parser, binder or catalog errors.
+        self.duck_conn = duckdb.connect(
+            database=":memory:",
+            config={
+                "lambda_syntax": "ENABLE_SINGLE_ARROW",
+                "default_transaction_invalidation_policy": "SYNTACTIC_ERRORS_DO_NOT_INVALIDATE",
+            },
+        )
 
         # create a "global" database for storing objects which span databases.
         self.duck_conn.execute(f"ATTACH IF NOT EXISTS ':memory:' AS {GLOBAL_DATABASE_NAME}")

tests/test_server.py

@@ -156,6 +156,25 @@ def test_server_connect_autocommit_false(server: dict) -> None:
         assert row and row[0] == 1
 
 
+def test_server_autocommit_false_failed_statement_keeps_transaction(server: dict) -> None:
+    # a failed statement rolls back only its own changes, the transaction stays active and can be committed
+    # see https://docs.snowflake.com/en/sql-reference/transactions#failed-statements-within-a-transaction
+    with (
+        snowflake.connector.connect(**server | {"autocommit": False}, database="db1", schema="schema1") as conn,
+        conn.cursor() as cur,
+    ):
+        cur.execute("create or replace table failed_statement (i int)")
+        cur.execute("insert into failed_statement values (1)")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("select * from missing_table")
+
+        conn.commit()
+        conn.rollback()
+        cur.execute("select count(*) from failed_statement")
+        assert cur.fetchone() == (1,)
+
+
 def test_server_client_session_keep_alive(server: dict) -> None:
     with snowflake.connector.connect(**server | {"client_session_keep_alive": True}):
         # shouldn't error
@@ -828,3 +847,48 @@ def test_server_fetch_multiple_arrow_batches(scur: snowflake.connector.cursor.Sn
 
     timestamp = datetime.datetime(2024, 1, 2, 3, 4, 5, 123456)
     assert scur.fetchall() == [(i, None if i % 2 == 0 else "héllo", Decimal("12.34"), timestamp) for i in range(5000)]
+
+
+@pytest.mark.parametrize("commit_via", ["connection", "sql", "autocommit"])
+def test_server_runtime_failure_commit_raises(server: dict, commit_via: str) -> None:
+    # Deviates from Snowflake: a runtime failure aborts DuckDB's transaction. The server must
+    # report the failed commit through the connector instead of returning success or HTTP 500.
+    with (
+        snowflake.connector.connect(**server | {"autocommit": False}, database="db1", schema="schema1") as conn,
+        conn.cursor() as cur,
+    ):
+        cur.execute("CREATE TABLE runtime_commit (id int)")
+        cur.execute("INSERT INTO runtime_commit VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_commit VALUES ('x')")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted") as excinfo:
+            if commit_via == "connection":
+                conn.commit()
+            elif commit_via == "sql":
+                cur.execute("COMMIT")
+            else:
+                cur.execute("ALTER SESSION SET AUTOCOMMIT = TRUE")
+        assert (excinfo.value.errno, excinfo.value.sqlstate) == (9999, "99999")
+        conn.rollback()
+        cur.execute("SELECT count(*) FROM runtime_commit")
+        assert cur.fetchone() == (0,)
+
+
+def test_server_runtime_failure_autocommit_keeps_commit_error(server: dict) -> None:
+    # The real connector suppresses ALTER SESSION errors in autocommit(). Keep the aborted
+    # transaction pending so that a subsequent commit still reports the lost writes.
+    # Deviates from Snowflake, which preserves earlier writes after the failed statement.
+    with (
+        snowflake.connector.connect(**server | {"autocommit": False}, database="db1", schema="schema1") as conn,
+        conn.cursor() as cur,
+    ):
+        cur.execute("CREATE TABLE runtime_autocommit (id int)")
+        cur.execute("INSERT INTO runtime_autocommit VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_autocommit VALUES ('x')")
+
+        conn.autocommit(True)
+        with pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted"):
+            conn.commit()
+        conn.rollback()

tests/test_sqlalchemy.py

@@ -47,3 +47,20 @@ def test_reflect(snowflake_engine: Engine):
 
         assert result
         assert result.fetchall() == [(1, "one")]
+
+
+def test_create_checkfirst_keeps_transaction(snowflake_engine: Engine):
+    # runs with autocommit=False, and checkfirst probes for the missing table with a statement that fails,
+    # which leaves the transaction and its earlier insert intact
+    metadata = MetaData()
+    foo = cast(Table, Table("foo", metadata, Column(types.Integer, name="id")))
+    bar = cast(Table, Table("bar", metadata, Column(types.Integer, name="id")))
+
+    with snowflake_engine.connect() as conn:
+        foo.create(conn)
+        conn.execute(foo.insert().values(id=1))
+        bar.create(conn, checkfirst=True)
+        conn.commit()
+
+    with snowflake_engine.connect() as conn:
+        assert conn.execute(foo.select()).fetchall() == [(1,)]

tests/test_transactions.py

@@ -221,3 +221,121 @@ def test_conn_context_exit_with_exception_rolls_back(_fakesnow: None):
 
     with snowflake.connector.connect(database="db1", schema="schema1") as conn:
         assert conn.cursor().execute("SELECT count(*) FROM context_exit").fetchone() == (0,)
+
+
+@pytest.mark.parametrize("failing_sql", ["DESC TABLE missing_table", "SELECT * FROM missing_table"])
+def test_autocommit_false_failed_statement_keeps_transaction(
+    conn: snowflake.connector.SnowflakeConnection, failing_sql: str
+):
+    # a failed statement rolls back only its own changes, the transaction stays active and can be committed
+    # see https://docs.snowflake.com/en/sql-reference/transactions#failed-statements-within-a-transaction
+    conn.autocommit(False)
+    with conn.cursor() as cur:
+        cur.execute("CREATE OR REPLACE TABLE failed_statement (id int)")
+        cur.execute("INSERT INTO failed_statement VALUES (1)")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute(failing_sql)
+
+        assert cur.execute("SELECT count(*) FROM failed_statement").fetchone() == (1,)
+        conn.commit()
+        conn.rollback()
+        assert cur.execute("SELECT count(*) FROM failed_statement").fetchone() == (1,)
+
+
+def test_autocommit_false_runtime_failure_commit_raises(conn: snowflake.connector.SnowflakeConnection):
+    # Deviates from Snowflake, which rolls back only the failed statement and commits the earlier insert, see
+    # https://docs.snowflake.com/en/sql-reference/transactions#failed-statements-within-a-transaction
+    # DuckDB has no savepoints, so a statement that fails at runtime aborts the whole transaction. Rather than
+    # commit silently dropping the earlier insert, fakesnow raises and requires an explicit rollback.
+    conn.autocommit(False)
+    with conn.cursor() as cur:
+        cur.execute("CREATE OR REPLACE TABLE runtime_failure (id int)")
+        cur.execute("INSERT INTO runtime_failure VALUES (1)")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_failure VALUES ('x')")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted"):
+            conn.commit()
+        conn.rollback()
+        assert cur.execute("SELECT count(*) FROM runtime_failure").fetchone() == (0,)
+
+
+@pytest.mark.parametrize(
+    "commit_sql",
+    [
+        "COMMIT",
+        "CREATE TABLE next_table (id int)",
+        "ALTER SESSION SET AUTOCOMMIT = TRUE",
+        "ALTER SESSION SET AUTOCOMMIT = FALSE",
+    ],
+)
+def test_runtime_failure_sql_commit_raises(conn: snowflake.connector.SnowflakeConnection, commit_sql: str):
+    # Deviates from Snowflake: DuckDB cannot preserve earlier writes after a runtime failure.
+    # Every explicit or implicit commit must report that loss instead of claiming success.
+    conn.autocommit(False)
+    with conn.cursor() as cur:
+        cur.execute("CREATE TABLE runtime_sql_commit (id int)")
+        cur.execute("INSERT INTO runtime_sql_commit VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_sql_commit VALUES ('x')")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted") as excinfo:
+            cur.execute(commit_sql)
+        assert (excinfo.value.errno, excinfo.value.sqlstate) == (9999, "99999")
+        conn.rollback()
+        assert cur.execute("SELECT count(*) FROM runtime_sql_commit").fetchone() == (0,)
+
+
+@pytest.mark.parametrize("mode", [False, True])
+def test_runtime_failure_autocommit_raises(conn: snowflake.connector.SnowflakeConnection, mode: bool):
+    # Deviates from Snowflake: changing AUTOCOMMIT cannot preserve writes after a DuckDB runtime failure.
+    conn.autocommit(False)
+    with conn.cursor() as cur:
+        cur.execute("CREATE TABLE runtime_autocommit (id int)")
+        cur.execute("INSERT INTO runtime_autocommit VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_autocommit VALUES ('x')")
+
+        with pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted"):
+            conn.autocommit(mode)
+        conn.rollback()
+        assert cur.execute("SELECT count(*) FROM runtime_autocommit").fetchone() == (0,)
+
+
+def test_runtime_failure_rollback(conn: snowflake.connector.SnowflakeConnection):
+    conn.autocommit(False)
+    with conn.cursor() as cur:
+        cur.execute("CREATE TABLE runtime_rollback (id int)")
+        cur.execute("INSERT INTO runtime_rollback VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_rollback VALUES ('x')")
+
+        conn.rollback()
+        cur.execute("INSERT INTO runtime_rollback VALUES (2)")
+        conn.commit()
+        assert cur.execute("SELECT * FROM runtime_rollback").fetchall() == [(2,)]
+
+
+def test_context_exit_failed_commit_closes_connection(_fakesnow: None):
+    # Deviates from Snowflake: the commit fails after a runtime error, but context exit must still close.
+    conn = snowflake.connector.connect(database="db1", schema="schema1", autocommit=False)
+    with (
+        pytest.raises(snowflake.connector.errors.ProgrammingError, match="transaction was aborted"),
+        conn,
+        conn.cursor() as cur,
+    ):
+        cur.execute("CREATE TABLE runtime_context_exit (id int)")
+        cur.execute("INSERT INTO runtime_context_exit VALUES (1)")
+        with pytest.raises(snowflake.connector.errors.ProgrammingError):
+            cur.execute("INSERT INTO runtime_context_exit VALUES ('x')")
+
+    assert conn.is_closed()
+
+
+def test_commit_closed_connection(conn: snowflake.connector.SnowflakeConnection):
+    conn.close()
+    with pytest.raises(snowflake.connector.errors.DatabaseError) as excinfo:
+        conn.commit()
+    assert (excinfo.value.errno, excinfo.value.sqlstate) == (250002, "08003")