giftig commented on code in PR #25897:
URL: https://github.com/apache/superset/pull/25897#discussion_r1385700521
##########
superset/db_engine_specs/trino.py:
##########
@@ -193,34 +201,32 @@ def execute_with_cursor(
in another thread and invoke `handle_cursor` to poll for the query ID
to appear on the cursor in parallel.
"""
+ # Fetch the query ID before hand, since it might fail inside the
thread due to
+ # how the SQLAlchemy session is handled.
+ query_id = query.id
+
execute_result: dict[str, Any] = {}
+ execute_event = threading.Event()
- def _execute(results: dict[str, Any]) -> None:
- logger.debug("Query %d: Running query: %s", query.id, sql)
+ def _execute(results: dict[str, Any], event: threading.Event) -> None:
+ logger.debug("Query %d: Running query: %s", query_id, sql)
- # Pass result / exception information back to the parent thread
try:
cls.execute(cursor, sql)
- results["complete"] = True
except Exception as ex: # pylint: disable=broad-except
- results["complete"] = True
results["error"] = ex
+ finally:
+ event.set()
- execute_thread = threading.Thread(target=_execute,
args=(execute_result,))
+ execute_thread = threading.Thread(
+ target=_execute,
+ args=(execute_result, execute_event),
+ )
execute_thread.start()
+ execute_event.wait()
- # Wait for a query ID to be available before handling the cursor, as
- # it's required by that method; it may never become available on error.
- while not cursor.query_id and not execute_result.get("complete"):
- time.sleep(0.1)
-
- logger.debug("Query %d: Handling cursor", query.id)
- cls.handle_cursor(cursor, query, session)
-
- # Block until the query completes; same behaviour as the client itself
- logger.debug("Query %d: Waiting for query to complete", query.id)
- while not execute_result.get("complete"):
- time.sleep(0.5)
+ logger.debug("Query %d: Handling cursor", query_id)
+ cls.handle_cursor_with_query_id(cursor, query, session, query_id)
Review Comment:
You've also passed `query_id` as `query.id`, i.e. the ID of the superset
`Query` object, when it should be `cursor.query_id`, the query ID within trino,
which we can use to cancel the query. I don't think the signature change is
necessary though; we're already passing in the cursor so we can take it from
there.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For additional commands, e-mail: notifications-h...@superset.apache.org
