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

zenml/event_sources/webhooks/__init__.py

@@ -0,0 +1,14 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Base Classes for Webhook Event Sources."""

zenml/event_sources/webhooks/base_webhook_event_source.py

@@ -0,0 +1,231 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Abstract BaseEvent class that all Event implementations must implement."""
+
+import hashlib
+import hmac
+import json
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar, Dict, Optional, Type
+
+from zenml.enums import PluginSubType
+from zenml.event_sources.base_event import BaseEvent
+from zenml.event_sources.base_event_source import (
+    BaseEventSourceFlavor,
+    BaseEventSourceHandler,
+    EventFilterConfig,
+    EventSourceConfig,
+)
+from zenml.exceptions import AuthorizationException
+from zenml.logger import get_logger
+from zenml.models import EventSourceResponse
+
+logger = get_logger(__name__)
+
+
+# -------------------- Event Models -----------------------------------
+
+
+class BaseWebhookEvent(BaseEvent):
+    """Base class for all inbound events."""
+
+
+# -------------------- Configuration Models ----------------------------------
+
+
+class WebhookEventSourceConfig(EventSourceConfig):
+    """The Event Source configuration."""
+
+
+class WebhookEventFilterConfig(EventFilterConfig):
+    """The Event Filter configuration."""
+
+
+# -------------------- Webhook Event Source -----------------------
+
+
+class BaseWebhookEventSourceHandler(BaseEventSourceHandler, ABC):
+    """Base implementation for all Webhook event sources."""
+
+    @property
+    @abstractmethod
+    def config_class(self) -> Type[WebhookEventSourceConfig]:
+        """Returns the webhook event source configuration class.
+
+        Returns:
+            The configuration.
+        """
+
+    @property
+    @abstractmethod
+    def filter_class(self) -> Type[WebhookEventFilterConfig]:
+        """Returns the webhook event filter configuration class.
+
+        Returns:
+            The event filter configuration class.
+        """
+
+    @property
+    @abstractmethod
+    def flavor_class(self) -> "Type[BaseWebhookEventSourceFlavor]":
+        """Returns the flavor class of the plugin.
+
+        Returns:
+            The flavor class of the plugin.
+        """
+
+    def is_valid_signature(
+        self, raw_body: bytes, secret_token: str, signature_header: str
+    ) -> bool:
+        """Verify the SHA256 signature of the payload.
+
+        Args:
+            raw_body: original request body to verify
+            secret_token: secret token used to generate the signature
+            signature_header: signature header to verify (x-hub-signature-256)
+
+        Returns:
+            Whether if the signature is valid.
+        """
+        hash_object = hmac.new(
+            secret_token.encode("utf-8"),
+            msg=raw_body,
+            digestmod=hashlib.sha256,
+        )
+        expected_signature = "sha256=" + hash_object.hexdigest()
+
+        if not hmac.compare_digest(expected_signature, signature_header):
+            return False
+        return True
+
+    @abstractmethod
+    def _interpret_event(self, event: Dict[str, Any]) -> BaseEvent:
+        """Converts the generic event body into a event-source specific pydantic model.
+
+        Args:
+            event: The generic event body
+
+        Return:
+            An instance of the event source specific pydantic model.
+        """
+
+    @abstractmethod
+    def _get_webhook_secret(
+        self, event_source: EventSourceResponse
+    ) -> Optional[str]:
+        """Get the webhook secret for the event source.
+
+        Inheriting classes should implement this method to retrieve the webhook
+        secret associated with an event source. If a webhook secret is not
+        applicable for the event source, this method should return None.
+
+        Args:
+            event_source: The event source to retrieve the secret for.
+
+        Return:
+            The webhook secret associated with the event source, or None if a
+            secret is not applicable.
+        """
+
+    def _validate_webhook_event_signature(
+        self, raw_body: bytes, headers: Dict[str, str], webhook_secret: str
+    ) -> None:
+        """Validate the signature of an incoming webhook event.
+
+        Args:
+            raw_body: The raw inbound webhook event.
+            headers: The headers of the inbound webhook event.
+            webhook_secret: The webhook secret to use for signature validation.
+
+        Raises:
+            AuthorizationException: If the signature validation fails.
+        """
+        signature_header = headers.get("x-hub-signature-256") or headers.get(
+            "x-hub-signature"
+        )
+        if not signature_header:
+            raise AuthorizationException(
+                "x-hub-signature-256 or x-hub-signature header is missing!"
+            )
+
+        if not self.is_valid_signature(
+            raw_body=raw_body,
+            secret_token=webhook_secret,
+            signature_header=signature_header,
+        ):
+            raise AuthorizationException(
+                "Webhook signature verification failed!"
+            )
+
+    def _load_payload(
+        self, raw_body: bytes, headers: Dict[str, str]
+    ) -> Dict[Any, Any]:
+        """Converts the raw body of the request into a python dictionary.
+
+        Args:
+            raw_body: The raw event body.
+            headers: The request headers.
+
+        Returns:
+            An instance of the event source specific pydantic model.
+
+        Raises:
+            ValueError: In case the body can not be parsed.
+        """
+        # For now assume all webhook events are json encoded and parse
+        # the body as such.
+        try:
+            body_dict: Dict[Any, Any] = json.loads(raw_body)
+            return body_dict
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON body received: {e}")
+
+    def process_webhook_event(
+        self,
+        event_source: EventSourceResponse,
+        raw_body: bytes,
+        headers: Dict[str, str],
+    ) -> None:
+        """Process an incoming webhook event.
+
+        Args:
+            event_source: The event source that the event belongs to.
+            raw_body: The raw inbound webhook event.
+            headers: The headers of the inbound webhook event.
+        """
+        json_body = self._load_payload(raw_body=raw_body, headers=headers)
+
+        webhook_secret = self._get_webhook_secret(event_source)
+        if webhook_secret:
+            self._validate_webhook_event_signature(
+                raw_body=raw_body,
+                headers=headers,
+                webhook_secret=webhook_secret,
+            )
+
+        event = self._interpret_event(json_body)
+
+        self.dispatch_event(
+            event=event,
+            event_source=event_source,
+        )
+
+
+# -------------------- Flavors ----------------------------------
+
+
+class BaseWebhookEventSourceFlavor(BaseEventSourceFlavor, ABC):
+    """Base Event Plugin Flavor to access an event plugin along with its configurations."""
+
+    SUBTYPE: ClassVar[PluginSubType] = PluginSubType.WEBHOOK

