frogml-core 0.0.116__py3-none-any.whl → 0.0.118__py3-none-any.whl

core/clients/alerts_registry/channel.py

@@ -0,0 +1,132 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from frogml._proto.qwak.monitoring.v0.alerting_channel_pb2 import (
+    AlertingChannelDescription,
+    AlertingChannelMetadata,
+    AlertingChannelSpec,
+)
+from frogml._proto.qwak.monitoring.v0.alerting_channel_pb2 import (
+    ChannelConfiguration as ChannelConfigurationProto,
+)
+from frogml._proto.qwak.monitoring.v0.alerting_channel_pb2 import (
+    OpsgenieType,
+    PagerDutyType,
+    SecretStringValue,
+    SlackType,
+)
+from frogml.core.exceptions import FrogmlException
+
+
+@dataclass
+class ChannelConfiguration(ABC):
+    @abstractmethod
+    def to_proto(self) -> "ChannelConfigurationProto":
+        pass
+
+    @classmethod
+    def from_proto(cls, proto: "ChannelConfigurationProto") -> "ChannelConfiguration":
+        channel_type: str = proto.WhichOneof("channel_type")
+
+        if channel_type == "slack":
+            return SlackChannel.from_proto(proto=proto.slack)
+        if channel_type == "pagerduty":
+            return PagerDutyChannel.from_proto(proto=proto.pagerduty)
+        if channel_type == "opsgenie":
+            return OpsGenieChannel.from_proto(proto=proto.opsgenie)
+
+        raise FrogmlException(f"Got unsupported channel type: {channel_type}")
+
+
+@dataclass
+class Channel:
+    name: str
+    channel_conf: ChannelConfiguration
+
+    def __post_init__(self):
+        import re
+
+        allowed_pattern = "^[a-z]+[a-z0-9_]*$"
+        result = re.match(allowed_pattern, self.name)
+        if not result:
+            raise FrogmlException(
+                "Channel name is limited to lower case letters and numbers, must start with a letter."
+            )
+
+    def to_proto(self) -> AlertingChannelDescription:
+        metadata = AlertingChannelMetadata(name=self.name, id="")
+        return AlertingChannelDescription(
+            metadata=metadata,
+            spec=AlertingChannelSpec(configuration=self.channel_conf.to_proto()),
+        )
+
+    @classmethod
+    def from_proto(cls, proto: "AlertingChannelDescription") -> "Channel":
+        channel_conf: ChannelConfiguration = ChannelConfiguration.from_proto(
+            proto.spec.configuration
+        )
+        return cls(
+            name=proto.metadata.name,
+            channel_conf=channel_conf,
+        )
+
+
+@dataclass
+class SlackChannel(ChannelConfiguration):
+    api_url: str
+
+    def to_proto(self) -> ChannelConfigurationProto:
+        return ChannelConfigurationProto(
+            slack=SlackType(api_url=SecretStringValue(value=self.api_url))
+        )
+
+    @classmethod
+    def from_proto(cls, proto: "SlackType") -> "SlackChannel":
+        return cls(api_url=proto.api_url.value)
+
+
+@dataclass
+class PagerDutyChannel(ChannelConfiguration):
+    url: str
+    routing_key: Optional[str] = None
+    service_key: Optional[str] = None
+
+    def __post_init__(self):
+        if (not self.routing_key and not self.service_key) or (
+            self.routing_key and self.service_key
+        ):
+            raise FrogmlException("Routing and service key are mutually exclusive.")
+
+    def to_proto(self) -> ChannelConfigurationProto:
+        return ChannelConfigurationProto(
+            pagerduty=PagerDutyType(
+                routing_key=SecretStringValue(value=self.routing_key),
+                service_key=SecretStringValue(value=self.service_key),
+                url=self.url,
+            )
+        )
+
+    @classmethod
+    def from_proto(cls, proto: "PagerDutyType") -> "PagerDutyChannel":
+        return cls(
+            routing_key=proto.routing_key.value,
+            service_key=proto.service_key.value,
+            url=proto.url,
+        )
+
+
+@dataclass
+class OpsGenieChannel(ChannelConfiguration):
+    api_key: str
+
+    def to_proto(self) -> ChannelConfigurationProto:
+        return ChannelConfigurationProto(
+            opsgenie=OpsgenieType(
+                api_key=SecretStringValue(value=self.api_key),
+            )
+        )
+
+    @classmethod
+    def from_proto(cls, proto: "OpsgenieType") -> "OpsGenieChannel":
+        return cls(api_key=proto.api_key.value)

