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.
  """The base interface to extend the ZenML artifact store."""
  
  import inspect
  import os
  import textwrap
  from abc import abstractmethod
  from pathlib import Path
  from typing import (
      Any,
      Callable,
      ClassVar,
      Dict,
      Iterable,
      List,
      Optional,
      Set,
      Tuple,
      Type,
      Union,
      cast,
  )
  
  from pydantic import Field, model_validator
  
  from zenml.constants import (
      ENV_ZENML_SERVER,
      ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS,
      IN_MEMORY_ARTIFACT_URI_PREFIX,
      handle_bool_env_var,
  )
  from zenml.enums import StackComponentType
  from zenml.exceptions import ArtifactStoreInterfaceError, IllegalOperationError
  from zenml.io import fileio
  from zenml.logger import get_logger
  from zenml.stack import Flavor, StackComponent, StackComponentConfig
  from zenml.utils import io_utils
  from zenml.utils.pydantic_utils import before_validator_handler
  
  logger = get_logger(__name__)
  
  PathType = Union[bytes, str]
  
  
  class _sanitize_paths:
      """Sanitizes path inputs before calling the original function.
  
      Extra decoration layer is needed to pass in fixed artifact store root
      path for static methods that are called on filesystems directly.
  
      Args:
          func: The function to decorate.
          fixed_root_path: The fixed artifact store root path.
          is_static: Whether the function is static or not.
  
      Returns:
          Function that calls the input function with sanitized path inputs.
      """
  
      def __init__(self, func: Callable[..., Any], fixed_root_path: str) -> None:
          """Initializes the decorator.
  
          Args:
              func: The function to decorate.
              fixed_root_path: The fixed artifact store root path.
          """
          self.func = func
          self.fixed_root_path = fixed_root_path
          if ENV_ZENML_SERVER in os.environ:
              self.allow_local_file_access = handle_bool_env_var(
                  ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS, False
              )
          else:
              self.allow_local_file_access = True
  
          self.path_args: List[int] = []
          self.path_kwargs: List[str] = []
          for i, param in enumerate(
              inspect.signature(self.func).parameters.values()
          ):
              if param.annotation == PathType:
                  self.path_kwargs.append(param.name)
                  if param.default == inspect.Parameter.empty:
                      self.path_args.append(i)
  
      def _validate_path(self, path: str) -> None:
          """Validates a path.
  
          Args:
              path: The path to validate.
  
          Raises:
              FileNotFoundError: If the path is outside of the artifact store
                  bounds.
              IllegalOperationError: If the path is a local file and the server
                  is not configured to allow local file access.
          """
          if path.startswith(IN_MEMORY_ARTIFACT_URI_PREFIX):
              # No need to validate in-memory URIs
              return
  
          if not self.allow_local_file_access and not io_utils.is_remote(path):
              raise IllegalOperationError(
                  "Files in a local artifact store cannot be accessed from the "
                  "server."
              )
  
          if not path.startswith(self.fixed_root_path):
              raise FileNotFoundError(
                  f"File `{path}` is outside of "
                  f"artifact store bounds `{self.fixed_root_path}`"
              )
  
      def _sanitize_potential_path(self, potential_path: Any) -> Any:
          """Sanitizes the input if it is a path.
  
          If the input is a **remote** path, this function replaces backslash path
          separators by forward slashes.
  
          Args:
              potential_path: Value that potentially refers to a (remote) path.
  
          Returns:
              The original input or a sanitized version of it in case of a remote
              path.
          """
          if isinstance(potential_path, bytes):
              path = fileio.convert_to_str(potential_path)
          elif isinstance(potential_path, str):
              path = potential_path
          else:
              # Neither string nor bytes, this is not a path
              return potential_path
  
          if path.startswith(IN_MEMORY_ARTIFACT_URI_PREFIX):
              return path
  
          if io_utils.is_remote(path):
              # If we have a remote path, replace windows path separators with
              # slashes
              import ntpath
              import posixpath
  
              path = path.replace(ntpath.sep, posixpath.sep)
              self._validate_path(path)
          else:
              self._validate_path(str(Path(path).absolute().resolve()))
  
          return path
  
      def __call__(self, *args: Any, **kwargs: Any) -> Any:
          """Decorator function that sanitizes paths before calling the original function.
  
          Args:
              *args: Positional args.
              **kwargs: Keyword args.
  
          Returns:
              Output of the input function called with sanitized paths.
          """
          # verify if `self` is part of the args
          has_self = bool(args and isinstance(args[0], BaseArtifactStore))
  
          # sanitize inputs for relevant args and kwargs, keep rest unchanged
          args = tuple(
              self._sanitize_potential_path(
                  arg,
              )
              if i + has_self in self.path_args
              else arg
              for i, arg in enumerate(args)
          )
          kwargs = {
              key: self._sanitize_potential_path(
                  value,
              )
              if key in self.path_kwargs
              else value
              for key, value in kwargs.items()
          }
  
          return self.func(*args, **kwargs)
  
  
  class BaseArtifactStoreConfig(StackComponentConfig):
      """Config class for `BaseArtifactStore`.
  
      Base configuration for artifact storage backends.
      Field descriptions are defined inline using Field() descriptors.
      """
  
      path: str = Field(
          description="Root path for artifact storage. Must be a valid URI supported by the "
          "specific artifact store implementation. Examples: 's3://my-bucket/artifacts', "
          "'/local/storage/path', 'gs://bucket-name/zenml-artifacts', 'azure://container/path'. "
          "Path must be accessible with the configured credentials and permissions"
      )
  
      SUPPORTED_SCHEMES: ClassVar[Set[str]]
      IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
  
      @model_validator(mode="before")
      @classmethod
      @before_validator_handler
      def _ensure_artifact_store(cls, data: Dict[str, Any]) -> Dict[str, Any]:
          """Validator function for the Artifact Stores.
  
          Checks whether supported schemes are defined and the given path is
          supported.
  
          Args:
              data: the input data to construct the artifact store.
  
          Returns:
              The validated values.
  
          Raises:
              ArtifactStoreInterfaceError: If the scheme is not supported.
          """
          try:
              getattr(cls, "SUPPORTED_SCHEMES")
          except AttributeError:
              raise ArtifactStoreInterfaceError(
                  textwrap.dedent(
                      """
                      When you are working with any classes which subclass from
                      zenml.artifact_store.BaseArtifactStore please make sure
                      that your class has a ClassVar named `SUPPORTED_SCHEMES`
                      which should hold a set of supported file schemes such
                      as {"s3://"} or {"gcs://"}.
  
                      Example:
  
                      class MyArtifactStoreConfig(BaseArtifactStoreConfig):
                          ...
                          # Class Variables
                          SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"s3://"}
                          ...
                      """
                  )
              )
  
          if "path" in data:
              data["path"] = data["path"].strip("'\"`")
              if not any(
                  data["path"].startswith(i) for i in cls.SUPPORTED_SCHEMES
              ):
                  raise ArtifactStoreInterfaceError(
                      f"The path: '{data['path']}' you defined for your "
                      f"artifact store is not supported by the implementation of "
                      f"{cls.schema()['title']}, because it does not start with "
                      f"one of its supported schemes: {cls.SUPPORTED_SCHEMES}."
                  )
  
          return data
  
  
  class BaseArtifactStore(StackComponent):
      """Base class for all ZenML artifact stores."""
  
      @property
      def config(self) -> BaseArtifactStoreConfig:
          """Returns the `BaseArtifactStoreConfig` config.
  
          Returns:
              The configuration.
          """
          return cast(BaseArtifactStoreConfig, self._config)
  
      @property
      def path(self) -> str:
          """The path to the artifact store.
  
          Returns:
              The path.
          """
          return self.config.path
  
      @property
      def custom_cache_key(self) -> Optional[bytes]:
          """Custom cache key.
  
          Any artifact store can override this property in case they need
          additional control over the caching behavior.
  
          Returns:
              Custom cache key.
          """
          return None
  
      # --- User interface ---
      @abstractmethod
      def open(self, path: PathType, mode: str = "r") -> Any:
          """Open a file at the given path.
  
          Args:
              path: The path of the file to open.
              mode: The mode to open the file.
  
          Returns:
              The file object.
          """
  
      @abstractmethod
      def copyfile(
          self, src: PathType, dst: PathType, overwrite: bool = False
      ) -> None:
          """Copy a file from the source to the destination.
  
          Args:
              src: The source path.
              dst: The destination path.
              overwrite: Whether to overwrite the destination file if it exists.
          """
  
      @abstractmethod
      def exists(self, path: PathType) -> bool:
          """Checks if a path exists.
  
          Args:
              path: The path to check.
  
          Returns:
              `True` if the path exists.
          """
  
      @abstractmethod
      def glob(self, pattern: PathType) -> List[PathType]:
          """Gets the paths that match a glob pattern.
  
          Args:
              pattern: The glob pattern.
  
          Returns:
              The list of paths that match the pattern.
          """
  
      @abstractmethod
      def isdir(self, path: PathType) -> bool:
          """Returns whether the given path points to a directory.
  
          Args:
              path: The path to check.
  
          Returns:
              `True` if the path points to a directory.
          """
  
      @abstractmethod
      def listdir(self, path: PathType) -> List[PathType]:
          """Returns a list of files under a given directory in the filesystem.
  
          Args:
              path: The path to list.
  
          Returns:
              The list of files under the given path.
          """
  
      @abstractmethod
      def makedirs(self, path: PathType) -> None:
          """Make a directory at the given path, recursively creating parents.
  
          Args:
              path: The path to create.
          """
  
      @abstractmethod
      def mkdir(self, path: PathType) -> None:
          """Make a directory at the given path; parent directory must exist.
  
          Args:
              path: The path to create.
          """
  
      @abstractmethod
      def remove(self, path: PathType) -> None:
          """Remove the file at the given path. Dangerous operation.
  
          Args:
              path: The path to remove.
          """
  
      @abstractmethod
      def rename(
          self, src: PathType, dst: PathType, overwrite: bool = False
      ) -> None:
          """Rename source file to destination file.
  
          Args:
              src: The source path.
              dst: The destination path.
              overwrite: Whether to overwrite the destination file if it exists.
          """
  
      @abstractmethod
      def rmtree(self, path: PathType) -> None:
          """Deletes dir recursively. Dangerous operation.
  
          Args:
              path: The path to delete.
          """
  
      @abstractmethod
      def stat(self, path: PathType) -> Any:
          """Return the stat descriptor for a given file path.
  
          Args:
              path: The path to check.
  
          Returns:
              The stat descriptor.
          """
  
      @abstractmethod
      def size(self, path: PathType) -> Optional[int]:
          """Get the size of a file in bytes.
  
          Args:
              path: The path to the file.
  
          Returns:
              The size of the file in bytes or `None` if the artifact store
              does not implement the `size` method.
          """
          logger.warning(
              "Cannot get size of file '%s' since the artifact store %s does not "
              "implement the `size` method.",
              path,
              self.__class__.__name__,
          )
          return None
  
      @abstractmethod
      def walk(
          self,
          top: PathType,
          topdown: bool = True,
          onerror: Optional[Callable[..., None]] = None,
      ) -> Iterable[Tuple[PathType, List[PathType], List[PathType]]]:
          """Return an iterator that walks the contents of the given directory.
  
          Args:
              top: The path to walk.
              topdown: Whether to walk the top-down or bottom-up.
              onerror: The error handler.
  
          Returns:
              The iterator that walks the contents of the given directory.
          """
  
      # --- Internal interface ---
      def __init__(self, *args: Any, **kwargs: Any) -> None:
          """Initiate the Pydantic object and register the corresponding filesystem.
  
          Args:
              *args: The positional arguments to pass to the Pydantic object.
              **kwargs: The keyword arguments to pass to the Pydantic object.
          """
          super(BaseArtifactStore, self).__init__(*args, **kwargs)
          self._add_path_sanitization()
  
          # If running in a ZenML server environment, we don't register
          # the filesystems. We always use the artifact stores directly.
          if ENV_ZENML_SERVER not in os.environ:
              self._register()
  
      def _add_path_sanitization(self) -> None:
          """Add path sanitization to the artifact store."""
          for method_name, method in inspect.getmembers(BaseArtifactStore):
              if getattr(method, "__isabstractmethod__", False):
                  method_implementation = getattr(self, method_name)
                  sanitized_method = _sanitize_paths(
                      method_implementation, self.path
                  )
                  setattr(self, method_name, sanitized_method)
  
      def _register(self) -> None:
          """Create and register a filesystem within the filesystem registry."""
          from zenml.io.filesystem import BaseFilesystem
          from zenml.io.filesystem_registry import default_filesystem_registry
          from zenml.io.local_filesystem import LocalFilesystem
  
          # Local filesystem is always registered, no point in doing it again.
          if isinstance(self, LocalFilesystem):
              return
  
          overloads: Dict[str, Any] = {
              "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
          }
          for method_name, method in inspect.getmembers(BaseArtifactStore):
              if getattr(method, "__isabstractmethod__", False):
                  method_implementation = getattr(self, method_name)
                  overloads[method_name] = staticmethod(method_implementation)
  
          filesystem_class = type(
              self.__class__.__name__, (BaseFilesystem,), overloads
          )
  
          default_filesystem_registry.register(filesystem_class)
  
      def _remove_previous_file_versions(self, path: PathType) -> None:
          """Remove all file versions but the latest in the given path.
  
          Method is useful for logs stored in versioned file systems
          like AWS S3.
  
          Args:
              path: The path to the file.
          """
          return
  
  
  class BaseArtifactStoreFlavor(Flavor):
      """Base class for artifact store flavors."""
  
      @property
      def type(self) -> StackComponentType:
          """Returns the flavor type.
  
          Returns:
              The flavor type.
          """
          return StackComponentType.ARTIFACT_STORE
  
      @property
      def config_class(self) -> Type[StackComponentConfig]:
          """Config class for this flavor.
  
          Returns:
              The config class.
          """
          return BaseArtifactStoreConfig
  
      @property
      @abstractmethod
      def implementation_class(self) -> Type["BaseArtifactStore"]:
          """Implementation class.
  
          Returns:
              The implementation class.
          """
  