diff --git a/setup.cfg b/setup.cfg index c6f3339eb8d..582523acd53 100755 --- a/setup.cfg +++ b/setup.cfg @@ -24,6 +24,7 @@ project_urls = packages = find: install_requires = aiorwlock==1.1.0 + anyio<4.0.0 appdirs>=1.4.4 appdirs-stubs>=0.1.0 async-generator>=1.10 diff --git a/src/firebolt/model/engine.py b/src/firebolt/model/engine.py index 1c830229ad9..5ed19c25f9e 100644 --- a/src/firebolt/model/engine.py +++ b/src/firebolt/model/engine.py @@ -57,6 +57,7 @@ class Engine(FireboltBaseModel): "AUTO_STOP", "RENAME_TO", "WARMUP", + "ENGINE_TYPE", ) DROP_SQL: ClassVar[str] = "DROP ENGINE {}" @@ -201,6 +202,7 @@ def stop(self) -> Engine: def update( self, name: Optional[str] = None, + engine_type: Union[EngineType, str, None] = None, scale: Optional[int] = None, spec: Union[InstanceType, str, None] = None, auto_stop: Optional[int] = None, @@ -211,7 +213,7 @@ def update( parameters are set to None, old engine parameter values remain. """ - if not any((name, scale, spec, auto_stop, warmup)): + if not any((name, scale, spec, auto_stop, warmup, engine_type)): # Nothing to be updated return self @@ -226,7 +228,8 @@ def update( sql = self.ALTER_PREFIX_SQL.format(self.name) parameters = [] for param, value in zip( - self.ALTER_PARAMETER_NAMES, (scale, spec, auto_stop, name, warmup) + self.ALTER_PARAMETER_NAMES, + (scale, spec, auto_stop, name, warmup, engine_type), ): if value: sql += f"{param} = ? " diff --git a/tests/integration/resource_manager/test_engine.py b/tests/integration/resource_manager/test_engine.py index f21bc4dd1f3..93b68ffc2ba 100644 --- a/tests/integration/resource_manager/test_engine.py +++ b/tests/integration/resource_manager/test_engine.py @@ -2,7 +2,7 @@ from firebolt.client.auth import Auth from firebolt.service.manager import ResourceManager -from firebolt.service.types import EngineStatus +from firebolt.service.types import EngineStatus, EngineType, WarmupMethod def make_engine_name(database_name: str, suffix: str) -> str: @@ -41,11 +41,11 @@ def test_create_start_stop_engine( ParamValue = namedtuple("ParamValue", "set expected") ENGINE_UPDATE_PARAMS = { - # commented parameters are not available yet - # "scale": ParamValue(23, 23), - # "spec": ParamValue("B1", "B1"), + "scale": ParamValue(3, 3), + "spec": ParamValue("B1", "B1"), "auto_stop": ParamValue(123, 7380), - # "warmup": ParamValue(WarmupMethod.PRELOAD_ALL_DATA, WarmupMethod.PRELOAD_ALL_DATA), + "warmup": ParamValue(WarmupMethod.PRELOAD_ALL_DATA, WarmupMethod.PRELOAD_ALL_DATA), + "engine_type": ParamValue(EngineType.DATA_ANALYTICS, EngineType.DATA_ANALYTICS), } @@ -69,7 +69,13 @@ def test_engine_update_single_parameter( engine.update(**{param: value.set}) engine_new = rm.engines.get(name) - assert getattr(engine_new, param) == value.expected, f"Invalid {param} value" + if param == "spec": + current_value = engine_new.spec.name + elif param == "engine_type": + current_value = engine_new.type + else: + current_value = getattr(engine_new, param) + assert current_value == value.expected, f"Invalid {param} value" engine.delete() @@ -97,6 +103,12 @@ def test_engine_update_multiple_parameters( engine_new = rm.engines.get(name) for param, value in ENGINE_UPDATE_PARAMS.items(): - assert getattr(engine_new, param) == value.expected, f"Invalid {param} value" + if param == "spec": + current_value = engine_new.spec.name + elif param == "engine_type": + current_value = engine_new.type + else: + current_value = getattr(engine_new, param) + assert current_value == value.expected, f"Invalid {param} value" engine.delete() diff --git a/tests/unit/service/conftest.py b/tests/unit/service/conftest.py index a639376929c..5dc058e1dd1 100644 --- a/tests/unit/service/conftest.py +++ b/tests/unit/service/conftest.py @@ -205,10 +205,15 @@ def get_obj_field(obj, f): value = getattr(obj, f.name) if isinstance(value, (InstanceType, EngineStatus)): return str(value) - if isinstance(value, (EngineType, WarmupMethod)): + if isinstance(value, WarmupMethod): return " ".join( map(lambda s: s.capitalize(), str(value).lower().split("_")) ) + if isinstance(value, EngineType): + return { + EngineType.GENERAL_PURPOSE: "General Purpose", + EngineType.DATA_ANALYTICS: "Analytics", + }[value] return value query_response = { @@ -272,11 +277,17 @@ def updated_engine_scale() -> int: return 10 +@fixture +def updated_engine_type() -> EngineType: + return EngineType.DATA_ANALYTICS + + @fixture def update_engine_callback( system_engine_no_db_query_url: str, mock_engine: Engine, updated_engine_scale: int, + updated_engine_type: EngineType, ) -> Callable: def do_mock( request: httpx.Request = None, @@ -284,6 +295,7 @@ def do_mock( ) -> Response: assert request.url == system_engine_no_db_query_url mock_engine.scale = updated_engine_scale + mock_engine.type = updated_engine_type return Response( status_code=httpx.codes.OK, json=empty_response, diff --git a/tests/unit/service/test_engine.py b/tests/unit/service/test_engine.py index 0097153c8bb..13db5e6947e 100644 --- a/tests/unit/service/test_engine.py +++ b/tests/unit/service/test_engine.py @@ -6,6 +6,7 @@ from firebolt.model.database import Database from firebolt.model.engine import Engine from firebolt.service.manager import ResourceManager +from firebolt.service.types import EngineType from firebolt.utils.exception import ( EngineNotFoundError, NoAttachedDatabaseError, @@ -136,6 +137,7 @@ def test_engine_update( update_engine_callback: Callable, system_engine_no_db_query_url: str, updated_engine_scale: int, + updated_engine_type: EngineType, ): httpx_mock.add_callback(instance_type_callback, url=instance_type_url) httpx_mock.add_callback(get_engine_callback, url=system_engine_no_db_query_url) @@ -143,6 +145,7 @@ def test_engine_update( httpx_mock.add_callback(get_engine_callback, url=system_engine_no_db_query_url) mock_engine._service = resource_manager.engines - mock_engine.update(scale=updated_engine_scale) + mock_engine.update(scale=updated_engine_scale, engine_type=updated_engine_type) assert mock_engine.scale == updated_engine_scale + assert mock_engine.type == updated_engine_type