zenml/exceptions.py

@@ -13,10 +13,7 @@
 #  permissions and limitations under the License.
 """ZenML specific exception definitions."""
 
-from typing import TYPE_CHECKING, Dict, List, Optional, Type
-
-if TYPE_CHECKING:
-    from zenml.steps import BaseParameters
+from typing import Dict, Optional
 
 
 class ZenMLBaseException(Exception):
@@ -50,6 +47,17 @@ class AuthorizationException(ZenMLBaseException):
     """Raised when an authorization error occurred while trying to access a ZenML resource ."""
 
 
+class CredentialsNotValid(AuthorizationException):
+    """Raised when the credentials provided are invalid.
+
+    This is a subclass of AuthorizationException and should only be raised when
+    the authentication credentials are invalid (e.g. expired API token, invalid
+    username/password, invalid signature). If caught by the ZenML client, it
+    will trigger an invalidation of the currently cached API token and a
+    re-authentication flow.
+    """
+
+
 class DoesNotExistException(ZenMLBaseException):
     """Raises exception when the entity does not exist in the system but an action is being done that requires it to be present."""
 
@@ -132,43 +140,6 @@ class PipelineConfigurationError(ZenMLBaseException):
     """Raises exceptions when a pipeline configuration contains invalid values."""
 
 
