Merge pull request #160 from opsmill/dga-20241201-cleanup-schema

Cleanup and restructure Pydantic models to manage the Schema

Author: dgarros

infrahub_sdk/checks.py

@@ -17,7 +17,7 @@
     from pathlib import Path
 
     from . import InfrahubClient
-    from .schema import InfrahubCheckDefinitionConfig
+    from .schema.repository import InfrahubCheckDefinitionConfig
 
 INFRAHUB_CHECK_VARIABLE_TO_IMPORT = "INFRAHUB_CHECKS"
 

infrahub_sdk/client.py

@@ -50,7 +50,7 @@
 from .protocols_base import CoreNode, CoreNodeSync
 from .queries import get_commit_update_mutation
 from .query_groups import InfrahubGroupContext, InfrahubGroupContextSync
-from .schema import InfrahubSchema, InfrahubSchemaSync, NodeSchema
+from .schema import InfrahubSchema, InfrahubSchemaSync, NodeSchemaAPI
 from .store import NodeStore, NodeStoreSync
 from .timestamp import Timestamp
 from .types import AsyncRequester, HTTPMethod, SyncRequester
@@ -448,12 +448,12 @@ async def get(
         filters: MutableMapping[str, Any] = {}
 
         if id:
-            if not is_valid_uuid(id) and isinstance(schema, NodeSchema) and schema.default_filter:
+            if not is_valid_uuid(id) and isinstance(schema, NodeSchemaAPI) and schema.default_filter:
                 filters[schema.default_filter] = id
             else:
                 filters["ids"] = [id]
         if hfid:
-            if isinstance(schema, NodeSchema) and schema.human_friendly_id:
+            if isinstance(schema, NodeSchemaAPI) and schema.human_friendly_id:
                 filters["hfid"] = hfid
             else:
                 raise ValueError("Cannot filter by HFID if the node doesn't have an HFID defined")
@@ -1916,12 +1916,12 @@ def get(
         filters: MutableMapping[str, Any] = {}
 
         if id:
-            if not is_valid_uuid(id) and isinstance(schema, NodeSchema) and schema.default_filter:
+            if not is_valid_uuid(id) and isinstance(schema, NodeSchemaAPI) and schema.default_filter:
                 filters[schema.default_filter] = id
             else:
                 filters["ids"] = [id]
         if hfid:
-            if isinstance(schema, NodeSchema) and schema.human_friendly_id:
+            if isinstance(schema, NodeSchemaAPI) and schema.human_friendly_id:
                 filters["hfid"] = hfid
             else:
                 raise ValueError("Cannot filter by HFID if the node doesn't have an HFID defined")

infrahub_sdk/code_generator.py

@@ -1,17 +1,19 @@
 from collections.abc import Mapping
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import jinja2
 
 from . import protocols as sdk_protocols
 from .ctl.constants import PROTOCOLS_TEMPLATE
 from .schema import (
-    AttributeSchema,
+    AttributeSchemaAPI,
     GenericSchema,
-    MainSchemaTypes,
+    GenericSchemaAPI,
+    MainSchemaTypesAll,
     NodeSchema,
-    ProfileSchema,
-    RelationshipSchema,
+    NodeSchemaAPI,
+    ProfileSchemaAPI,
+    RelationshipSchemaAPI,
 )
 
 ATTRIBUTE_KIND_MAP = {
@@ -40,17 +42,17 @@
 
 
 class CodeGenerator:
-    def __init__(self, schema: dict[str, MainSchemaTypes]):
-        self.generics: dict[str, GenericSchema] = {}
-        self.nodes: dict[str, NodeSchema] = {}
-        self.profiles: dict[str, ProfileSchema] = {}
+    def __init__(self, schema: dict[str, MainSchemaTypesAll]):
+        self.generics: dict[str, Union[GenericSchemaAPI, GenericSchema]] = {}
+        self.nodes: dict[str, Union[NodeSchemaAPI, NodeSchema]] = {}
+        self.profiles: dict[str, ProfileSchemaAPI] = {}
 
         for name, schema_type in schema.items():
-            if isinstance(schema_type, GenericSchema):
+            if isinstance(schema_type, (GenericSchemaAPI, GenericSchema)):
                 self.generics[name] = schema_type
-            if isinstance(schema_type, NodeSchema):
+            if isinstance(schema_type, (NodeSchemaAPI, NodeSchema)):
                 self.nodes[name] = schema_type
-            if isinstance(schema_type, ProfileSchema):
+            if isinstance(schema_type, ProfileSchemaAPI):
                 self.profiles[name] = schema_type
 
         self.base_protocols = [
@@ -92,7 +94,7 @@ def _jinja2_filter_inheritance(value: dict[str, Any]) -> str:
         return ", ".join(inherit_from)
 
     @staticmethod
-    def _jinja2_filter_render_attribute(value: AttributeSchema) -> str:
+    def _jinja2_filter_render_attribute(value: AttributeSchemaAPI) -> str:
         attribute_kind: str = ATTRIBUTE_KIND_MAP[value.kind]
 
         if value.optional:
@@ -101,7 +103,7 @@ def _jinja2_filter_render_attribute(value: AttributeSchema) -> str:
         return f"{value.name}: {attribute_kind}"
 
     @staticmethod
-    def _jinja2_filter_render_relationship(value: RelationshipSchema, sync: bool = False) -> str:
+    def _jinja2_filter_render_relationship(value: RelationshipSchemaAPI, sync: bool = False) -> str:
         name = value.name
         cardinality = value.cardinality
 
@@ -116,12 +118,12 @@ def _jinja2_filter_render_relationship(value: RelationshipSchema, sync: bool = F
 
     @staticmethod
     def _sort_and_filter_models(
-        models: Mapping[str, MainSchemaTypes], filters: Optional[list[str]] = None
-    ) -> list[MainSchemaTypes]:
+        models: Mapping[str, MainSchemaTypesAll], filters: Optional[list[str]] = None
+    ) -> list[MainSchemaTypesAll]:
         if filters is None:
             filters = ["CoreNode"]
 
-        filtered: list[MainSchemaTypes] = []
+        filtered: list[MainSchemaTypesAll] = []
         for name, model in models.items():
             if name in filters:
                 continue

infrahub_sdk/ctl/check.py

@@ -17,7 +17,7 @@
 from ..ctl.repository import get_repository_config
 from ..ctl.utils import catch_exception, execute_graphql_query
 from ..exceptions import ModuleImportError
-from ..schema import InfrahubCheckDefinitionConfig, InfrahubRepositoryConfig
+from ..schema.repository import InfrahubCheckDefinitionConfig, InfrahubRepositoryConfig
 
 app = typer.Typer()
 console = Console()

infrahub_sdk/ctl/cli_commands.py

@@ -38,11 +38,8 @@
 from ..ctl.validate import app as validate_app
 from ..exceptions import GraphQLError, ModuleImportError
 from ..jinja2 import identify_faulty_jinja_code
-from ..schema import (
-    InfrahubRepositoryConfig,
-    MainSchemaTypes,
-    SchemaRoot,
-)
+from ..schema import MainSchemaTypesAll, SchemaRoot
+from ..schema.repository import InfrahubRepositoryConfig
 from ..utils import get_branch, write_to_file
 from ..yaml import SchemaFile
 from .exporter import dump
@@ -364,7 +361,7 @@ def protocols(
 ) -> None:
     """Export Python protocols corresponding to a schema."""
 
-    schema: dict[str, MainSchemaTypes] = {}
+    schema: dict[str, MainSchemaTypesAll] = {}
 
     if schemas:
         schemas_data = load_yamlfile_from_disk_and_exit(paths=schemas, file_type=SchemaFile, console=console)

infrahub_sdk/ctl/generator.py

@@ -10,7 +10,7 @@
 from ..ctl.utils import execute_graphql_query, parse_cli_vars
 from ..exceptions import ModuleImportError
 from ..node import InfrahubNode
-from ..schema import InfrahubRepositoryConfig
+from ..schema.repository import InfrahubRepositoryConfig
 
 
 async def run(

infrahub_sdk/ctl/render.py

@@ -1,6 +1,6 @@
 from rich.console import Console
 
-from ..schema import InfrahubRepositoryConfig
+from ..schema.repository import InfrahubRepositoryConfig
 
 
 def list_jinja2_transforms(config: InfrahubRepositoryConfig) -> None:

infrahub_sdk/ctl/repository.py

@@ -12,7 +12,7 @@
 from ..ctl.exceptions import FileNotValidError
 from ..ctl.utils import init_logging
 from ..graphql import Mutation
-from ..schema import InfrahubRepositoryConfig
+from ..schema.repository import InfrahubRepositoryConfig
 from ._file import read_file
 from .parameters import CONFIG_PARAM
 

infrahub_sdk/ctl/transform.py

@@ -1,6 +1,6 @@
 from rich.console import Console
 
-from ..schema import InfrahubRepositoryConfig
+from ..schema.repository import InfrahubRepositoryConfig
 
 
 def list_transforms(config: InfrahubRepositoryConfig) -> None:

infrahub_sdk/ctl/utils.py

@@ -26,7 +26,7 @@
     ServerNotReachableError,
     ServerNotResponsiveError,
 )
-from ..schema import InfrahubRepositoryConfig
+from ..schema.repository import InfrahubRepositoryConfig
 from ..yaml import YamlFile
 from .client import initialize_client_sync