Prototype to generate Infrahub Schema from Pydantic models

[dgarros]

This is an initial Prototype to explore how we could better integrate Pydantic within the SDK

The end goal is to leverage Pydantic to define the schema and to support all CRUD operations, similar to how SQLmodel is working for a Relational Database.

In the current prototype, it's possible to: Define the schema and read objects, the other operations (Create, Update, Delete) aren't available yet

Define Infrahub Schema using Pydantic Models

Below a simple example with 2 models connected together with a relationship (One to Many)

# This snippet is a simplified version of the example script in `docs/docs/python-sdk/examples/schema_pydantic.py`

from pydantic import Field
from infrahub_sdk.schema import (
     InfrahubAttributeParam as AttrParam,  
     InfrahubRelationshipParam as RelParam, 
    AttributeKind,
    from_pydantic,
    NodeModel
)

class TestCar(NodeModel):
    name: str = Field(description="The name of the car")
    owner: Annotated[Person, RelParam(identifier="car__person")]

class Person(NodeModel):
    model_config = ConfigDict(
        node_schema=NodeSchema(name="Person", namespace="Test", human_readable_fields=["name__value"])
    )

    name: str
    cars: Annotated[list[TestCar] | None, RelParam(identifier="car__person")] = None
    
async def main():
    client = InfrahubClient()
    schema = from_pydantic(models=[Person, TestCar])
    response = await client.schema.load(schemas=[schema.to_schema_dict()], wait_until_converged=True)

The integration is based on the following principles

• All models must inherit from either NodeModel or GenericModel
• Node specific parameters must be provided using model_config
• Attribute or Relationship specific parameters that are not supported by Pydantic, must be provided via an Annotation (InfrahubAttributeParam and InfrahubRelationshipParam)
• Attribute or Relationship parameters that are natively supported by Pydantic (default value, type, regex etc ..) can be provided using pydantic Field

In the example below, the annotation is used on cars to define the identifier for this relationship

class Person(NodeModel):
    name: str
    cars: Annotated[list[TestCar] | None, RelParam(identifier="car__person")] = None

This implementation is slightly different than SQLModel that is using a more Integrated approach with a custom Field method, would be good to understand which approach is preferable.

Query data using Pydantic

The main methods to query objects get, all & filters have been updated to access a Pydantic Models

In the example, below the model Site that was used to define the schema can also be used to query data.
The client.all() will automatically format the data returned by the API using the model Site, in Pydantic format

class Site(NodeModel):
    model_config = ConfigDict(
        node_schema=NodeSchema(
            name="Site", namespace="Infra", human_friendly_id=["name__value"], display_labels=["name__value"]
        )
    )

    name: Annotated[str, AttrParam(unique=True)] = Field(description="The name of the site")

async def main() -> None:
    client = InfrahubClient()
    sites = await client.all(kind=Site)
    rprint(sites)

Create / Update objects

This part hasn't been implemented yet but here is an overview of what it could look like following something similar to SQLModel.

async def main() -> None:
    client = InfrahubClient()
    jfk = Site(name="JFK")

    # objects to create or update will be added to the client and later, the changes will be `saved` / `committed` 
    client.add(jfk)

    await client.sync() # Could be client.save() or client.commit() too

[codecov]

Codecov Report

Attention: Patch coverage is 0.52083% with 191 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
infrahub_sdk/schema/pydantic_utils.py	0.00%	160 Missing ⚠️
infrahub_sdk/client.py	0.00%	25 Missing and 1 partial ⚠️
infrahub_sdk/schema/__init__.py	0.00%	4 Missing and 1 partial ⚠️

❗ There is a different number of reports uploaded between BASE (5bee453) and HEAD (8a8a768). Click for more details.

HEAD has 13 uploads less than BASE

Flag	BASE (5bee453)	HEAD (8a8a768)
python-3.11	2	0
python-3.9	2	0
python-3.10	2	0
python-3.12	2	0
python-filler-3.12	2	0
integration-tests	2	1
python-3.13	2	0

@@             Coverage Diff              @@
##           develop     #306       +/-   ##
============================================
- Coverage    72.45%   21.77%   -50.68%     
============================================
  Files           91       92        +1     
  Lines         8165     8353      +188     
  Branches      1572     1537       -35     
============================================
- Hits          5916     1819     -4097     
- Misses        1838     6345     +4507     
+ Partials       411      189      -222     

Flag	Coverage Δ
integration-tests	21.77% <0.52%> (-0.51%)	⬇️
python-3.10	?
python-3.11	?
python-3.12	?
python-3.13	?
python-3.9	?
python-filler-3.12	?

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
infrahub_sdk/schema/main.py	2.33% <100.00%> (-81.72%)	⬇️
infrahub_sdk/schema/__init__.py	14.19% <0.00%> (-55.87%)	⬇️
infrahub_sdk/client.py	16.97% <0.00%> (-50.93%)	⬇️
infrahub_sdk/schema/pydantic_utils.py	0.00% <0.00%> (ø)

... and 81 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.