fix: Re-adding bindings to resource manager (#325)

src/firebolt/model/V1/binding.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from firebolt.model.V1 import FireboltBaseModel
+
+
+class BindingKey(BaseModel):
+    account_id: str
+    database_id: str
+    engine_id: str
+
+
+class Binding(FireboltBaseModel):
+    """A binding between an engine and a database."""
+
+    binding_key: BindingKey = Field(alias="id")
+    is_default_engine: bool = Field(alias="engine_is_default")
+
+    # optional
+    current_status: Optional[str]
+    health_status: Optional[str]
+    create_time: Optional[datetime]
+    create_actor: Optional[str]
+    last_update_time: Optional[datetime]
+    last_update_actor: Optional[str]
+    desired_status: Optional[str]
+
+    @property
+    def database_id(self) -> str:
+        return self.binding_key.database_id
+
+    @property
+    def engine_id(self) -> str:
+        return self.binding_key.engine_id

src/firebolt/model/engine.py

@@ -0,0 +1,3 @@
+# Temporary fix for airflow-provider-firebolt
+class Engine:
+    pass

src/firebolt/service/V1/binding.py

@@ -0,0 +1,49 @@
+import logging
+from typing import List, Optional
+
+from firebolt.model.V1.binding import Binding
+from firebolt.service.V1.base import BaseService
+from firebolt.utils.urls import ACCOUNT_BINDINGS_URL
+from firebolt.utils.util import prune_dict
+
+logger = logging.getLogger(__name__)
+
+
+class BindingService(BaseService):
+    def get_many(
+        self,
+        database_id: Optional[str] = None,
+        engine_id: Optional[str] = None,
+        is_system_database: Optional[bool] = None,
+    ) -> List[Binding]:
+        """
+        List bindings on Firebolt, optionally filtering by database and engine.
+
+        Args:
+            database_id:
+                Return bindings matching the database_id.
+                If None, match any databases.
+            engine_id:
+                Return bindings matching the engine_id.
+                If None, match any engines.
+            is_system_database:
+                If True, return only system databases.
+                If False, return only non-system databases.
+                If None, do not filter on this parameter.
+
+        Returns:
+            List of bindings matching the filter parameters
+        """
+
+        response = self.client.get(
+            url=ACCOUNT_BINDINGS_URL.format(account_id=self.account_id),
+            params=prune_dict(
+                {
+                    "page.first": 5000,  # FUTURE: pagination support w/ generator
+                    "filter.id_database_id_eq": database_id,
+                    "filter.id_engine_id_eq": engine_id,
+                    "filter.is_system_database_eq": is_system_database,
+                }
+            ),
+        )
+        return [Binding.parse_obj(i["node"]) for i in response.json()["edges"]]

src/firebolt/service/manager.py

@@ -133,8 +133,10 @@ def _init_services_v2(self) -> None:
 
     def _init_services_v1(self) -> None:
         # avoid circular import
+        from firebolt.service.V1.binding import BindingService
         from firebolt.service.V1.engine import EngineService
 
+        self.bindings = BindingService(resource_manager=self)  # type: ignore
         self.engines = EngineService(resource_manager=self)  # type: ignore
 
     def __del__(self) -> None:

src/firebolt/utils/urls.py

@@ -12,8 +12,6 @@
 ACCOUNT_DATABASE_BINDING_URL = ACCOUNT_DATABASE_URL + "/bindings/{engine_id}"
 ACCOUNT_DATABASE_BY_NAME_URL = ACCOUNT_DATABASES_URL + ":getIdByName"
 
-ACCOUNT_BINDINGS_URL = "/core/v1/accounts/{account_id}/bindings"
-
 ACCOUNT_INSTANCE_TYPES_URL = "/aws/v2/accounts/{account_id}/instanceTypes"
 
 PROVIDERS_URL = "/compute/v1/providers"
@@ -39,3 +37,4 @@
 ACCOUNT_ENGINE_ID_BY_NAME_URL = ACCOUNT_LIST_ENGINES_URL + ":getIdByName"
 ACCOUNT_URL = "/iam/v2/account"
 ACCOUNT_BY_NAME_URL_V1 = "/iam/v2/accounts:getIdByName"
+ACCOUNT_BINDINGS_URL = "/core/v1/accounts/{account_id}/bindings"

tests/unit/service/V1/conftest.py

@@ -8,6 +8,7 @@
 from pytest import fixture
 
 from firebolt.client.auth.base import Auth
+from firebolt.model.V1.binding import Binding, BindingKey
 from firebolt.model.V1.engine import Engine, EngineKey, EngineSettings
 from firebolt.model.V1.region import Region, RegionKey
 from firebolt.utils.exception import AccountNotFoundError
@@ -25,7 +26,9 @@
     PROVIDERS_URL,
     REGIONS_URL,
 )