core/clients/alerts_registry/client.py

@@ -0,0 +1,137 @@
+from typing import List, Optional, Tuple
+
+import grpc
+from dependency_injector.wiring import Provide
+from google.protobuf.wrappers_pb2 import StringValue
+
+from frogml._proto.qwak.monitoring.v0.alerting_channel_management_service_pb2 import (
+    CreateAlertingChannelRequest,
+    DeleteAlertingChannelRequest,
+    GetAlertingChannelRequestByName,
+    GetAlertingChannelResponseByName,
+    ListAlertingChannelRequest,
+    ListAlertingChannelResponse,
+    UpdateAlertingChannelRequest,
+)
+from frogml._proto.qwak.monitoring.v0.alerting_channel_management_service_pb2_grpc import (
+    AlertingChannelManagementServiceStub,
+)
+from frogml._proto.qwak.monitoring.v0.alerting_channel_pb2 import (
+    AlertingChannelCreateOptions,
+    AlertingChannelSpec,
+    AlertingChannelUpdateOptions,
+)
+from frogml.core.clients.alerts_registry.channel import Channel
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.di_configuration import FrogmlContainer
+
+
+class AlertingRegistryClient:
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self.alerting_channel_mgmt_service = AlertingChannelManagementServiceStub(
+            grpc_channel
+        )
+
+    def create_alerting_channel(self, channel: Channel) -> None:
+        try:
+            request = CreateAlertingChannelRequest(
+                options=AlertingChannelCreateOptions(
+                    name=channel.name,
+                    spec=AlertingChannelSpec(
+                        configuration=channel.channel_conf.to_proto()
+                    ),
+                )
+            )
+            self.alerting_channel_mgmt_service.CreateAlertingChannel(request)
+        except grpc.RpcError as rpc_error:
+            raise FrogmlException(
+                f"Failed to create alerting channel named: {channel.name}. {rpc_error.details()}"
+            )
+        except Exception as e:
+            raise FrogmlException(
+                f"Failed to create alerting channel named: {channel.name}"
+            ) from e
+
+    def update_alerting_channel(self, channel_id: str, channel: Channel) -> None:
+        try:
+            request = UpdateAlertingChannelRequest(
+                options=AlertingChannelUpdateOptions(
+                    id=channel_id,
+                    name=StringValue(value=channel.name),
+                    spec=AlertingChannelSpec(
+                        configuration=channel.channel_conf.to_proto()
+                    ),
+                )
+            )
+            self.alerting_channel_mgmt_service.UpdateAlertingChannel(request)
+        except grpc.RpcError as rpc_error:
+            raise FrogmlException(
+                f"Failed to update alerting channel named: {channel.name}. \n{rpc_error.details()}"
+            )
+        except Exception as e:
+            raise FrogmlException(
+                f"Failed to update alerting channel named: {channel.name}"
+            ) from e
+
+    def get_alerting_channel(self, channel_name: str) -> Optional[Tuple[str, Channel]]:
+        try:
+            request = GetAlertingChannelRequestByName(name=channel_name)
+            response: GetAlertingChannelResponseByName = (
+                self.alerting_channel_mgmt_service.GetAlertingChannelByName(request)
+            )
+
+            return response.description.metadata.id, Channel.from_proto(
+                response.description
+            )
+        except grpc.RpcError as rpc_error:
+            code, name = rpc_error.code().value
+            if code == 5:
+                return None, None
+            else:
+                raise FrogmlException(
+                    f"Failed to get alerting channel name: {channel_name}.\n {rpc_error.details()}"
+                )
+        except Exception as e:
+            raise FrogmlException(
+                f"Failed to get alerting channel name: {channel_name}"
+            ) from e
+
+    def list_alerting_channel(self) -> List[Tuple[str, Channel]]:
+        try:
+            response: ListAlertingChannelResponse = (
+                self.list_alerting_channel_from_client()
+            )
+            return [
+                (d.metadata.id, Channel.from_proto(d)) for d in response.description
+            ]
+        except grpc.RpcError as rpc_error:
+            raise FrogmlException(
+                f"Failed to list alerting channels\n {rpc_error.details()}"
+            )
+        except Exception as e:
+            raise FrogmlException("Failed to list alerting channels") from e
+
+    def delete_alerting_channel(self, channel_name: str):
+        try:
+            self.alerting_channel_mgmt_service.DeleteAlertingChannel(
+                DeleteAlertingChannelRequest(name=channel_name)
+            )
+        except grpc.RpcError as rpc_error:
+            raise FrogmlException(
+                f"Failed to delete an alerting channel named: {channel_name}\n {rpc_error.details()}"
+            )
+        except Exception as e:
+            raise FrogmlException(
+                f"Failed to delete an alerting channel named: {channel_name}"
+            ) from e
+
+    def list_alerting_channel_from_client(self) -> ListAlertingChannelResponse:
+        """
+        Get list of alerting channels from the server
+        """
+        response: ListAlertingChannelResponse = (
+            self.alerting_channel_mgmt_service.ListAlertingChannel(
+                ListAlertingChannelRequest()
+            )
+        )
+        return response

