fix: Fir 31676 python sdk resource manager errors when referencing engines (#366)

Author: stepansergeevitch

src/firebolt/common/constants.py

@@ -3,4 +3,4 @@
 DEFAULT_TIMEOUT_SECONDS: int = 60
 
 # Running statuses in infromation schema
-ENGINE_STATUS_RUNNING_LIST = ["Running", "ENGINE_STATE_RUNNING"]
+ENGINE_STATUS_RUNNING_LIST = ["RUNNING", "Running", "ENGINE_STATE_RUNNING"]

src/firebolt/service/V2/types.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import Optional
 
 
 class EngineType(Enum):
@@ -50,9 +51,20 @@ class EngineStatus(Enum):
     FAILED = "Failed"
     DELETING = "Deleting"
 
+    @classmethod
+    def _missing_(cls, value: object) -> Optional["EngineStatus"]:
+        """Handle case-insensitive string values."""
+        if isinstance(value, str):
+            return cls(value.capitalize())
+        return super()._missing_(value)
+
     def __str__(self) -> str:
         return self.value
 
+    @classmethod
+    def from_string(cls, value: str) -> "EngineStatus":
+        return cls(value.capitalize())
+
 
 class DatabaseOrder(Enum):
     DATABASE_ORDER_UNSPECIFIED = "DATABASE_ORDER_UNSPECIFIED"

tests/unit/async_db/V2/test_connection.py

@@ -136,6 +136,40 @@ async def test_connect_engine_name(
         assert await connection.cursor().execute("select*") == len(python_query_data)
 
 
+async def test_connect_engine_running(
+    httpx_mock: HTTPXMock,
+    db_name: str,
+    auth_url: str,
+    server: str,
+    auth: Auth,
+    account_name: str,
+    get_system_engine_url: str,
+    get_system_engine_callback: Callable,
+    account_id_url: str,
+    account_id_callback: Callable,
+    check_credentials_callback: Callable,
+    engine_name: str,
+    system_engine_query_url: str,
+    get_engine_url_callback_test_status: Callable,
+) -> None:
+    httpx_mock.add_callback(check_credentials_callback, url=auth_url)
+    httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url)
+    httpx_mock.add_callback(account_id_callback, url=account_id_url)
+    httpx_mock.add_callback(
+        get_engine_url_callback_test_status, url=system_engine_query_url
+    )
+
+    async with await connect(
+        engine_name=engine_name,
+        database=db_name,
+        auth=auth,
+        account_name=account_name,
+        api_endpoint=server,
+    ):
+        # Engine is running, no exception should be raised
+        pass
+
+
 async def test_connect_database(
     db_name: str,
     auth_url: str,

tests/unit/db/V2/test_connection.py

@@ -141,6 +141,40 @@ def test_connect_engine_name(
         assert connection.cursor().execute("select*") == len(python_query_data)
 
 
+def test_connect_engine_running(
+    httpx_mock: HTTPXMock,
+    db_name: str,
+    auth_url: str,
+    server: str,
+    auth: Auth,
+    account_name: str,
+    get_system_engine_url: str,
+    get_system_engine_callback: Callable,
+    account_id_url: str,
+    account_id_callback: Callable,
+    check_credentials_callback: Callable,
+    engine_name: str,
+    system_engine_query_url: str,
+    get_engine_url_callback_test_status: Callable,
+) -> None:
+    httpx_mock.add_callback(check_credentials_callback, url=auth_url)
+    httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url)
+    httpx_mock.add_callback(account_id_callback, url=account_id_url)
+    httpx_mock.add_callback(
+        get_engine_url_callback_test_status, url=system_engine_query_url
+    )
+
+    with connect(
+        engine_name=engine_name,
+        database=db_name,
+        auth=auth,
+        account_name=account_name,
+        api_endpoint=server,
+    ):
+        # Engine is running, no exception should be raised
+        pass
+
+
 def test_connect_database(
     db_name: str,
     auth_url: str,

tests/unit/db_conftest.py

@@ -303,6 +303,13 @@ def get_engine_url_not_running_callback(
     return _get_engine_url_callback(engine_name, db_name, query_statistics, "Stopped")
 
 
+@fixture(params=["Running", "RUNNING", "ENGINE_STATE_RUNNING"])
+def get_engine_url_callback_test_status(
+    server: str, db_name: str, query_statistics: Dict[str, Any], request: Any
+) -> Callable:
+    return _get_engine_url_callback(server, db_name, query_statistics, request.param)
+
+
 @fixture
 def get_engine_url_invalid_db_callback(
     engine_name,

tests/unit/service/V2/test_engine.py

@@ -1,7 +1,7 @@
 from typing import Callable
 
 from httpx import Request
-from pytest import raises
+from pytest import mark, raises
 from pytest_httpx import HTTPXMock
 
 from firebolt.model.V2.database import Database
@@ -228,3 +228,28 @@ def test_engine_deleteting(
     engine = resource_manager.engines.get_by_name(mock_engine.name)
 
     assert engine.current_status == EngineStatus.DELETING
+
+
+@mark.parametrize(
+    "engine_status, expected_status",
+    [(status.value.upper(), status) for status in EngineStatus],
+)
+def test_engine_new_status(
+    engine_status: str,
+    expected_status: EngineStatus,
+    httpx_mock: HTTPXMock,
+    resource_manager: ResourceManager,
+    instance_type_callback: Callable,
+    instance_type_url: str,
+    system_engine_no_db_query_url: str,
+    mock_engine: Engine,
+):
+    mock_engine.current_status = engine_status
+    get_engine_callback = get_objects_from_db_callback([mock_engine])
+
+    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)
+
+    engine = resource_manager.engines.get_by_name(mock_engine.name)
+
+    assert engine.current_status == expected_status