-from tests.unit.util import list_to_paginated_response
+from tests.unit.util import (
+    list_to_paginated_response_v1 as list_to_paginated_response,
+)
 
 
 @fixture
@@ -48,6 +51,11 @@ def region_key() -> RegionKey:
     return RegionKey(provider_id="pid", region_id="rid")
 
 
+@fixture
+def db_id() -> str:
+    return "db_id"
+
+
 @fixture
 def mock_engine(engine_name, region_key, engine_settings, account_id, server) -> Engine:
     return Engine(
@@ -392,6 +400,33 @@ def do_mock(
     return do_mock
 
 
+@fixture
+def binding(account_id, mock_engine, db_id) -> Binding:
+    return Binding(
+        binding_key=BindingKey(
+            account_id=account_id,
+            database_id=db_id,
+            engine_id=mock_engine.engine_id,
+        ),
+        is_default_engine=True,
+    )
+
+
+@fixture
+def bindings_callback(bindings_url: str, binding: Binding) -> Callable:
+    def do_mock(
+        request: httpx.Request = None,
+        **kwargs,
+    ) -> Response:
+        assert request.url == bindings_url
+        return Response(
+            status_code=httpx.codes.OK,
+            json=list_to_paginated_response([binding]),
+        )
+
+    return do_mock
+
+
 @fixture
 def auth_url(server: str) -> str:
     return f"https://{server}{AUTH_URL}"

tests/unit/service/V1/test_bindings.py

@@ -0,0 +1,29 @@
+from re import Pattern
+from typing import Callable
+
+from pytest_httpx import HTTPXMock
+
+from firebolt.common.settings import Settings
+from firebolt.model.V1.engine import Engine
+from firebolt.service.manager import ResourceManager
+
+
+def test_get_many_bindings(
+    httpx_mock: HTTPXMock,
+    auth_callback: Callable,
+    auth_url: str,
+    account_id_callback: Callable,
+    account_id_url: Pattern,
+    bindings_url: str,
+    bindings_callback: Callable,
+    settings: Settings,
+    mock_engine: Engine,
+):
+    httpx_mock.add_callback(auth_callback, url=auth_url)
+    httpx_mock.add_callback(account_id_callback, url=account_id_url)
+    httpx_mock.add_callback(bindings_callback, url=bindings_url)
+
+    resource_manager = ResourceManager(settings=settings)
+    bindings = resource_manager.bindings.get_many(engine_id=mock_engine.engine_id)
+    assert len(bindings) > 0
+    assert any(binding.is_default_engine for binding in bindings)

tests/unit/util.py

@@ -6,6 +6,7 @@
 from firebolt.client import AsyncClientV2 as AsyncClient
 from firebolt.client import ClientV2
 from firebolt.client.auth import Auth
+from firebolt.model.V1 import FireboltBaseModel as FireboltBaseModelV1
 from firebolt.model.V2 import FireboltBaseModel
 
 
@@ -21,6 +22,10 @@ def list_to_paginated_response(items: List[FireboltBaseModel]) -> Dict:
     return {"edges": [{"node": to_dict(i)} for i in items]}
 
 
+def list_to_paginated_response_v1(items: List[FireboltBaseModelV1]) -> Dict:
+    return {"edges": [{"node": i.dict()} for i in items]}
+
+
 def execute_generator_requests(
     requests: Generator[Request, Response, None], api_endpoint: str = ""
 ) -> None: