From 9daee18fa69dcb31bcde23a09fb58f5f7a454960 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Mon, 18 Oct 2021 15:12:54 +0300 Subject: [PATCH 1/9] migrate to python3.7 --- src/firebolt/client/client.py | 5 +++-- src/firebolt/db/_types.py | 6 ++++-- src/firebolt/model/database.py | 13 +----------- src/firebolt/model/engine.py | 6 +----- src/firebolt/service/engine.py | 30 +++++++++++++-------------- src/firebolt/service/instance_type.py | 10 ++++++--- src/firebolt/service/region.py | 13 ++++++++---- 7 files changed, 39 insertions(+), 44 deletions(-) diff --git a/src/firebolt/client/client.py b/src/firebolt/client/client.py index 8e7852af6b2..be0f83de6fd 100644 --- a/src/firebolt/client/client.py +++ b/src/firebolt/client/client.py @@ -1,6 +1,6 @@ import time import typing -from functools import wraps +from functools import lru_cache, wraps from inspect import cleandoc from json import JSONDecodeError from typing import Any, Optional, Tuple, Type @@ -154,6 +154,7 @@ def send(self, *args: Any, **kwargs: Any) -> httpx.Response: resp = super().send(*args, **kwargs) return resp - @cached_property + @property + @lru_cache() def account_id(self) -> str: return self.get(url="/iam/v2/account").json()["account"]["id"] diff --git a/src/firebolt/db/_types.py b/src/firebolt/db/_types.py index 65059e0a6bf..6a4384b82e4 100644 --- a/src/firebolt/db/_types.py +++ b/src/firebolt/db/_types.py @@ -2,6 +2,7 @@ from datetime import date, datetime from enum import Enum +from functools import lru_cache from typing import Union from ciso8601 import parse_datetime @@ -55,7 +56,7 @@ class ARRAY: _prefix = "Array(" def __init__(self, subtype: Union[type, ARRAY]): - assert (subtype in _col_types and subtype is not list) or isinstance( + assert (subtype in ColType.__args__ and subtype is not list) or isinstance( subtype, ARRAY ), f"Invalid array subtype: {str(subtype)}" self.subtype = subtype @@ -103,7 +104,8 @@ class _InternalType(Enum): # Nullable(Nothing) Nothing = "Nothing" - @cached_property + @property + @lru_cache() def python_type(self) -> type: """Convert internal type to python type.""" types = { diff --git a/src/firebolt/model/database.py b/src/firebolt/model/database.py index 69e4023ebe2..5c1898e4115 100644 --- a/src/firebolt/model/database.py +++ b/src/firebolt/model/database.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, List, Optional +from typing import Optional from pydantic import Field, PrivateAttr @@ -22,17 +22,6 @@ class DatabaseKey(FireboltBaseModel): class Database(FireboltBaseModel): - """ - A Firebolt database. - - Databases belong to a region and have a description, - but otherwise are not configurable. - """ - - # internal - _service: DatabaseService = PrivateAttr() - - # required name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$") compute_region_key: RegionKey = Field(alias="compute_region_id") diff --git a/src/firebolt/model/engine.py b/src/firebolt/model/engine.py index 34795ef8982..c9b94bb36a3 100644 --- a/src/firebolt/model/engine.py +++ b/src/firebolt/model/engine.py @@ -4,7 +4,7 @@ import logging import time from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import Optional from pydantic import Field, PrivateAttr @@ -87,10 +87,6 @@ class Engine(FireboltBaseModel): Engines are configured in Settings and in EngineRevisions. """ - # internal - _service: EngineService = PrivateAttr() - - # required name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$") compute_region_key: RegionKey = Field(alias="compute_region_id") settings: EngineSettings diff --git a/src/firebolt/service/engine.py b/src/firebolt/service/engine.py index 5eebb1a814b..39e3d65ffad 100644 --- a/src/firebolt/service/engine.py +++ b/src/firebolt/service/engine.py @@ -1,5 +1,7 @@ -from logging import getLogger -from typing import List, Optional, Union +import json +import logging +import time +from typing import List, Optional from firebolt.common.util import prune_dict from firebolt.model import FireboltBaseModel @@ -147,20 +149,16 @@ def create( ), ) - instance_type_key = self.resource_manager.instance_types.get_by_name( - instance_type_name=spec - ).key - - engine_revision = EngineRevision( - specification=EngineRevisionSpecification( - db_compute_instances_type_key=instance_type_key, - db_compute_instances_count=scale, - db_compute_instances_use_spot=False, - db_version="", - proxy_instances_type_key=instance_type_key, - proxy_instances_count=1, - proxy_version="", - ) + def get_engines_by_ids(self, engine_ids: List[str]) -> List[Engine]: + """Get multiple Engines from Firebolt by their ids.""" + response = self.client.post( + url=f"/core/v1/engines:getByIds", + json={ + "engine_ids": [ + {"account_id": self.account_id, "engine_id": engine_id} + for engine_id in engine_ids + ] + }, ) return self._send_create_engine(engine=engine, engine_revision=engine_revision) diff --git a/src/firebolt/service/instance_type.py b/src/firebolt/service/instance_type.py index 25bd8458214..3ad7ee37d83 100644 --- a/src/firebolt/service/instance_type.py +++ b/src/firebolt/service/instance_type.py @@ -1,3 +1,4 @@ +from functools import lru_cache from typing import Dict, List, NamedTuple, Optional from firebolt.common.utils import cached_property @@ -13,7 +14,8 @@ class InstanceTypeLookup(NamedTuple): class InstanceTypeService(BaseService): - @cached_property + @property + @lru_cache() def instance_types(self) -> List[InstanceType]: """List of instance types available on Firebolt.""" response = self.client.get( @@ -21,12 +23,14 @@ def instance_types(self) -> List[InstanceType]: ) return [InstanceType.parse_obj(i["node"]) for i in response.json()["edges"]] - @cached_property + @property + @lru_cache() def instance_types_by_key(self) -> Dict[InstanceTypeKey, InstanceType]: """Dict of {InstanceTypeKey: InstanceType}""" return {i.key: i for i in self.instance_types} - @cached_property + @property + @lru_cache() def instance_types_by_name(self) -> Dict[InstanceTypeLookup, InstanceType]: """Dict of {InstanceTypeLookup: InstanceType}""" return { diff --git a/src/firebolt/service/region.py b/src/firebolt/service/region.py index 74660869da0..e7ff94ee2d4 100644 --- a/src/firebolt/service/region.py +++ b/src/firebolt/service/region.py @@ -1,3 +1,4 @@ +from functools import lru_cache from typing import Dict, List from firebolt.common.utils import cached_property @@ -20,7 +21,8 @@ def __init__( self.default_region_name = default_region_name super().__init__(resource_manager=resource_manager) - @cached_property + @property + @lru_cache() def regions(self) -> List[Region]: """List of available AWS Regions on Firebolt.""" response = self.client.get( @@ -28,17 +30,20 @@ def regions(self) -> List[Region]: ) return [Region.parse_obj(i["node"]) for i in response.json()["edges"]] - @cached_property + @property + @lru_cache() def regions_by_name(self) -> Dict[str, Region]: """Dict of {RegionLookup: Region}""" return {r.name: r for r in self.regions} - @cached_property + @property + @lru_cache() def regions_by_key(self) -> Dict[RegionKey, Region]: """Dict of {RegionKey: Region}""" return {r.key: r for r in self.regions} - @cached_property + @property + @lru_cache() def default_region(self) -> Region: """Default AWS Region, could be provided from environment.""" From db53926e95758aad0b79270a9b809518dbc15faa Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Mon, 18 Oct 2021 16:36:43 +0300 Subject: [PATCH 2/9] fix mypy issues --- src/firebolt/client/client.py | 5 ++--- src/firebolt/db/_types.py | 6 ++---- src/firebolt/service/instance_type.py | 10 +++------- src/firebolt/service/region.py | 13 ++++--------- 4 files changed, 11 insertions(+), 23 deletions(-) diff --git a/src/firebolt/client/client.py b/src/firebolt/client/client.py index be0f83de6fd..8e7852af6b2 100644 --- a/src/firebolt/client/client.py +++ b/src/firebolt/client/client.py @@ -1,6 +1,6 @@ import time import typing -from functools import lru_cache, wraps +from functools import wraps from inspect import cleandoc from json import JSONDecodeError from typing import Any, Optional, Tuple, Type @@ -154,7 +154,6 @@ def send(self, *args: Any, **kwargs: Any) -> httpx.Response: resp = super().send(*args, **kwargs) return resp - @property - @lru_cache() + @cached_property def account_id(self) -> str: return self.get(url="/iam/v2/account").json()["account"]["id"] diff --git a/src/firebolt/db/_types.py b/src/firebolt/db/_types.py index 6a4384b82e4..65059e0a6bf 100644 --- a/src/firebolt/db/_types.py +++ b/src/firebolt/db/_types.py @@ -2,7 +2,6 @@ from datetime import date, datetime from enum import Enum -from functools import lru_cache from typing import Union from ciso8601 import parse_datetime @@ -56,7 +55,7 @@ class ARRAY: _prefix = "Array(" def __init__(self, subtype: Union[type, ARRAY]): - assert (subtype in ColType.__args__ and subtype is not list) or isinstance( + assert (subtype in _col_types and subtype is not list) or isinstance( subtype, ARRAY ), f"Invalid array subtype: {str(subtype)}" self.subtype = subtype @@ -104,8 +103,7 @@ class _InternalType(Enum): # Nullable(Nothing) Nothing = "Nothing" - @property - @lru_cache() + @cached_property def python_type(self) -> type: """Convert internal type to python type.""" types = { diff --git a/src/firebolt/service/instance_type.py b/src/firebolt/service/instance_type.py index 3ad7ee37d83..25bd8458214 100644 --- a/src/firebolt/service/instance_type.py +++ b/src/firebolt/service/instance_type.py @@ -1,4 +1,3 @@ -from functools import lru_cache from typing import Dict, List, NamedTuple, Optional from firebolt.common.utils import cached_property @@ -14,8 +13,7 @@ class InstanceTypeLookup(NamedTuple): class InstanceTypeService(BaseService): - @property - @lru_cache() + @cached_property def instance_types(self) -> List[InstanceType]: """List of instance types available on Firebolt.""" response = self.client.get( @@ -23,14 +21,12 @@ def instance_types(self) -> List[InstanceType]: ) return [InstanceType.parse_obj(i["node"]) for i in response.json()["edges"]] - @property - @lru_cache() + @cached_property def instance_types_by_key(self) -> Dict[InstanceTypeKey, InstanceType]: """Dict of {InstanceTypeKey: InstanceType}""" return {i.key: i for i in self.instance_types} - @property - @lru_cache() + @cached_property def instance_types_by_name(self) -> Dict[InstanceTypeLookup, InstanceType]: """Dict of {InstanceTypeLookup: InstanceType}""" return { diff --git a/src/firebolt/service/region.py b/src/firebolt/service/region.py index e7ff94ee2d4..74660869da0 100644 --- a/src/firebolt/service/region.py +++ b/src/firebolt/service/region.py @@ -1,4 +1,3 @@ -from functools import lru_cache from typing import Dict, List from firebolt.common.utils import cached_property @@ -21,8 +20,7 @@ def __init__( self.default_region_name = default_region_name super().__init__(resource_manager=resource_manager) - @property - @lru_cache() + @cached_property def regions(self) -> List[Region]: """List of available AWS Regions on Firebolt.""" response = self.client.get( @@ -30,20 +28,17 @@ def regions(self) -> List[Region]: ) return [Region.parse_obj(i["node"]) for i in response.json()["edges"]] - @property - @lru_cache() + @cached_property def regions_by_name(self) -> Dict[str, Region]: """Dict of {RegionLookup: Region}""" return {r.name: r for r in self.regions} - @property - @lru_cache() + @cached_property def regions_by_key(self) -> Dict[RegionKey, Region]: """Dict of {RegionKey: Region}""" return {r.key: r for r in self.regions} - @property - @lru_cache() + @cached_property def default_region(self) -> Region: """Default AWS Region, could be provided from environment.""" From cfee09a2525c9896b1c2f363e78ea774cfb3ae17 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 20 Oct 2021 11:37:26 +0300 Subject: [PATCH 3/9] implement connect function --- src/firebolt/db/__init__.py | 2 +- src/firebolt/db/connection.py | 66 +++++++++++++++++++++++++++++++---- 2 files changed, 60 insertions(+), 8 deletions(-) diff --git a/src/firebolt/db/__init__.py b/src/firebolt/db/__init__.py index 9a20b4636b8..7c7af17cbdb 100644 --- a/src/firebolt/db/__init__.py +++ b/src/firebolt/db/__init__.py @@ -13,7 +13,7 @@ Timestamp, TimestampFromTicks, ) -from firebolt.db.connection import Connection +from firebolt.db.connection import Connection, connect from firebolt.db.cursor import Cursor apilevel = "2.0" diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 66b91bd6439..89c29f32218 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -2,31 +2,86 @@ from inspect import cleandoc from types import TracebackType -from typing import List +from typing import List, Optional from httpx import Timeout from readerwriterlock.rwlock import RWLockWrite from firebolt.client import DEFAULT_API_URL, Client -from firebolt.common.exception import ConnectionClosedError +from firebolt.common.exception import ConnectionClosedError, InterfaceError +from firebolt.common.settings import Settings from firebolt.db.cursor import Cursor +from firebolt.service.manager import ResourceManager DEFAULT_TIMEOUT_SECONDS: int = 5 +def connect( + database: str = None, + username: str = None, + password: str = None, + engine_name: Optional[str] = None, + engine_url: Optional[str] = None, + api_endpoint: str = DEFAULT_API_URL, +) -> Connection: + cleandoc( + """ + Connect to Firebolt database. + + Connection parameters: + database - name of the database to connect + username - user name to use for authentication + password - password to use for authentication + engine_name - name of the engine to connect to + engine_url - engine endpoint to use + note: either engine_name or engine_url should be provided, but not both + """ + ) + if engine_name and engine_url: + raise InterfaceError( + "Both engine_name and engine_url are provided. Provide only one to connect" + ) + if not engine_name and not engine_url: + raise InterfaceError( + "Neither engine_name nor engine_url are provided. Provide one to connect" + ) + # This parameters are optional in function signature, but are required to connect. + # It's recomended to make them kwargs by PEP 249 + assert database, "database_name required" + assert username, "username required" + assert password, "password required" + + if engine_name is not None: + rm = ResourceManager( + Settings(user=username, password=password, server=api_endpoint) + ) + endpoint = rm.engines.get_engine_by_name(engine_name).endpoint + if endpoint is None: + raise InterfaceError("unable to retrieve engine endpoint") + else: + engine_url = endpoint + + assert engine_url is not None + engine_url = ( + engine_url if engine_url.startswith("http") else f"https://{engine_url}" + ) + return Connection(engine_url, database, username, password, api_endpoint) + + class Connection: cleandoc( """ Firebolt database connection class. Implements PEP-249. Parameters: + engine_url - Firebolt database engine REST API url + database - Firebolt database name username - Firebolt account username password - Firebolt account password - engine_url - Firebolt database engine REST API url api_endpoint(optional) - Firebolt API endpoint. Used for authentication Methods: - cursor - created new Cursor object + cursor - create new Cursor object close - close the Connection and all it's cursors Firebolt currenly doesn't support transactions so commit and rollback methods @@ -43,9 +98,6 @@ def __init__( password: str, api_endpoint: str = DEFAULT_API_URL, ): - engine_url = ( - engine_url if engine_url.startswith("http") else f"https://{engine_url}" - ) self._client = Client( auth=(username, password), base_url=engine_url, From c72cf966963ce05590bcfaca251d33539151bde1 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 20 Oct 2021 15:21:35 +0300 Subject: [PATCH 4/9] improve examples.ipynb --- src/firebolt/db/connection.py | 4 +++- src/firebolt/db/examples.ipynb | 19 +++++++++++++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 89c29f32218..13920c3e395 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -53,7 +53,9 @@ def connect( if engine_name is not None: rm = ResourceManager( - Settings(user=username, password=password, server=api_endpoint) + Settings( + user=username, password=password, server=api_endpoint, default_region="" + ) ) endpoint = rm.engines.get_engine_by_name(engine_name).endpoint if endpoint is None: diff --git a/src/firebolt/db/examples.ipynb b/src/firebolt/db/examples.ipynb index 301283725e4..22ef690bec6 100644 --- a/src/firebolt/db/examples.ipynb +++ b/src/firebolt/db/examples.ipynb @@ -15,7 +15,7 @@ "metadata": {}, "outputs": [], "source": [ - "from firebolt.db import Connection\n", + "from firebolt.db import connect\n", "from firebolt.client import DEFAULT_API_URL" ] }, @@ -34,8 +34,14 @@ "metadata": {}, "outputs": [], "source": [ - "database_name = \"\"\n", + "# Only one of these two parameters should be specified\n", "engine_url = \"\"\n", + "engine_name = \"\"\n", + "assert bool(engine_url) != bool(\n", + " engine_name\n", + "), \"Specify only one of engine_name and engine_url\"\n", + "\n", + "database_name = \"\"\n", "username = \"\"\n", "password = \"\"\n", "api_endpoint = DEFAULT_API_URL" @@ -57,8 +63,13 @@ "outputs": [], "source": [ "# create a connection based on provided credentials\n", - "connection = Connection(\n", - " engine_url, database_name, username, password, api_endpoint=api_endpoint\n", + "connection = connect(\n", + " engine_url=engine_url,\n", + " engine_name=engine_name,\n", + " database=database_name,\n", + " username=username,\n", + " password=password,\n", + " api_endpoint=api_endpoint,\n", ")\n", "\n", "# create a cursor for connection\n", From 23051df456882a8c7d33992fe03e5f7c23dca3ac Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 20 Oct 2021 17:09:12 +0300 Subject: [PATCH 5/9] fix tests, implement unit tests for connect --- src/firebolt/db/connection.py | 2 +- tests/integration/dbapi/conftest.py | 10 ++- tests/integration/dbapi/test_errors.py | 26 +++++-- tests/unit/conftest.py | 84 ++++++++++++++++++++++ tests/unit/db/conftest.py | 10 ++- tests/unit/db/test_connection.py | 98 ++++++++++++++++++++++++-- 6 files changed, 210 insertions(+), 20 deletions(-) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 13920c3e395..4102ec09760 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -47,7 +47,7 @@ def connect( ) # This parameters are optional in function signature, but are required to connect. # It's recomended to make them kwargs by PEP 249 - assert database, "database_name required" + assert database, "database required" assert username, "username required" assert password, "password required" diff --git a/tests/integration/dbapi/conftest.py b/tests/integration/dbapi/conftest.py index 5e35b317606..1f336e0ce66 100644 --- a/tests/integration/dbapi/conftest.py +++ b/tests/integration/dbapi/conftest.py @@ -5,7 +5,7 @@ from pytest import fixture -from firebolt.db import ARRAY, Connection +from firebolt.db import ARRAY, Connection, connect from firebolt.db._types import ColType from firebolt.db.cursor import Column @@ -53,8 +53,12 @@ def api_endpoint() -> str: def connection( engine_url: str, database_name: str, username: str, password: str, api_endpoint: str ) -> Connection: - return Connection( - engine_url, database_name, username, password, api_endpoint=api_endpoint + return connect( + engine_url=engine_url, + database=database_name, + username=username, + password=password, + api_endpoint=api_endpoint, ) diff --git a/tests/integration/dbapi/test_errors.py b/tests/integration/dbapi/test_errors.py index 39bce4f7f27..2fb9d3e3f5f 100644 --- a/tests/integration/dbapi/test_errors.py +++ b/tests/integration/dbapi/test_errors.py @@ -6,15 +6,19 @@ OperationalError, ProgrammingError, ) -from firebolt.db import Connection +from firebolt.db import Connection, connect def test_invalid_credentials( engine_url: str, database_name: str, username: str, password: str, api_endpoint: str ) -> None: """Connection properly reacts to invalid credentials error""" - connection = Connection( - engine_url, database_name, username + "_", password + "_", api_endpoint + connection = connect( + engine_url=engine_url, + database=database_name, + username=username + "_", + password=password + "_", + api_endpoint=api_endpoint, ) with raises(AuthenticationError) as exc_info: connection.cursor().execute("show tables") @@ -28,8 +32,12 @@ def test_engine_not_exists( engine_url: str, database_name: str, username: str, password: str, api_endpoint: str ) -> None: """Connection properly reacts to invalid engine url error""" - connection = Connection( - engine_url + "_", database_name, username, password, api_endpoint + connection = connect( + engine_url=engine_url + "_", + database=database_name, + username=username, + password=password, + api_endpoint=api_endpoint, ) with raises(ConnectError): connection.cursor().execute("show tables") @@ -40,7 +48,13 @@ def test_database_not_exists( ) -> None: """Connection properly reacts to invalid database error""" new_db_name = database_name + "_" - connection = Connection(engine_url, new_db_name, username, password, api_endpoint) + connection = connect( + engine_url=engine_url, + database=new_db_name, + username=username, + password=password, + api_endpoint=api_endpoint, + ) with raises(ProgrammingError) as exc_info: connection.cursor().execute("show tables") diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 6c7af704fb1..f24c11f033d 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -150,3 +150,87 @@ def provider_url(settings: Settings) -> str: @pytest.fixture def db_name() -> str: return "database" + + +@pytest.fixture +def account_id_url(settings: Settings) -> str: + return f"https://{settings.server}/iam/v2/account" + + +@pytest.fixture +def account_id_callback(account_id: str, account_id_url: str) -> Callable: + def do_mock( + request: httpx.Request = None, + **kwargs, + ) -> Response: + assert request.url == account_id_url + return to_response( + status_code=httpx.codes.OK, json={"account": {"id": account_id}} + ) + + return do_mock + + +@pytest.fixture +def engine_id() -> str: + return "engine_id" + + +@pytest.fixture +def get_engine_url(settings: Settings, account_id: str, engine_id: str) -> str: + return ( + f"https://{settings.server}/core/v1/accounts/{account_id}/engines/{engine_id}" + ) + + +@pytest.fixture +def get_engine_callback( + get_engine_url: str, engine_id: str, settings: Settings +) -> Callable: + def do_mock( + request: httpx.Request = None, + **kwargs, + ) -> Response: + assert request.url == get_engine_url + return to_response( + status_code=httpx.codes.OK, + json={ + "engine": { + "name": "name", + "compute_region_id": { + "provider_id": "provider", + "region_id": "region", + }, + "settings": { + "preset": "", + "auto_stop_delay_duration": "", + "minimum_logging_level": "", + "is_read_only": False, + "warm_up": "", + }, + "endpoint": f"https://{settings.server}", + } + }, + ) + + return do_mock + + +@pytest.fixture +def get_providers_url(settings: Settings, account_id: str, engine_id: str) -> str: + return f"https://{settings.server}/compute/v1/providers" + + +@pytest.fixture +def get_providers_callback(get_providers_url: str, provider: Provider) -> Callable: + def do_mock( + request: httpx.Request = None, + **kwargs, + ) -> Response: + assert request.url == get_providers_url + return to_response( + status_code=httpx.codes.OK, + json=list_to_paginated_response([provider]), + ) + + return do_mock diff --git a/tests/unit/db/conftest.py b/tests/unit/db/conftest.py index 870573eef02..fdff4b6131e 100644 --- a/tests/unit/db/conftest.py +++ b/tests/unit/db/conftest.py @@ -6,7 +6,7 @@ from pytest_httpx import to_response from firebolt.common.settings import Settings -from firebolt.db import ARRAY, Connection, Cursor +from firebolt.db import ARRAY, Connection, Cursor, connect from firebolt.db.cursor import JSON_OUTPUT_FORMAT, ColType, Column QUERY_ROW_COUNT: int = 10 @@ -139,8 +139,12 @@ def query_url(settings: Settings, db_name: str) -> str: @fixture def connection(settings: Settings, db_name: str) -> Connection: - return Connection( - f"https://{settings.server}", db_name, "u", "p", api_endpoint=settings.server + return connect( + engine_url=settings.server, + database=db_name, + username="u", + password="p", + api_endpoint=settings.server, ) diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index 8420d6060a3..ff2cad0ac97 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -1,11 +1,12 @@ from typing import Callable, List +from httpx import codes from pytest import raises from pytest_httpx import HTTPXMock -from firebolt.common.exception import ConnectionClosedError +from firebolt.common.exception import ConnectionClosedError, InterfaceError from firebolt.common.settings import Settings -from firebolt.db import Connection +from firebolt.db import Connection, connect from firebolt.db._types import ColType @@ -52,11 +53,11 @@ def test_cursor_initialized( httpx_mock.add_callback(query_callback, url=query_url) for url in (settings.server, f"https://{settings.server}"): - connection = Connection( - url, - db_name, - "u", - "p", + connection = connect( + engine_url=url, + database=db_name, + username="u", + password="p", api_endpoint=settings.server, ) @@ -70,3 +71,86 @@ def test_cursor_initialized( assert ( cursor not in connection._cursors ), "Cursor wasn't removed from connection after close" + + +def test_connect_empty_parameters(): + params = ("database", "username", "password") + kwargs = {"engine_url": "engine_url", **{p: p for p in params}} + + for param in params: + with raises(AssertionError) as exc_info: + kwargs = { + "engine_url": "engine_url", + **{p: p for p in params if p != param}, + } + connect(**kwargs) + assert str(exc_info.value) == f"{param} required" + + +def test_connect_engine_name( + 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, + get_providers_url: str, + get_providers_callback: Callable, + python_query_data: List[List[ColType]], + account_id: str, +): + """connect properly handles engine_name""" + + with raises(InterfaceError) as exc_info: + connect( + engine_url="engine_url", + engine_name="engine_name", + database="db", + username="username", + password="password", + ) + assert str(exc_info.value).startswith( + "Both engine_name and engine_url are provided" + ) + + with raises(InterfaceError) as exc_info: + connect( + database="db", + username="username", + password="password", + ) + assert str(exc_info.value).startswith( + "Neither engine_name nor engine_url are provided" + ) + + 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(get_engine_callback, url=get_engine_url) + httpx_mock.add_callback(get_providers_callback, url=get_providers_url) + + engine_name = settings.server.split(".")[0] + + # Mock engine id lookup by name + httpx_mock.add_response( + url=f"https://{settings.server}/core/v1/account/engines:getIdByName?" + f"engine_name={engine_name}", + status_code=codes.OK, + json={"engine_id": {"engine_id": engine_id}}, + ) + + cursor = connect( + engine_name=engine_name, + database=db_name, + username="u", + password="p", + api_endpoint=settings.server, + ).cursor() + + assert cursor.execute("select*") == len(python_query_data) From 1c7f2bc890dbca328caec009cb3c610f806fc47e Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 20 Oct 2021 17:24:16 +0300 Subject: [PATCH 6/9] fix merge issues --- src/firebolt/db/connection.py | 2 +- src/firebolt/model/database.py | 13 ++++++++++++- src/firebolt/model/engine.py | 6 +++++- src/firebolt/service/engine.py | 30 ++++++++++++++++-------------- tests/unit/conftest.py | 2 +- 5 files changed, 35 insertions(+), 18 deletions(-) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 4102ec09760..0d16a85ec89 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -57,7 +57,7 @@ def connect( user=username, password=password, server=api_endpoint, default_region="" ) ) - endpoint = rm.engines.get_engine_by_name(engine_name).endpoint + endpoint = rm.engines.get_by_name(engine_name).endpoint if endpoint is None: raise InterfaceError("unable to retrieve engine endpoint") else: diff --git a/src/firebolt/model/database.py b/src/firebolt/model/database.py index 5c1898e4115..69e4023ebe2 100644 --- a/src/firebolt/model/database.py +++ b/src/firebolt/model/database.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import datetime -from typing import Optional +from typing import TYPE_CHECKING, Any, List, Optional from pydantic import Field, PrivateAttr @@ -22,6 +22,17 @@ class DatabaseKey(FireboltBaseModel): class Database(FireboltBaseModel): + """ + A Firebolt database. + + Databases belong to a region and have a description, + but otherwise are not configurable. + """ + + # internal + _service: DatabaseService = PrivateAttr() + + # required name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$") compute_region_key: RegionKey = Field(alias="compute_region_id") diff --git a/src/firebolt/model/engine.py b/src/firebolt/model/engine.py index c9b94bb36a3..34795ef8982 100644 --- a/src/firebolt/model/engine.py +++ b/src/firebolt/model/engine.py @@ -4,7 +4,7 @@ import logging import time from datetime import datetime -from typing import Optional +from typing import TYPE_CHECKING, Any, Callable, Optional from pydantic import Field, PrivateAttr @@ -87,6 +87,10 @@ class Engine(FireboltBaseModel): Engines are configured in Settings and in EngineRevisions. """ + # internal + _service: EngineService = PrivateAttr() + + # required name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$") compute_region_key: RegionKey = Field(alias="compute_region_id") settings: EngineSettings diff --git a/src/firebolt/service/engine.py b/src/firebolt/service/engine.py index 39e3d65ffad..5eebb1a814b 100644 --- a/src/firebolt/service/engine.py +++ b/src/firebolt/service/engine.py @@ -1,7 +1,5 @@ -import json -import logging -import time -from typing import List, Optional +from logging import getLogger +from typing import List, Optional, Union from firebolt.common.util import prune_dict from firebolt.model import FireboltBaseModel @@ -149,16 +147,20 @@ def create( ), ) - def get_engines_by_ids(self, engine_ids: List[str]) -> List[Engine]: - """Get multiple Engines from Firebolt by their ids.""" - response = self.client.post( - url=f"/core/v1/engines:getByIds", - json={ - "engine_ids": [ - {"account_id": self.account_id, "engine_id": engine_id} - for engine_id in engine_ids - ] - }, + instance_type_key = self.resource_manager.instance_types.get_by_name( + instance_type_name=spec + ).key + + engine_revision = EngineRevision( + specification=EngineRevisionSpecification( + db_compute_instances_type_key=instance_type_key, + db_compute_instances_count=scale, + db_compute_instances_use_spot=False, + db_version="", + proxy_instances_type_key=instance_type_key, + proxy_instances_count=1, + proxy_version="", + ) ) return self._send_create_engine(engine=engine, engine_revision=engine_revision) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index f24c11f033d..61238d2cba2 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -203,7 +203,7 @@ def do_mock( }, "settings": { "preset": "", - "auto_stop_delay_duration": "", + "auto_stop_delay_duration": "1s", "minimum_logging_level": "", "is_read_only": False, "warm_up": "", From c6d9298ff3e34a98b7b8f0b4c21014d87cb185c2 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Wed, 20 Oct 2021 18:05:05 +0300 Subject: [PATCH 7/9] add integration tests for connect --- tests/integration/dbapi/conftest.py | 23 +++++++++++++++++++++++ tests/integration/dbapi/test_errors.py | 21 ++++++++++++++++++++- tests/integration/dbapi/test_queries.py | 16 ++++++++++++++++ 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/integration/dbapi/conftest.py b/tests/integration/dbapi/conftest.py index 1f336e0ce66..362b957d704 100644 --- a/tests/integration/dbapi/conftest.py +++ b/tests/integration/dbapi/conftest.py @@ -12,6 +12,7 @@ LOGGER = getLogger(__name__) ENGINE_URL_ENV = "ENGINE_URL" +ENGINE_NAME_ENV = "ENGINE_NAME" DATABASE_NAME_ENV = "DATABASE_NAME" USERNAME_ENV = "USERNAME" PASSWORD_ENV = "PASSWORD" @@ -29,6 +30,11 @@ def engine_url() -> str: return must_env(ENGINE_URL_ENV) +@fixture(scope="session") +def engine_name() -> str: + return must_env(ENGINE_NAME_ENV) + + @fixture(scope="session") def database_name() -> str: return must_env(DATABASE_NAME_ENV) @@ -62,6 +68,23 @@ def connection( ) +@fixture +def connection_engine_name( + engine_name: str, + database_name: str, + username: str, + password: str, + api_endpoint: str, +) -> Connection: + return connect( + engine_name=engine_name, + database=database_name, + username=username, + password=password, + api_endpoint=api_endpoint, + ) + + @fixture def all_types_query() -> str: return ( diff --git a/tests/integration/dbapi/test_errors.py b/tests/integration/dbapi/test_errors.py index 2fb9d3e3f5f..b2ae3d7b2c6 100644 --- a/tests/integration/dbapi/test_errors.py +++ b/tests/integration/dbapi/test_errors.py @@ -28,7 +28,7 @@ def test_invalid_credentials( ), "Invalid authentication error message" -def test_engine_not_exists( +def test_engine_url_not_exists( engine_url: str, database_name: str, username: str, password: str, api_endpoint: str ) -> None: """Connection properly reacts to invalid engine url error""" @@ -43,6 +43,25 @@ def test_engine_not_exists( connection.cursor().execute("show tables") +def test_engine_name_not_exists( + engine_name: str, + database_name: str, + username: str, + password: str, + api_endpoint: str, +) -> None: + """Connection properly reacts to invalid engine url error""" + connection = connect( + engine_url=engine_name + "_________", + database=database_name, + username=username, + password=password, + api_endpoint=api_endpoint, + ) + with raises(ConnectError): + connection.cursor().execute("show tables") + + def test_database_not_exists( engine_url: str, database_name: str, username: str, password: str, api_endpoint: str ) -> None: diff --git a/tests/integration/dbapi/test_queries.py b/tests/integration/dbapi/test_queries.py index 2eb8f6d1e6f..2bed6d5b407 100644 --- a/tests/integration/dbapi/test_queries.py +++ b/tests/integration/dbapi/test_queries.py @@ -14,6 +14,22 @@ def assert_deep_eq(got: Any, expected: Any, msg: str) -> bool: ), f"{msg}: {got}(got) != {expected}(expected)" +def test_connect_engine_name( + connection_engine_name: 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.""" + # Basically run selects on connection with engine name + test_select( + connection_engine_name, + all_types_query, + all_types_query_description, + all_types_query_response, + ) + + def test_select( connection: Connection, all_types_query: str, From 0fc6da35e13ef81c22e377cec500e0e1e83da341 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 21 Oct 2021 11:21:35 +0300 Subject: [PATCH 8/9] address comments --- src/firebolt/db/connection.py | 16 ++++++++++------ tests/integration/dbapi/test_errors.py | 3 ++- tests/integration/dbapi/test_queries.py | 1 - tests/unit/db/test_connection.py | 4 ++-- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 0d16a85ec89..cb0a4bda98a 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -39,17 +39,21 @@ def connect( ) if engine_name and engine_url: raise InterfaceError( - "Both engine_name and engine_url are provided. Provide only one to connect" + "Both engine_name and engine_url are provided. Provide only one to connect." ) if not engine_name and not engine_url: raise InterfaceError( - "Neither engine_name nor engine_url are provided. Provide one to connect" + "Neither engine_name nor engine_url are provided. Provide one to connect." ) # This parameters are optional in function signature, but are required to connect. # It's recomended to make them kwargs by PEP 249 - assert database, "database required" - assert username, "username required" - assert password, "password required" + for param, name in ( + (database, "database"), + (username, "username"), + (password, "password"), + ): + if not param: + raise InterfaceError(f"{name} is required to connect.") if engine_name is not None: rm = ResourceManager( @@ -59,7 +63,7 @@ def connect( ) endpoint = rm.engines.get_by_name(engine_name).endpoint if endpoint is None: - raise InterfaceError("unable to retrieve engine endpoint") + raise InterfaceError("unable to retrieve engine endpoint.") else: engine_url = endpoint diff --git a/tests/integration/dbapi/test_errors.py b/tests/integration/dbapi/test_errors.py index b2ae3d7b2c6..f50cb9114a4 100644 --- a/tests/integration/dbapi/test_errors.py +++ b/tests/integration/dbapi/test_errors.py @@ -50,7 +50,7 @@ def test_engine_name_not_exists( password: str, api_endpoint: str, ) -> None: - """Connection properly reacts to invalid engine url error""" + """Connection properly reacts to invalid engine name error""" connection = connect( engine_url=engine_name + "_________", database=database_name, @@ -83,6 +83,7 @@ def test_database_not_exists( def test_sql_error(connection: Connection) -> None: + """Connection properly reacts to sql execution error""" with connection.cursor() as c: with raises(OperationalError) as exc_info: c.execute("select ]") diff --git a/tests/integration/dbapi/test_queries.py b/tests/integration/dbapi/test_queries.py index 2bed6d5b407..4eaa356bacb 100644 --- a/tests/integration/dbapi/test_queries.py +++ b/tests/integration/dbapi/test_queries.py @@ -21,7 +21,6 @@ def test_connect_engine_name( all_types_query_response: List[ColType], ) -> None: """Connecting with engine name is handled properly.""" - # Basically run selects on connection with engine name test_select( connection_engine_name, all_types_query, diff --git a/tests/unit/db/test_connection.py b/tests/unit/db/test_connection.py index ff2cad0ac97..f9dc409d292 100644 --- a/tests/unit/db/test_connection.py +++ b/tests/unit/db/test_connection.py @@ -78,13 +78,13 @@ def test_connect_empty_parameters(): kwargs = {"engine_url": "engine_url", **{p: p for p in params}} for param in params: - with raises(AssertionError) as exc_info: + with raises(InterfaceError) as exc_info: kwargs = { "engine_url": "engine_url", **{p: p for p in params if p != param}, } connect(**kwargs) - assert str(exc_info.value) == f"{param} required" + assert str(exc_info.value) == f"{param} is required to connect." def test_connect_engine_name( From 34ae24e1e0aa5fa73c03957c644ddcd80338868c Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 21 Oct 2021 11:24:51 +0300 Subject: [PATCH 9/9] fix mypy --- src/firebolt/db/connection.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index cb0a4bda98a..8b1fbf88abd 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -67,7 +67,12 @@ def connect( else: engine_url = endpoint + # Mypy checks, this should never happen assert engine_url is not None + assert database is not None + assert username is not None + assert password is not None + engine_url = ( engine_url if engine_url.startswith("http") else f"https://{engine_url}" )