diff --git a/setup.cfg b/setup.cfg index f1fd9b6e343..cc7dae3c4af 100755 --- a/setup.cfg +++ b/setup.cfg @@ -29,6 +29,7 @@ install_requires = async-generator>=1.10 async-property>=0.2.1 cryptography>=3.4.0 + httpcore<0.17.3 httpx[http2]==0.24.0 python-dateutil>=2.8.2 readerwriterlock>=1.0.9 diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index f73fb28627a..b659f533ca7 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -1,6 +1,5 @@ from __future__ import annotations -import logging import socket from types import TracebackType from typing import Any, Dict, List, Optional @@ -29,9 +28,7 @@ InterfaceError, ) from firebolt.utils.usage_tracker import get_user_agent_header -from firebolt.utils.util import fix_url_schema - -logger = logging.getLogger(__name__) +from firebolt.utils.util import Timer, fix_url_schema class OverriddenHttpBackend(AutoBackend): @@ -231,9 +228,10 @@ async def connect( else: try: - engine_url, status, attached_db = await _get_engine_url_status_db( - system_engine_connection, engine_name - ) + with Timer("[PERFORMANCE] Resolving engine name "): + engine_url, status, attached_db = await _get_engine_url_status_db( + system_engine_connection, engine_name + ) if status != "Running": raise EngineNotRunningError(engine_name) diff --git a/src/firebolt/async_db/cursor.py b/src/firebolt/async_db/cursor.py index 850039ac451..0ded9783f2c 100644 --- a/src/firebolt/async_db/cursor.py +++ b/src/firebolt/async_db/cursor.py @@ -49,6 +49,7 @@ if TYPE_CHECKING: from firebolt.async_db.connection import Connection +from firebolt.utils.util import Timer logger = logging.getLogger(__name__) @@ -199,14 +200,20 @@ async def _do_execute( skip_parsing, async_execution, ) - response = await self._api_request( - query, - { - "async_execution": 1, - "advanced_mode": 1, - "output_format": JSON_OUTPUT_FORMAT, - }, - ) + + with Timer( + f"[PERFORMANCE] Running query {query[:50]} " + f"{'... ' if len(query) > 50 else ''}" + ): + response = await self._api_request( + query, + { + "async_execution": 1, + "advanced_mode": 1, + "output_format": JSON_OUTPUT_FORMAT, + }, + ) + await self._raise_if_error(response) if response.headers.get("content-length", "") == "0": raise OperationalError("No response to asynchronous query.") diff --git a/src/firebolt/client/auth/base.py b/src/firebolt/client/auth/base.py index 75d6b29d816..cacd286913e 100644 --- a/src/firebolt/client/auth/base.py +++ b/src/firebolt/client/auth/base.py @@ -5,7 +5,7 @@ from httpx import Request, Response, codes from firebolt.utils.token_storage import TokenSecureStorage -from firebolt.utils.util import cached_property +from firebolt.utils.util import Timer, cached_property class AuthRequest(Request): @@ -107,15 +107,16 @@ def auth_flow(self, request: Request) -> Generator[Request, Response, None]: Yields: Request: Request required for auth flow """ - if not self.token or self.expired: - yield from self.get_new_token_generator() - self._cache_token() + with Timer("[PERFORMANCE] Authentication "): + if not self.token or self.expired: + yield from self.get_new_token_generator() + self._cache_token() - request.headers["Authorization"] = f"Bearer {self.token}" + request.headers["Authorization"] = f"Bearer {self.token}" - response = yield request + response = yield request - if response.status_code == codes.UNAUTHORIZED: - yield from self.get_new_token_generator() - request.headers["Authorization"] = f"Bearer {self.token}" - yield request + if response.status_code == codes.UNAUTHORIZED: + yield from self.get_new_token_generator() + request.headers["Authorization"] = f"Bearer {self.token}" + yield request diff --git a/src/firebolt/common/base_cursor.py b/src/firebolt/common/base_cursor.py index b6c690bfce3..9b78145cf11 100644 --- a/src/firebolt/common/base_cursor.py +++ b/src/firebolt/common/base_cursor.py @@ -24,6 +24,7 @@ DataError, QueryNotRunError, ) +from firebolt.utils.util import Timer logger = logging.getLogger(__name__) @@ -345,7 +346,9 @@ def fetchone(self) -> Optional[List[ColType]]: # We are out of elements return None assert self._rows is not None - return self._parse_row(self._rows[left]) + with Timer("[PERFORMANCE] Parsing query output into native Python types "): + result = self._parse_row(self._rows[left]) + return result @check_not_closed @check_query_executed @@ -358,7 +361,9 @@ def fetchmany(self, size: Optional[int] = None) -> List[List[ColType]]: left, right = self._get_next_range(size) assert self._rows is not None rows = self._rows[left:right] - return [self._parse_row(row) for row in rows] + with Timer("[PERFORMANCE] Parsing query output into native Python types "): + result = [self._parse_row(row) for row in rows] + return result @check_not_closed @check_query_executed @@ -367,7 +372,9 @@ def fetchall(self) -> List[List[ColType]]: left, right = self._get_next_range(self.rowcount) assert self._rows is not None rows = self._rows[left:right] - return [self._parse_row(row) for row in rows] + with Timer("[PERFORMANCE] Parsing query output into native Python types "): + result = [self._parse_row(row) for row in rows] + return result @check_not_closed def setinputsizes(self, sizes: List[int]) -> None: diff --git a/src/firebolt/utils/util.py b/src/firebolt/utils/util.py index c7b7533b847..9efcdceed8e 100644 --- a/src/firebolt/utils/util.py +++ b/src/firebolt/utils/util.py @@ -1,9 +1,14 @@ +import logging from functools import lru_cache +from os import environ +from time import time +from types import TracebackType from typing import TYPE_CHECKING, Callable, Type, TypeVar from httpx import URL T = TypeVar("T") +logger = logging.getLogger(__name__) def cached_property(func: Callable[..., T]) -> T: @@ -99,3 +104,26 @@ def merge_urls(base: URL, merge: URL) -> URL: merge_raw_path = base.raw_path + merge.raw_path.lstrip(b"/") return base.copy_with(raw_path=merge_raw_path) return merge + + +class Timer: + def __init__(self, message: str = ""): + self._message = message + + def __enter__(self) -> "Timer": + self._start_time: float = time() + return self + + def __exit__( + self, + exc_type: Type[BaseException], + exc_val: BaseException, + exc_tb: TracebackType, + ) -> None: + self.elapsed_time: str = "{:.2f}".format(round((time() - self._start_time), 2)) + if ( + environ.get("FIREBOLT_SDK_PERFORMANCE_DEBUG", "0") == "1" + and self._message != "" + ): + log_message = self._message + self.elapsed_time + "s" + logger.debug(log_message)