diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index f05bd67ec35a..8d8bbae4cb2b 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -134,6 +134,19 @@ def get_url_for_impersonation( def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool: return True + @classmethod + def execute(cls, cursor: Cursor, query: str, **kwargs: Any) -> None: + opts = {} + if "async_" in kwargs: + opts["deferred_fetch"] = kwargs["async_"] + if cls.arraysize: + cursor.arraysize = cls.arraysize + + try: + cursor.execute(query, **opts) + except Exception as ex: + raise cls.get_dbapi_mapped_exception(ex) + @classmethod def get_tracking_url(cls, cursor: Cursor) -> str | None: try: @@ -160,8 +173,8 @@ def handle_cursor(cls, cursor: Cursor, query: Query, session: Session) -> None: session.commit() - # if query cancelation was requested prior to the handle_cursor call, but - # the query was still executed, trigger the actual query cancelation now + # if query cancellation was requested prior to the handle_cursor call, but + # the query was still executed, trigger the actual query cancellation now if query.extra.get(QUERY_EARLY_CANCEL_KEY): cls.cancel_query( cursor=cursor, @@ -178,7 +191,7 @@ def prepare_cancel_query(cls, query: Query, session: Session) -> None: session.commit() @classmethod - def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: + def cancel_query(cls, cursor: Cursor, query: Query, cancel_query_id: str) -> bool: """ Cancel query in the underlying database. @@ -188,11 +201,12 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: :return: True if query cancelled successfully, False otherwise """ try: + # Can't use cursor.cancel() because + # cursor is new object created in sql_lab.cancel_query cursor.execute( f"CALL system.runtime.kill_query(query_id => '{cancel_query_id}'," "message => 'Query cancelled by Superset')" ) - cursor.fetchall() # needed to trigger the call except Exception: # pylint: disable=broad-except return False