zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/zen_server/rbac/models.py

@@ -16,7 +16,7 @@
 from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.utils.enum_utils import StrEnum
 
@@ -89,7 +89,4 @@ class Resource(BaseModel):
 
         return representation
 
-    class Config:
-        """Pydantic configuration class."""
-
-        frozen = True
+    model_config = ConfigDict(frozen=True)

zenml/zen_server/rbac/utils.py

@@ -72,7 +72,7 @@ def dehydrate_page(page: Page[AnyResponse]) -> Page[AnyResponse]:
         for item in page.items
     ]
 
-    return page.copy(update={"items": new_items})
+    return page.model_copy(update={"items": new_items})
 
 
 def dehydrate_response_model(
@@ -103,15 +103,12 @@ def dehydrate_response_model(
         )
 
     dehydrated_values = {}
-    for key, value in model.__dict__.items():
-        if key in model.__private_attributes__:
-            dehydrated_values[key] = value
-        else:
-            dehydrated_values[key] = _dehydrate_value(
-                value, permissions=permissions
-            )
+    for key, value in dict(model).items():
+        dehydrated_values[key] = _dehydrate_value(
+            value, permissions=permissions
+        )
 
-    return type(model).parse_obj(dehydrated_values)
+    return type(model).model_validate(dehydrated_values)
 
 
 def _dehydrate_value(
@@ -144,6 +141,8 @@ def _dehydrate_value(
             return dehydrate_response_model(value, permissions=permissions)
         else:
             return get_permission_denied_model(value)
+    elif isinstance(value, Page):
+        return dehydrate_page(page=value)
     elif isinstance(value, BaseModel):
         return dehydrate_response_model(value, permissions=permissions)
     elif isinstance(value, Dict):
@@ -189,8 +188,8 @@ def get_permission_denied_model(model: AnyResponse) -> AnyResponse:
     Returns:
         The permission denied model.
     """
-    return model.copy(
-        exclude={"body", "metadata"}, update={"permission_denied": True}
+    return model.model_copy(
+        update={"body": None, "metadata": None, "permission_denied": True}
     )
 
 
@@ -472,9 +471,7 @@ def get_subresources_for_model(
     """
     resources = set()
 
-    for key, value in model.__dict__.items():
-        if key in model.__private_attributes__:
-            continue
+    for value in dict(model).values():
         resources.update(_get_subresources_for_value(value))
 
     return resources

zenml/zen_server/routers/auth_endpoints.py

@@ -436,7 +436,7 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             )
         )
     else:
@@ -455,7 +455,7 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             ),
         )
 

zenml/zen_server/routers/pipeline_builds_endpoints.py

@@ -177,7 +177,7 @@ if server_config().workload_manager_enabled:
         deployment = zen_store().get_deployment(
             deployment_id=build.template_deployment_id, hydrate=True
         )
