dify/api/services/trigger/subscription_gateway.py

"""Adapt plugin provider discovery and dispatch for subscription management and renewal."""

from collections.abc import Mapping
from typing import Any

from core.plugin.entities.plugin_daemon import CredentialType
from core.trigger.entities.entities import Subscription, UnsubscribeResult
from core.trigger.errors import TriggerProviderCredentialValidationError
from models.provider_ids import TriggerProviderID
from services.trigger.credential_encryption import TriggerProviderPropertiesCache
from services.trigger.plugin_provider import PluginTriggerProvider
from services.trigger.provider_loader import PluginTriggerProviderLoader
from services.trigger.subscription_query_service import TriggerSubscriptionSchemas
from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob
from services.trigger.subscription_service import (
    TriggerSubscriptionInvalidError,
    TriggerSubscriptionProviderError,
)


class PluginTriggerSubscriptionGateway:
    """Reuse PluginTriggerProviderLoader's context cache for a synchronous subscription operation."""

    def __init__(self, *, providers: PluginTriggerProviderLoader) -> None:
        self._providers = providers

    def get_schemas(
        self, *, tenant_id: str, provider_id: TriggerProviderID, credential_type: CredentialType
    ) -> TriggerSubscriptionSchemas:
        provider = self._provider(tenant_id=tenant_id, provider_id=provider_id)
        return TriggerSubscriptionSchemas(
            credentials={credential_type: provider.get_credential_schema_config(credential_type)},
            properties=provider.get_properties_schema(),
        )

    def validate_credentials(
        self, *, tenant_id: str, provider_id: TriggerProviderID, user_id: str, credentials: Mapping[str, Any]
    ) -> None:
        provider = self._provider(tenant_id=tenant_id, provider_id=provider_id)
        try:
            provider.validate_credentials(user_id=user_id, credentials=credentials)
        except TriggerProviderCredentialValidationError as error:
            raise TriggerSubscriptionInvalidError(
                f"Invalid credentials for trigger provider {provider_id}: {error}"
            ) from error
        except Exception as error:
            raise TriggerSubscriptionProviderError(
                f"Failed to verify credentials for trigger provider {provider_id}."
            ) from error

    def subscribe(
        self,
        *,
        tenant_id: str,
        provider_id: TriggerProviderID,
        user_id: str,
        endpoint: str,
        parameters: Mapping[str, Any],
        credentials: Mapping[str, object],
        credential_type: CredentialType,
    ) -> Subscription:
        provider = self._provider(tenant_id=tenant_id, provider_id=provider_id)
        return provider.subscribe_trigger(
            user_id=user_id,
            endpoint=endpoint,
            parameters=parameters,
            credentials=credentials,
            credential_type=credential_type,
        )

    def unsubscribe(
        self,
        *,
        tenant_id: str,
        provider_id: TriggerProviderID,
        user_id: str,
        subscription: Subscription,
        credentials: Mapping[str, object],
        credential_type: CredentialType,
    ) -> UnsubscribeResult:
        provider = self._provider(tenant_id=tenant_id, provider_id=provider_id)
        return provider.unsubscribe_trigger(
            user_id=user_id,
            subscription=subscription,
            credentials=credentials,
            credential_type=credential_type,
        )

    def refresh(
        self,
        *,
        tenant_id: str,
        provider_id: TriggerProviderID,
        subscription: Subscription,
        credentials: Mapping[str, object],
        credential_type: CredentialType,
    ) -> Subscription:
        provider = self._provider(tenant_id=tenant_id, provider_id=provider_id)
        return provider.refresh_trigger(
            subscription=subscription, credentials=credentials, credential_type=credential_type
        )

    def _provider(self, *, tenant_id: str, provider_id: TriggerProviderID) -> PluginTriggerProvider:
        try:
            return self._providers.get_trigger_provider(tenant_id=tenant_id, provider_id=provider_id)
        except Exception as error:
            raise TriggerSubscriptionProviderError(f"Failed to load trigger provider {provider_id}.") from error


def invalidate_subscription_properties(*, tenant_id: str, provider_id: str, subscription_id: str) -> None:
    TriggerProviderPropertiesCache(
        tenant_id=tenant_id, provider_id=provider_id, subscription_id=subscription_id
    ).delete()


def enqueue_subscription_refresh(job: TriggerSubscriptionRefreshJob) -> None:
    from tasks.trigger_subscription_refresh_tasks import trigger_subscription_refresh

    # Keep the two-argument task contract for older workers during rolling upgrades.
    trigger_subscription_refresh.apply_async(
        kwargs={"tenant_id": job.subscription.tenant_id, "subscription_id": job.subscription.subscription_id},
        headers={"reservation_token": job.token},
    )