fix(sqllab): reinstate "Force trino client async execution" (#25680)

This commit is contained in:
Rob Moore
2023-10-19 14:38:13 +01:00
committed by GitHub
parent e7cdfeeb2c
commit 4e94fc4cff
9 changed files with 125 additions and 24 deletions

View File

@@ -1066,6 +1066,24 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
query object"""
# TODO: Fix circular import error caused by importing sql_lab.Query
@classmethod
def execute_with_cursor(
cls, cursor: Any, sql: str, query: Query, session: Session
) -> None:
"""
Trigger execution of a query and handle the resulting cursor.
For most implementations this just makes calls to `execute` and
`handle_cursor` consecutively, but in some engines (e.g. Trino) we may
need to handle client limitations such as lack of async support and
perform a more complicated operation to get information from the cursor
in a timely manner and facilitate operations such as query stop
"""
logger.debug("Query %d: Running query: %s", query.id, sql)
cls.execute(cursor, sql, async_=True)
logger.debug("Query %d: Handling cursor", query.id)
cls.handle_cursor(cursor, query, session)
@classmethod
def extract_error_message(cls, ex: Exception) -> str:
return f"{cls.engine} error: {cls._extract_error_message(ex)}"