impprovments

superset/config.py

@@ -204,7 +204,9 @@ def _try_json_readsha(filepath: str, length: int) -> str | None:
 # `SQLALCHEMY_ENGINE_OPTIONS = {"isolation_level": "READ COMMITTED"}`
 # Also note that we recommend READ COMMITTED for regular operation.
 # Find out more here https://flask-sqlalchemy.palletsprojects.com/en/3.1.x/config/
-SQLALCHEMY_ENGINE_OPTIONS = {}
+from sqlalchemy.pool import NullPool
+
+SQLALCHEMY_ENGINE_OPTIONS = {"poolclass": NullPool}
 
 # In order to hook up a custom password store for all SQLALCHEMY connections
 # implement a function that takes a single argument of type 'sqla.engine.url',

superset/models/core.py

@@ -95,6 +95,17 @@
 DB_CONNECTION_MUTATOR = config["DB_CONNECTION_MUTATOR"]
 
 
+def print_connection_status():
+    logger.info("Pinging database")
+    logger.info(f"STATUS: {db.session.connection().closed}")
+
+    try:
+        db.session.execute("SELECT 1")
+        logger.info("success")
+    except Exception:
+        logger.error("connection failed as expected")
+
+
 @contextmanager
 def temporarily_disconnect_metadata_db():  # type: ignore
     """
@@ -108,20 +119,28 @@ def temporarily_disconnect_metadata_db():  # type: ignore
     Only has an effect if feature flag DISABLE_METADATA_DB_DURING_ANALYTICS is on
     """
     do_it = is_feature_enabled("DISABLE_METADATA_DB_DURING_ANALYTICS")
-    print("YO: temporarily_disconnect_metadata_db")
+    connection = db.session.connection()
     try:
         if do_it:
-            print("YO: Disconnecting")
+            logger.info("Disconnecting metadata database temporarily")
             db.session.close()
+            print_connection_status()
+            connection.close()
+            logger.info(f"STATUS: {connection.closed}")
         yield None
     finally:
         if do_it:
-            print("YO: Reconnecting to metadata database")
-            conn = db.engine.connect()
-            print(f"YO: conn: {conn.closed}")
-            print(conn)
-            conn.open()
-            db.session = db._make_scoped_session()
+            logger.info("Reconnecting to metadata database")
+            logger.info("BEFORE")
+            print_connection_status()
+
+            # NOTE: Interface changes in flask-sqlalchemy 3.0
+
+            db.session.connection()
+            db.session = db.create_scoped_session()
+
+            logger.info("AFTER")
+            print_connection_status()
 
 
 class KeyValue(Model):  # pylint: disable=too-few-public-methods

superset/sql_lab.py

@@ -50,7 +50,7 @@
     SupersetParseError,
 )
 from superset.extensions import celery_app, event_logger
-from superset.models.core import Database, temporarily_disconnect_metadata_db
+from superset.models.core import Database
 from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
 from superset.sql.parse import SQLStatement, Table
@@ -304,8 +304,7 @@ def execute_sql_statement(  # pylint: disable=too-many-statements, too-many-loca
             object_ref=__name__,
         ):
             with stats_timing("sqllab.query.time_executing_query", stats_logger):
-                with temporarily_disconnect_metadata_db():
-                    db_engine_spec.execute_with_cursor(cursor, sql, query)
+                db_engine_spec.execute_with_cursor(cursor, sql, query)
 
             with stats_timing("sqllab.query.time_fetching_results", stats_logger):
                 logger.debug(

superset/sqllab/sql_json_executer.py

@@ -31,6 +31,7 @@
     SupersetGenericDBErrorException,
     SupersetTimeoutException,
 )
+from superset.models.core import temporarily_disconnect_metadata_db
 from superset.sqllab.command_status import SqlJsonExecutionStatus
 from superset.utils import core as utils
 from superset.utils.core import get_username
@@ -100,7 +101,9 @@ def execute(
         except SupersetTimeoutException:
             raise
         except Exception as ex:
+            raise ex
             logger.exception("Query %i failed unexpectedly", query_id)
+            logger.error(ex)
             raise SupersetGenericDBErrorException(
                 utils.error_msg_from_exception(ex)
             ) from ex
@@ -127,7 +130,10 @@ def _get_sql_results_with_timeout(
             seconds=self._timeout_duration_in_seconds,
             error_message=self._get_timeout_error_msg(),
         ):
-            return self._get_sql_results(execution_context, rendered_query, log_params)
+            with temporarily_disconnect_metadata_db():
+                return self._get_sql_results(
+                    execution_context, rendered_query, log_params
+                )
 
     def _get_sql_results(
         self,