zenml-orm-schema

ZenML ORM Schemas

These are common patterns, rules and conventions used in defining ORM schemas within the ZenML project. Make sure to understand these patterns and adhere to these practices when working with the ORM schema.

All ORM schemas are located in the @zenml/src/zenml/zen_stores/schemas directory and referenced in the __init__.py file. The files are named following the convention <concept>_schemas.py (e.g. stack_schemas.py and component_schemas.py) and usually contain only one schema ORM class, unless multiple classes are required to define the same concept.

ORM Model Definition

Base Class Usage

ORM models typically inherit from the BaseSchema or NamedSchema class to ensure consistency and reuse common functionality. BaseSchema is defined in @zenml/src/zenml/zen_stores/schemas/base_schemas.py:

class BaseSchema(SQLModel):
    """Base SQL Model for ZenML entities."""

    id: UUID = Field(default_factory=uuid4, primary_key=True)
    created: datetime = Field(default_factory=datetime.utcnow)
    updated: datetime = Field(default_factory=datetime.utcnow)

NamedSchema is defined in @zenml/src/zenml/zen_stores/schemas/base_schemas.py, inheriting from BaseSchema and adding a name field:

class NamedSchema(BaseSchema):
    name: str

Declarative Style

Models are defined using a declarative style, where classes represent tables and attributes represent columns. For instance, in @zenml/src/zenml/zen_stores/schemas/component_schemas.py:

class StackComponentSchema(NamedSchema, table=True):
    """SQL Model for stack components."""

    __tablename__ = "stack_component"

    type: str
    flavor: str
    configuration: bytes
    labels: Optional[bytes]
    component_spec_path: Optional[str]
    ...

Note that the table=True argument is used to indicate that the model should be a table in the database and the __tablename__ attribute is used to specify the name of the table.

There is a tight relationship between ORM schemas and their associated domain models, especially concerning the names and types of attributes. Make sure to keep them in sync. Read zenml-domain-models.mdc to understand what ZenML domain models are and how they work.