-class MissingStepParameterError(ZenMLBaseException):
-    """Raises exceptions when a step parameter is missing when running a pipeline."""
-
-    def __init__(
-        self,
-        step_name: str,
-        missing_parameters: List[str],
-        parameters_class: Type["BaseParameters"],
-    ):
-        """Initializes a MissingStepParameterError object.
-
-        Args:
-            step_name: Name of the step for which one or more parameters
-                are missing.
-            missing_parameters: Names of all parameters which are missing.
-            parameters_class: Class of the parameters object for which
-                the parameters are missing.
-        """
-        import textwrap
-
-        message = textwrap.fill(
-            textwrap.dedent(
-                f"""
-            Missing parameters {missing_parameters} for '{step_name}' step.
-            There are three ways to solve this issue:
-            (1) Specify a default value in the parameters class
-            `{parameters_class.__name__}`
-            (2) Specify the parameters in code when creating the pipeline:
-            `my_pipeline({step_name}(params={parameters_class.__name__}(...))`
-            (3) Specify the parameters in a yaml configuration file and pass
-            it to the pipeline: `my_pipeline(...).run(config_path='path_to_yaml')`
-            """
-            )
-        )
-        super().__init__(message)
-
-
 class IntegrationError(ZenMLBaseException):
     """Raises exceptions when a requested integration can not be activated."""
 
@@ -197,6 +168,22 @@ class EntityExistsError(ZenMLBaseException):
     """Raised when trying to register an entity that already exists."""
 
 
+class EntityCreationError(ZenMLBaseException, RuntimeError):
+    """Raised when failing to create an entity."""
+
+
+class ActionExistsError(EntityExistsError):
+    """Raised when registering an action with a name that already exists."""
+
+
+class TriggerExistsError(EntityExistsError):
+    """Raised when registering a trigger with name that already exists."""
+
+
+class WebhookInactiveError(ZenMLBaseException):
+    """Raised when source is inactive."""
+
+
 class StackExistsError(EntityExistsError):
     """Raised when trying to register a stack with name that already exists."""
 
@@ -205,6 +192,10 @@ class StackComponentExistsError(EntityExistsError):
     """Raised when trying to register a stack component with existing name."""
 
 
+class EventSourceExistsError(EntityExistsError):
+    """Raised when trying to register an event source with existing name."""
+
+
 class SecretExistsError(EntityExistsError):
     """Raised when trying to register a secret with existing name."""
 
@@ -233,6 +224,10 @@ class IllegalOperationError(ZenMLBaseException):
     """Raised when an illegal operation is attempted."""
 
 
+class MethodNotAllowedError(ZenMLBaseException):
+    """Raised when the server does not allow a request method."""
+
+
 class SettingsResolvingError(ZenMLBaseException):
     """Raised when resolving settings failed."""
 
@@ -241,6 +236,10 @@ class InputResolutionError(ZenMLBaseException):
     """Raised when step input resolving failed."""
 
 
+class SubscriptionUpgradeRequiredError(ZenMLBaseException):
+    """Raised when user tries to perform an action outside their current subscription tier."""
+
+
 class HydrationError(ZenMLBaseException):
     """Raised when the model hydration failed."""
 

zenml/feature_stores/base_feature_store.py

@@ -14,9 +14,7 @@
 """The base class for feature stores."""
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Type, Union, cast
-
-import pandas as pd
+from typing import Any, Dict, List, Type, cast
 
 from zenml.enums import StackComponentType
 from zenml.stack import Flavor, StackComponent
@@ -42,10 +40,10 @@ class BaseFeatureStore(StackComponent, ABC):
     @abstractmethod
     def get_historical_features(
         self,
-        entity_df: Union[pd.DataFrame, str],
+        entity_df: Any,
         features: List[str],
         full_feature_names: bool = False,
-    ) -> pd.DataFrame:
+    ) -> Any:
         """Returns the historical features for training or batch scoring.
 
         Args:
@@ -54,7 +52,7 @@ class BaseFeatureStore(StackComponent, ABC):
             full_feature_names: Whether to return the full feature names.
 
         Returns:
-            The historical features as a Pandas DataFrame.
+            The historical features.
         """
 
     @abstractmethod

