From 1978759e1bf0bcd001f699af3b80fcd27adf3ebb Mon Sep 17 00:00:00 2001 From: Maxime Beauchemin Date: Thu, 5 Dec 2024 09:08:14 -0800 Subject: [PATCH] impprovments --- superset-frontend/package-lock.json | 2 +- superset/config.py | 4 +++- superset/models/core.py | 35 +++++++++++++++++++++------- superset/sql_lab.py | 5 ++-- superset/sqllab/sql_json_executer.py | 8 ++++++- 5 files changed, 40 insertions(+), 14 deletions(-) diff --git a/superset-frontend/package-lock.json b/superset-frontend/package-lock.json index 99e08f8d59f1a..f4bef440366ab 100644 --- a/superset-frontend/package-lock.json +++ b/superset-frontend/package-lock.json @@ -69182,7 +69182,7 @@ "@types/jest": "^29.5.12", "@types/lodash": "^4.17.7", "handlebars": "^4.7.8", - "handlebars-group-by": "*", + "handlebars-group-by": "^1.0.1", "jest": "^29.7.0", "just-handlebars-helpers": "^1.0.19" }, diff --git a/superset/config.py b/superset/config.py index a4dbe5a877e40..f9d10aca18124 100644 --- a/superset/config.py +++ b/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', diff --git a/superset/models/core.py b/superset/models/core.py index 94e915f0d7c67..4e60ba0cb8279 100755 --- a/superset/models/core.py +++ b/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 diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 537fe9ca1df0a..88e1bc1aad858 100644 --- a/superset/sql_lab.py +++ b/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( diff --git a/superset/sqllab/sql_json_executer.py b/superset/sqllab/sql_json_executer.py index 27483fb31cb01..9ad70ca6ab0d8 100644 --- a/superset/sqllab/sql_json_executer.py +++ b/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,