From 8bdd01e984f3d6a69018f9b8869c07012d720119 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 28 Jan 2021 14:32:00 -0800
Subject: [PATCH] fix: don't close cursor before closing connection (#12821)

---
 superset/sql_lab.py | 72 ++++++++++++++++++++++-----------------------
 1 file changed, 35 insertions(+), 37 deletions(-)

diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 9acb02097841c..dd8282652c5cf 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -379,44 +379,42 @@ def execute_sql_statements(  # pylint: disable=too-many-arguments, too-many-loca
     # Sharing a single connection and cursor across the
     # execution of all statements (if many)
     with closing(engine.raw_connection()) as conn:
-        with closing(conn.cursor()) as cursor:
-            statement_count = len(statements)
-            for i, statement in enumerate(statements):
-                # Check if stopped
-                query = get_query(query_id, session)
-                if query.status == QueryStatus.STOPPED:
-                    return None
-
-                # For CTAS we create the table only on the last statement
-                apply_ctas = query.select_as_cta and (
-                    query.ctas_method == CtasMethod.VIEW
-                    or (
-                        query.ctas_method == CtasMethod.TABLE
-                        and i == len(statements) - 1
-                    )
-                )
+        # closing the connection closes the cursor as well
+        cursor = conn.cursor()
+        statement_count = len(statements)
+        for i, statement in enumerate(statements):
+            # Check if stopped
+            query = get_query(query_id, session)
+            if query.status == QueryStatus.STOPPED:
+                return None
 
-                # Run statement
-                msg = f"Running statement {i+1} out of {statement_count}"
-                logger.info("Query %s: %s", str(query_id), msg)
-                query.set_extra_json_key("progress", msg)
-                session.commit()
-                try:
-                    result_set = execute_sql_statement(
-                        statement,
-                        query,
-                        user_name,
-                        session,
-                        cursor,
-                        log_params,
-                        apply_ctas,
-                    )
-                except Exception as ex:  # pylint: disable=broad-except
-                    msg = str(ex)
-                    if statement_count > 1:
-                        msg = f"[Statement {i+1} out of {statement_count}] " + msg
-                    payload = handle_query_error(msg, query, session, payload)
-                    return payload
+            # For CTAS we create the table only on the last statement
+            apply_ctas = query.select_as_cta and (
+                query.ctas_method == CtasMethod.VIEW
+                or (query.ctas_method == CtasMethod.TABLE and i == len(statements) - 1)
+            )
+
+            # Run statement
+            msg = f"Running statement {i+1} out of {statement_count}"
+            logger.info("Query %s: %s", str(query_id), msg)
+            query.set_extra_json_key("progress", msg)
+            session.commit()
+            try:
+                result_set = execute_sql_statement(
+                    statement,
+                    query,
+                    user_name,
+                    session,
+                    cursor,
+                    log_params,
+                    apply_ctas,
+                )
+            except Exception as ex:  # pylint: disable=broad-except
+                msg = str(ex)
+                if statement_count > 1:
+                    msg = f"[Statement {i+1} out of {statement_count}] " + msg
+                payload = handle_query_error(msg, query, session, payload)
+                return payload
 
         # Commit the connection so CTA queries will create the table.
         conn.commit()