-        deployment.metadata.build = build  # type: ignore[union-attr]
+        deployment.get_metadata().build = build
 
         return run_pipeline(
             deployment=deployment,

zenml/zen_server/routers/runs_endpoints.py

@@ -242,7 +242,7 @@ def get_pipeline_configuration(
     run = verify_permissions_and_get_entity(
         id=run_id, get_method=zen_store().get_run, hydrate=True
     )
-    return run.config.dict()
+    return run.config.model_dump()
 
 
 @router.get(

zenml/zen_server/routers/secrets_endpoints.py

@@ -175,8 +175,9 @@ def update_secret(
         # the update to None, so that they are deleted.
         secret = zen_store().get_secret(secret_id=secret_id)
         for key in secret.values.keys():
-            if key not in secret_update.values:
-                secret_update.values[key] = None
+            if secret_update.values is not None:
+                if key not in secret_update.values:
+                    secret_update.values[key] = None
 
     return verify_permissions_and_update_entity(
         id=secret_id,

zenml/zen_server/routers/server_endpoints.py

@@ -122,7 +122,7 @@ if server_config().external_server_id is None:
         """
         if not server_config().rbac_enabled:
             will_update_admin_properties = bool(
-                settings_update.dict(
+                settings_update.model_dump(
                     exclude_none=True, exclude={"onboarding_state"}
                 )
             )

zenml/zen_server/routers/steps_endpoints.py

@@ -200,7 +200,7 @@ def get_step_configuration(
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
     verify_permission_for_model(pipeline_run, action=Action.READ)
 
-    return step.config.dict()
+    return step.config.model_dump()
 
 
 @router.get(

zenml/zen_server/routers/workspaces_endpoints.py

@@ -1067,7 +1067,7 @@ def create_code_repository(
 
 @router.get(
     WORKSPACES + "/{workspace_name_or_id}" + STATISTICS,
-    response_model=Dict[str, str],
+    response_model=Dict[str, int],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_stores/base_zen_store.py

@@ -18,6 +18,7 @@ from abc import ABC
 from typing import (
     Any,
     ClassVar,
+    Dict,
     Optional,
     Tuple,
     Type,
@@ -26,7 +27,7 @@ from typing import (
 from urllib.parse import urlparse
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, model_validator
 from requests import ConnectionError
 
 import zenml
@@ -54,6 +55,7 @@ from zenml.models import (
     UserResponse,
     WorkspaceResponse,
 )
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.sql_secrets_store import (
     SqlSecretsStoreConfiguration,
 )
@@ -78,6 +80,45 @@ class BaseZenStore(
     TYPE: ClassVar[StoreType]
     CONFIG_TYPE: ClassVar[Type[StoreConfiguration]]
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def convert_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Method to infer the correct type of the config and convert.
+
+        Args:
+            data: The provided configuration object, can potentially be a
+                generic object
+
+        Raises:
+            ValueError: If the provided config object's type does not match
+                any of the current implementations.
+
+        Returns:
+            The converted configuration object.
+        """
+        if data["config"].type == StoreType.SQL:
+            from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration
+
+            data["config"] = SqlZenStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == StoreType.REST:
+            from zenml.zen_stores.rest_zen_store import (
+                RestZenStoreConfiguration,
+            )
+
+            data["config"] = RestZenStoreConfiguration(
+                **data["config"].model_dump()
+            )
+        else:
+            raise ValueError(
+                f"Unknown type '{data['config'].type}' for the configuration."
+            )
+
+        return data
+
     # ---------------------------------
     # Initialization and configuration
     # ---------------------------------
@@ -477,14 +518,10 @@ class BaseZenStore(
             raise KeyError(f"User with external ID '{user_id}' not found.")
         return users.items[0]
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from configs of previous ZenML versions
-        extra = "ignore"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        extra="ignore",
+    )

zenml/zen_stores/migrations/utils.py

@@ -28,8 +28,7 @@ from typing import (
 )
 
 import pymysql
-from pydantic import BaseModel
-from pydantic.json import pydantic_encoder
+from pydantic import BaseModel, ConfigDict
 from sqlalchemy import MetaData, func, text
 from sqlalchemy.engine import URL, Engine
 from sqlalchemy.exc import (
@@ -42,6 +41,7 @@ from sqlmodel import (
 )
 
 from zenml.logger import get_logger
+from zenml.utils.json_utils import pydantic_encoder
 
 logger = get_logger(__name__)
 
@@ -300,25 +300,26 @@ class MigrationUtils(BaseModel):
 
                 # Fetch the number of rows in the table
                 row_count = conn.scalar(
-                    select([func.count("*")]).select_from(table)
+                    select(func.count()).select_from(table)
                 )
 
                 # Fetch the data from the table in batches
-                batch_size = 50
-                for i in range(0, row_count, batch_size):
-                    rows = conn.execute(
-                        table.select()
-                        .order_by(*order_by)
-                        .limit(batch_size)
-                        .offset(i)
-                    ).fetchall()
-
-                    store_db_info(
-                        dict(
-                            table=table.name,
-                            data=[row._asdict() for row in rows],
-                        ),
-                    )
+                if row_count is not None:
+                    batch_size = 50
+                    for i in range(0, row_count, batch_size):
+                        rows = conn.execute(
+                            table.select()
+                            .order_by(*order_by)
+                            .limit(batch_size)
+                            .offset(i)
+                        ).fetchall()
+
+                        store_db_info(
+                            dict(
+                                table=table.name,
+                                data=[row._asdict() for row in rows],
+                            ),
+                        )
 
     def restore_database_from_storage(
         self, load_db_info: Callable[[], Generator[Dict[str, Any], None, None]]
@@ -342,11 +343,9 @@ class MigrationUtils(BaseModel):
                 information.
         """
         # Drop and re-create the primary database
-        self.create_database(
-            drop=True,
-        )
+        self.create_database(drop=True)
 
-        metadata = MetaData(bind=self.engine)
+        metadata = MetaData()
 
         with self.engine.begin() as connection:
             # read the DB information one JSON object at a time
@@ -356,7 +355,7 @@ class MigrationUtils(BaseModel):
                     # execute the table creation statement
                     connection.execute(text(table_dump["create_stmt"]))
                     # Reload the database metadata after creating the table
-                    metadata.reflect()
+                    metadata.reflect(bind=self.engine)
 
                 if "data" in table_dump:
                     # insert the data into the database
@@ -604,11 +603,11 @@ class MigrationUtils(BaseModel):
             src_engine: The source SQLAlchemy engine.
             dst_engine: The destination SQLAlchemy engine.
         """
-        src_metadata = MetaData(bind=src_engine)
-        src_metadata.reflect()
+        src_metadata = MetaData()
+        src_metadata.reflect(bind=src_engine)
 
-        dst_metadata = MetaData(bind=dst_engine)
-        dst_metadata.reflect()
+        dst_metadata = MetaData()
+        dst_metadata.reflect(bind=dst_engine)
 
         # @event.listens_for(src_metadata, "column_reflect")
         # def generalize_datatypes(inspector, tablename, column_dict):
@@ -620,7 +619,7 @@ class MigrationUtils(BaseModel):
 
         # Refresh target metadata after creating the tables
         dst_metadata.clear()
-        dst_metadata.reflect()
+        dst_metadata.reflect(bind=dst_engine)
 
         # Copy all data from the source database to the destination database
         with src_engine.begin() as src_conn:
@@ -648,22 +647,23 @@ class MigrationUtils(BaseModel):
                         order_by.append(src_table.columns["id"])
 
                     row_count = src_conn.scalar(
-                        select([func.count("*")]).select_from(src_table)
+                        select(func.count()).select_from(src_table)
                     )
 
                     # Copy rows in batches
-                    batch_size = 50
-                    for i in range(0, row_count, batch_size):
-                        rows = src_conn.execute(
-                            src_table.select()
-                            .order_by(*order_by)
-                            .limit(batch_size)
-                            .offset(i)
-                        ).fetchall()
-
-                        dst_conn.execute(
-                            insert, [row._asdict() for row in rows]
-                        )
+                    if row_count is not None:
+                        batch_size = 50
+                        for i in range(0, row_count, batch_size):
+                            rows = src_conn.execute(
+                                src_table.select()
+                                .order_by(*order_by)
+                                .limit(batch_size)
+                                .offset(i)
+                            ).fetchall()
+
+                            dst_conn.execute(
+                                insert, [row._asdict() for row in rows]
+                            )
 
     def backup_database_to_db(self, backup_db_name: str) -> None:
         """Backup the database to a backup database.
@@ -713,8 +713,4 @@ class MigrationUtils(BaseModel):
             "backup database."
         )
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # all attributes with leading underscore are private
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)

zenml/zen_stores/migrations/versions/0701da9951a0_added_service_table.py

@@ -22,7 +22,7 @@ def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     # If the tables already exist, skip this migration.
     conn = op.get_bind()
-    inspector = Inspector.from_engine(conn)
+    inspector = Inspector.from_engine(conn.engine)
     tables = inspector.get_table_names()
     if "service" in tables:
         return

zenml/zen_stores/migrations/versions/1041bc644e0d_remove_secrets_manager.py

@@ -19,14 +19,16 @@ depends_on = None
 def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("stack_component", "stack_composition"))
+    meta = sa.MetaData()
+    meta.reflect(
+        only=("stack_component", "stack_composition"), bind=op.get_bind()
+    )
     components = sa.Table("stack_component", meta)
     compositions = sa.Table("stack_composition", meta)
 
     # Find all secrets manager components
     secrets_manager_components = conn.execute(
-        sa.select([components.c.id]).where(
+        sa.select(components.c.id).where(
             components.c.type == "secrets_manager"
         )
     ).all()

zenml/zen_stores/migrations/versions/10a907dad202_delete_mlmd_tables.py

@@ -20,7 +20,7 @@ def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     # Bye bye MLMD :wave: :wave: :wave:
     conn = op.get_bind()
-    inspector = Inspector.from_engine(conn)
+    inspector = Inspector.from_engine(conn.engine)
     tables = inspector.get_table_names()
     for table in (
         "Artifact",

zenml/zen_stores/migrations/versions/26b776ad583e_redesign_artifacts.py

@@ -6,7 +6,7 @@ Create Date: 2022-11-17 08:00:24.936750
 
 """
 
-from typing import TYPE_CHECKING, Dict
+from typing import Any, Dict
 
 import sqlalchemy as sa
 import sqlmodel
@@ -14,9 +14,6 @@ from alembic import op
 from sqlalchemy import select
 from sqlalchemy.sql.expression import false, true
 
-if TYPE_CHECKING:
-    from sqlalchemy.engine.row import Row
-
 # revision identifiers, used by Alembic.
 revision = "26b776ad583e"
 down_revision = "0.23.0"
@@ -53,13 +50,14 @@ def upgrade() -> None:
     # Migrate data
     # ------------
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
     meta.reflect(
+        bind=op.get_bind(),
         only=(
             "artifacts",
             "step_run_output_artifact",
             "step_run_input_artifact",
-        )
+        ),
     )
     artifacts = sa.Table("artifacts", meta)
     step_run_output_artifact = sa.Table("step_run_output_artifact", meta)
@@ -85,7 +83,7 @@ def upgrade() -> None:
         ).where(artifacts.c.is_cached == true())
     ).fetchall()
 
-    def _find_produced_artifact(cached_artifact: "Row") -> "Row":
+    def _find_produced_artifact(cached_artifact: Any) -> Any:
         """For a given cached artifact, find the original produced artifact.
 
         Args:
@@ -115,9 +113,9 @@ def upgrade() -> None:
         produced_artifact = _find_produced_artifact(cached_artifact)
         cached_to_produced_mapping[cached_artifact.id] = produced_artifact.id
         conn.execute(
-            step_run_input_artifact.update(
-                step_run_input_artifact.c.artifact_id == cached_artifact.id
-            ).values({"artifact_id": produced_artifact.id})
+            step_run_input_artifact.update()
+            .where(step_run_input_artifact.c.artifact_id == cached_artifact.id)
+            .values({"artifact_id": produced_artifact.id})
         )
 
     # Delete all cached artifacts from the artifacts table

zenml/zen_stores/migrations/versions/37835ce041d2_optimizing_database.py

@@ -62,7 +62,7 @@ def upgrade() -> None:
             "pipeline_configuration",
             existing_type=sa.TEXT(),
             type_=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             existing_nullable=False,
         )
@@ -121,7 +121,7 @@ def upgrade() -> None:
         batch_op.alter_column(
             "step_configuration",
             existing_type=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             nullable=True,
         )
@@ -272,7 +272,7 @@ def downgrade() -> None:
         batch_op.alter_column(
             "pipeline_configuration",
             existing_type=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             type_=sa.TEXT(),
             existing_nullable=False,

zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py

@@ -43,18 +43,17 @@ def upgrade() -> None:
     )
 
     # Get metadata from current connection
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
 
     # Pass in tuple with tables we want to reflect, otherwise whole database
     # will get reflected
-    meta.reflect(only=("identity", "server_settings"))
+    meta.reflect(only=("identity", "server_settings"), bind=op.get_bind())
 
     # Fetch the deployment id from the identity table
     deployment_id = (
-        sa.select([meta.tables["identity"].c.id])
-        .limit(1)
-        .execute()
-        .fetchone()[0]
+        op.get_bind()
+        .execute(sa.select(meta.tables["identity"].c.id).limit(1))
+        .scalar_one()
     )
 
     # Prefill the settings table with a single row that contains the deployment
@@ -94,18 +93,17 @@ def downgrade() -> None:
     )
 
     # Get metadata from current connection
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
 
     # Pass in tuple with tables we want to reflect, otherwise whole database
     # will get reflected
-    meta.reflect(only=("identity", "server_settings"))
+    meta.reflect(only=("identity", "server_settings"), bind=op.get_bind())
 
     # Fetch the deployment id from the settings table
     deployment_id = (
-        sa.select([meta.tables["server_settings"].c.id])
-        .limit(1)
-        .execute()
-        .fetchone()[0]
+        op.get_bind()
+        .execute(sa.select(meta.tables["server_settings"].c.id).limit(1))
+        .scalar_one()
     )
 
     # Prefill the identity table with a single row that contains the deployment

zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py

@@ -36,18 +36,19 @@ def upgrade() -> None:
         sa.PrimaryKeyConstraint("name", "role_id"),
     )
     # get metadata from current connection
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
 
     # pass in tuple with tables we want to reflect, otherwise whole database
     #  will get reflected
     meta.reflect(
+        bind=op.get_bind(),
         only=(
             "rolepermissionschema",
             "roleschema",
             "userroleassignmentschema",
             "teamroleassignmentschema",
             "userschema",
-        )
+        ),
     )
 
     # In order to ensure unique names on roles delete potential admin/guest role

zenml/zen_stores/migrations/versions/6917bce75069_add_pipeline_run_unique_constraint.py

@@ -24,8 +24,8 @@ def add_orchestrator_run_id_for_old_runs() -> None:
     this constraint. This is not the case for old pipeline runs which existed
     before we had deployments, so we add a dummy value for those runs.
     """
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("pipeline_run",))
+    meta = sa.MetaData()
+    meta.reflect(bind=op.get_bind(), only=("pipeline_run",))
     run_table = sa.Table("pipeline_run", meta)
     connection = op.get_bind()
 
@@ -48,8 +48,8 @@ def verify_unique_constraint_satisfied() -> None:
         RuntimeError: If there are rows which have identical values for the
             `deployment_id` and `orchestrator_run_id` columns.
     """
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("pipeline_run",))
+    meta = sa.MetaData()
+    meta.reflect(bind=op.get_bind(), only=("pipeline_run",))
     run_table = sa.Table("pipeline_run", meta)
     connection = op.get_bind()
 

zenml/zen_stores/migrations/versions/728c6369cfaa_add_name_column_to_input_artifact_pk.py

@@ -131,7 +131,8 @@ def _disable_primary_key_requirement_if_necessary() -> None:
     Raises:
         NotImplementedError: If the database engine is not MySQL or mariadb.
     """
-    engine = op.get_bind().engine
+    conn = op.get_bind()
+    engine = conn.engine
     engine_name = engine.name.lower()
 
     if engine_name == "mysql":
@@ -141,7 +142,7 @@ def _disable_primary_key_requirement_if_necessary() -> None:
             0,
             13,
         ):
-            potential_session_var = engine.execute(
+            potential_session_var = conn.execute(
                 text('SHOW SESSION VARIABLES LIKE "sql_require_primary_key";')
             ).fetchone()
             if potential_session_var and potential_session_var[1] == "ON":

zenml/zen_stores/migrations/versions/743ec82b1b3c_update_size_of_build_images.py

@@ -25,7 +25,7 @@ def upgrade() -> None:
             "images",
             existing_type=sa.TEXT(),
             type_=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             existing_nullable=False,
         )
@@ -40,7 +40,7 @@ def downgrade() -> None:
         batch_op.alter_column(
             "images",
             existing_type=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             type_=sa.TEXT(),
             existing_nullable=False,