Merge pull request #155 from opsmill/pog-wait-for-schema-converge

Add function to wait until the schema has converged

Author: ogenstad

changelog/+0e61a54f.added.md

@@ -0,0 +1 @@
+Added InfrahubClient.schema.wait_until_converged() which allowes you to wait until the schema has converged across all Infrahub workers before proceeding with an operation. The InfrahubClient.schema.load() method has also been updated with a new parameter "wait_until_converged".

infrahub_sdk/config.py

@@ -54,6 +54,9 @@ class ConfigBase(BaseSettings):
     pagination_size: int = Field(default=50, description="Page size for queries to the server")
     retry_delay: int = Field(default=5, description="Number of seconds to wait until attempting a retry.")
     retry_on_failure: bool = Field(default=False, description="Retry operation in case of failure")
+    schema_converge_timeout: int = Field(
+        default=60, description="Number of seconds to wait for schema to have converged"
+    )
     timeout: int = Field(default=60, description="Default connection timeout in seconds")
     transport: RequesterTransport = Field(
         default=RequesterTransport.HTTPX, description="Set an alternate transport using a predefined option"

infrahub_sdk/schema.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 from collections import defaultdict
 from collections.abc import MutableMapping
 from enum import Enum
 from pathlib import Path
+from time import sleep
 from typing import TYPE_CHECKING, Any, Optional, TypedDict, TypeVar, Union
 from urllib.parse import urlencode
 
@@ -22,6 +24,7 @@
 )
 from .generator import InfrahubGenerator
 from .graphql import Mutation
+from .queries import SCHEMA_HASH_SYNC_STATUS
 from .transforms import InfrahubTransform
 from .utils import duplicates
 
@@ -616,15 +619,36 @@ async def all(
 
         return self.cache[branch]
 
-    async def load(self, schemas: list[dict], branch: Optional[str] = None) -> SchemaLoadResponse:
+    async def load(
+        self, schemas: list[dict], branch: Optional[str] = None, wait_until_converged: bool = False
+    ) -> SchemaLoadResponse:
         branch = branch or self.client.default_branch
         url = f"{self.client.address}/api/schema/load?branch={branch}"
         response = await self.client._post(
             url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
         )
 
+        if wait_until_converged:
+            await self.wait_until_converged(branch=branch)
+
         return self._validate_load_schema_response(response=response)
 
+    async def wait_until_converged(self, branch: Optional[str] = None) -> None:
+        """Wait until the schema has converged on the selected branch or the timeout has been reached"""
+        waited = 0
+        while True:
+            status = await self.client.execute_graphql(query=SCHEMA_HASH_SYNC_STATUS, branch_name=branch)
+            if status["InfrahubStatus"]["summary"]["schema_hash_synced"]:
+                self.client.log.info(f"Schema successfully converged after {waited} seconds")
+                return
+
+            if waited >= self.client.config.schema_converge_timeout:
+                self.client.log.warning(f"Schema not converged after {waited} seconds, proceeding regardless")
+                return
+
+            waited += 1
+            await asyncio.sleep(delay=1)
+
     async def check(self, schemas: list[dict], branch: Optional[str] = None) -> tuple[bool, Optional[dict]]:
         branch = branch or self.client.default_branch
         url = f"{self.client.address}/api/schema/check?branch={branch}"
@@ -999,15 +1023,36 @@ def fetch(
 
         return nodes
 
-    def load(self, schemas: list[dict], branch: Optional[str] = None) -> SchemaLoadResponse:
+    def load(
+        self, schemas: list[dict], branch: Optional[str] = None, wait_until_converged: bool = False
+    ) -> SchemaLoadResponse:
         branch = branch or self.client.default_branch
         url = f"{self.client.address}/api/schema/load?branch={branch}"
         response = self.client._post(
             url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
         )
 
+        if wait_until_converged:
+            self.wait_until_converged(branch=branch)
+
         return self._validate_load_schema_response(response=response)
 
+    def wait_until_converged(self, branch: Optional[str] = None) -> None:
+        """Wait until the schema has converged on the selected branch or the timeout has been reached"""
+        waited = 0
+        while True:
+            status = self.client.execute_graphql(query=SCHEMA_HASH_SYNC_STATUS, branch_name=branch)
+            if status["InfrahubStatus"]["summary"]["schema_hash_synced"]:
+                self.client.log.info(f"Schema successfully converged after {waited} seconds")
+                return
+
+            if waited >= self.client.config.schema_converge_timeout:
+                self.client.log.warning(f"Schema not converged after {waited} seconds, proceeding regardless")
+                return
+
+            waited += 1
+            sleep(1)
+
     def check(self, schemas: list[dict], branch: Optional[str] = None) -> tuple[bool, Optional[dict]]:
         branch = branch or self.client.default_branch
         url = f"{self.client.address}/api/schema/check?branch={branch}"

tests/unit/sdk/test_schema.py

@@ -3,6 +3,7 @@
 from unittest import mock
 
 import pytest
+from pytest_httpx import HTTPXMock
 from rich.console import Console
 
 from infrahub_sdk import Config, InfrahubClient, InfrahubClientSync
@@ -18,6 +19,7 @@
     InfrahubSchemaSync,
     NodeSchema,
 )
+from tests.unit.sdk.conftest import BothClients
 
 async_schema_methods = [method for method in dir(InfrahubSchema) if not method.startswith("_")]
 sync_schema_methods = [method for method in dir(InfrahubSchemaSync) if not method.startswith("_")]
@@ -170,6 +172,27 @@ async def test_remove_enum_option_raises(clients, client_type, mock_schema_query
             clients.sync.schema.add_enum_option("BuiltinTag", "attribute", "option")
 
 
+@pytest.mark.parametrize("client_type", client_types)
+async def test_schema_wait_happy_path(clients: BothClients, client_type: list[str], httpx_mock: HTTPXMock) -> None:
+    """Simplistic unittest that can be removed once we have the integration tests running again."""
+    httpx_mock.add_response(
+        method="POST",
+        url="http://mock/graphql/branch1",
+        json={"data": {"InfrahubStatus": {"summary": {"schema_hash_synced": False}}}},
+    )
+    httpx_mock.add_response(
+        method="POST",
+        url="http://mock/graphql/branch1",
+        json={"data": {"InfrahubStatus": {"summary": {"schema_hash_synced": True}}}},
+    )
+    if client_type == "standard":
+        await clients.standard.schema.wait_until_converged(branch="branch1")
+    else:
+        clients.sync.schema.wait_until_converged(branch="branch1")
+
+    assert len(httpx_mock.get_requests()) == 2
+
+
 async def test_infrahub_repository_config_getters():
     repo_config = InfrahubRepositoryConfig(
         jinja2_transforms=[