From 559f52619968a69c19ba245234b6e9d9eb0734c8 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Fri, 23 Aug 2024 15:17:28 +0300 Subject: [PATCH] refactor: FIR-35468 remove account resolve endpoint from python sdk (#401) --- src/firebolt/async_db/connection.py | 6 +- src/firebolt/async_db/cursor.py | 8 -- src/firebolt/client/client.py | 110 +----------------- src/firebolt/common/cache.py | 3 - src/firebolt/db/connection.py | 6 +- src/firebolt/db/cursor.py | 8 -- src/firebolt/utils/urls.py | 2 - .../async/V2/test_system_engine_async.py | 9 -- .../dbapi/sync/V2/test_system_engine.py | 9 -- tests/unit/V1/client/conftest.py | 4 +- tests/unit/V1/client/test_client.py | 1 - tests/unit/V1/client/test_client_async.py | 1 - tests/unit/V1/conftest.py | 4 +- tests/unit/V1/service/conftest.py | 4 +- tests/unit/async_db/test_connection.py | 65 +---------- tests/unit/client/test_client.py | 26 ----- tests/unit/client/test_client_async.py | 27 +---- tests/unit/conftest.py | 46 +------- tests/unit/db/test_connection.py | 82 +------------ tests/unit/db_conftest.py | 38 +----- tests/unit/service/test_resource_manager.py | 29 +---- 21 files changed, 23 insertions(+), 465 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index f20f980283e..e98b7d1e237 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -11,10 +11,7 @@ from firebolt.client.auth import Auth from firebolt.client.client import AsyncClient, AsyncClientV1, AsyncClientV2 from firebolt.common.base_connection import BaseConnection -from firebolt.common.cache import ( - _firebolt_account_info_cache, - _firebolt_system_engine_cache, -) +from firebolt.common.cache import _firebolt_system_engine_cache from firebolt.common.constants import DEFAULT_TIMEOUT_SECONDS from firebolt.utils.exception import ConfigurationError, ConnectionClosedError from firebolt.utils.usage_tracker import get_user_agent_header @@ -135,7 +132,6 @@ async def connect( user_agent_header = get_user_agent_header(user_drivers, user_clients) if disable_cache: _firebolt_system_engine_cache.disable() - _firebolt_account_info_cache.disable() # Use v2 if auth is ClientCredentials # Use v1 if auth is ServiceAccount or UsernamePassword auth_version = auth.get_firebolt_version() diff --git a/src/firebolt/async_db/cursor.py b/src/firebolt/async_db/cursor.py index 6d1e1da3204..c652be67999 100644 --- a/src/firebolt/async_db/cursor.py +++ b/src/firebolt/async_db/cursor.py @@ -136,14 +136,6 @@ async def _parse_response_headers(self, headers: Headers) -> None: endpoint, params = _parse_update_endpoint( headers.get(UPDATE_ENDPOINT_HEADER) ) - if ( - params.get("account_id", await self._client.account_id) - != await self._client.account_id - ): - raise OperationalError( - "USE ENGINE command failed. Account parameter mismatch. " - "Contact support" - ) self._update_set_parameters(params) self.engine_url = endpoint diff --git a/src/firebolt/client/client.py b/src/firebolt/client/client.py index 97edee53a79..dd6cb620943 100644 --- a/src/firebolt/client/client.py +++ b/src/firebolt/client/client.py @@ -20,16 +20,12 @@ AsyncKeepaliveTransport, KeepaliveTransport, ) -from firebolt.common._types import _AccountInfo -from firebolt.common.cache import _firebolt_account_info_cache from firebolt.utils.exception import ( AccountNotFoundError, - AccountNotFoundOrNoAccessError, FireboltEngineError, InterfaceError, ) from firebolt.utils.urls import ( - ACCOUNT_BY_NAME_URL, ACCOUNT_BY_NAME_URL_V1, ACCOUNT_ENGINE_ID_BY_NAME_URL, ACCOUNT_ENGINE_URL, @@ -47,13 +43,6 @@ FireboltClientMixinBase = mixin_for(HttpxClient) # type: Any -def parse_response_for_account_info(response: Response) -> _AccountInfo: - """Construct account info object from the API response.""" - account_id = response.json()["id"] - account_version = int(response.json().get("infraVersion", 2)) - return _AccountInfo(id=account_id, version=account_version) - - class FireboltClientMixin(FireboltClientMixinBase): """HttpxAsyncClient mixin with Firebolt authentication functionality.""" @@ -101,7 +90,7 @@ def _merge_auth_request(self, request: Request) -> Request: return request def _enforce_trailing_slash(self, url: URL) -> URL: - """Don't automatically append trailing slach to a base url""" + """Don't automatically append trailing slash to a base url""" return url def clone(self) -> "Client": @@ -124,11 +113,6 @@ def __init__(self, *args: Any, **kwargs: Any): def account_id(self) -> str: ... - @property - @abstractmethod - def _account_version(self) -> int: - ... - def _send_handling_redirects( self, request: Request, *args: Any, **kwargs: Any ) -> Response: @@ -161,40 +145,6 @@ def __init__( **kwargs, ) - @property - def _account_info(self) -> _AccountInfo: - if account_info := _firebolt_account_info_cache.get( - [self.account_name, self._api_endpoint.host] - ): - return account_info - response = self.get( - url=self._api_endpoint.copy_with( - path=ACCOUNT_BY_NAME_URL.format(account_name=self.account_name) - ) - ) - if response.status_code == HttpxCodes.NOT_FOUND: - assert self.account_name is not None - raise AccountNotFoundOrNoAccessError(self.account_name) - # process all other status codes - response.raise_for_status() - account_info = parse_response_for_account_info(response) - _firebolt_account_info_cache.set( - key=[self.account_name, self._api_endpoint.host], value=account_info - ) - return account_info - - @property - def _account_version(self) -> int: - """User account version. 2 means both database and engine v2 are supported. - - Returns: - int: Account version - - Raises: - AccountNotFoundError: No account found with provided name - """ - return self._account_info.version - @property def account_id(self) -> str: """User account ID. @@ -208,7 +158,7 @@ def account_id(self) -> str: Raises: AccountNotFoundError: No account found with provided name """ - return self._account_info.id + return "" class ClientV1(Client): @@ -236,13 +186,6 @@ def __init__( ) self._auth_endpoint = URL(fix_url_schema(api_endpoint)) - @property - def _account_version(self) -> int: - """User account version. Hardcoded since it's not returned - by the backend for V1. - """ - return 1 - @cached_property def account_id(self) -> str: """User account ID. @@ -333,11 +276,6 @@ def __init__(self, *args: Any, **kwargs: Any): async def account_id(self) -> str: ... - @property - @abstractmethod - async def _account_version(self) -> int: - ... - async def _send_handling_redirects( self, request: Request, *args: Any, **kwargs: Any ) -> Response: @@ -370,30 +308,6 @@ def __init__( **kwargs, ) - async def _account_info(self) -> _AccountInfo: - # manual caching to avoid async_cached_property issues - if account_info := _firebolt_account_info_cache.get( - [self.account_name, self._api_endpoint.host] - ): - return account_info - - response = await self.get( - url=self._api_endpoint.copy_with( - path=ACCOUNT_BY_NAME_URL.format(account_name=self.account_name) - ) - ) - if response.status_code == HttpxCodes.NOT_FOUND: - assert self.account_name is not None - raise AccountNotFoundOrNoAccessError(self.account_name) - # process all other status codes - response.raise_for_status() - account_info = parse_response_for_account_info(response) - # cache for future use - _firebolt_account_info_cache.set( - key=[self.account_name, self._api_endpoint.host], value=account_info - ) - return account_info - @property async def account_id(self) -> str: """User account ID. @@ -407,19 +321,7 @@ async def account_id(self) -> str: Raises: AccountNotFoundError: No account found with provided name """ - return (await self._account_info()).id - - @property - async def _account_version(self) -> int: - """User account version. 2 means both database and engine v2 are supported. - - Returns: - int: Account version - - Raises: - AccountNotFoundError: No account found with provided name - """ - return (await self._account_info()).version + return "" class AsyncClientV1(AsyncClient): @@ -448,12 +350,6 @@ def __init__( self.account_id_cache: Dict[str, str] = {} self._auth_endpoint = URL(fix_url_schema(api_endpoint)) - @property - async def _account_version(self) -> int: - """User account version. Hardcoded since it's not returned - by the backend for V1.""" - return 1 - @property async def account_id(self) -> str: """User account ID. diff --git a/src/firebolt/common/cache.py b/src/firebolt/common/cache.py index ce7552d39ea..66a13df0ede 100644 --- a/src/firebolt/common/cache.py +++ b/src/firebolt/common/cache.py @@ -10,8 +10,6 @@ TypeVar, ) -from firebolt.common._types import _AccountInfo - T = TypeVar("T") @@ -85,4 +83,3 @@ def __contains__(self, key: str) -> bool: _firebolt_system_engine_cache = UtilCache[Tuple[str, Dict[str, str]]]( cache_name="system_engine" ) -_firebolt_account_info_cache = UtilCache[_AccountInfo](cache_name="account_info") diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 8f774bc3867..9914e3174a8 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -10,10 +10,7 @@ from firebolt.client import DEFAULT_API_URL, Client, ClientV1, ClientV2 from firebolt.client.auth import Auth from firebolt.common.base_connection import BaseConnection -from firebolt.common.cache import ( - _firebolt_account_info_cache, - _firebolt_system_engine_cache, -) +from firebolt.common.cache import _firebolt_system_engine_cache from firebolt.common.constants import DEFAULT_TIMEOUT_SECONDS from firebolt.db.cursor import Cursor, CursorV1, CursorV2 from firebolt.db.util import _get_system_engine_url_and_params @@ -48,7 +45,6 @@ def connect( auth_version = auth.get_firebolt_version() if disable_cache: _firebolt_system_engine_cache.disable() - _firebolt_account_info_cache.disable() # Use v2 if auth is ClientCredentials # Use v1 if auth is ServiceAccount or UsernamePassword if auth_version == 2: diff --git a/src/firebolt/db/cursor.py b/src/firebolt/db/cursor.py index 6c994f92d50..57c4679f6b2 100644 --- a/src/firebolt/db/cursor.py +++ b/src/firebolt/db/cursor.py @@ -135,14 +135,6 @@ def _parse_response_headers(self, headers: Headers) -> None: endpoint, params = _parse_update_endpoint( headers.get(UPDATE_ENDPOINT_HEADER) ) - if ( - params.get("account_id", self._client.account_id) - != self._client.account_id - ): - raise OperationalError( - "USE ENGINE command failed. Account parameter mismatch. " - "Contact support" - ) self._update_set_parameters(params) self.engine_url = endpoint diff --git a/src/firebolt/utils/urls.py b/src/firebolt/utils/urls.py index 1bb93fd3466..a556a10803b 100644 --- a/src/firebolt/utils/urls.py +++ b/src/firebolt/utils/urls.py @@ -5,8 +5,6 @@ ENGINES_URL = "/core/v1/account/engines" ENGINES_BY_IDS_URL = "/core/v1/engines:getByIds" -ACCOUNT_BY_NAME_URL = "/web/v3/account/{account_name}/resolve" - ACCOUNT_DATABASES_URL = "/core/v1/accounts/{account_id}/databases" ACCOUNT_DATABASE_URL = "/core/v1/accounts/{account_id}/databases/{database_id}" ACCOUNT_DATABASE_BINDING_URL = ACCOUNT_DATABASE_URL + "/bindings/{engine_id}" diff --git a/tests/integration/dbapi/async/V2/test_system_engine_async.py b/tests/integration/dbapi/async/V2/test_system_engine_async.py index b711bed626b..02e50318dc9 100644 --- a/tests/integration/dbapi/async/V2/test_system_engine_async.py +++ b/tests/integration/dbapi/async/V2/test_system_engine_async.py @@ -77,15 +77,6 @@ async def test_system_engine_no_db( ) -async def test_system_engine_account(connection_system_engine: Connection): - assert ( - await connection_system_engine._client.account_id - ), "Can't get account id explicitly" - assert ( - await connection_system_engine._client._account_version - ) == 2, "Invalid account version" - - async def test_system_engine_use_engine( connection_system_engine: Connection, database_name: str, engine_name: str ): diff --git a/tests/integration/dbapi/sync/V2/test_system_engine.py b/tests/integration/dbapi/sync/V2/test_system_engine.py index 409671f05bc..62d455ecea6 100644 --- a/tests/integration/dbapi/sync/V2/test_system_engine.py +++ b/tests/integration/dbapi/sync/V2/test_system_engine.py @@ -78,15 +78,6 @@ def test_system_engine_no_db( ) -def test_system_engine_account(connection_system_engine: Connection): - assert ( - connection_system_engine._client.account_id - ), "Can't get account id explicitly" - assert ( - connection_system_engine._client._account_version == 2 - ), "Invalid account version" - - def test_system_engine_use_engine( connection_system_engine: Connection, database_name: str, engine_name: str ): diff --git a/tests/unit/V1/client/conftest.py b/tests/unit/V1/client/conftest.py index 3ad226155a5..03ee5146f40 100644 --- a/tests/unit/V1/client/conftest.py +++ b/tests/unit/V1/client/conftest.py @@ -7,7 +7,7 @@ from pytest import fixture from firebolt.utils.exception import AccountNotFoundError -from firebolt.utils.urls import ACCOUNT_BY_NAME_URL, ACCOUNT_URL, AUTH_URL +from firebolt.utils.urls import ACCOUNT_BY_NAME_URL_V1, ACCOUNT_URL, AUTH_URL from tests.unit.response import Response @@ -91,7 +91,7 @@ def check_credentials( @fixture def account_id_url(api_endpoint: str) -> Pattern: - base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL}?account_name=" + base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL_V1}?account_name=" default_base = f"https://{api_endpoint}{ACCOUNT_URL}" base = base.replace("/", "\\/").replace("?", "\\?") default_base = default_base.replace("/", "\\/").replace("?", "\\?") diff --git a/tests/unit/V1/client/test_client.py b/tests/unit/V1/client/test_client.py index 1bca2ce440b..425e9e4094f 100644 --- a/tests/unit/V1/client/test_client.py +++ b/tests/unit/V1/client/test_client.py @@ -112,7 +112,6 @@ def test_client_account_id( api_endpoint=api_endpoint, ) as c: assert c.account_id == account_id, "Invalid account id returned" - assert c._account_version == 1, "Invalid account version returned" # FIR-14945 diff --git a/tests/unit/V1/client/test_client_async.py b/tests/unit/V1/client/test_client_async.py index 2dbc67efe6f..dec8e1fc123 100644 --- a/tests/unit/V1/client/test_client_async.py +++ b/tests/unit/V1/client/test_client_async.py @@ -122,7 +122,6 @@ async def test_client_account_id( api_endpoint=api_endpoint, ) as c: assert await c.account_id == account_id, "Invalid account id returned." - assert await c._account_version == 1, "Invalid account version returned." async def test_concurent_auth_lock( diff --git a/tests/unit/V1/conftest.py b/tests/unit/V1/conftest.py index 93693a73bf7..15b64d67e6e 100644 --- a/tests/unit/V1/conftest.py +++ b/tests/unit/V1/conftest.py @@ -8,7 +8,7 @@ from firebolt.client.auth import Auth, UsernamePassword from firebolt.utils.exception import AccountNotFoundError from firebolt.utils.urls import ( - ACCOUNT_BY_NAME_URL, + ACCOUNT_BY_NAME_URL_V1, ACCOUNT_DATABASE_BY_NAME_URL, ACCOUNT_ENGINE_URL, ACCOUNT_ENGINE_URL_BY_DATABASE_NAME_V1, @@ -148,7 +148,7 @@ def do_mock( @fixture def account_id_url(api_endpoint: str) -> Pattern: - base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL}?account_name=" + base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL_V1}?account_name=" default_base = f"https://{api_endpoint}{ACCOUNT_URL}" base = base.replace("/", "\\/").replace("?", "\\?") default_base = default_base.replace("/", "\\/").replace("?", "\\?") diff --git a/tests/unit/V1/service/conftest.py b/tests/unit/V1/service/conftest.py index add800dae31..17682980d1e 100644 --- a/tests/unit/V1/service/conftest.py +++ b/tests/unit/V1/service/conftest.py @@ -16,7 +16,7 @@ from firebolt.utils.exception import AccountNotFoundError from firebolt.utils.urls import ( ACCOUNT_BINDINGS_URL, - ACCOUNT_BY_NAME_URL, + ACCOUNT_BY_NAME_URL_V1, ACCOUNT_DATABASE_BINDING_URL, ACCOUNT_DATABASE_BY_NAME_URL, ACCOUNT_DATABASE_URL, @@ -599,7 +599,7 @@ def check_credentials( @fixture def account_id_url(api_endpoint: str) -> Pattern: - base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL}?account_name=" + base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL_V1}?account_name=" default_base = f"https://{api_endpoint}{ACCOUNT_URL}" base = base.replace("/", "\\/").replace("?", "\\?") default_base = default_base.replace("/", "\\/").replace("?", "\\?") diff --git a/tests/unit/async_db/test_connection.py b/tests/unit/async_db/test_connection.py index 5c12112862b..183eb67621c 100644 --- a/tests/unit/async_db/test_connection.py +++ b/tests/unit/async_db/test_connection.py @@ -8,10 +8,7 @@ from firebolt.async_db.connection import Connection, connect from firebolt.client.auth import Auth, ClientCredentials from firebolt.common._types import ColType -from firebolt.common.cache import ( - _firebolt_account_info_cache, - _firebolt_system_engine_cache, -) +from firebolt.common.cache import _firebolt_system_engine_cache from firebolt.utils.exception import ( AccountNotFoundOrNoAccessError, ConfigurationError, @@ -188,50 +185,6 @@ async def test_connect_engine_failed( httpx_mock.reset(False) -async def test_connect_invalid_account( - db_name: str, - engine_name: str, - engine_url: str, - auth_url: str, - api_endpoint: str, - auth: Auth, - account_name: str, - httpx_mock: HTTPXMock, - check_credentials_callback: Callable, - system_engine_no_db_query_url: str, - system_engine_query_url: str, - get_system_engine_url: str, - get_system_engine_callback: Callable, - use_database_callback: Callable, - use_engine_with_account_id_callback: Callable, - account_id_url: str, - account_id_invalid_callback: Callable, -): - httpx_mock.add_callback(check_credentials_callback, url=auth_url) - httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) - httpx_mock.add_callback(account_id_invalid_callback, url=account_id_url) - httpx_mock.add_callback( - use_database_callback, - url=system_engine_no_db_query_url, - match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"), - ) - - httpx_mock.add_callback( - use_engine_with_account_id_callback, - url=system_engine_query_url, - match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"), - ) - with raises(AccountNotFoundOrNoAccessError): - async with await connect( - engine_name=engine_name, - database=db_name, - auth=auth, - account_name=account_name, - api_endpoint=api_endpoint, - ) as connection: - await connection.cursor().execute("select *") - - @mark.parametrize("cache_enabled", [True, False]) async def test_connect_caching( db_name: str, @@ -244,32 +197,23 @@ async def test_connect_caching( check_credentials_callback: Callable, get_system_engine_url: str, get_system_engine_callback: Callable, - account_id_url: str, - account_id_callback: Callable, system_engine_query_url: str, system_engine_no_db_query_url: str, query_url: str, use_database_callback: Callable, - use_engine_with_account_id_callback: Callable, + use_engine_callback: Callable, query_callback: Callable, cache_enabled: bool, ): system_engine_call_counter = 0 - account_id_call_counter = 0 def system_engine_callback_counter(request, **kwargs): nonlocal system_engine_call_counter system_engine_call_counter += 1 return get_system_engine_callback(request, **kwargs) - def account_id_callback_counter(request, **kwargs): - nonlocal account_id_call_counter - account_id_call_counter += 1 - return account_id_callback(request, **kwargs) - httpx_mock.add_callback(check_credentials_callback, url=auth_url) httpx_mock.add_callback(system_engine_callback_counter, url=get_system_engine_url) - httpx_mock.add_callback(account_id_callback_counter, url=account_id_url) httpx_mock.add_callback( use_database_callback, url=system_engine_no_db_query_url, @@ -277,7 +221,7 @@ def account_id_callback_counter(request, **kwargs): ) httpx_mock.add_callback( - use_engine_with_account_id_callback, + use_engine_callback, url=system_engine_query_url, match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"), ) @@ -296,14 +240,11 @@ def account_id_callback_counter(request, **kwargs): if cache_enabled: assert system_engine_call_counter == 1, "System engine URL was not cached" - assert account_id_call_counter == 1, "Account ID was not cached" else: assert system_engine_call_counter != 1, "System engine URL was cached" - assert account_id_call_counter != 1, "Account ID was cached" # Reset caches for the next test iteration _firebolt_system_engine_cache.enable() - _firebolt_account_info_cache.enable() async def test_connect_system_engine_404( diff --git a/tests/unit/client/test_client.py b/tests/unit/client/test_client.py index 8d5304f3613..42a34911dde 100644 --- a/tests/unit/client/test_client.py +++ b/tests/unit/client/test_client.py @@ -1,4 +1,3 @@ -from re import Pattern from typing import Callable from httpx import Request, Timeout, codes @@ -11,7 +10,6 @@ from firebolt.client.resource_manager_hooks import raise_on_4xx_5xx from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL -from firebolt.utils.util import fix_url_schema from tests.unit.conftest import Response @@ -88,30 +86,6 @@ def test_client_different_auths( ), "invalid auth validation error message" -def test_client_account_id( - httpx_mock: HTTPXMock, - auth: Auth, - account_name: str, - account_id: str, - account_id_url: Pattern, - account_id_callback: Callable, - auth_url: str, - auth_callback: Callable, - api_endpoint: str, -): - httpx_mock.add_callback(account_id_callback, url=account_id_url) - httpx_mock.add_callback(auth_callback, url=auth_url) - - with Client( - account_name=account_name, - auth=auth, - base_url=fix_url_schema(api_endpoint), - api_endpoint=api_endpoint, - ) as cursor: - assert cursor.account_id == account_id, "Invalid account id returned" - assert cursor._account_version == 2, "Invalid account version returned" - - # FIR-14945 def test_refresh_with_hooks( fs: FakeFilesystem, diff --git a/tests/unit/client/test_client_async.py b/tests/unit/client/test_client_async.py index d75bff7a582..b24100f67f5 100644 --- a/tests/unit/client/test_client_async.py +++ b/tests/unit/client/test_client_async.py @@ -1,6 +1,6 @@ import random from queue import Queue -from re import Pattern, compile +from re import compile from types import MethodType from typing import Any, Callable @@ -12,7 +12,6 @@ from firebolt.client import AsyncClientV2 as AsyncClient from firebolt.client.auth import Auth, ClientCredentials from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL -from firebolt.utils.util import fix_url_schema from tests.unit.conftest import Response, retry_if_failed @@ -94,30 +93,6 @@ async def test_client_different_auths( ), "invalid auth validation error message" -async def test_client_account_id( - httpx_mock: HTTPXMock, - auth: Auth, - account_name: str, - account_id: str, - account_id_url: Pattern, - account_id_callback: Callable, - auth_url: str, - auth_callback: Callable, - api_endpoint: str, -): - httpx_mock.add_callback(account_id_callback, url=account_id_url) - httpx_mock.add_callback(auth_callback, url=auth_url) - - async with AsyncClient( - account_name=account_name, - auth=auth, - base_url=fix_url_schema(api_endpoint), - api_endpoint=api_endpoint, - ) as c: - assert await c.account_id == account_id, "Invalid account id returned." - assert await c._account_version == 2, "Invalid account version returned" - - async def test_concurent_auth_lock( httpx_mock: HTTPXMock, account_name: str, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4b4880b2bfa..7d77ca906a9 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,5 +1,4 @@ import functools -from re import Pattern, compile from typing import Callable import httpx @@ -8,11 +7,10 @@ from pytest import fixture from firebolt.client.auth import Auth, ClientCredentials -from firebolt.client.client import ClientV2, _firebolt_account_info_cache +from firebolt.client.client import ClientV2 from firebolt.common.cache import _firebolt_system_engine_cache from firebolt.common.settings import Settings from firebolt.utils.exception import ( - AccountNotFoundError, DatabaseError, DataError, Error, @@ -24,7 +22,7 @@ ProgrammingError, Warning, ) -from firebolt.utils.urls import ACCOUNT_BY_NAME_URL, AUTH_SERVICE_ACCOUNT_URL +from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL from tests.unit.db_conftest import * # noqa from tests.unit.response import Response @@ -46,7 +44,6 @@ def global_fake_fs(request) -> None: @fixture(autouse=True) def clear_cache() -> None: _firebolt_system_engine_cache.clear() - _firebolt_account_info_cache.clear() @fixture @@ -147,45 +144,6 @@ def db_name_updated() -> str: return "database_new" -@fixture -def account_id_url(api_endpoint: str, account_name: str) -> Pattern: - account_name_re = r"[^\\\\]*" - base = f"https://{api_endpoint}{ACCOUNT_BY_NAME_URL}" - base = base.replace("/", "\\/").replace("?", "\\?") - base = base.format(account_name=account_name_re) - return compile(base) - - -@fixture -def account_id_callback( - account_id: str, - account_name: str, -) -> Callable: - def do_mock( - request: Request, - **kwargs, - ) -> Response: - if request.url.path.split("/")[-2] != account_name: - raise AccountNotFoundError(request.url.path.split("/")[-2]) - return Response( - status_code=httpx.codes.OK, - json={"id": account_id, "infraVersion": 2}, - ) - - return do_mock - - -@fixture -def account_id_invalid_callback() -> Callable: - def do_mock( - request: Request, - **kwargs, - ) -> Response: - return Response(status_code=httpx.codes.NOT_FOUND, json={"error": "not found"}) - - return do_mock - - @fixture def engine_name() -> str: return "mock_engine_name" diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index 406d81b9508..d874079aff8 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -3,7 +3,6 @@ from typing import Callable, List from unittest.mock import patch -from httpx import codes from pyfakefs.fake_filesystem_unittest import Patcher from pytest import mark, raises, warns from pytest_httpx import HTTPXMock @@ -11,10 +10,7 @@ from firebolt.client.auth import Auth, ClientCredentials from firebolt.client.client import ClientV2 from firebolt.common._types import ColType -from firebolt.common.cache import ( - _firebolt_account_info_cache, - _firebolt_system_engine_cache, -) +from firebolt.common.cache import _firebolt_system_engine_cache from firebolt.db import Connection, connect from firebolt.db.cursor import CursorV2 from firebolt.utils.exception import ( @@ -24,7 +20,6 @@ FireboltError, ) from firebolt.utils.token_storage import TokenSecureStorage -from tests.unit.response import Response def test_connection_attributes(connection: Connection) -> None: @@ -194,65 +189,6 @@ def test_connect_engine_failed( httpx_mock.reset(False) -def test_connect_invalid_account( - db_name: str, - engine_name: str, - engine_url: str, - auth_url: str, - api_endpoint: str, - auth: Auth, - account_name: str, - httpx_mock: HTTPXMock, - check_credentials_callback: Callable, - system_engine_no_db_query_url: str, - system_engine_query_url: str, - get_system_engine_url: str, - get_system_engine_callback: Callable, - use_database_callback: Callable, - account_id_url: str, - account_id_invalid_callback: Callable, -): - httpx_mock.add_callback(check_credentials_callback, url=auth_url) - httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) - httpx_mock.add_callback(account_id_invalid_callback, url=account_id_url) - httpx_mock.add_callback( - use_database_callback, - url=system_engine_no_db_query_url, - match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"), - ) - - def use_engine_callback(*args, **kwargs): - query_response = { - "meta": [], - "data": [], - "rows": 0, - "statistics": {}, - } - - return Response( - status_code=codes.OK, - json=query_response, - headers={ - "Firebolt-Update-Endpoint": engine_url + "?account_id=2", - }, - ) - - httpx_mock.add_callback( - use_engine_callback, - url=system_engine_query_url, - match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"), - ) - with raises(AccountNotFoundOrNoAccessError): - with connect( - engine_name=engine_name, - database=db_name, - auth=auth, - account_name=account_name, - api_endpoint=api_endpoint, - ) as connection: - connection.cursor().execute("select *") - - @mark.parametrize("cache_enabled", [True, False]) def test_connect_caching( db_name: str, @@ -265,32 +201,23 @@ def test_connect_caching( check_credentials_callback: Callable, get_system_engine_url: str, get_system_engine_callback: Callable, - account_id_url: str, - account_id_callback: Callable, system_engine_query_url: str, system_engine_no_db_query_url: str, query_url: str, use_database_callback: Callable, - use_engine_with_account_id_callback: Callable, + use_engine_callback: Callable, query_callback: Callable, cache_enabled: bool, ): system_engine_call_counter = 0 - account_id_call_counter = 0 def system_engine_callback_counter(request, **kwargs): nonlocal system_engine_call_counter system_engine_call_counter += 1 return get_system_engine_callback(request, **kwargs) - def account_id_callback_counter(request, **kwargs): - nonlocal account_id_call_counter - account_id_call_counter += 1 - return account_id_callback(request, **kwargs) - httpx_mock.add_callback(check_credentials_callback, url=auth_url) httpx_mock.add_callback(system_engine_callback_counter, url=get_system_engine_url) - httpx_mock.add_callback(account_id_callback_counter, url=account_id_url) httpx_mock.add_callback( use_database_callback, url=system_engine_no_db_query_url, @@ -298,7 +225,7 @@ def account_id_callback_counter(request, **kwargs): ) httpx_mock.add_callback( - use_engine_with_account_id_callback, + use_engine_callback, url=system_engine_query_url, match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"), ) @@ -317,14 +244,11 @@ def account_id_callback_counter(request, **kwargs): if cache_enabled: assert system_engine_call_counter == 1, "System engine URL was not cached" - assert account_id_call_counter == 1, "Account ID was not cached" else: assert system_engine_call_counter != 1, "System engine URL was cached" - assert account_id_call_counter != 1, "Account ID was cached" # Reset caches for the next test iteration _firebolt_system_engine_cache.enable() - _firebolt_account_info_cache.enable() def test_connect_system_engine_404( diff --git a/tests/unit/db_conftest.py b/tests/unit/db_conftest.py index 7a0b48d94bd..29fa4eae861 100644 --- a/tests/unit/db_conftest.py +++ b/tests/unit/db_conftest.py @@ -260,14 +260,12 @@ def system_engine_url() -> str: @fixture -def system_engine_query_url( - system_engine_url: str, db_name: str, account_id: str -) -> str: +def system_engine_query_url(system_engine_url: str, db_name: str) -> str: return f"{system_engine_url}/?output_format=JSON_Compact&database={db_name}" @fixture -def system_engine_no_db_query_url(system_engine_url: str, account_id: str) -> str: +def system_engine_no_db_query_url(system_engine_url: str) -> str: return f"{system_engine_url}/?output_format=JSON_Compact" @@ -381,35 +379,6 @@ def inner( return inner -@fixture -def use_engine_with_account_id_callback( - engine_url: str, query_statistics: Dict[str, Any], account_id: int -) -> Callable: - def inner( - request: Request = None, - **kwargs, - ) -> Response: - assert request, "empty request" - assert request.method == "POST", "invalid request method" - - query_response = { - "meta": [], - "data": [], - "rows": 0, - "statistics": query_statistics, - } - - return Response( - status_code=codes.OK, - json=query_response, - headers={ - "Firebolt-Update-Endpoint": engine_url + f"?account_id={account_id}" - }, - ) - - return inner - - @fixture def use_engine_failed_callback( engine_url, db_name, query_statistics: Dict[str, Any] @@ -436,13 +405,10 @@ def mock_system_engine_connection_flow( check_credentials_callback: Callable, get_system_engine_url: str, get_system_engine_callback: Callable, - account_id_url: str, - account_id_callback: Callable, ) -> Callable: def inner() -> None: httpx_mock.add_callback(check_credentials_callback, url=auth_url) httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) - httpx_mock.add_callback(account_id_callback, url=account_id_url) return inner diff --git a/tests/unit/service/test_resource_manager.py b/tests/unit/service/test_resource_manager.py index 0a18309d031..229e9cc7227 100644 --- a/tests/unit/service/test_resource_manager.py +++ b/tests/unit/service/test_resource_manager.py @@ -1,15 +1,12 @@ -from re import Pattern from typing import Callable from pyfakefs.fake_filesystem_unittest import Patcher -from pytest import mark, raises +from pytest import mark from pytest_httpx import HTTPXMock from firebolt.client.auth import Auth, ClientCredentials from firebolt.service.manager import ResourceManager -from firebolt.utils.exception import AccountNotFoundError from firebolt.utils.token_storage import TokenSecureStorage -from firebolt.utils.urls import GATEWAY_HOST_BY_ACCOUNT_NAME def test_rm_credentials( @@ -76,27 +73,3 @@ def test_rm_token_cache( assert ( ts.get_cached_token() is None ), "Token is cached even though caching is disabled" - - -def test_rm_invalid_account_name( - httpx_mock: HTTPXMock, - auth: Auth, - api_endpoint: str, - auth_url: str, - check_credentials_callback: Callable, - account_id_url: Pattern, - account_id_callback: Callable, - get_system_engine_callback: Callable, -) -> None: - """Resource manager raises an error on invalid account name.""" - get_system_engine_url = ( - f"https://{api_endpoint}" - f"{GATEWAY_HOST_BY_ACCOUNT_NAME.format(account_name='invalid')}" - ) - - httpx_mock.add_callback(check_credentials_callback, url=auth_url) - httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) - httpx_mock.add_callback(account_id_callback, url=account_id_url) - - with raises(AccountNotFoundError): - ResourceManager(auth=auth, account_name="invalid", api_endpoint=api_endpoint)