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

zenml/event_sources/base_event_source.py

@@ -0,0 +1,695 @@
+#  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 implementation for event sources."""
+
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    Optional,
+    Type,
+)
+
+from pydantic import BaseModel
+
+from zenml.enums import PluginType
+from zenml.event_hub.base_event_hub import BaseEventHub
+from zenml.event_sources.base_event import BaseEvent
+from zenml.logger import get_logger
+from zenml.models import (
+    EventSourceFlavorResponse,
+    EventSourceFlavorResponseBody,
+    EventSourceFlavorResponseMetadata,
+    EventSourceRequest,
+    EventSourceResponse,
+    EventSourceUpdate,
+)
+from zenml.plugins.base_plugin_flavor import (
+    BasePlugin,
+    BasePluginConfig,
+    BasePluginFlavor,
+)
+
+logger = get_logger(__name__)
+
+# -------------------- Configuration Models ----------------------------------
+
+
+class EventSourceConfig(BasePluginConfig):
+    """The Event Source configuration."""
+
+
+class EventFilterConfig(BaseModel, ABC):
+    """The Event Filter configuration."""
+
+    @abstractmethod
+    def event_matches_filter(self, event: BaseEvent) -> bool:
+        """All implementations need to implement this check.
+
+        If the filter matches the inbound event instance, this should
+        return True, else False.
+
+        Args:
+            event: The inbound event instance.
+
+        Returns:
+            Whether the filter matches the event.
+        """
+
+
+# -------------------- Event Source -----------------------------
+
+
+class BaseEventSourceHandler(BasePlugin, ABC):
+    """Base event source handler implementation.
+
+    This class provides a base implementation for event source handlers.
+
+    The base event source handler acts as an intermediary between the REST API
+    endpoints and the ZenML store for all operations related to event sources.
+    It implements all operations related creating, updating, deleting, and
+    fetching event sources from the database. It also provides a set of methods
+    that can be overridden by concrete event source handlers that need to
+    react to or participate in the lifecycle of an event source:
+
+    * `_validate_event_source_request`: validate and/or modify an event source
+    before it is created (before it is persisted in the database)
+    * `_process_event_source_request`: react to a new event source being created
+    (after it is persisted in the database)
+    * `_validate_event_source_update`: validate and/or modify an event source
+    update before it is applied (before the update is saved in the database)
+    * `_process_event_source_update`: react to an event source being updated
+    (after the update is saved in the database)
+    * `_process_event_source_delete`: react to an event source being deleted
+    (before it is deleted from the database)
+    * `_process_event_source_response`: modify an event source before it is
+    returned to the user
+
+    In addition to optionally overriding these methods, every event source
+    handler must define and provide configuration classes for the event source
+    and the event filter. These are used to validate and instantiate the
+    configuration from the user requests.
+
+    Finally, since event source handlers are also sources of events, the base
+    class provides methods that implementations can use to dispatch events to
+    the central event hub where they can be processed and eventually trigger
+    actions.
+    """
+
+    _event_hub: Optional[BaseEventHub] = None
+
+    def __init__(self, event_hub: Optional[BaseEventHub] = None) -> None:
+        """Event source handler initialization.
+
+        Args:
+            event_hub: Optional event hub to use to initialize the event source
+                handler. If not set during initialization, it may be set
+                at a later time by calling `set_event_hub`. An event hub must
+                be configured before the event handler needs to dispatch events.
+        """
+        super().__init__()
+        if event_hub is None:
+            from zenml.event_hub.event_hub import (
+                event_hub as default_event_hub,
+            )
+
+            # TODO: for now, we use the default internal event hub. In
+            # the future, this should be configurable.
+            event_hub = default_event_hub
+
+        self.set_event_hub(event_hub)
+
+    @property
+    @abstractmethod
+    def config_class(self) -> Type[EventSourceConfig]:
+        """Returns the event source configuration class.
+
+        Returns:
+            The configuration.
+        """
+
+    @property
+    @abstractmethod
+    def filter_class(self) -> Type[EventFilterConfig]:
+        """Returns the event filter configuration class.
+
+        Returns:
+            The event filter configuration class.
+        """
+
+    @property
+    @abstractmethod
+    def flavor_class(self) -> "Type[BaseEventSourceFlavor]":
+        """Returns the flavor class of the plugin.
+
+        Returns:
+            The flavor class of the plugin.
+        """
+
+    @property
+    def event_hub(self) -> BaseEventHub:
+        """Get the event hub used to dispatch events.
+
+        Returns:
+            The event hub.
+
+        Raises:
+            RuntimeError: if an event hub isn't configured.
+        """
+        if self._event_hub is None:
+            raise RuntimeError(
+                f"An event hub is not configured for the "
+                f"{self.flavor_class.FLAVOR} {self.flavor_class.SUBTYPE} "
+                f"event source handler."
+            )
+
+        return self._event_hub
+
+    def set_event_hub(self, event_hub: BaseEventHub) -> None:
+        """Configure an event hub for this event source plugin.
+
+        Args:
+            event_hub: Event hub to be used by this event handler to dispatch
+                events.
+        """
+        self._event_hub = event_hub
+
+    def create_event_source(
+        self, event_source: EventSourceRequest
+    ) -> EventSourceResponse:
+        """Process an event source request and create the event source in the database.
+
+        Args:
+            event_source: Event source request.
+
+        Returns:
+            The created event source.
+
+        # noqa: DAR401
+        """
+        # Validate and instantiate the configuration from the request
+        config = self.validate_event_source_configuration(
+            event_source.configuration
+        )
+        # Call the implementation specific method to validate the request
+        # before it is sent to the database
+        self._validate_event_source_request(
+            event_source=event_source, config=config
+        )
+        # Serialize the configuration back into the request
+        event_source.configuration = config.model_dump(exclude_none=True)
+        # Create the event source in the database
+        event_source_response = self.zen_store.create_event_source(
+            event_source=event_source
+        )
+        try:
+            # Instantiate the configuration from the response
+            config = self.validate_event_source_configuration(
+                event_source.configuration
+            )
+            # Call the implementation specific method to process the created
+            # event source
+            self._process_event_source_request(
+                event_source=event_source_response, config=config
+            )
+        except Exception:
+            # If the event source creation fails, delete the event source from
+            # the database
+            logger.exception(
+                f"Failed to create event source {event_source_response}. "
+                f"Deleting the event source."
+            )
+            self.zen_store.delete_event_source(
+                event_source_id=event_source_response.id
+            )
+            raise
+
+        # Serialize the configuration back into the response
+        event_source_response.set_configuration(
+            config.model_dump(exclude_none=True)
+        )
+
+        # Return the response to the user
+        return event_source_response
+
+    def update_event_source(
+        self,
+        event_source: EventSourceResponse,
+        event_source_update: EventSourceUpdate,
+    ) -> EventSourceResponse:
+        """Process an event source update request and update the event source in the database.
+
+        Args:
+            event_source: The event source to update.
+            event_source_update: The update to be applied to the event source.
+
+        Returns:
+            The updated event source.
+
+        # noqa: DAR401
+        """
+        # Validate and instantiate the configuration from the original event
+        # source
+        config = self.validate_event_source_configuration(
+            event_source.configuration
+        )
+        # Validate and instantiate the configuration from the update request
+        # NOTE: if supplied, the configuration update is a full replacement
+        # of the original configuration
+        config_update = config
+        if event_source_update.configuration is not None:
+            config_update = self.validate_event_source_configuration(
+                event_source_update.configuration
+            )
+        # Call the implementation specific method to validate the update request
+        # before it is sent to the database
+        self._validate_event_source_update(
+            event_source=event_source,
+            config=config,
+            event_source_update=event_source_update,
+            config_update=config_update,
+        )
+        # Serialize the configuration update back into the update request
+        event_source_update.configuration = config_update.model_dump(
+            exclude_none=True
+        )
+
+        # Update the event source in the database
+        event_source_response = self.zen_store.update_event_source(
+            event_source_id=event_source.id,
+            event_source_update=event_source_update,
+        )
+        try:
+            # Instantiate the configuration from the response
+            response_config = self.validate_event_source_configuration(
+                event_source_response.configuration
+            )
+            # Call the implementation specific method to process the update
+            # request before it is sent to the database
+            self._process_event_source_update(
+                event_source=event_source_response,
+                config=response_config,
+                previous_event_source=event_source,
+                previous_config=config,
+            )
+        except Exception:
+            # If the event source update fails, roll back the event source in
+            # the database to the original state
+            logger.exception(
+                f"Failed to update event source {event_source}. "
+                f"Rolling back the event source to the previous state."
+            )
+            self.zen_store.update_event_source(
+                event_source_id=event_source.id,
+                event_source_update=EventSourceUpdate.from_response(
+                    event_source
+                ),
+            )
+            raise
+
+        # Serialize the configuration back into the response
+        event_source_response.set_configuration(
+            response_config.model_dump(exclude_none=True)
+        )
+        # Return the response to the user
+        return event_source_response
+
+    def delete_event_source(
+        self,
+        event_source: EventSourceResponse,
+        force: bool = False,
+    ) -> None:
+        """Process an event source delete request and delete the event source in the database.
+
+        Args:
+            event_source: The event source to delete.
+            force: Whether to force delete the event source from the database
+                even if the event source handler fails to delete the event
+                source.
+
+        # noqa: DAR401
+        """
+        # Validate and instantiate the configuration from the original event
+        # source
+        config = self.validate_event_source_configuration(
+            event_source.configuration
+        )
+        try:
+            # Call the implementation specific method to process the deleted
+            # event source before it is deleted from the database
+            self._process_event_source_delete(
+                event_source=event_source,
+                config=config,
+                force=force,
+            )
+        except Exception:
+            logger.exception(f"Failed to delete event source {event_source}. ")
+            if not force:
+                raise
+
+            logger.warning(f"Force deleting event source {event_source}.")
+
+        # Delete the event source from the database
+        self.zen_store.delete_event_source(
+            event_source_id=event_source.id,
+        )
+
+    def get_event_source(
+        self, event_source: EventSourceResponse, hydrate: bool = False
+    ) -> EventSourceResponse:
+        """Process an event source response before it is returned to the user.
+
+        Args:
+            event_source: The event source fetched from the database.
+            hydrate: Whether to hydrate the event source.
+
+        Returns:
+            The event source.
+        """
+        if hydrate:
+            # Instantiate the configuration from the response
+            config = self.validate_event_source_configuration(
+                event_source.configuration
+            )
+            # Call the implementation specific method to process the response
+            self._process_event_source_response(
+                event_source=event_source, config=config
+            )
+            # Serialize the configuration back into the response
+            event_source.set_configuration(
+                config.model_dump(exclude_none=True)
+            )
+
+        # Return the response to the user
+        return event_source
+
+    def validate_event_source_configuration(
+        self, event_source_config: Dict[str, Any]
+    ) -> EventSourceConfig:
+        """Validate and return the event source configuration.
+
+        Args:
+            event_source_config: The event source configuration to validate.
+
+        Returns:
+            The validated event source configuration.
+
+        Raises:
+            ValueError: if the configuration is invalid.
+        """
+        try:
+            return self.config_class(**event_source_config)
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid configuration for event source: {e}."
+            ) from e
+
+    def validate_event_filter_configuration(
+        self,
+        configuration: Dict[str, Any],
+    ) -> EventFilterConfig:
+        """Validate and return the configuration of an event filter.
+
+        Args:
+            configuration: The configuration to validate.
+
+        Returns:
+            Instantiated event filter config.
+
+        Raises:
+            ValueError: if the configuration is invalid.
+        """
+        try:
+            return self.filter_class(**configuration)
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid configuration for event filter: {e}."
+            ) from e
+
+    def dispatch_event(
+        self,
+        event: BaseEvent,
+        event_source: EventSourceResponse,
+    ) -> None:
+        """Dispatch an event to all active triggers that match the event.
+
+        Args:
+            event: The event to dispatch.
+            event_source: The event source that produced the event.
+        """
+        if not event_source.is_active:
+            logger.debug(
+                f"Event source {event_source.id} is not active. Skipping event "
+                f"dispatch."
+            )
+            return
+
+        self.event_hub.publish_event(
+            event=event,
+            event_source=event_source,
+        )
+
+    def _validate_event_source_request(
+        self, event_source: EventSourceRequest, config: EventSourceConfig
+    ) -> None:
+        """Validate an event source request before it is created in the database.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to the validation of
+        a new event source. The implementation may also modify the event source
+        request and/or configuration in place to apply implementation specific
+        changes before the request is stored in the database.
+
+        If validation is required, the implementation should raise a
+        ValueError if the configuration is invalid. If any exceptions are raised
+        during the validation, the event source will not be created in the
+        database.
+
+        The resulted configuration is serialized back into the event source
+        request before it is stored in the database.
+
+        The implementation should not use this method to provision any external
+        resources, as the event source may not be created in the database if
+        the database level validation fails.
+
+        Args:
+            event_source: Event source request.
+            config: Event source configuration instantiated from the request.
+        """
+        pass
+
+    def _process_event_source_request(
+        self, event_source: EventSourceResponse, config: EventSourceConfig
+    ) -> None:
+        """Process an event source request after it is created in the database.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to the creation of
+        a new event source. The implementation may also modify the event source
+        response and/or configuration in place to apply implementation specific
+        changes before the response is returned to the user.
+
+        The resulted configuration is serialized back into the event source
+        response before it is returned to the user.
+
+        The implementation should use this method to provision any external
+        resources required for the event source. If any of the provisioning
+        fails, the implementation should raise an exception and the event source
+        will be deleted from the database.
+
+        Args:
+            event_source: Newly created event source
+            config: Event source configuration instantiated from the response.
+        """
+        pass
+
+    def _validate_event_source_update(
+        self,
+        event_source: EventSourceResponse,
+        config: EventSourceConfig,
+        event_source_update: EventSourceUpdate,
+        config_update: EventSourceConfig,
+    ) -> None:
+        """Validate an event source update before it is reflected in the database.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to validation of an
+        event source update request. The implementation may also modify the
+        event source update and/or configuration update in place to apply
+        implementation specific changes.
+
+        If validation is required, the implementation should raise a
+        ValueError if the configuration update is invalid. If any exceptions are
+        raised during the validation, the event source will not be updated in
+        the database.
+
+        The resulted configuration update is serialized back into the event
+        source update before it is stored in the database.
+
+        The implementation should not use this method to provision any external
+        resources, as the event source may not be updated in the database if
+        the database level validation fails.
+
+        Args:
+            event_source: Original event source before the update.
+            config: Event source configuration instantiated from the original
+                event source.
+            event_source_update: Event source update request.
+            config_update: Event source configuration instantiated from the
+                updated event source.
+        """
+        pass
+
+    def _process_event_source_update(
+        self,
+        event_source: EventSourceResponse,
+        config: EventSourceConfig,
+        previous_event_source: EventSourceResponse,
+        previous_config: EventSourceConfig,
+    ) -> None:
+        """Process an event source after it is updated in the database.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to updating an existing
+        event source. The implementation may also modify the event source
+        and/or configuration in place to apply implementation specific
+        changes before the response is returned to the user.
+
+        The resulted configuration is serialized back into the event source
+        response before it is returned to the user.
+
+        The implementation should use this method to provision any external
+        resources required for the event source update. If any of the
+        provisioning fails, the implementation should raise an exception and the
+        event source will be rolled back to the previous state in the database.
+
+        Args:
+            event_source: Event source after the update.
+            config: Event source configuration instantiated from the updated
+                event source.
+            previous_event_source: Original event source before the update.
+            previous_config: Event source configuration instantiated from the
+                original event source.
+        """
+        pass
+
+    def _process_event_source_delete(
+        self,
+        event_source: EventSourceResponse,
+        config: EventSourceConfig,
+        force: Optional[bool] = False,
+    ) -> None:
+        """Process an event source before it is deleted from the database.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to the deletion of an
+        event source.
+
+        The implementation should use this method to deprovision any external
+        resources required for the event source. If any of the deprovisioning
+        fails, the implementation should raise an exception and the
+        event source will kept in the database (unless force is True, in which
+        case the event source will be deleted from the database regardless of
+        any deprovisioning failures).
+
+        Args:
+            event_source: Event source before the deletion.
+            config: Validated instantiated event source configuration before
+                the deletion.
+            force: Whether to force deprovision the event source.
+        """
+        pass
+
+    def _process_event_source_response(
+        self, event_source: EventSourceResponse, config: EventSourceConfig
+    ) -> None:
+        """Process an event source response before it is returned to the user.
+
+        Concrete event source handlers should override this method to add
+        implementation specific functionality pertaining to how the event source
+        response is returned to the user. The implementation may modify the
+        the event source response and/or configuration in place.
+
+        The resulted configuration is serialized back into the event source
+        response before it is returned to the user.
+
+        This method is applied to all event source responses fetched from the
+        database before they are returned to the user, with the exception of
+        those returned from the `create_event_source`, `update_event_source`,
+        and `delete_event_source` methods, which have their own specific
+        processing methods.
+
+        Args:
+            event_source: Event source response.
+            config: Event source configuration instantiated from the response.
+        """
+        pass
+
+
+# -------------------- Flavors ----------------------------------
+
+
+class BaseEventSourceFlavor(BasePluginFlavor, ABC):
+    """Base Event Plugin Flavor to access an event plugin along with its configurations."""
+
+    TYPE: ClassVar[PluginType] = PluginType.EVENT_SOURCE
+
+    # EventPlugin specific
+    EVENT_SOURCE_CONFIG_CLASS: ClassVar[Type[EventSourceConfig]]
+    EVENT_FILTER_CONFIG_CLASS: ClassVar[Type[EventFilterConfig]]
+
+    @classmethod
+    def get_event_filter_config_schema(cls) -> Dict[str, Any]:
+        """The config schema for a flavor.
+
+        Returns:
+            The config schema.
+        """
+        return cls.EVENT_SOURCE_CONFIG_CLASS.model_json_schema()
+
+    @classmethod
+    def get_event_source_config_schema(cls) -> Dict[str, Any]:
+        """The config schema for a flavor.
+
+        Returns:
+            The config schema.
+        """
+        return cls.EVENT_FILTER_CONFIG_CLASS.model_json_schema()
+
+    @classmethod
+    def get_flavor_response_model(
+        cls, hydrate: bool
+    ) -> EventSourceFlavorResponse:
+        """Convert the Flavor into a Response Model.
+
+        Args:
+            hydrate: Whether the model should be hydrated.
+
+        Returns:
+            The Flavor Response model for the Event Source implementation
+        """
+        metadata = None
+        if hydrate:
+            metadata = EventSourceFlavorResponseMetadata(
+                source_config_schema=cls.get_event_source_config_schema(),
+                filter_config_schema=cls.get_event_filter_config_schema(),
+            )
+        return EventSourceFlavorResponse(
+            body=EventSourceFlavorResponseBody(),
+            metadata=metadata,
+            name=cls.FLAVOR,
+            type=cls.TYPE,
+            subtype=cls.SUBTYPE,
+        )