zenml/hooks/hook_validators.py

@@ -14,14 +14,13 @@
 """Validation functions for hooks."""
 
 import inspect
-from types import FunctionType
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from zenml.config.source import Source
 from zenml.utils import source_utils
 
 if TYPE_CHECKING:
-    HookSpecification = Union[str, Source, FunctionType]
+    from zenml.types import HookSpecification
 
 
 def resolve_and_validate_hook(hook: "HookSpecification") -> Source:
@@ -44,9 +43,6 @@ def resolve_and_validate_hook(hook: "HookSpecification") -> Source:
     if not callable(func):
         raise ValueError(f"{func} is not a valid function.")
 
-    from zenml.new.steps.step_context import StepContext
-    from zenml.steps.base_parameters import BaseParameters
-
     sig = inspect.getfullargspec(inspect.unwrap(func))
     sig_annotations = sig.annotations
     if "return" in sig_annotations:
@@ -63,11 +59,7 @@ def resolve_and_validate_hook(hook: "HookSpecification") -> Source:
         seen_annotations = set()
         for annotation in annotations:
             if annotation:
-                if annotation not in (
-                    BaseException,
-                    BaseParameters,
-                    StepContext,
-                ):
+                if annotation not in (BaseException,):
                     raise ValueError(
                         "Hook arguments must be of type `BaseException`, not "
                         f"`{annotation}`."

zenml/image_builders/base_image_builder.py

@@ -25,6 +25,7 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent
 from zenml.stack.stack_component import StackComponentConfig
+from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
     from zenml.container_registries import BaseContainerRegistry
@@ -100,6 +101,7 @@ class BaseImageBuilder(StackComponent, ABC):
     def _upload_build_context(
         build_context: "BuildContext",
         parent_path_directory_name: str,
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> str:
         """Uploads a Docker image build context to a remote location.
 
@@ -109,6 +111,7 @@ class BaseImageBuilder(StackComponent, ABC):
                 the build context to. It will be appended to the artifact
                 store path to create the parent path where the build context
                 will be uploaded to.
+            archive_type: The type of archive to create.
 
         Returns:
             The path to the uploaded build context.
@@ -119,7 +122,7 @@ class BaseImageBuilder(StackComponent, ABC):
 
         hash_ = hashlib.sha1()  # nosec
         with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as f:
-            build_context.write_archive(f, gzip=True)
+            build_context.write_archive(f, archive_type)
 
             while True:
                 data = f.read(64 * 1024)
@@ -127,7 +130,7 @@ class BaseImageBuilder(StackComponent, ABC):
                     break
                 hash_.update(data)
 
-            filename = f"{hash_.hexdigest()}.tar.gz"
+            filename = f"{hash_.hexdigest()}.{archive_type.value}"
             filepath = f"{parent_path}/{filename}"
             if not fileio.exists(filepath):
                 logger.info("Uploading build context to `%s`.", filepath)

zenml/image_builders/build_context.py

@@ -14,18 +14,18 @@
 """Image build context."""
 
 import os
-from pathlib import Path
-from typing import IO, Dict, List, Optional, Set, Tuple, cast
+from typing import IO, Dict, List, Optional, Set, cast
 
 from zenml.constants import REPOSITORY_DIRECTORY_NAME
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import io_utils, string_utils
+from zenml.utils.archivable import Archivable, ArchiveType
 
 logger = get_logger(__name__)
 
 
-class BuildContext:
+class BuildContext(Archivable):
     """Image build context.
 
     This class is responsible for creating an archive of the files needed to
@@ -45,9 +45,9 @@ class BuildContext:
                 given, a file called `.dockerignore` in the build context root
                 directory will be used instead if it exists.
         """
+        super().__init__()
         self._root = root
         self._dockerignore_file = dockerignore_file
-        self._extra_files: Dict[str, str] = {}
 
     @property
     def dockerignore_file(self) -> Optional[str]:
@@ -68,73 +68,20 @@ class BuildContext:
 
         return None
 
-    def add_file(self, source: str, destination: str) -> None:
-        """Adds a file to the build context.
-
-        Args:
-            source: The source of the file to add. This can either be a path
-                or the file content.
-            destination: The path inside the build context where the file
-                should be added.
-        """
-        if fileio.exists(source):
-            with fileio.open(source) as f:
-                self._extra_files[destination] = f.read()
-        else:
-            self._extra_files[destination] = source
-
-    def add_directory(self, source: str, destination: str) -> None:
-        """Adds a directory to the build context.
-
-        Args:
-            source: Path to the directory.
-            destination: The path inside the build context where the directory
-                should be added.
-
-        Raises:
-            ValueError: If `source` does not point to a directory.
-        """
-        if not fileio.isdir(source):
-            raise ValueError(
-                f"Can't add directory {source} to the build context as it "
-                "does not exist or is not a directory."
-            )
-
-        for dir, _, files in fileio.walk(source):
-            dir_path = Path(fileio.convert_to_str(dir))
-            for file_name in files:
-                file_name = fileio.convert_to_str(file_name)
-                file_source = dir_path / file_name
-                file_destination = (
-                    Path(destination)
-                    / dir_path.relative_to(source)
-                    / file_name
-                )
-
-                with file_source.open("r") as f:
-                    self._extra_files[file_destination.as_posix()] = f.read()
-
-    def write_archive(self, output_file: IO[bytes], gzip: bool = True) -> None:
+    def write_archive(
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
+    ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            gzip: Whether to use `gzip` to compress the file.
+            archive_type: The type of archive to create.
         """
-        from docker.utils import build as docker_build_utils
-
-        files = self._get_files()
-        extra_files = self._get_extra_files()
+        super().write_archive(output_file, archive_type)
 
-        context_archive = docker_build_utils.create_archive(
-            fileobj=output_file,
-            root=self._root,
-            files=sorted(files),
-            gzip=gzip,
-            extra_files=extra_files,
-        )
-
-        build_context_size = os.path.getsize(context_archive.name)
+        build_context_size = os.path.getsize(output_file.name)
         if (
             self._root
             and build_context_size > 50 * 1024 * 1024
@@ -151,33 +98,30 @@ class BuildContext:
                 os.path.join(self._root, ".dockerignore"),
             )
 
-    def _get_files(self) -> Set[str]:
-        """Gets all non-ignored files in the build context root directory.
+    def get_files(self) -> Dict[str, str]:
+        """Gets all regular files that should be included in the archive.
 
         Returns:
-            All build context files.
+            A dict {path_in_archive: path_on_filesystem} for all regular files
+            in the archive.
         """
         if self._root:
-            exclude_patterns = self._get_exclude_patterns()
             from docker.utils import build as docker_build_utils
 
-            return cast(
+            exclude_patterns = self._get_exclude_patterns()
+
+            archive_paths = cast(
                 Set[str],
                 docker_build_utils.exclude_paths(
                     self._root, patterns=exclude_patterns
                 ),
             )
+            return {
+                archive_path: os.path.join(self._root, archive_path)
+                for archive_path in archive_paths
+            }
         else:
-            return set()
-
-    def _get_extra_files(self) -> List[Tuple[str, str]]:
-        """Gets all extra files of the build context.
-
-        Returns:
-            A tuple (path, file_content) for all extra files in the build
-            context.
-        """
-        return list(self._extra_files.items())
+            return {}
 
     def _get_exclude_patterns(self) -> List[str]:
         """Gets all exclude patterns from the dockerignore file.