From 221b69bd75939e3ec52019f73957104435e37f9c Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Tue, 29 Mar 2022 13:45:20 +0200 Subject: [PATCH 1/8] allow not passing engine to connection --- src/firebolt/async_db/connection.py | 91 ++++++++++++++++++++++---- tests/unit/async_db/test_connection.py | 9 --- tests/unit/db/test_connection.py | 7 -- 3 files changed, 78 insertions(+), 29 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 8bd4ec17066..e3eaada0d4c 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -15,10 +15,16 @@ from firebolt.common.exception import ( ConfigurationError, ConnectionClosedError, + FireboltDatabaseError, FireboltEngineError, InterfaceError, ) -from firebolt.common.urls import ACCOUNT_ENGINE_BY_NAME_URL, ACCOUNT_ENGINE_URL +from firebolt.common.urls import ( + ACCOUNT_BINDINGS_URL, + ACCOUNT_DATABASE_BY_NAME_URL, + ACCOUNT_ENGINE_BY_NAME_URL, + ACCOUNT_ENGINE_URL, +) from firebolt.common.util import fix_url_schema DEFAULT_TIMEOUT_SECONDS: int = 5 @@ -26,6 +32,16 @@ KEEPIDLE_RATE: int = 60 # seconds +async def _get_engine_endpoint(client: AsyncClient, engine_id: int) -> str: + response = await client.get( + url=ACCOUNT_ENGINE_URL.format( + account_id=(await client.account_id), engine_id=engine_id + ), + ) + response.raise_for_status() + return response.json()["engine"]["endpoint"] + + async def _resolve_engine_url( engine_name: str, auth: AuthTypes, @@ -46,13 +62,7 @@ async def _resolve_engine_url( ) response.raise_for_status() engine_id = response.json()["engine_id"]["engine_id"] - response = await client.get( - url=ACCOUNT_ENGINE_URL.format( - account_id=account_id, engine_id=engine_id - ), - ) - response.raise_for_status() - return response.json()["engine"]["endpoint"] + return await _get_engine_endpoint(client, engine_id) except HTTPStatusError as e: # Engine error would be 404. if e.response.status_code != 404: @@ -65,6 +75,57 @@ async def _resolve_engine_url( raise InterfaceError(f"Unable to retrieve engine endpoint: {e}.") +async def _get_database_default_engine_url( + database: str, + auth: AuthTypes, + api_endpoint: str, + account_name: Optional[str] = None, +) -> str: + async with AsyncClient( + auth=auth, + base_url=api_endpoint, + account_name=account_name, + api_endpoint=api_endpoint, + ) as client: + try: + account_id = await client.account_id + # Get database id by name + response = await client.get( + url=ACCOUNT_DATABASE_BY_NAME_URL.format(account_id=account_id), + params={"database_name": database}, + ) + response.raise_for_status() + database_id = response.json()["database_id"]["database_id"] + + # Get attachend engines to a database + response = await client.get( + url=ACCOUNT_BINDINGS_URL.format(account_id=account_id), + params={ + "filter.id_database_id_eq": database_id, + }, + ) + response.raise_for_status() + default_engines_bindings = [ + b for b in response.json()["edges"] if b["engine_is_default"] + ] + + if len(default_engines_bindings) == 0: + raise FireboltDatabaseError( + f"Database {database} has no default engines" + ) + engine_id = default_engines_bindings[0].binding_key.engine_id + + return await _get_engine_endpoint(client, engine_id) + except ( + JSONDecodeError, + RequestError, + RuntimeError, + HTTPStatusError, + KeyError, + ) as e: + raise InterfaceError(f"Unable to retrieve default engine endpoint: {e}.") + + def _validate_engine_name_and_url( engine_name: Optional[str], engine_url: Optional[str] ) -> None: @@ -72,10 +133,6 @@ def _validate_engine_name_and_url( raise ConfigurationError( "Both engine_name and engine_url are provided. Provide only one to connect." ) - if not engine_name and not engine_url: - raise ConfigurationError( - "Neither engine_name nor engine_url is provided. Provide one to connect." - ) def _get_auth( @@ -136,7 +193,15 @@ async def connect_inner( # Mypy checks, this should never happen assert database is not None - if engine_name: + if not engine_name and not engine_url: + engine_url = await _get_database_default_engine_url( + database=database, + auth=auth, + account_name=account_name, + api_endpoint=api_endpoint, + ) + + elif engine_name: engine_url = await _resolve_engine_url( engine_name=engine_name, auth=auth, diff --git a/tests/unit/async_db/test_connection.py b/tests/unit/async_db/test_connection.py index 938ab85c743..c534fef647c 100644 --- a/tests/unit/async_db/test_connection.py +++ b/tests/unit/async_db/test_connection.py @@ -167,15 +167,6 @@ async def test_connect_engine_name( ): pass - with raises(ConfigurationError): - async with await connect( - database="db", - username="username", - password="password", - account_name="account", - ): - pass - httpx_mock.add_callback(auth_callback, url=auth_url) httpx_mock.add_callback(query_callback, url=query_url) httpx_mock.add_callback(account_id_callback, url=account_id_url) diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index 8f3779ace9e..b1b38714d6e 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -152,13 +152,6 @@ def test_connect_engine_name( password="password", ) - with raises(ConfigurationError): - connect( - database="db", - username="username", - password="password", - ) - httpx_mock.add_callback(auth_callback, url=auth_url) httpx_mock.add_callback(query_callback, url=query_url) httpx_mock.add_callback(account_id_callback, url=account_id_url) From 0d406c99f9f80b4803de0f48d1527ccacc51605d Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Tue, 29 Mar 2022 15:44:42 +0200 Subject: [PATCH 2/8] add unit tests --- src/firebolt/async_db/connection.py | 2 +- tests/unit/async_db/test_connection.py | 97 ++++++++++++++++++++++++- tests/unit/conftest.py | 62 ++++++++++++++++ tests/unit/db/test_connection.py | 99 +++++++++++++++++++++++++- 4 files changed, 256 insertions(+), 4 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index e3eaada0d4c..ca467785fd2 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -113,7 +113,7 @@ async def _get_database_default_engine_url( raise FireboltDatabaseError( f"Database {database} has no default engines" ) - engine_id = default_engines_bindings[0].binding_key.engine_id + engine_id = default_engines_bindings[0]["id"]["engine_id"] return await _get_engine_endpoint(client, engine_id) except ( diff --git a/tests/unit/async_db/test_connection.py b/tests/unit/async_db/test_connection.py index c534fef647c..fd5790ed803 100644 --- a/tests/unit/async_db/test_connection.py +++ b/tests/unit/async_db/test_connection.py @@ -9,6 +9,7 @@ from firebolt.common.exception import ( ConfigurationError, ConnectionClosedError, + FireboltDatabaseError, FireboltEngineError, ) from firebolt.common.settings import Settings @@ -149,8 +150,6 @@ async def test_connect_engine_name( engine_id: str, get_engine_url: str, get_engine_callback: Callable, - get_providers_url: str, - get_providers_callback: Callable, python_query_data: List[List[ColType]], account_id: str, ): @@ -213,6 +212,100 @@ async def test_connect_engine_name( assert await connection.cursor().execute("select*") == len(python_query_data) +@mark.asyncio +async def test_connect_default_engine( + settings: Settings, + db_name: str, + httpx_mock: HTTPXMock, + auth_callback: Callable, + auth_url: str, + query_callback: Callable, + query_url: str, + account_id_url: str, + account_id_callback: Callable, + engine_id: str, + get_engine_url: str, + get_engine_callback: Callable, + database_by_name_url: str, + database_by_name_callback: Callable, + database_id: str, + bindings_url: str, + python_query_data: List[List[ColType]], + account_id: str, +): + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(query_callback, url=query_url) + httpx_mock.add_callback(account_id_callback, url=account_id_url) + httpx_mock.add_callback(database_by_name_callback, url=database_by_name_url) + bindings_url = f"{bindings_url}?filter.id_database_id_eq={database_id}" + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={"edges": []}, + ) + + with raises(FireboltDatabaseError): + async with await connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ): + pass + + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={ + "edges": [ + { + "engine_is_default": False, + "id": {"engine_id": engine_id}, + } + ] + }, + ) + + with raises(FireboltDatabaseError): + async with await connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ): + pass + + non_default_engine_id = "non_default_engine_id" + + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={ + "edges": [ + { + "engine_is_default": False, + "id": {"engine_id": non_default_engine_id}, + }, + { + "engine_is_default": True, + "id": {"engine_id": engine_id}, + }, + ] + }, + ) + httpx_mock.add_callback(get_engine_callback, url=get_engine_url) + async with await connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ) as connection: + assert await connection.cursor().execute("select*") == len(python_query_data) + + @mark.asyncio async def test_connection_commit(connection: Connection): # nothing happens diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 8c5849ace9e..1a5e56bcddf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -21,7 +21,9 @@ ) from firebolt.common.settings import Settings from firebolt.common.urls import ( + ACCOUNT_BINDINGS_URL, ACCOUNT_BY_NAME_URL, + ACCOUNT_DATABASE_BY_NAME_URL, ACCOUNT_ENGINE_URL, ACCOUNT_URL, AUTH_URL, @@ -242,6 +244,66 @@ def get_databases_url(settings: Settings) -> str: return f"https://{settings.server}{DATABASES_URL}" +@fixture +def database_id() -> str: + return "database_id" + + +@fixture +def database_by_name_url(settings: Settings, account_id: str, db_name: str) -> str: + return ( + f"https://{settings.server}" + f"{ACCOUNT_DATABASE_BY_NAME_URL.format(account_id=account_id)}" + f"?database_name={db_name}" + ) + + +@fixture +def database_by_name_callback(account_id: str, database_id: str) -> str: + def do_mock( + request: httpx.Request = None, + **kwargs, + ) -> Response: + return Response( + status_code=httpx.codes.OK, + json={ + "database_id": { + "database_id": database_id, + "account_id": account_id, + } + }, + ) + + return do_mock + + +@fixture +def bindings_url(settings: Settings, account_id: str) -> str: + return ( + f"https://{settings.server}{ACCOUNT_BINDINGS_URL.format(account_id=account_id)}" + ) + + +@fixture +def bindings_callback(account_id: str, database_id: str) -> str: + def do_mock( + request: httpx.Request = None, + **kwargs, + ) -> Response: + assert request.url == get_providers_url + return Response( + status_code=httpx.codes.OK, + json={ + "database_id": { + "database_id": database_id, + "account_id": account_id, + } + }, + ) + + return do_mock + + @fixture def db_api_exceptions(): exceptions = { diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index b1b38714d6e..354b38aaa4e 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -5,7 +5,11 @@ from pytest_httpx import HTTPXMock from firebolt.async_db._types import ColType -from firebolt.common.exception import ConfigurationError, ConnectionClosedError +from firebolt.common.exception import ( + ConfigurationError, + ConnectionClosedError, + FireboltDatabaseError, +) from firebolt.common.settings import Settings from firebolt.common.urls import ACCOUNT_ENGINE_BY_NAME_URL from firebolt.db import Connection, connect @@ -179,6 +183,99 @@ def test_connect_engine_name( assert connection.cursor().execute("select*") == len(python_query_data) +def test_connect_default_engine( + settings: Settings, + db_name: str, + httpx_mock: HTTPXMock, + auth_callback: Callable, + auth_url: str, + query_callback: Callable, + query_url: str, + account_id_url: str, + account_id_callback: Callable, + engine_id: str, + get_engine_url: str, + get_engine_callback: Callable, + database_by_name_url: str, + database_by_name_callback: Callable, + database_id: str, + bindings_url: str, + python_query_data: List[List[ColType]], + account_id: str, +): + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(query_callback, url=query_url) + httpx_mock.add_callback(account_id_callback, url=account_id_url) + httpx_mock.add_callback(database_by_name_callback, url=database_by_name_url) + bindings_url = f"{bindings_url}?filter.id_database_id_eq={database_id}" + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={"edges": []}, + ) + + with raises(FireboltDatabaseError): + with connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ): + pass + + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={ + "edges": [ + { + "engine_is_default": False, + "id": {"engine_id": engine_id}, + } + ] + }, + ) + + with raises(FireboltDatabaseError): + with connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ): + pass + + non_default_engine_id = "non_default_engine_id" + + httpx_mock.add_response( + url=bindings_url, + status_code=codes.OK, + json={ + "edges": [ + { + "engine_is_default": False, + "id": {"engine_id": non_default_engine_id}, + }, + { + "engine_is_default": True, + "id": {"engine_id": engine_id}, + }, + ] + }, + ) + httpx_mock.add_callback(get_engine_callback, url=get_engine_url) + with connect( + database=db_name, + username="u", + password="p", + account_name=settings.account_name, + api_endpoint=settings.server, + ) as connection: + assert connection.cursor().execute("select*") == len(python_query_data) + + def test_connection_unclosed_warnings(): c = Connection("", "", ("", ""), "") with warns(UserWarning) as winfo: From 069f31a677dc074d46ba2b61d589cb968a898059 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 10:45:04 +0200 Subject: [PATCH 3/8] add integration tests --- src/firebolt/async_db/connection.py | 5 ++++- tests/integration/dbapi/async/conftest.py | 19 +++++++++++++++++++ .../dbapi/async/test_queries_async.py | 16 ++++++++++++++++ tests/integration/dbapi/sync/conftest.py | 19 +++++++++++++++++++ tests/integration/dbapi/sync/test_queries.py | 15 +++++++++++++++ tests/unit/async_db/test_connection.py | 18 ++++++++++++------ tests/unit/db/test_connection.py | 18 ++++++++++++------ 7 files changed, 97 insertions(+), 13 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index ca467785fd2..09f0df82674 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -105,8 +105,11 @@ async def _get_database_default_engine_url( }, ) response.raise_for_status() + print(response.json()["edges"]) default_engines_bindings = [ - b for b in response.json()["edges"] if b["engine_is_default"] + b["node"] + for b in response.json()["edges"] + if b["node"]["engine_is_default"] ] if len(default_engines_bindings) == 0: diff --git a/tests/integration/dbapi/async/conftest.py b/tests/integration/dbapi/async/conftest.py index 4354d03d98e..61282b6bcd7 100644 --- a/tests/integration/dbapi/async/conftest.py +++ b/tests/integration/dbapi/async/conftest.py @@ -42,3 +42,22 @@ async def connection_engine_name( api_endpoint=api_endpoint, ) as connection: yield connection + + +@fixture +async def connection_no_engine( + database_name: str, + username: str, + password: str, + account_name: str, + api_endpoint: str, +) -> Connection: + + async with await connect( + database=database_name, + username=username, + password=password, + account_name=account_name, + api_endpoint=api_endpoint, + ) as connection: + yield connection diff --git a/tests/integration/dbapi/async/test_queries_async.py b/tests/integration/dbapi/async/test_queries_async.py index 3292a09f192..fbe5ee329ff 100644 --- a/tests/integration/dbapi/async/test_queries_async.py +++ b/tests/integration/dbapi/async/test_queries_async.py @@ -32,6 +32,22 @@ async def test_connect_engine_name( ) +@mark.asyncio +async def test_connect_no_engine( + connection_no_engine: Connection, + all_types_query: str, + all_types_query_description: List[Column], + all_types_query_response: List[ColType], +) -> None: + """Connecting with engine name is handled properly.""" + await test_select( + connection_no_engine, + all_types_query, + all_types_query_description, + all_types_query_response, + ) + + @mark.asyncio async def test_select( connection: Connection, diff --git a/tests/integration/dbapi/sync/conftest.py b/tests/integration/dbapi/sync/conftest.py index c518bd2c83e..b55eb46c30d 100644 --- a/tests/integration/dbapi/sync/conftest.py +++ b/tests/integration/dbapi/sync/conftest.py @@ -43,3 +43,22 @@ def connection_engine_name( ) yield connection connection.close() + + +@fixture +def connection_no_engine( + database_name: str, + username: str, + password: str, + account_name: str, + api_endpoint: str, +) -> Connection: + connection = connect( + database=database_name, + username=username, + password=password, + account_name=account_name, + api_endpoint=api_endpoint, + ) + yield connection + connection.close() diff --git a/tests/integration/dbapi/sync/test_queries.py b/tests/integration/dbapi/sync/test_queries.py index 8e4ed16c44e..4900a94ef16 100644 --- a/tests/integration/dbapi/sync/test_queries.py +++ b/tests/integration/dbapi/sync/test_queries.py @@ -31,6 +31,21 @@ def test_connect_engine_name( ) +def test_connect_no_engine( + connection_no_engine: Connection, + all_types_query: str, + all_types_query_description: List[Column], + all_types_query_response: List[ColType], +) -> None: + """Connecting with engine name is handled properly.""" + test_select( + connection_no_engine, + all_types_query, + all_types_query_description, + all_types_query_response, + ) + + def test_select( connection: Connection, all_types_query: str, diff --git a/tests/unit/async_db/test_connection.py b/tests/unit/async_db/test_connection.py index fd5790ed803..af8edd4937f 100644 --- a/tests/unit/async_db/test_connection.py +++ b/tests/unit/async_db/test_connection.py @@ -260,8 +260,10 @@ async def test_connect_default_engine( json={ "edges": [ { - "engine_is_default": False, - "id": {"engine_id": engine_id}, + "node": { + "engine_is_default": False, + "id": {"engine_id": engine_id}, + } } ] }, @@ -285,12 +287,16 @@ async def test_connect_default_engine( json={ "edges": [ { - "engine_is_default": False, - "id": {"engine_id": non_default_engine_id}, + "node": { + "engine_is_default": False, + "id": {"engine_id": non_default_engine_id}, + } }, { - "engine_is_default": True, - "id": {"engine_id": engine_id}, + "node": { + "engine_is_default": True, + "id": {"engine_id": engine_id}, + } }, ] }, diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index 354b38aaa4e..842556bbe02 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -230,8 +230,10 @@ def test_connect_default_engine( json={ "edges": [ { - "engine_is_default": False, - "id": {"engine_id": engine_id}, + "node": { + "engine_is_default": False, + "id": {"engine_id": engine_id}, + } } ] }, @@ -255,12 +257,16 @@ def test_connect_default_engine( json={ "edges": [ { - "engine_is_default": False, - "id": {"engine_id": non_default_engine_id}, + "node": { + "engine_is_default": False, + "id": {"engine_id": non_default_engine_id}, + } }, { - "engine_is_default": True, - "id": {"engine_id": engine_id}, + "node": { + "engine_is_default": True, + "id": {"engine_id": engine_id}, + } }, ] }, From 98d49b824f27708080db575e6178978199dc2c5b Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 10:52:23 +0200 Subject: [PATCH 4/8] remove debug --- src/firebolt/async_db/connection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 09f0df82674..f8f79d92a6b 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -105,7 +105,6 @@ async def _get_database_default_engine_url( }, ) response.raise_for_status() - print(response.json()["edges"]) default_engines_bindings = [ b["node"] for b in response.json()["edges"] From f5ca9a203ded4625d7d52d203918190467d52e97 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 13:09:54 +0200 Subject: [PATCH 5/8] use single endpoint for default engine --- src/firebolt/async_db/connection.py | 49 ++++-------------- src/firebolt/common/urls.py | 1 + tests/unit/async_db/test_connection.py | 66 ++---------------------- tests/unit/conftest.py | 27 ++-------- tests/unit/db/test_connection.py | 71 ++------------------------ 5 files changed, 25 insertions(+), 189 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index f8f79d92a6b..687e9fe5513 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -15,15 +15,13 @@ from firebolt.common.exception import ( ConfigurationError, ConnectionClosedError, - FireboltDatabaseError, FireboltEngineError, InterfaceError, ) from firebolt.common.urls import ( - ACCOUNT_BINDINGS_URL, - ACCOUNT_DATABASE_BY_NAME_URL, ACCOUNT_ENGINE_BY_NAME_URL, ACCOUNT_ENGINE_URL, + ACCOUNT_ENGINE_URL_BY_DATABASE_NAME, ) from firebolt.common.util import fix_url_schema @@ -32,16 +30,6 @@ KEEPIDLE_RATE: int = 60 # seconds -async def _get_engine_endpoint(client: AsyncClient, engine_id: int) -> str: - response = await client.get( - url=ACCOUNT_ENGINE_URL.format( - account_id=(await client.account_id), engine_id=engine_id - ), - ) - response.raise_for_status() - return response.json()["engine"]["endpoint"] - - async def _resolve_engine_url( engine_name: str, auth: AuthTypes, @@ -62,7 +50,14 @@ async def _resolve_engine_url( ) response.raise_for_status() engine_id = response.json()["engine_id"]["engine_id"] - return await _get_engine_endpoint(client, engine_id) + + response = await client.get( + url=ACCOUNT_ENGINE_URL.format( + account_id=(await client.account_id), engine_id=engine_id + ), + ) + response.raise_for_status() + return response.json()["engine"]["endpoint"] except HTTPStatusError as e: # Engine error would be 404. if e.response.status_code != 404: @@ -91,33 +86,11 @@ async def _get_database_default_engine_url( account_id = await client.account_id # Get database id by name response = await client.get( - url=ACCOUNT_DATABASE_BY_NAME_URL.format(account_id=account_id), + url=ACCOUNT_ENGINE_URL_BY_DATABASE_NAME.format(account_id=account_id), params={"database_name": database}, ) response.raise_for_status() - database_id = response.json()["database_id"]["database_id"] - - # Get attachend engines to a database - response = await client.get( - url=ACCOUNT_BINDINGS_URL.format(account_id=account_id), - params={ - "filter.id_database_id_eq": database_id, - }, - ) - response.raise_for_status() - default_engines_bindings = [ - b["node"] - for b in response.json()["edges"] - if b["node"]["engine_is_default"] - ] - - if len(default_engines_bindings) == 0: - raise FireboltDatabaseError( - f"Database {database} has no default engines" - ) - engine_id = default_engines_bindings[0]["id"]["engine_id"] - - return await _get_engine_endpoint(client, engine_id) + return response.json()["engine_url"] except ( JSONDecodeError, RequestError, diff --git a/src/firebolt/common/urls.py b/src/firebolt/common/urls.py index 07a29d74d89..470b5e81aca 100644 --- a/src/firebolt/common/urls.py +++ b/src/firebolt/common/urls.py @@ -15,6 +15,7 @@ ACCOUNT_ENGINES_URL = "/core/v1/accounts/{account_id}/engines" ACCOUNT_ENGINE_BY_NAME_URL = ACCOUNT_ENGINES_URL + ":getIdByName" ACCOUNT_ENGINE_REVISION_URL = ACCOUNT_ENGINE_URL + "/engineRevisions/{revision_id}" +ACCOUNT_ENGINE_URL_BY_DATABASE_NAME = ACCOUNT_ENGINES_URL + ":getURLByDatabaseName" ACCOUNT_DATABASES_URL = "/core/v1/accounts/{account_id}/databases" ACCOUNT_DATABASE_URL = "/core/v1/accounts/{account_id}/databases/{database_id}" diff --git a/tests/unit/async_db/test_connection.py b/tests/unit/async_db/test_connection.py index af8edd4937f..314f1a74337 100644 --- a/tests/unit/async_db/test_connection.py +++ b/tests/unit/async_db/test_connection.py @@ -9,7 +9,6 @@ from firebolt.common.exception import ( ConfigurationError, ConnectionClosedError, - FireboltDatabaseError, FireboltEngineError, ) from firebolt.common.settings import Settings @@ -229,79 +228,22 @@ async def test_connect_default_engine( database_by_name_url: str, database_by_name_callback: Callable, database_id: str, - bindings_url: str, + engine_by_db_url: str, python_query_data: List[List[ColType]], account_id: str, ): httpx_mock.add_callback(auth_callback, url=auth_url) httpx_mock.add_callback(query_callback, url=query_url) httpx_mock.add_callback(account_id_callback, url=account_id_url) - httpx_mock.add_callback(database_by_name_callback, url=database_by_name_url) - bindings_url = f"{bindings_url}?filter.id_database_id_eq={database_id}" - httpx_mock.add_response( - url=bindings_url, - status_code=codes.OK, - json={"edges": []}, - ) - - with raises(FireboltDatabaseError): - async with await connect( - database=db_name, - username="u", - password="p", - account_name=settings.account_name, - api_endpoint=settings.server, - ): - pass + engine_by_db_url = f"{engine_by_db_url}?database_name={db_name}" httpx_mock.add_response( - url=bindings_url, + url=engine_by_db_url, status_code=codes.OK, json={ - "edges": [ - { - "node": { - "engine_is_default": False, - "id": {"engine_id": engine_id}, - } - } - ] + "engine_url": settings.server, }, ) - - with raises(FireboltDatabaseError): - async with await connect( - database=db_name, - username="u", - password="p", - account_name=settings.account_name, - api_endpoint=settings.server, - ): - pass - - non_default_engine_id = "non_default_engine_id" - - httpx_mock.add_response( - url=bindings_url, - status_code=codes.OK, - json={ - "edges": [ - { - "node": { - "engine_is_default": False, - "id": {"engine_id": non_default_engine_id}, - } - }, - { - "node": { - "engine_is_default": True, - "id": {"engine_id": engine_id}, - } - }, - ] - }, - ) - httpx_mock.add_callback(get_engine_callback, url=get_engine_url) async with await connect( database=db_name, username="u", diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 1a5e56bcddf..6d37ceebb53 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -21,10 +21,10 @@ ) from firebolt.common.settings import Settings from firebolt.common.urls import ( - ACCOUNT_BINDINGS_URL, ACCOUNT_BY_NAME_URL, ACCOUNT_DATABASE_BY_NAME_URL, ACCOUNT_ENGINE_URL, + ACCOUNT_ENGINE_URL_BY_DATABASE_NAME, ACCOUNT_URL, AUTH_URL, DATABASES_URL, @@ -278,32 +278,13 @@ def do_mock( @fixture -def bindings_url(settings: Settings, account_id: str) -> str: +def engine_by_db_url(settings: Settings, account_id: str) -> str: return ( - f"https://{settings.server}{ACCOUNT_BINDINGS_URL.format(account_id=account_id)}" + f"https://{settings.server}" + f"{ACCOUNT_ENGINE_URL_BY_DATABASE_NAME.format(account_id=account_id)}" ) -@fixture -def bindings_callback(account_id: str, database_id: str) -> str: - def do_mock( - request: httpx.Request = None, - **kwargs, - ) -> Response: - assert request.url == get_providers_url - return Response( - status_code=httpx.codes.OK, - json={ - "database_id": { - "database_id": database_id, - "account_id": account_id, - } - }, - ) - - return do_mock - - @fixture def db_api_exceptions(): exceptions = { diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index 842556bbe02..780b8692aca 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -5,11 +5,7 @@ from pytest_httpx import HTTPXMock from firebolt.async_db._types import ColType -from firebolt.common.exception import ( - ConfigurationError, - ConnectionClosedError, - FireboltDatabaseError, -) +from firebolt.common.exception import ConfigurationError, ConnectionClosedError from firebolt.common.settings import Settings from firebolt.common.urls import ACCOUNT_ENGINE_BY_NAME_URL from firebolt.db import Connection, connect @@ -199,79 +195,22 @@ def test_connect_default_engine( database_by_name_url: str, database_by_name_callback: Callable, database_id: str, - bindings_url: str, + engine_by_db_url: str, python_query_data: List[List[ColType]], account_id: str, ): httpx_mock.add_callback(auth_callback, url=auth_url) httpx_mock.add_callback(query_callback, url=query_url) httpx_mock.add_callback(account_id_callback, url=account_id_url) - httpx_mock.add_callback(database_by_name_callback, url=database_by_name_url) - bindings_url = f"{bindings_url}?filter.id_database_id_eq={database_id}" - httpx_mock.add_response( - url=bindings_url, - status_code=codes.OK, - json={"edges": []}, - ) - - with raises(FireboltDatabaseError): - with connect( - database=db_name, - username="u", - password="p", - account_name=settings.account_name, - api_endpoint=settings.server, - ): - pass + engine_by_db_url = f"{engine_by_db_url}?database_name={db_name}" httpx_mock.add_response( - url=bindings_url, + url=engine_by_db_url, status_code=codes.OK, json={ - "edges": [ - { - "node": { - "engine_is_default": False, - "id": {"engine_id": engine_id}, - } - } - ] + "engine_url": settings.server, }, ) - - with raises(FireboltDatabaseError): - with connect( - database=db_name, - username="u", - password="p", - account_name=settings.account_name, - api_endpoint=settings.server, - ): - pass - - non_default_engine_id = "non_default_engine_id" - - httpx_mock.add_response( - url=bindings_url, - status_code=codes.OK, - json={ - "edges": [ - { - "node": { - "engine_is_default": False, - "id": {"engine_id": non_default_engine_id}, - } - }, - { - "node": { - "engine_is_default": True, - "id": {"engine_id": engine_id}, - } - }, - ] - }, - ) - httpx_mock.add_callback(get_engine_callback, url=get_engine_url) with connect( database=db_name, username="u", From 93b83724b090bdcda94384f00e8dbc46cdfcef04 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 13:10:49 +0200 Subject: [PATCH 6/8] revert redundant changes --- src/firebolt/async_db/connection.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 687e9fe5513..465254d1938 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -50,10 +50,9 @@ async def _resolve_engine_url( ) response.raise_for_status() engine_id = response.json()["engine_id"]["engine_id"] - response = await client.get( url=ACCOUNT_ENGINE_URL.format( - account_id=(await client.account_id), engine_id=engine_id + account_id=account_id, engine_id=engine_id ), ) response.raise_for_status() From 9207bad0efee28679e039aaa51c63797dbf5b99b Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 13:11:48 +0200 Subject: [PATCH 7/8] remove redundant comment --- src/firebolt/async_db/connection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 465254d1938..4d06e3d5b2c 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -83,7 +83,6 @@ async def _get_database_default_engine_url( ) as client: try: account_id = await client.account_id - # Get database id by name response = await client.get( url=ACCOUNT_ENGINE_URL_BY_DATABASE_NAME.format(account_id=account_id), params={"database_name": database}, From 26692240712ea3d0bf6d2afeb33e3667062b114b Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 30 Mar 2022 13:49:36 +0200 Subject: [PATCH 8/8] updated docstring --- src/firebolt/async_db/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 4d06e3d5b2c..be517e57fc6 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -150,7 +150,7 @@ async def connect_inner( api_endpoint(optional): Firebolt API endpoint. Used for authentication. Note: - Either `engine_name` or `engine_url` should be provided, but not both. + Providing both `engine_name` and `engine_url` would result in an error. """ # These parameters are optional in function signature