content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Amazon S3 artifact store flavor."""
  
  import re
  from typing import (
      TYPE_CHECKING,
      Any,
      ClassVar,
      Dict,
      Optional,
      Set,
      Type,
  )
  
  from pydantic import Field, field_validator
  
  from zenml.artifact_stores import (
      BaseArtifactStoreConfig,
      BaseArtifactStoreFlavor,
  )
  from zenml.integrations.s3 import S3_ARTIFACT_STORE_FLAVOR
  from zenml.integrations.s3.utils import split_s3_path
  from zenml.models import ServiceConnectorRequirements
  from zenml.stack.authentication_mixin import AuthenticationConfigMixin
  from zenml.utils.networking_utils import (
      replace_localhost_with_internal_hostname,
  )
  from zenml.utils.secret_utils import SecretField
  
  if TYPE_CHECKING:
      from zenml.integrations.s3.artifact_stores import S3ArtifactStore
  
  
  class S3ArtifactStoreConfig(
      BaseArtifactStoreConfig, AuthenticationConfigMixin
  ):
      """Configuration for the S3 Artifact Store.
  
      All attributes of this class except `path` will be passed to the
      `s3fs.S3FileSystem` initialization. See
      [here](https://s3fs.readthedocs.io/en/latest/) for more information on how
      to use those configuration options to connect to any S3-compatible storage.
  
      When you want to register an S3ArtifactStore from the CLI and need to pass
      `client_kwargs`, `config_kwargs` or `s3_additional_kwargs`, you should pass
      them as a json string:
      ```
      zenml artifact-store register my_s3_store --flavor=s3 \
      --path=s3://my_bucket --client_kwargs='{"endpoint_url": "http://my-s3-endpoint"}'
      ```
      """
  
      SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"s3://"}
  
      key: Optional[str] = SecretField(
          default=None,
          description="AWS access key ID for authentication. "
          "If not provided, credentials will be inferred from the environment.",
      )
      secret: Optional[str] = SecretField(
          default=None,
          description="AWS secret access key for authentication. "
          "If not provided, credentials will be inferred from the environment.",
      )
      token: Optional[str] = SecretField(
          default=None,
          description="AWS session token for temporary credentials. "
          "If not provided, credentials will be inferred from the environment.",
      )
      client_kwargs: Optional[Dict[str, Any]] = Field(
          None,
          description="Additional keyword arguments to pass to the S3 client. "
          "For example, to connect to a custom S3-compatible endpoint: "
          "{'endpoint_url': 'http://minio:9000'}",
      )
      config_kwargs: Optional[Dict[str, Any]] = Field(
          None,
          description="Additional keyword arguments to pass to the S3 client configuration. "
          "For example: {'region_name': 'us-west-2', 'signature_version': 's3v4'}",
      )
      s3_additional_kwargs: Optional[Dict[str, Any]] = Field(
          None,
          description="Additional keyword arguments for S3 operations. "
          "For example: {'ACL': 'bucket-owner-full-control'}",
      )
  
      _bucket: Optional[str] = None
  
      @field_validator("client_kwargs")
      @classmethod
      def _validate_client_kwargs(
          cls, value: Optional[Dict[str, Any]]
      ) -> Optional[Dict[str, Any]]:
          """Validates the `client_kwargs` attribute.
  
          Args:
              value: The value to validate.
  
          Raises:
              ValueError: If the value is not a valid URL.
  
          Returns:
              The validated value.
          """
          if value is None:
              return value
  
          if "endpoint_url" in value and value["endpoint_url"]:
              url = value["endpoint_url"].rstrip("/")
              scheme = re.search("^([a-z0-9]+://)", url)
              if scheme is None or scheme.group() not in ("https://", "http://"):
                  raise ValueError(
                      "Invalid URL for endpoint url: {url}. Should be in the form "
                      "https://hostname[:port] or http://hostname[:port]."
                  )
  
              # When running inside a container, if the URL uses localhost, the
              # target service will not be available. We try to replace localhost
              # with one of the special Docker or K3D internal hostnames.
              value["endpoint_url"] = replace_localhost_with_internal_hostname(
                  url
              )
          return value
  
      @property
      def bucket(self) -> str:
          """The bucket name of the artifact store.
  
          Returns:
              The bucket name of the artifact store.
          """
          if self._bucket is None:
              self._bucket, _ = split_s3_path(self.path)
          return self._bucket
  
  
  class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
      """Flavor of the S3 artifact store."""
  
      @property
      def name(self) -> str:
          """Name of the flavor.
  
          Returns:
              The name of the flavor.
          """
          return S3_ARTIFACT_STORE_FLAVOR
  
      @property
      def display_name(self) -> str:
          """Display name of the flavor.
  
          Returns:
              The display name of the flavor.
          """
          return "S3"
  
      @property
      def service_connector_requirements(
          self,
      ) -> Optional[ServiceConnectorRequirements]:
          """Service connector resource requirements for service connectors.
  
          Specifies resource requirements that are used to filter the available
          service connector types that are compatible with this flavor.
  
          Returns:
              Requirements for compatible service connectors, if a service
              connector is required for this flavor.
          """
          return ServiceConnectorRequirements(
              resource_type="s3-bucket",
              resource_id_attr="path",
          )
  
      @property
      def docs_url(self) -> Optional[str]:
          """A URL to point at docs explaining this flavor.
  
          Returns:
              A flavor docs url.
          """
          return self.generate_default_docs_url()
  
      @property
      def sdk_docs_url(self) -> Optional[str]:
          """A URL to point at SDK docs explaining this flavor.
  
          Returns:
              A flavor SDK docs url.
          """
          return self.generate_default_sdk_docs_url()
  
      @property
      def logo_url(self) -> str:
          """A URL to represent the flavor in the dashboard.
  
          Returns:
              The flavor logo.
          """
          return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/artifact_store/aws.png"
  
      @property
      def config_class(self) -> Type[S3ArtifactStoreConfig]:
          """The config class of the flavor.
  
          Returns:
              The config class of the flavor.
          """
          return S3ArtifactStoreConfig
  
      @property
      def implementation_class(self) -> Type["S3ArtifactStore"]:
          """Implementation class for this flavor.
  
          Returns:
              The implementation class for this flavor.
          """
          from zenml.integrations.s3.artifact_stores import S3ArtifactStore
  
          return S3ArtifactStore
  