From 69c5bf0de2061e329597309638efec49dd2b2a97 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Mon, 22 Jul 2024 10:00:03 -0700
Subject: [PATCH] fix(mssql): avoid calling `.commit()` unless a DDL operation
 is being performed (#9658)

This PR attempts to address flaky behavior of MS SQL observed in CI.

It seems like calling `commit()`/`rollback()` with a `SELECT` statement,
AKA autocommit behavior
can cause this issue.

The origin is not 100% clear to me, but there are a number of places
where the
same problem (function sequence error) with the same solution show up
(disabling autocommit or not calling commit with a `SELECT` statement):

-
https://stackoverflow.com/questions/25769043/function-sequence-error-in-pyodbc
- https://github.com/ansible-collections/community.general/issues/1137
-
https://www.reddit.com/r/learnpython/comments/x8568m/function_sequence_error_in_sqlalchemy/
- https://github.com/explorerhq/django-sql-explorer/issues/423


Here I avoid calling `cur.commit()` unless a DDL statement is being
executed.

The remaining case is `raw_sql()`, which I opted to avoid calling
`commit` in
since that would have this problem when calling it with a `SELECT`
statement.

Users of `raw_sql` are therefore responsible for handling
commit/rollback when
invoking it for DDL statements.

Closes #9654.
---
 ibis/backends/mssql/__init__.py | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/ibis/backends/mssql/__init__.py b/ibis/backends/mssql/__init__.py
index 79a86611bb8c..a259850767a0 100644
--- a/ibis/backends/mssql/__init__.py
+++ b/ibis/backends/mssql/__init__.py
@@ -251,6 +251,11 @@ def current_database(self) -> str:
 
     @contextlib.contextmanager
     def begin(self):
+        with contextlib.closing(self.con.cursor()) as cur:
+            yield cur
+
+    @contextlib.contextmanager
+    def _ddl_begin(self):
         con = self.con
         cur = con.cursor()
         try:
@@ -272,6 +277,15 @@ def _safe_raw_sql(self, query, *args, **kwargs):
             cur.execute(query, *args, **kwargs)
             yield cur
 
+    @contextlib.contextmanager
+    def _safe_ddl(self, query, *args, **kwargs):
+        with contextlib.suppress(AttributeError):
+            query = query.sql(self.dialect)
+
+        with self._ddl_begin() as cur:
+            cur.execute(query, *args, **kwargs)
+            yield cur
+
     def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         with contextlib.suppress(AttributeError):
             query = query.sql(self.dialect)
@@ -279,15 +293,8 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         con = self.con
         cursor = con.cursor()
 
-        try:
-            cursor.execute(query, **kwargs)
-        except Exception:
-            con.rollback()
-            cursor.close()
-            raise
-        else:
-            con.commit()
-            return cursor
+        cursor.execute(query, **kwargs)
+        return cursor
 
     def create_catalog(self, name: str, force: bool = False) -> None:
         expr = (
@@ -308,11 +315,11 @@ def create_catalog(self, name: str, force: bool = False) -> None:
             if force
             else stmt
         )
-        with self._safe_raw_sql(create_stmt):
+        with self._safe_ddl(create_stmt):
             pass
 
     def drop_catalog(self, name: str, force: bool = False) -> None:
-        with self._safe_raw_sql(
+        with self._safe_ddl(
             sge.Drop(
                 kind="DATABASE",
                 this=sg.to_identifier(name, quoted=self.compiler.quoted),
@@ -583,7 +590,7 @@ def create_table(
 
         this = sg.table(name, catalog=catalog, db=db, quoted=self.compiler.quoted)
         raw_this = sg.table(name, catalog=catalog, db=db, quoted=False)
-        with self._safe_raw_sql(create_stmt) as cur:
+        with self._safe_ddl(create_stmt) as cur:
             if query is not None:
                 # You can specify that a table is temporary for the sqlglot `Create` but not
                 # for the subsequent `Insert`, so we need to shove a `#` in
@@ -665,7 +672,7 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
             data = df.itertuples(index=False)
 
             insert_stmt = self._build_insert_template(name, schema=schema, columns=True)
-            with self._safe_raw_sql(create_stmt) as cur:
+            with self._safe_ddl(create_stmt) as cur:
                 if not df.empty:
                     cur.executemany(insert_stmt, data)