core/clients/analytics/client.py

@@ -0,0 +1,108 @@
+from datetime import timedelta
+from time import sleep, time
+from typing import Optional
+
+import grpc
+from frogml.core.clients.location_discovery import LocationDiscoveryClient
+from frogml.core.inner.di_configuration import FrogmlContainer
+from frogml.core.inner.tool.grpc.grpc_tools import create_grpc_channel
+from frogml._proto.qwak.analytics.analytics_pb2 import (
+    QueryResultDownloadURLParams,
+    QueryStatus,
+)
+from frogml._proto.qwak.analytics.analytics_service_pb2 import (
+    GetQueryResultDownloadURLRequest,
+    GetQueryResultsPageRequest,
+    GetQueryResultsPageResponse,
+    QueryRequest,
+)
+from frogml._proto.qwak.analytics.analytics_service_pb2_grpc import (
+    AnalyticsQueryServiceStub,
+)
+
+
+class AnalyticsEngineError(RuntimeError):
+    def __init__(self, failure_reason: str):
+        super().__init__(
+            f"Cannot retrieve results from JFrog ML Analytics Engine\n{failure_reason}"
+        )
+
+
+class AnalyticsEngineClient:
+    def __init__(
+        self, grpc_channel: Optional[FrogmlContainer.core_grpc_channel] = None
+    ):
+        self.grpc_client = self._create_grpc_client(grpc_channel)
+
+    def get_analytics_data(self, query: str, timeout: timedelta = None) -> str:
+        """
+        Sends a given query to JFrog ML Analytics Engine and returns the URL of the results file.
+
+        Args:
+            query (str): SQL query to be sent to JFrog ML Analytics Engine
+            timeout (timedelta): maximum time to wait for the query to complete. If None, the function will wait indefinitely.
+
+        Returns:
+            str: a path from which the results can be downloaded
+
+        Raises:
+            AnalyticsEngineError: if the query failed
+            TimeoutError: if the query timed out
+        """
+        max_time = time() + timeout.total_seconds() if timeout else None
+        query_id = self._send_query(query)
+
+        while True:
+            status_response = self._check_status(query_id)
+
+            if status_response.status == QueryStatus.SUCCESS:
+                download_url = self._get_download_url(query_id)
+                return download_url
+            elif status_response.status == QueryStatus.FAILED:
+                raise AnalyticsEngineError(status_response.failure_reason)
+            elif status_response.status == QueryStatus.CANCELED:
+                raise AnalyticsEngineError("Query was cancelled")
+            elif max_time and time() > max_time:
+                raise TimeoutError()
+            else:
+                sleep(1)
+
+    def _send_query(self, query: str) -> str:
+        request = QueryRequest(query=query)
+        response = self.grpc_client.Query(request)
+        return response.query_id
+
+    def _check_status(self, query_id: str) -> GetQueryResultsPageResponse:
+        request = GetQueryResultsPageRequest(
+            query_id=query_id, page_id=None, max_results=1
+        )
+        response = self.grpc_client.GetQueryResultsPage(request)
+        return response
+
+    def _get_download_url(self, query_id: str) -> str:
+        request_params = QueryResultDownloadURLParams(query_id=query_id)
+        request = GetQueryResultDownloadURLRequest(params=request_params)
+        response = self.grpc_client.GetQueryResultDownloadURL(request)
+        return response.download_url
+
+    def _create_grpc_client(
+        self, grpc_channel: Optional[grpc.Channel] = None
+    ) -> AnalyticsQueryServiceStub:
+        if grpc_channel:
+            return AnalyticsQueryServiceStub(grpc_channel)
+
+        channel = create_grpc_channel(
+            url=self._get_analytics_engine_url(),
+            status_for_retry=(
+                grpc.StatusCode.UNAVAILABLE,
+                grpc.StatusCode.CANCELLED,
+            ),
+        )
+        return AnalyticsQueryServiceStub(channel)
+
+    @staticmethod
+    def _get_analytics_engine_url() -> str:
+        """
+        Fetches the analytics engine service URL from the LocationDiscoveryService.
+        """
+        return LocationDiscoveryClient().get_analytics_engine().service_url

