From 77947fe95ea53c23e9c702f38eb07e34729a89f7 Mon Sep 17 00:00:00 2001 From: ymaayan <57786539+ymaayan@users.noreply.github.com> Date: Thu, 8 May 2025 12:04:15 +0300 Subject: [PATCH 1/2] expose _get_async_query_info in db --- src/firebolt/db/connection.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 8f740560dce..e4e7bc5a4bf 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -230,7 +230,7 @@ def close(self) -> None: # Server-side async methods - def _get_async_query_info(self, token: str) -> AsyncQueryInfo: + def get_async_query_info(self, token: str) -> AsyncQueryInfo: if self.cursor_type != CursorV2: raise FireboltError( "This method is only supported for connection with service account." @@ -268,7 +268,7 @@ def is_async_query_running(self, token: str) -> bool: Returns: bool: True if async query is still running, False otherwise """ - return self._get_async_query_info(token).status == ASYNC_QUERY_STATUS_RUNNING + return self.get_async_query_info(token).status == ASYNC_QUERY_STATUS_RUNNING def is_async_query_successful(self, token: str) -> Optional[bool]: """ @@ -281,7 +281,7 @@ def is_async_query_successful(self, token: str) -> Optional[bool]: bool: None if the query is still running, True if successful, False otherwise """ - async_query_info = self._get_async_query_info(token) + async_query_info = self.get_async_query_info(token) if async_query_info.status == ASYNC_QUERY_STATUS_RUNNING: return None return async_query_info.status == ASYNC_QUERY_STATUS_SUCCESSFUL @@ -293,7 +293,7 @@ def cancel_async_query(self, token: str) -> None: Args: token: Async query token. Can be obtained from Cursor.async_query_token. """ - async_query_id = self._get_async_query_info(token).query_id + async_query_id = self.get_async_query_info(token).query_id cursor = self.cursor() cursor.execute(ASYNC_QUERY_CANCEL, [async_query_id]) From 1cedcf67de313adcc283ebc6b1ae659f8c948cb1 Mon Sep 17 00:00:00 2001 From: ymaayan <57786539+ymaayan@users.noreply.github.com> Date: Thu, 8 May 2025 12:05:13 +0300 Subject: [PATCH 2/2] expose _get_async_query_info in async_db --- src/firebolt/async_db/connection.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 05603bbe10d..c083025e162 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -92,7 +92,7 @@ def cursor(self, **kwargs: Any) -> Cursor: return c # Server-side async methods - async def _get_async_query_info(self, token: str) -> AsyncQueryInfo: + async def get_async_query_info(self, token: str) -> AsyncQueryInfo: if self.cursor_type != CursorV2: raise FireboltError( "This method is only supported for connection with service account." @@ -130,7 +130,7 @@ async def is_async_query_running(self, token: str) -> bool: Returns: bool: True if async query is still running, False otherwise """ - async_query_details = await self._get_async_query_info(token) + async_query_details = await self.get_async_query_info(token) return async_query_details.status == ASYNC_QUERY_STATUS_RUNNING async def is_async_query_successful(self, token: str) -> Optional[bool]: @@ -144,7 +144,7 @@ async def is_async_query_successful(self, token: str) -> Optional[bool]: bool: None if the query is still running, True if successful, False otherwise """ - async_query_details = await self._get_async_query_info(token) + async_query_details = await self.get_async_query_info(token) if async_query_details.status == ASYNC_QUERY_STATUS_RUNNING: return None return async_query_details.status == ASYNC_QUERY_STATUS_SUCCESSFUL @@ -156,7 +156,7 @@ async def cancel_async_query(self, token: str) -> None: Args: token: Async query token. Can be obtained from Cursor.async_query_token. """ - async_query_details = await self._get_async_query_info(token) + async_query_details = await self.get_async_query_info(token) async_query_id = async_query_details.query_id cursor = self.cursor() await cursor.execute(ASYNC_QUERY_CANCEL, [async_query_id])