zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/models/v2/core/service_connector.py

@@ -18,21 +18,23 @@ 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,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     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__)
 
@@ -48,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="",
@@ -98,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.",
     )
@@ -203,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
@@ -231,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,
+    )
+    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,
     )
-    configuration: Optional[Dict[str, Any]] = 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 configuration, not including secrets.",
     )
-    secrets: Optional[Dict[str, Optional[SecretStr]]] = 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,
-        title="The service connector secrets.",
     )
-    labels: Optional[Dict[str, str]] = Field(  # type: ignore[assignment]
+    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 ------------------
 
@@ -261,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 "
@@ -318,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.",
     )
@@ -328,18 +475,45 @@ class ServiceConnectorResponseMetadata(WorkspaceScopedResponseMetadata):
     )
 
 
+class ServiceConnectorResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the service connector entity."""
+
+
 class ServiceConnectorResponse(
     WorkspaceScopedResponse[
-        ServiceConnectorResponseBody, ServiceConnectorResponseMetadata
+        ServiceConnectorResponseBody,
+        ServiceConnectorResponseMetadata,
+        ServiceConnectorResponseResources,
     ]
 ):
     """Response model for service connectors."""
 
+    # Disable the warning for updating responses, because we update the
+    # service connector type in place
+    _warn_on_response_updates: bool = False
+
     name: str = Field(
         title="The service connector name.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the service connector labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        metadata.update(
+            {
+                label[6:]: value
+                for label, value in self.labels.items()
+                if label.startswith("zenml:")
+            }
+        )
+        return metadata
+
     def get_hydrated_version(self) -> "ServiceConnectorResponse":
         """Get the hydrated version of this service connector.
 
@@ -627,7 +801,6 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         default=None,
         description="The type to scope this query to.",
     )
-
     name: Optional[str] = Field(
         default=None,
         description="The name to filter by",
@@ -636,12 +809,6 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         default=None,
         description="The type of service connector to filter by",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace to filter by"
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User to filter by"
-    )
     auth_method: Optional[str] = Field(
         default=None,
         title="Filter by the authentication method configured for the "
@@ -671,56 +838,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,
@@ -750,17 +910,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.
@@ -785,6 +952,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,
@@ -828,7 +996,27 @@ def _validate_and_configure_resources(
         required = attr_name in auth_method_spec.config_schema.get(
             "required", []
         )
-        secret = attr_schema.get("format", "") == "password"
+
+        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:
@@ -846,6 +1034,11 @@ def _validate_and_configure_resources(
             else:
                 update_connector_metadata.secrets[attr_name] = SecretStr(value)
         else:
+            if attr_type == "array" and isinstance(value, str):
+                try:
+                    value = json.loads(value)
+                except json.decoder.JSONDecodeError:
+                    value = value.split(",")
             update_connector_metadata.configuration[attr_name] = value
 
     # Warn about attributes that are not part of the configuration schema
@@ -856,7 +1049,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: "