|  | 
| 13 | 13 | from typing_extensions import TypeAlias | 
| 14 | 14 | 
 | 
| 15 | 15 | from ..exceptions import ( | 
|  | 16 | +    BranchNotFoundError, | 
| 16 | 17 |     InvalidResponseError, | 
| 17 | 18 |     JsonDecodeError, | 
| 18 | 19 |     SchemaNotFoundError, | 
| @@ -167,6 +168,25 @@ def _get_schema_name(schema: type[SchemaType | SchemaTypeSync] | str) -> str: | 
| 167 | 168 | 
 | 
| 168 | 169 |         raise ValueError("schema must be a protocol or a string") | 
| 169 | 170 | 
 | 
|  | 171 | +    @staticmethod | 
|  | 172 | +    def _parse_schema_response(response: httpx.Response, branch: str) -> MutableMapping[str, Any]: | 
|  | 173 | +        if response.status_code == 400: | 
|  | 174 | +            raise BranchNotFoundError( | 
|  | 175 | +                identifier=branch, message=f"The requested branch was not found on the server [{branch}]" | 
|  | 176 | +            ) | 
|  | 177 | +        response.raise_for_status() | 
|  | 178 | + | 
|  | 179 | +        try: | 
|  | 180 | +            data: MutableMapping[str, Any] = response.json() | 
|  | 181 | +        except json.decoder.JSONDecodeError as exc: | 
|  | 182 | +            raise JsonDecodeError( | 
|  | 183 | +                message=f"Invalid Schema response received from the server at {response.url}: {response.text} [{response.status_code}] ", | 
|  | 184 | +                content=response.text, | 
|  | 185 | +                url=str(response.url), | 
|  | 186 | +            ) from exc | 
|  | 187 | + | 
|  | 188 | +        return data | 
|  | 189 | + | 
| 170 | 190 | 
 | 
| 171 | 191 | class InfrahubSchema(InfrahubSchemaBase): | 
| 172 | 192 |     def __init__(self, client: InfrahubClient): | 
| @@ -420,16 +440,8 @@ async def _fetch( | 
| 420 | 440 |         url = f"{self.client.address}/api/schema?{query_params}" | 
| 421 | 441 | 
 | 
| 422 | 442 |         response = await self.client._get(url=url, timeout=timeout) | 
| 423 |  | -        response.raise_for_status() | 
| 424 | 443 | 
 | 
| 425 |  | -        try: | 
| 426 |  | -            data: MutableMapping[str, Any] = response.json() | 
| 427 |  | -        except json.decoder.JSONDecodeError as exc: | 
| 428 |  | -            raise JsonDecodeError( | 
| 429 |  | -                message=f"Invalid Schema response received from the server at {response.url}: {response.text} [{response.status_code}] ", | 
| 430 |  | -                content=response.text, | 
| 431 |  | -                url=response.url, | 
| 432 |  | -            ) from exc | 
|  | 444 | +        data = self._parse_schema_response(response=response, branch=branch) | 
| 433 | 445 | 
 | 
| 434 | 446 |         nodes: MutableMapping[str, MainSchemaTypesAPI] = {} | 
| 435 | 447 |         for node_schema in data.get("nodes", []): | 
| @@ -648,9 +660,7 @@ def _fetch(self, branch: str, namespaces: list[str] | None = None, timeout: int | 
| 648 | 660 |         query_params = urlencode(url_parts) | 
| 649 | 661 |         url = f"{self.client.address}/api/schema?{query_params}" | 
| 650 | 662 |         response = self.client._get(url=url, timeout=timeout) | 
| 651 |  | -        response.raise_for_status() | 
| 652 |  | - | 
| 653 |  | -        data: MutableMapping[str, Any] = response.json() | 
|  | 663 | +        data = self._parse_schema_response(response=response, branch=branch) | 
| 654 | 664 | 
 | 
| 655 | 665 |         nodes: MutableMapping[str, MainSchemaTypesAPI] = {} | 
| 656 | 666 |         for node_schema in data.get("nodes", []): | 
|  | 
0 commit comments