feat: Add rm support for engine type update

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

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} = ? "

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()

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,18 +277,25 @@ 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,
         **kwargs,
     ) -> 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,

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,13 +137,15 @@ 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)
     httpx_mock.add_callback(update_engine_callback, url=system_engine_no_db_query_url)
     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