From d79c47a320a4d9e70abdb0a9674764c66987f811 Mon Sep 17 00:00:00 2001 From: Petro Tiurin Date: Wed, 9 Feb 2022 17:18:22 +0000 Subject: [PATCH] feat: Internal args for engine --- src/firebolt/model/engine.py | 11 +++- src/firebolt/service/engine.py | 20 ++---- tests/unit/service/conftest.py | 25 ++++++++ tests/unit/service/test_engine.py | 101 +++++++++++++++++++++++++++++- 4 files changed, 141 insertions(+), 16 deletions(-) diff --git a/src/firebolt/model/engine.py b/src/firebolt/model/engine.py index e1c64cca566..67b6e17282a 100644 --- a/src/firebolt/model/engine.py +++ b/src/firebolt/model/engine.py @@ -69,6 +69,7 @@ def default( engine_type: EngineType = EngineType.GENERAL_PURPOSE, auto_stop_delay_duration: str = "1200s", warm_up: WarmupMethod = WarmupMethod.PRELOAD_INDEXES, + minimum_logging_level: str = "ENGINE_SETTINGS_LOGGING_LEVEL_INFO", ) -> EngineSettings: if engine_type == EngineType.GENERAL_PURPOSE: preset = engine_type.GENERAL_PURPOSE.api_settings_preset_name # type: ignore # noqa: E501 @@ -80,7 +81,7 @@ def default( return cls( preset=preset, auto_stop_delay_duration=auto_stop_delay_duration, - minimum_logging_level="ENGINE_SETTINGS_LOGGING_LEVEL_INFO", + minimum_logging_level=minimum_logging_level, is_read_only=is_read_only, warm_up=warm_up.api_name, ) @@ -459,3 +460,11 @@ def _send_engine_request(self, url: str) -> Engine: return Engine.parse_obj_with_service( obj=response.json()["engine"], engine_service=self._service ) + + +class _EngineCreateRequest(FireboltBaseModel): + """Helper model for sending Engine create requests.""" + + account_id: str + engine: Engine + engine_revision: Optional[EngineRevision] diff --git a/src/firebolt/service/engine.py b/src/firebolt/service/engine.py index 843a03e3b9c..c18529148f8 100644 --- a/src/firebolt/service/engine.py +++ b/src/firebolt/service/engine.py @@ -1,5 +1,5 @@ from logging import getLogger -from typing import List, Optional, Union +from typing import Any, Dict, List, Optional, Union from firebolt.common.exception import FireboltError from firebolt.common.urls import ( @@ -9,8 +9,7 @@ ENGINES_BY_IDS_URL, ) from firebolt.common.util import prune_dict -from firebolt.model import FireboltBaseModel -from firebolt.model.engine import Engine, EngineSettings +from firebolt.model.engine import Engine, EngineSettings, _EngineCreateRequest from firebolt.model.engine_revision import ( EngineRevision, EngineRevisionSpecification, @@ -116,6 +115,8 @@ def create( auto_stop: int = 20, warmup: Union[str, WarmupMethod] = WarmupMethod.PRELOAD_INDEXES, description: str = "", + engine_settings_kwargs: Dict[str, Any] = {}, + revision_spec_kwargs: Dict[str, Any] = {}, ) -> Engine: """ Create a new Engine. @@ -165,6 +166,7 @@ def create( engine_type=engine_type, auto_stop_delay_duration=f"{auto_stop * 60}s", warm_up=warmup, + **engine_settings_kwargs, ), ) @@ -186,11 +188,8 @@ def create( 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="", + **revision_spec_kwargs, ) ) @@ -210,13 +209,6 @@ def _send_create_engine( The newly created engine. """ - class _EngineCreateRequest(FireboltBaseModel): - """Helper model for sending Engine create requests.""" - - account_id: str - engine: Engine - engine_revision: Optional[EngineRevision] - response = self.client.post( url=ACCOUNT_ENGINES_URL.format(account_id=self.account_id), headers={"Content-type": "application/json"}, diff --git a/tests/unit/service/conftest.py b/tests/unit/service/conftest.py index eaf885c05c1..667f5e88e87 100644 --- a/tests/unit/service/conftest.py +++ b/tests/unit/service/conftest.py @@ -22,6 +22,10 @@ from firebolt.model.binding import Binding, BindingKey from firebolt.model.database import Database, DatabaseKey from firebolt.model.engine import Engine, EngineKey, EngineSettings +from firebolt.model.engine_revision import ( + EngineRevision, + EngineRevisionSpecification, +) from firebolt.model.instance_type import InstanceType, InstanceTypeKey from firebolt.model.region import Region from tests.unit.util import list_to_paginated_response @@ -32,6 +36,11 @@ def engine_name() -> str: return "my_engine" +@pytest.fixture +def engine_scale() -> int: + return 2 + + @pytest.fixture def engine_settings() -> EngineSettings: return EngineSettings.default() @@ -48,6 +57,22 @@ def mock_engine(engine_name, region_1, engine_settings, account_id, settings) -> ) +@pytest.fixture +def mock_engine_revision_spec( + instance_type_2, engine_scale +) -> EngineRevisionSpecification: + return EngineRevisionSpecification( + db_compute_instances_type_key=instance_type_2.key, + db_compute_instances_count=engine_scale, + proxy_instances_type_key=instance_type_2.key, + ) + + +@pytest.fixture +def mock_engine_revision(mock_engine_revision_spec) -> EngineRevision: + return EngineRevision(specification=mock_engine_revision_spec) + + @pytest.fixture def instance_type_1(provider, region_1) -> InstanceType: return InstanceType( diff --git a/tests/unit/service/test_engine.py b/tests/unit/service/test_engine.py index bc3c2d238cc..a700c28a759 100644 --- a/tests/unit/service/test_engine.py +++ b/tests/unit/service/test_engine.py @@ -1,11 +1,13 @@ from typing import Callable, List +from pydantic import ValidationError from pytest import raises from pytest_httpx import HTTPXMock from firebolt.common import Settings from firebolt.common.exception import FireboltError, NoAttachedDatabaseError -from firebolt.model.engine import Engine +from firebolt.model.engine import Engine, _EngineCreateRequest +from firebolt.model.engine_revision import EngineRevision from firebolt.model.instance_type import InstanceType from firebolt.model.region import Region from firebolt.service.manager import ResourceManager @@ -47,6 +49,103 @@ def test_engine_create( assert engine.name == engine_name +def test_engine_create_with_kwargs( + httpx_mock: HTTPXMock, + auth_callback: Callable, + auth_url: str, + provider_callback: Callable, + provider_url: str, + instance_type_region_1_callback: Callable, + instance_type_region_1_url: str, + region_callback: Callable, + region_url: str, + settings: Settings, + mock_engine: Engine, + engine_name: str, + account_id_callback: Callable, + account_id_url: str, + engine_callback: Callable, + engine_url: str, + account_id: str, + mock_engine_revision: EngineRevision, +): + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(provider_callback, url=provider_url) + httpx_mock.add_callback( + instance_type_region_1_callback, url=instance_type_region_1_url + ) + httpx_mock.add_callback(account_id_callback, url=account_id_url) + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(region_callback, url=region_url) + # Setting to manager.engines.create defaults + mock_engine.key = None + mock_engine.description = "" + mock_engine.endpoint = None + # Testing kwargs + mock_engine.settings.minimum_logging_level = "ENGINE_SETTINGS_LOGGING_LEVEL_DEBUG" + mock_engine_revision.specification.proxy_version = "0.2.3" + engine_content = _EngineCreateRequest( + account_id=account_id, engine=mock_engine, engine_revision=mock_engine_revision + ) + httpx_mock.add_callback( + engine_callback, + url=engine_url, + method="POST", + match_content=engine_content.json(by_alias=True).encode("ascii"), + ) + + manager = ResourceManager(settings=settings) + engine_settings_kwargs = { + "minimum_logging_level": "ENGINE_SETTINGS_LOGGING_LEVEL_DEBUG" + } + revision_spec_kwargs = {"proxy_version": "0.2.3"} + engine = manager.engines.create( + name=engine_name, + engine_settings_kwargs=engine_settings_kwargs, + revision_spec_kwargs=revision_spec_kwargs, + ) + + assert engine.name == engine_name + + +def test_engine_create_with_kwargs_fail( + httpx_mock: HTTPXMock, + auth_callback: Callable, + auth_url: str, + provider_callback: Callable, + provider_url: str, + instance_type_region_1_callback: Callable, + instance_type_region_1_url: str, + region_callback: Callable, + region_url: str, + settings: Settings, + engine_name: str, + account_id_callback: Callable, + account_id_url: str, +): + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(provider_callback, url=provider_url) + httpx_mock.add_callback( + instance_type_region_1_callback, url=instance_type_region_1_url + ) + httpx_mock.add_callback(account_id_callback, url=account_id_url) + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(region_callback, url=region_url) + + manager = ResourceManager(settings=settings) + revision_spec_kwargs = {"incorrect_kwarg": "val"} + with raises(ValidationError): + manager.engines.create( + name=engine_name, revision_spec_kwargs=revision_spec_kwargs + ) + + engine_settings_kwargs = {"incorrect_kwarg": "val"} + with raises(TypeError): + manager.engines.create( + name=engine_name, engine_settings_kwargs=engine_settings_kwargs + ) + + def test_engine_create_no_available_types( httpx_mock: HTTPXMock, auth_callback: Callable,