fix(backend): use internal schema cache for task workers

Signed-off-by: Fatih Acar <[email protected]>

Author: fatih-acar

backend/infrahub/core/schema/schema_branch.py

@@ -162,6 +162,14 @@ def to_dict_schema_object(self, duplicate: bool = False) -> dict[str, dict[str,
             "templates": {name: self.get(name, duplicate=duplicate) for name in self.templates},
         }
 
+    def to_dict_api_schema_object(self, duplicate: bool = False) -> dict[str, list[dict]]:
+        return {
+            "nodes": [self.get(name, duplicate=duplicate).model_dump() for name in self.nodes],
+            "profiles": [self.get(name, duplicate=duplicate).model_dump() for name in self.profiles],
+            "generics": [self.get(name, duplicate=duplicate).model_dump() for name in self.generics],
+            "templates": [self.get(name, duplicate=duplicate).model_dump() for name in self.templates],
+        }
+
     @classmethod
     def from_dict_schema_object(cls, data: dict) -> Self:
         type_mapping = {

backend/infrahub/git/tasks.py

@@ -378,7 +378,7 @@ async def generate_request_artifact_definition(
 ) -> None:
     await add_tags(branches=[model.branch])
 
-    client = get_client()
+    client = get_client(branch=model.branch)
 
     # Needs to be fetched before fetching group members otherwise `object` relationship would override
     # existing node in client store without the `name` attribute due to #521

backend/infrahub/groups/tasks.py

@@ -12,7 +12,7 @@
 async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate) -> None:
     """Create or Update a GraphQLQueryGroup."""
 
-    client = get_client()
+    client = get_client(branch=model.branch)
 
     # If there is only one subscriber, associate the task to it
     # If there are more than one, for now we can't associate all of them

backend/infrahub/workers/dependencies.py

@@ -7,6 +7,7 @@
 from infrahub import config
 from infrahub.components import ComponentType
 from infrahub.constants.environment import INSTALLATION_TYPE
+from infrahub.core.registry import registry
 from infrahub.database import InfrahubDatabase, get_db
 from infrahub.services.adapters.cache import InfrahubCache
 from infrahub.services.adapters.event import InfrahubEventService
@@ -33,12 +34,21 @@ def get_component_type() -> ComponentType:
         raise ValueError("Component type is not set. It needs to be initialized before working with services.") from exc
 
 
-def build_client() -> InfrahubClient:
-    return InfrahubClient(config=Config(address=config.SETTINGS.main.internal_address, retry_on_failure=True))
+def build_client(branch: str | None = None) -> InfrahubClient:
+    client = InfrahubClient(config=Config(address=config.SETTINGS.main.internal_address, retry_on_failure=True))
+
+    # Populate client schema cache using our internal schema cache
+    _branch = branch or registry.default_branch
+    schema_branch = registry.schema.get_schema_branch(_branch)
+    client.schema.set_cache(schema=schema_branch.to_dict_api_schema_object(), branch=_branch)
+    return client
 
 
 @inject
-def get_client(client: InfrahubClient = Depends(build_client)) -> InfrahubClient:  # noqa: B008
+def get_client(
+    branch: str | None = None,  # noqa: ARG001
+    client: InfrahubClient = Depends(build_client),  # noqa: B008
+) -> InfrahubClient:
     return client