core/clients/audience/client.py

@@ -0,0 +1,53 @@
+from typing import List
+
+from dependency_injector.wiring import Provide
+
+from frogml._proto.qwak.audience.v1.audience_api_pb2 import (
+    CreateAudienceRequest,
+    CreateAudienceResponse,
+    DeleteAudienceRequest,
+    GetAudienceRequest,
+    GetAudienceResponse,
+    ListAudienceRequest,
+    ListAudienceResponse,
+    UpdateAudienceRequest,
+)
+from frogml._proto.qwak.audience.v1.audience_api_pb2_grpc import AudienceAPIStub
+from frogml._proto.qwak.audience.v1.audience_pb2 import Audience, AudienceEntry
+from frogml.core.inner.di_configuration import FrogmlContainer
+
+
+class AudienceClient:
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._service_stub = AudienceAPIStub(grpc_channel)
+
+    def create_audience(self, audience: Audience) -> str:
+        create_audience_request = CreateAudienceRequest(audience=audience)
+        create_audience_response: CreateAudienceResponse = (
+            self._service_stub.CreateAudience(create_audience_request)
+        )
+        return create_audience_response.audience_id
+
+    def get_audience(self, audience_id: str) -> Audience:
+        get_audience_request = GetAudienceRequest(audience_id=audience_id)
+        get_audience_response: GetAudienceResponse = self._service_stub.GetAudience(
+            get_audience_request
+        )
+        return get_audience_response.audience
+
+    def delete_audience(self, audience_id: str) -> None:
+        delete_audience_request = DeleteAudienceRequest(audience_id=audience_id)
+        self._service_stub.DeleteAudience(delete_audience_request)
+
+    def list_audience(self) -> List[AudienceEntry]:
+        list_audience_request = ListAudienceRequest()
+        list_audience_response: ListAudienceResponse = self._service_stub.ListAudience(
+            list_audience_request
+        )
+        return list_audience_response.audience_entries
+
+    def update_audience(self, audience_id: str, audience: Audience) -> None:
+        update_audience_request = UpdateAudienceRequest(
+            id=audience_id, audience=audience
+        )
+        self._service_stub.UpdateAudience(update_audience_request)

