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

zenml/zen_stores/secrets_stores/aws_secrets_store.py

@@ -25,7 +25,7 @@ from uuid import UUID
 
 import boto3
 from botocore.exceptions import ClientError
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -38,6 +38,7 @@ from zenml.integrations.aws.service_connectors.aws_service_connector import (
     AWSAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -74,20 +75,22 @@ class AWSSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
 
         raise ValueError("AWS `region` must be specified in the auth_config.")
 
-    @root_validator(pre=True)
-    def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
         """
         # Search for legacy attributes and populate the connector configuration
         # from them, if they exist.
-        if values.get("region_name"):
-            if not values.get("aws_access_key_id") or not values.get(
+        if data.get("region_name"):
+            if not data.get("aws_access_key_id") or not data.get(
                 "aws_secret_access_key"
             ):
                 logger.warning(
@@ -96,9 +99,9 @@ class AWSSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
                     "the `auth_method` and `auth_config` attributes instead. "
                     "Using an implicit authentication method for AWS Secrets."
                 )
-                values["auth_method"] = AWSAuthenticationMethods.IMPLICIT
-                values["auth_config"] = dict(
-                    region=values.get("region_name"),
+                data["auth_method"] = AWSAuthenticationMethods.IMPLICIT
+                data["auth_config"] = dict(
+                    region=data.get("region_name"),
                 )
             else:
                 logger.warning(
@@ -107,20 +110,16 @@ class AWSSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
                     "will be removed in a future version of ZenML. Please use the "
                     "`auth_method` and `auth_config` attributes instead."
                 )
-                values["auth_method"] = AWSAuthenticationMethods.SECRET_KEY
-                values["auth_config"] = dict(
-                    aws_access_key_id=values.get("aws_access_key_id"),
-                    aws_secret_access_key=values.get("aws_secret_access_key"),
-                    region=values.get("region_name"),
+                data["auth_method"] = AWSAuthenticationMethods.SECRET_KEY
+                data["auth_config"] = dict(
+                    aws_access_key_id=data.get("aws_access_key_id"),
+                    aws_secret_access_key=data.get("aws_secret_access_key"),
+                    region=data.get("region_name"),
                 )
 
-        return values
+        return data
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # Allow extra attributes set in the class.
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class AWSSecretsStore(ServiceConnectorSecretsStore):

zenml/zen_stores/secrets_stores/azure_secrets_store.py

@@ -27,7 +27,7 @@ from uuid import UUID
 from azure.core.credentials import TokenCredential
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azure.keyvault.secrets import SecretClient
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -40,6 +40,7 @@ from zenml.integrations.azure.service_connectors.azure_service_connector import
     AzureAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -65,12 +66,14 @@ class AzureSecretsStoreConfiguration(
     type: SecretsStoreType = SecretsStoreType.AZURE
     key_vault_name: str
 
-    @root_validator(pre=True)
-    def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
@@ -78,9 +81,9 @@ class AzureSecretsStoreConfiguration(
         # Search for legacy attributes and populate the connector configuration
         # from them, if they exist.
         if (
-            values.get("azure_client_id")
-            and values.get("azure_client_secret")
-            and values.get("azure_tenant_id")
+            data.get("azure_client_id")
+            and data.get("azure_client_secret")
+            and data.get("azure_tenant_id")
         ):
             logger.warning(
                 "The `azure_client_id`, `azure_client_secret` and "
@@ -88,22 +91,16 @@ class AzureSecretsStoreConfiguration(
                 "removed in a future version or ZenML. Please use the "
                 "`auth_method` and `auth_config` attributes instead."
             )
-            values["auth_method"] = (
-                AzureAuthenticationMethods.SERVICE_PRINCIPAL
+            data["auth_method"] = AzureAuthenticationMethods.SERVICE_PRINCIPAL
+            data["auth_config"] = dict(
+                client_id=data.get("azure_client_id"),
+                client_secret=data.get("azure_client_secret"),
+                tenant_id=data.get("azure_tenant_id"),
             )
-            values["auth_config"] = dict(
-                client_id=values.get("azure_client_id"),
-                client_secret=values.get("azure_client_secret"),
-                tenant_id=values.get("azure_tenant_id"),
-            )
-
-        return values
 
-    class Config:
-        """Pydantic configuration class."""
+        return data
 
-        # Forbid extra attributes set in the class.
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class AzureSecretsStore(ServiceConnectorSecretsStore):

zenml/zen_stores/secrets_stores/base_secrets_store.py

@@ -24,12 +24,13 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.enums import SecretsStoreType
 from zenml.logger import get_logger
 from zenml.utils import source_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.secrets_store_interface import (
     SecretsStoreInterface,
 )
@@ -58,6 +59,79 @@ class BaseSecretsStore(BaseModel, SecretsStoreInterface, ABC):
     TYPE: ClassVar[SecretsStoreType]
     CONFIG_TYPE: ClassVar[Type[SecretsStoreConfiguration]]
 
+    @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 == SecretsStoreType.SQL:
+            from zenml.zen_stores.secrets_stores.sql_secrets_store import (
+                SqlSecretsStoreConfiguration,
+            )
+
+            data["config"] = SqlSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.GCP:
+            from zenml.zen_stores.secrets_stores.gcp_secrets_store import (
+                GCPSecretsStoreConfiguration,
+            )
+
+            data["config"] = GCPSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.AWS:
+            from zenml.zen_stores.secrets_stores.aws_secrets_store import (
+                AWSSecretsStoreConfiguration,
+            )
+
+            data["config"] = AWSSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.AZURE:
+            from zenml.zen_stores.secrets_stores.azure_secrets_store import (
+                AzureSecretsStoreConfiguration,
+            )
+
+            data["config"] = AzureSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.HASHICORP:
+            from zenml.zen_stores.secrets_stores.hashicorp_secrets_store import (
+                HashiCorpVaultSecretsStoreConfiguration,
+            )
+
+            data["config"] = HashiCorpVaultSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+        elif (
+            data["config"].type == SecretsStoreType.CUSTOM
+            or data["config"].type == SecretsStoreType.NONE
+        ):
+            pass
+        else:
+            raise ValueError(
+                f"Unknown type '{data['config'].type}' for the configuration."
+            )
+
+        return data
+
     # ---------------------------------
     # Initialization and configuration
     # ---------------------------------
@@ -297,14 +371,7 @@ class BaseSecretsStore(BaseModel, SecretsStoreInterface, ABC):
                 f"expected {secret_id}, got {stored_secret_id}"
             )
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # Validate attributes when assigning them. We need to set this in order
-        # to have a mix of mutable and immutable attributes
-        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
+    model_config = ConfigDict(
+        validate_assignment=True,
+        extra="ignore",
+    )

zenml/zen_stores/secrets_stores/gcp_secrets_store.py

@@ -27,7 +27,7 @@ from uuid import UUID
 
 from google.api_core import exceptions as google_exceptions
 from google.cloud.secretmanager import SecretManagerServiceClient
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -40,6 +40,7 @@ from zenml.integrations.gcp.service_connectors.gcp_service_connector import (
     GCPAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -78,19 +79,21 @@ class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
 
         raise ValueError("GCP `project_id` must be specified in auth_config.")
 
-    @root_validator(pre=True)
-    def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
         """
         # Search for legacy attributes and populate the connector configuration
         # from them, if they exist.
-        if values.get("project_id"):
+        if data.get("project_id"):
             if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
                 logger.warning(
                     "The `project_id` GCP secrets store attribute is "
@@ -99,9 +102,9 @@ class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
                     "instead. Using an implicit GCP authentication to access "
                     "the GCP Secrets Manager API."
                 )
-                values["auth_method"] = GCPAuthenticationMethods.IMPLICIT
-                values["auth_config"] = dict(
-                    project_id=values.get("project_id"),
+                data["auth_method"] = GCPAuthenticationMethods.IMPLICIT
+                data["auth_config"] = dict(
+                    project_id=data.get("project_id"),
                 )
             else:
                 logger.warning(
@@ -111,23 +114,17 @@ class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
                     "Please use the `auth_method` and `auth_config` attributes "
                     "instead."
                 )
-                values["auth_method"] = (
-                    GCPAuthenticationMethods.SERVICE_ACCOUNT
-                )
-                values["auth_config"] = dict(
-                    project_id=values.get("project_id"),
+                data["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT
+                data["auth_config"] = dict(
+                    project_id=data.get("project_id"),
                 )
                 # Load the service account credentials from the file
                 with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f:
-                    values["auth_config"]["service_account_json"] = f.read()
-
-        return values
+                    data["auth_config"]["service_account_json"] = f.read()
 
-    class Config:
-        """Pydantic configuration class."""
+        return data
 
-        # Forbid extra attributes set in the class.
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class GCPSecretsStore(ServiceConnectorSecretsStore):

zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

@@ -26,13 +26,14 @@ from hvac.exceptions import (  # type: ignore[import-untyped]
     InvalidPath,
     VaultError,
 )
-from pydantic import SecretStr
+from pydantic import ConfigDict
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.enums import (
     SecretsStoreType,
 )
 from zenml.logger import get_logger
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 from zenml.zen_stores.secrets_stores.base_secrets_store import (
     BaseSecretsStore,
 )
@@ -63,16 +64,11 @@ class HashiCorpVaultSecretsStoreConfiguration(SecretsStoreConfiguration):
     type: SecretsStoreType = SecretsStoreType.HASHICORP
 
     vault_addr: str
-    vault_token: Optional[SecretStr] = None
+    vault_token: Optional[PlainSerializedSecretStr] = None
     vault_namespace: Optional[str] = None
     mount_point: Optional[str] = None
     max_versions: int = 1
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Forbid extra attributes set in the class.
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class HashiCorpVaultSecretsStore(BaseSecretsStore):

zenml/zen_stores/secrets_stores/service_connector_secrets_store.py

@@ -25,7 +25,7 @@ from typing import (
 )
 from uuid import uuid4
 
-from pydantic import Field, root_validator
+from pydantic import Field, model_validator
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.logger import get_logger
@@ -36,6 +36,7 @@ from zenml.service_connectors.service_connector import ServiceConnector
 from zenml.service_connectors.service_connector_registry import (
     service_connector_registry,
 )
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.base_secrets_store import (
     BaseSecretsStore,
 )
@@ -54,12 +55,14 @@ class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration):
     auth_method: str
     auth_config: Dict[str, Any] = Field(default_factory=dict)
 
-    @root_validator(pre=True)
-    def validate_auth_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def validate_auth_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Convert the authentication configuration if given in JSON format.
 
         Args:
-            values: The configuration values.
+            data: The configuration values.
 
         Returns:
             The validated configuration values.
@@ -68,15 +71,15 @@ class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration):
             ValueError: If the authentication configuration is not a valid
                 JSON object.
         """
-        if isinstance(values.get("auth_config"), str):
+        if isinstance(data.get("auth_config"), str):
             try:
-                values["auth_config"] = json.loads(values["auth_config"])
+                data["auth_config"] = json.loads(data["auth_config"])
             except json.JSONDecodeError as e:
                 raise ValueError(
                     f"The authentication configuration is not a valid JSON "
                     f"object: {e}"
                 )
-        return values
+        return data
 
 
 class ServiceConnectorSecretsStore(BaseSecretsStore):

zenml/zen_stores/secrets_stores/sql_secrets_store.py

@@ -23,6 +23,7 @@ from typing import (
 )
 from uuid import UUID
 
+from pydantic import ConfigDict
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import NoResultFound
 from sqlalchemy_utils.types.encrypted.encrypted_type import AesGcmEngine
@@ -62,16 +63,14 @@ class SqlSecretsStoreConfiguration(SecretsStoreConfiguration):
 
     type: SecretsStoreType = SecretsStoreType.SQL
     encryption_key: Optional[str] = None
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Don't validate attributes when assigning them. This is necessary
         # because the certificate attributes can be expanded to the contents
         # of the certificate files.
-        validate_assignment = False
+        validate_assignment=False,
         # Forbid extra attributes set in the class.
-        extra = "forbid"
+        extra="forbid",
+    )
 
 
 class SqlSecretsStore(BaseSecretsStore):