zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240619__py3-none-any.whl

zenml/models/v2/core/service_connector.py

@@ -18,10 +18,11 @@ from datetime import datetime
 from typing import Any, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import Field, SecretStr, root_validator
+from pydantic import Field, SecretStr, ValidationError, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -30,10 +31,10 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import update_model
 from zenml.models.v2.misc.service_connector_type import (
     ServiceConnectorTypeModel,
 )
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
@@ -49,6 +50,7 @@ class ServiceConnectorRequest(WorkspaceScopedRequest):
     )
     connector_type: Union[str, "ServiceConnectorTypeModel"] = Field(
         title="The type of service connector.",
+        union_mode="left_to_right",
     )
     description: str = Field(
         default="",
@@ -99,7 +101,7 @@ class ServiceConnectorRequest(WorkspaceScopedRequest):
         default_factory=dict,
         title="The service connector configuration, not including secrets.",
     )
-    secrets: Dict[str, Optional[SecretStr]] = Field(
+    secrets: Dict[str, Optional[PlainSerializedSecretStr]] = Field(
         default_factory=dict,
         title="The service connector secrets.",
     )
@@ -204,8 +206,7 @@ class ServiceConnectorRequest(WorkspaceScopedRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class ServiceConnectorUpdate(ServiceConnectorRequest):
+class ServiceConnectorUpdate(BaseUpdate):
     """Model used for service connector updates.
 
     Most fields in the update model are optional and will not be updated if
@@ -232,24 +233,169 @@ class ServiceConnectorUpdate(ServiceConnectorRequest):
     have a None default value.
     """
 
-    resource_types: Optional[List[str]] = Field(  # type: ignore[assignment]
+    name: Optional[str] = Field(
+        title="The service connector name.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    connector_type: Optional[Union[str, "ServiceConnectorTypeModel"]] = Field(
+        title="The type of service connector.",
+        default=None,
+        union_mode="left_to_right",
+    )
+    description: Optional[str] = Field(
+        title="The service connector instance description.",
         default=None,
+    )
+    auth_method: Optional[str] = Field(
+        title="The authentication method that the connector instance uses to "
+        "access the resources.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    resource_types: Optional[List[str]] = Field(
         title="The type(s) of resource that the connector instance can be used "
         "to gain access to.",
+        default=None,
     )
-    configuration: Optional[Dict[str, Any]] = Field(  # type: ignore[assignment]
+    resource_id: Optional[str] = Field(
+        title="Uniquely identifies a specific resource instance that the "
+        "connector instance can be used to access. If omitted, the "
+        "connector instance can be used to access any and all resource "
+        "instances that the authentication method and resource type(s) "
+        "allow.",
+        max_length=STR_FIELD_MAX_LENGTH,
         default=None,
-        title="The service connector configuration, not including secrets.",
     )
-    secrets: Optional[Dict[str, Optional[SecretStr]]] = Field(  # type: ignore[assignment]
+    supports_instances: Optional[bool] = Field(
+        title="Indicates whether the connector instance can be used to access "
+        "multiple instances of the configured resource type.",
+        default=None,
+    )
+    expires_at: Optional[datetime] = Field(
+        title="Time when the authentication credentials configured for the "
+        "connector expire. If omitted, the credentials do not expire.",
         default=None,
-        title="The service connector secrets.",
     )
-    labels: Optional[Dict[str, str]] = Field(  # type: ignore[assignment]
+    expires_skew_tolerance: Optional[int] = Field(
+        title="The number of seconds of tolerance to apply when checking "
+        "whether the authentication credentials configured for the "
+        "connector have expired. If omitted, no tolerance is applied.",
         default=None,
+    )
+    expiration_seconds: Optional[int] = Field(
+        title="The duration, in seconds, that the temporary credentials "
+        "generated by this connector should remain valid. Only "
+        "applicable for connectors and authentication methods that "
+        "involve generating temporary credentials from the ones "
+        "configured in the connector.",
+        default=None,
+    )
+    configuration: Optional[Dict[str, Any]] = Field(
+        title="The service connector configuration, not including secrets.",
+        default=None,
+    )
+    secrets: Optional[Dict[str, Optional[PlainSerializedSecretStr]]] = Field(
+        title="The service connector secrets.",
+        default=None,
+    )
+    labels: Optional[Dict[str, str]] = Field(
         title="Service connector labels.",
+        default=None,
     )
 
+    # Analytics
+    ANALYTICS_FIELDS: ClassVar[List[str]] = [
+        "connector_type",
+        "auth_method",
+        "resource_types",
+    ]
+
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Format the resource types in the analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        if self.resource_types is not None:
+            if len(self.resource_types) == 1:
+                metadata["resource_types"] = self.resource_types[0]
+            else:
+                metadata["resource_types"] = ", ".join(self.resource_types)
+
+        if self.connector_type is not None:
+            metadata["connector_type"] = self.type
+
+        return metadata
+
+    # Helper methods
+    @property
+    def type(self) -> Optional[str]:
+        """Get the connector type.
+
+        Returns:
+            The connector type.
+        """
+        if self.connector_type is not None:
+            if isinstance(self.connector_type, str):
+                return self.connector_type
+            return self.connector_type.connector_type
+        return None
+
+    def validate_and_configure_resources(
+        self,
+        connector_type: "ServiceConnectorTypeModel",
+        resource_types: Optional[Union[str, List[str]]] = None,
+        resource_id: Optional[str] = None,
+        configuration: Optional[Dict[str, Any]] = None,
+        secrets: Optional[Dict[str, Optional[SecretStr]]] = None,
+    ) -> None:
+        """Validate and configure the resources that the connector can be used to access.
+
+        Args:
+            connector_type: The connector type specification used to validate
+                the connector configuration.
+            resource_types: The type(s) of resource that the connector instance
+                can be used to access. If omitted, a multi-type connector is
+                configured.
+            resource_id: Uniquely identifies a specific resource instance that
+                the connector instance can be used to access.
+            configuration: The connector configuration.
+            secrets: The connector secrets.
+        """
+        _validate_and_configure_resources(
+            connector=self,
+            connector_type=connector_type,
+            resource_types=resource_types,
+            resource_id=resource_id,
+            configuration=configuration,
+            secrets=secrets,
+        )
+
+    def convert_to_request(self) -> "ServiceConnectorRequest":
+        """Method to generate a service connector request object from self.
+
+        For certain operations, the service connector update model need to
+        adhere to the limitations set by the request model. In order to use
+        update models in such situations, we need to be able to convert an
+        update model into a request model.
+
+        Returns:
+            The equivalent request model
+
+        Raises:
+            RuntimeError: if the model can not be converted to a request model.
+        """
+        try:
+            return ServiceConnectorRequest.model_validate(self.model_dump())
+        except ValidationError as e:
+            raise RuntimeError(
+                "The service connector update model can not be converted into "
+                f"an equivalent request model: {e}"
+            )
+
 
 # ------------------ Response Model ------------------
 
@@ -262,7 +408,7 @@ class ServiceConnectorResponseBody(WorkspaceScopedResponseBody):
         title="The service connector instance description.",
     )
     connector_type: Union[str, "ServiceConnectorTypeModel"] = Field(
-        title="The type of service connector.",
+        title="The type of service connector.", union_mode="left_to_right"
     )
     auth_method: str = Field(
         title="The authentication method that the connector instance uses to "
@@ -319,7 +465,7 @@ class ServiceConnectorResponseMetadata(WorkspaceScopedResponseMetadata):
         "connectors and authentication methods that involve generating "
         "temporary credentials from the ones configured in the connector.",
     )
-    secrets: Dict[str, Optional[SecretStr]] = Field(
+    secrets: Dict[str, Optional[PlainSerializedSecretStr]] = Field(
         default_factory=dict,
         title="The service connector secrets.",
     )
@@ -644,10 +790,14 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         description="The type of service connector to filter by",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace to filter by"
+        default=None,
+        description="Workspace to filter by",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User to filter by"
+        default=None,
+        description="User to filter by",
+        union_mode="left_to_right",
     )
     auth_method: Optional[str] = Field(
         default=None,
@@ -678,56 +828,49 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         default=None,
         title="Filter by the ID of the secret that contains the service "
         "connector's credentials",
+        union_mode="left_to_right",
     )
 
     # Use this internally to configure and access the labels as a dictionary
     labels: Optional[Dict[str, Optional[str]]] = Field(
         default=None,
         title="The labels to filter by, as a dictionary",
+        exclude=True,
     )
 
-    @root_validator
-    def validate_labels(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def validate_labels(self) -> "ServiceConnectorFilter":
         """Parse the labels string into a label dictionary and vice-versa.
 
-        Args:
-            values: The values to validate.
-
         Returns:
             The validated values.
         """
-        labels_str = values.get("labels_str")
-        labels = values.get("labels")
-        if labels_str is not None:
+        if self.labels_str is not None:
             try:
-                values["labels"] = json.loads(labels_str)
+                self.labels = json.loads(self.labels_str)
             except json.JSONDecodeError:
                 # Interpret as comma-separated values instead
-                values["labels"] = {
+                self.labels = {
                     label.split("=", 1)[0]: label.split("=", 1)[1]
                     if "=" in label
                     else None
-                    for label in labels_str.split(",")
+                    for label in self.labels_str.split(",")
                 }
-        elif labels is not None:
-            values["labels_str"] = json.dumps(values["labels"])
-
-        return values
+        elif self.labels is not None:
+            self.labels_str = json.dumps(self.labels)
 
-    class Config:
-        """Pydantic config class."""
-
-        # Exclude the labels field from the serialized response
-        # (it is only used internally). The labels_str field is a string
-        # representation of the labels that can be used in the API.
-        exclude = ["labels"]
+        return self
 
 
 # ------------------ Helper Functions ------------------
 
 
 def _validate_and_configure_resources(
-    connector: Union[ServiceConnectorRequest, ServiceConnectorResponse],
+    connector: Union[
+        ServiceConnectorRequest,
+        ServiceConnectorUpdate,
+        ServiceConnectorResponse,
+    ],
     connector_type: "ServiceConnectorTypeModel",
     resource_types: Optional[Union[str, List[str]]] = None,
     resource_id: Optional[str] = None,
@@ -757,17 +900,24 @@ def _validate_and_configure_resources(
     # connector model itself, while for the response model, they are in the
     # metadata field.
     update_connector_metadata: Union[
-        ServiceConnectorRequest, ServiceConnectorResponseMetadata
+        ServiceConnectorRequest,
+        ServiceConnectorUpdate,
+        ServiceConnectorResponseMetadata,
     ]
     update_connector_body: Union[
-        ServiceConnectorRequest, ServiceConnectorResponseBody
+        ServiceConnectorRequest,
+        ServiceConnectorUpdate,
+        ServiceConnectorResponseBody,
     ]
     if isinstance(connector, ServiceConnectorRequest):
         update_connector_metadata = connector
         update_connector_body = connector
+    elif isinstance(connector, ServiceConnectorUpdate):
+        update_connector_metadata = connector
+        update_connector_body = connector
     else:
         # Updating service connector responses must only be done on hydrated
-        # instances, otherwise the metadata will be missing and we risk calling
+        # instances, otherwise the metadata will be missing, and we risk calling
         # the ZenML store to update the connector with additional information.
         # This is just a safety measure, but it will never happen because
         # this method will always be called on a hydrated response.
@@ -792,6 +942,7 @@ def _validate_and_configure_resources(
     try:
         # Validate the connector configuration and retrieve the resource
         # specification
+        assert connector.auth_method is not None
         (
             auth_method_spec,
             resource_spec,
@@ -835,8 +986,27 @@ def _validate_and_configure_resources(
         required = attr_name in auth_method_spec.config_schema.get(
             "required", []
         )
-        secret = attr_schema.get("format", "") == "password"
-        attr_type = attr_schema.get("type", "string")
+
+        attr_any_of = attr_schema.get("anyOf", [])
+
+        if attr_any_of:
+            no_null_attr_any_of = [
+                a for a in attr_any_of if a.get("type", "string") != "null"
+            ]
+
+            if len(no_null_attr_any_of) == 1:
+                secret = no_null_attr_any_of[0].get("format", "") == "password"
+                attr_type = no_null_attr_any_of[0].get("type", "string")
+            else:
+                # TODO: Still not sure what needs to happen here. We will
+                #   only end up here if the auth method config schema has a
+                #   field with a Union[SecretStr, int] or something.
+                raise RuntimeError("Service connector schema error.")
+
+        else:
+            secret = attr_schema.get("format", "") == "password"
+            attr_type = attr_schema.get("type", "string")
+
         value = configuration.get(attr_name, secrets.get(attr_name))
         if required:
             if value is None:
@@ -869,7 +1039,12 @@ def _validate_and_configure_resources(
             f"{supported_attrs}",
         )
     # Warn about secrets that are not part of the configuration schema
-    for attr_name in set(secrets.keys()) - connector.secrets.keys():
+    connector_secrets = (
+        set(connector.secrets.keys())
+        if connector.secrets is not None
+        else set()
+    )
+    for attr_name in set(secrets.keys()) - connector_secrets:
         logger.warning(
             f"Ignoring unknown attribute in connector '{connector.name}' "
             f"configuration {attr_name}. Supported attributes are: "

zenml/models/v2/core/stack.py

@@ -14,16 +14,15 @@
 """Models representing stacks."""
 
 import json
-from typing import Any, ClassVar, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import Field
-from sqlalchemy import and_
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from sqlmodel import and_
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import StackComponentType
-from zenml.models.v2.base.internal import server_owned_request_model
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -32,9 +31,11 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import update_model
 from zenml.models.v2.core.component import ComponentResponse
 
+if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
 # ------------------ Request Model ------------------
 
 
@@ -74,20 +75,41 @@ class StackRequest(WorkspaceScopedRequest):
         )
 
 
-@server_owned_request_model
 class InternalStackRequest(StackRequest):
     """Internal stack request model."""
 
-    pass
+    user: Optional[UUID] = Field(  # type: ignore[assignment]
+        title="The id of the user that created this resource.",
+        default=None,
+    )
 
 
 # ------------------ Update Model ------------------
 
 
-@update_model
-class StackUpdate(StackRequest):
+class StackUpdate(BaseUpdate):
     """Update model for stacks."""
 
+    name: Optional[str] = Field(
+        title="The name of the stack.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    description: Optional[str] = Field(
+        title="The description of the stack",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    stack_spec_path: Optional[str] = Field(
+        title="The path to the stack spec used for mlstacks deployments.",
+        default=None,
+    )
+    components: Optional[Dict[StackComponentType, List[UUID]]] = Field(
+        title="A mapping of stack component types to the actual"
+        "instances of components of this type.",
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 
@@ -167,7 +189,9 @@ class StackResponse(
                 flavor=component.flavor,
             )
             configuration = json.loads(
-                component.get_metadata().json(include={"configuration"})
+                component.get_metadata().model_dump_json(
+                    include={"configuration"}
+                )
             )
             component_dict.update(configuration)
 
@@ -248,18 +272,22 @@ class StackFilter(WorkspaceScopedFilter):
         default=None, description="Description of the stack"
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the stack"
+        default=None,
+        description="Workspace of the stack",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the stack"
+        default=None,
+        description="User of the stack",
+        union_mode="left_to_right",
     )
     component_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Component in the stack"
+        default=None,
+        description="Component in the stack",
+        union_mode="left_to_right",
     )
 
-    def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
         Returns:
@@ -273,7 +301,7 @@ class StackFilter(WorkspaceScopedFilter):
         )
 
         if self.component_id:
-            component_id_filter = and_(  # type: ignore[type-var]
+            component_id_filter = and_(
                 StackCompositionSchema.stack_id == StackSchema.id,
                 StackCompositionSchema.component_id == self.component_id,
             )

zenml/models/v2/core/step_run.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.config.step_configurations import StepConfiguration, StepSpec
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
@@ -220,7 +220,15 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
 class StepRunResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the step run entity."""
 
-    model_version: Optional[ModelVersionResponse]
+    model_version: Optional[ModelVersionResponse] = None
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class StepRunResponse(
@@ -471,20 +479,32 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Status of the Step Run",
     )
     start_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="Start time for this run"
+        default=None,
+        description="Start time for this run",
+        union_mode="left_to_right",
     )
     end_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="End time for this run"
+        default=None,
+        description="End time for this run",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Pipeline run of this step run"
+        default=None,
+        description="Pipeline run of this step run",
+        union_mode="left_to_right",
     )
     original_step_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Original id for this step run"
+        default=None,
+        description="Original id for this step run",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that produced this step run"
+        default=None,
+        description="User that produced this step run",
+        union_mode="left_to_right",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of this step run"
+        default=None,
+        description="Workspace of this step run",
+        union_mode="left_to_right",
     )

zenml/models/v2/core/tag.py

@@ -51,8 +51,8 @@ class TagRequest(BaseRequest):
 class TagUpdate(BaseModel):
     """Update model for tags."""
 
-    name: Optional[str]
-    color: Optional[ColorVariants]
+    name: Optional[str] = None
+    color: Optional[ColorVariants] = None
 
 
 # ------------------ Response Model ------------------
@@ -121,5 +121,9 @@ class TagResponse(
 class TagFilter(BaseFilter):
     """Model to enable advanced filtering of all tags."""
 
-    name: Optional[str]
-    color: Optional[ColorVariants]
+    name: Optional[str] = Field(
+        description="The unique title of the tag.", default=None
+    )
+    color: Optional[ColorVariants] = Field(
+        description="The color variant assigned to the tag.", default=None
+    )

zenml/models/v2/core/trigger.py

@@ -68,7 +68,6 @@ class TriggerBase(BaseModel):
     )
     action_subtype: PluginSubType = Field(
         title="The subtype of the action that is executed by this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
     )
     service_account_id: UUID = Field(
         title="The service account that is used to execute the action.",
@@ -164,7 +163,6 @@ class TriggerResponseBody(WorkspaceScopedResponseBody):
     )
     action_subtype: PluginSubType = Field(
         title="The subtype of the action that is executed by this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
     )
     is_active: bool = Field(
         title="Whether the trigger is active.",
@@ -341,6 +339,7 @@ class TriggerFilter(WorkspaceScopedFilter):
     event_source_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="By the event source this trigger is attached to.",
+        union_mode="left_to_right",
     )
     is_active: Optional[bool] = Field(
         default=None,
@@ -358,6 +357,7 @@ class TriggerFilter(WorkspaceScopedFilter):
     resource_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="By the resource this trigger references.",
+        union_mode="left_to_right",
     )
     resource_type: Optional[str] = Field(
         default=None,

zenml/models/v2/core/trigger_execution.py

@@ -115,4 +115,5 @@ class TriggerExecutionFilter(WorkspaceScopedFilter):
     trigger_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="ID of the trigger of the execution.",
+        union_mode="left_to_right",
     )