core/clients/automation_management/client.py

@@ -0,0 +1,274 @@
+import uuid
+from typing import List, Optional
+
+from dependency_injector.wiring import Provide
+from grpc import RpcError, StatusCode
+
+from frogml._proto.qwak.automation.v1.automation_management_service_pb2 import (
+    CreateAutomationRequest,
+    DeleteAutomationRequest,
+    GetAutomationByNameRequest,
+    GetAutomationRequest,
+    ListAutomationExecutionsRequest,
+    ListAutomationsRequest,
+    RegisterAutomationExecutionRequest,
+    RunAutomationRequest,
+    UpdateAutomationExecutionRequest,
+    UpdateAutomationRequest,
+)
+from frogml._proto.qwak.automation.v1.automation_management_service_pb2_grpc import (
+    AutomationManagementServiceStub,
+)
+from frogml._proto.qwak.automation.v1.automation_pb2 import (
+    Automation as AutomationProto,
+)
+from frogml.core.automations.automation_executions import (
+    AutomationExecution,
+    ExecutionRunDetails,
+)
+from frogml.core.automations.automations import Automation
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.di_configuration import FrogmlContainer
+
+
+class AutomationsManagementClient:
+    """
+    Used for interacting with Automations Registry endpoints
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._service_stub = AutomationManagementServiceStub(grpc_channel)
+
+    def get_automation_by_name(self, automation_name: str) -> Optional[Automation]:
+        """
+
+        Args:
+            automation_name: automation name to get
+        Returns:
+            Automation
+        """
+        try:
+            automation: AutomationProto = self._service_stub.GetAutomationByName(
+                GetAutomationByNameRequest(automation_name=automation_name)
+            ).automation
+            return Automation.from_proto(automation)
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                return None
+            else:
+                raise FrogmlException(
+                    f"Failed to fetch automation by name [{automation_name}], error code: [{e.details()}]"
+                )
+
+    def get_automation_by_id(self, automation_id: str) -> Optional[Automation]:
+        """
+
+        Args:
+            automation_id: automation name to get
+        Returns:
+            Automation
+        """
+        try:
+            automation: AutomationProto = self._service_stub.GetAutomation(
+                GetAutomationRequest(automation_id=automation_id)
+            ).automation
+            return Automation.from_proto(automation)
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                return None
+            else:
+                raise FrogmlException(
+                    f"Failed to fetch automation with ID [{automation_id}], error code: [{e.details()}]"
+                )
+
+    def update_automation(
+        self, automation_id: str, automation: AutomationProto
+    ) -> None:
+        """
+
+        Args:
+            automation: automation to update
+        """
+        automation.automation_spec.is_sdk_v1 = True
+        try:
+            self._service_stub.UpdateAutomation(
+                UpdateAutomationRequest(
+                    automation_id=automation_id,
+                    automation_spec=automation.automation_spec,
+                )
+            )
+            return Automation.from_proto(automation)
+        except RpcError as e:
+            raise FrogmlException(
+                f"Failed to update automation with ID [{automation.automation_id}], error code: [{e.details()}]"
+            )
+
+    def create_automation(self, automation: AutomationProto) -> str:
+        """
+
+        Args:
+            automation: automation to create
+        Returns:
+            Automation id of created automation
+        """
+        automation.automation_spec.is_sdk_v1 = True
+        try:
+            return self._service_stub.CreateAutomation(
+                CreateAutomationRequest(automation_spec=automation.automation_spec)
+            ).automation_id
+        except RpcError as e:
+            raise FrogmlException(
+                f"Failed to update automation, error code: [{e.details()}]"
+            )
+
+    def run_automation(self, automation_id: str) -> str:
+        """
+        Trigger an automation to run
+
+        Args:
+            automation_id: automation id to trigger
+        """
+        try:
+            return self._service_stub.RunAutomation(
+                RunAutomationRequest(automation_id=automation_id)
+            )
+        except RpcError as e:
+            raise FrogmlException(
+                f"Failed to execute a run automation command, error code: [{e.details()}]"
+            )
+
+    def list_automations(
+        self,
+    ) -> Optional[List[Automation]]:
+        """
+
+        Returns:
+            List of client automations
+        """
+        try:
+            automations = self.get_list_automations_from_server()
+            return list(
+                map(lambda automation: Automation.from_proto(automation), automations)
+            )
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                return None
+            else:
+                raise FrogmlException(
+                    f"Failed to fetch automations, error code: [{e.details()}]"
+                )
+
+    def get_list_automations_from_server(self):
+        """
+        Get list of automations from server
+        """
+        automations: List[AutomationProto] = self._service_stub.ListAutomations(
+            ListAutomationsRequest()
+        ).automations
+        return automations
+
+    def delete_automation(self, automation_id: str) -> bool:
+        """
+
+        Args:
+            automation_id: The automation id to delete
+
+        Returns:
+            boolean whether the automation was deleted or not
+        """
+        try:
+            self._service_stub.DeleteAutomation(
+                DeleteAutomationRequest(automation_id=automation_id)
+            )
+            return True
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                return False
+            else:
+                raise FrogmlException(
+                    f"Failed to fetch automations, error code: [{e.details()}]"
+                )
+
+    def register_execution(self, automation_id: str) -> str:
+        """
+
+        Args:
+            automation_id: The automation id to register
+
+        Returns:
+            The execution id for the automation execution
+
+        """
+        try:
+            execution_id = str(uuid.uuid4())
+            self._service_stub.RegisterAutomationExecution(
+                RegisterAutomationExecutionRequest(
+                    automation_id=automation_id, execution_id=execution_id
+                )
+            )
+            return execution_id
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                raise FrogmlException(
+                    f"Failed to register automation execution, automation {automation_id} was not found"
+                )
+            else:
+                raise FrogmlException(
+                    f"Failed to register automation execution, error code: [{e.details()}]"
+                )
+
+    def update_execution(
+        self, execution_id: str, run_details: ExecutionRunDetails
+    ) -> None:
+        """
+
+        Args:
+            execution_id: The execution id to update
+            run_details: The details to update on the execution
+
+        """
+        try:
+            self._service_stub.UpdateAutomationExecution(
+                UpdateAutomationExecutionRequest(
+                    execution_id=execution_id, run_details=run_details.to_proto()
+                )
+            )
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                raise FrogmlException(
+                    f"Failed to update execution, execution {execution_id} was not found"
+                )
+            else:
+                raise FrogmlException(
+                    f"Failed to update execution, error code: [{e.details()}]"
+                )
+
+    def list_executions(self, automation_id: str) -> List[AutomationExecution]:
+        """
+
+        Args:
+            automation_id: The automation to list its executions
+
+        Returns:
+            A list of executions for the specified automation
+
+        """
+        try:
+            executions = self._service_stub.ListAutomationExecutions(
+                ListAutomationExecutionsRequest(automation_id=automation_id)
+            ).automation_executions
+            return list(
+                map(
+                    lambda execution: AutomationExecution.from_proto(execution),
+                    executions,
+                )
+            )
+        except RpcError as e:
+            if e.code() == StatusCode.NOT_FOUND:
+                raise FrogmlException(
+                    f"Failed to list executions for automation, automation {automation_id} was not found"
+                )
+            else:
+                raise FrogmlException(
+                    f"Failed to list executions for automation, error code: [{e.details()}]"
+                )