mirror of
https://github.com/langgenius/dify.git
synced 2025-12-22 15:27:32 +00:00
feat: introduce trigger functionality (#27644)
Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com> Co-authored-by: Stream <Stream_2@qq.com> Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com> Co-authored-by: zhsama <torvalds@linux.do> Co-authored-by: Harry <xh001x@hotmail.com> Co-authored-by: lyzno1 <yuanyouhuilyz@gmail.com> Co-authored-by: yessenia <yessenia.contact@gmail.com> Co-authored-by: hjlarry <hjlarry@163.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: WTW0313 <twwu@dify.ai> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
115
api/tasks/trigger_subscription_refresh_tasks.py
Normal file
115
api/tasks/trigger_subscription_refresh_tasks.py
Normal file
@@ -0,0 +1,115 @@
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from celery import shared_task
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.plugin.entities.plugin_daemon import CredentialType
|
||||
from core.trigger.utils.locks import build_trigger_refresh_lock_key
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from models.trigger import TriggerSubscription
|
||||
from services.trigger.trigger_provider_service import TriggerProviderService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _now_ts() -> int:
|
||||
return int(time.time())
|
||||
|
||||
|
||||
def _load_subscription(session: Session, tenant_id: str, subscription_id: str) -> TriggerSubscription | None:
|
||||
return session.query(TriggerSubscription).filter_by(tenant_id=tenant_id, id=subscription_id).first()
|
||||
|
||||
|
||||
def _refresh_oauth_if_expired(tenant_id: str, subscription: TriggerSubscription, now: int) -> None:
|
||||
if (
|
||||
subscription.credential_expires_at != -1
|
||||
and int(subscription.credential_expires_at) <= now
|
||||
and CredentialType.of(subscription.credential_type) == CredentialType.OAUTH2
|
||||
):
|
||||
logger.info(
|
||||
"Refreshing OAuth token: tenant=%s subscription_id=%s expires_at=%s now=%s",
|
||||
tenant_id,
|
||||
subscription.id,
|
||||
subscription.credential_expires_at,
|
||||
now,
|
||||
)
|
||||
try:
|
||||
result: Mapping[str, Any] = TriggerProviderService.refresh_oauth_token(
|
||||
tenant_id=tenant_id, subscription_id=subscription.id
|
||||
)
|
||||
logger.info(
|
||||
"OAuth token refreshed: tenant=%s subscription_id=%s result=%s", tenant_id, subscription.id, result
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("OAuth refresh failed: tenant=%s subscription_id=%s", tenant_id, subscription.id)
|
||||
|
||||
|
||||
def _refresh_subscription_if_expired(
|
||||
tenant_id: str,
|
||||
subscription: TriggerSubscription,
|
||||
now: int,
|
||||
) -> None:
|
||||
if subscription.expires_at == -1 or int(subscription.expires_at) > now:
|
||||
logger.debug(
|
||||
"Subscription not due: tenant=%s subscription_id=%s expires_at=%s now=%s",
|
||||
tenant_id,
|
||||
subscription.id,
|
||||
subscription.expires_at,
|
||||
now,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
result: Mapping[str, Any] = TriggerProviderService.refresh_subscription(
|
||||
tenant_id=tenant_id, subscription_id=subscription.id, now=now
|
||||
)
|
||||
logger.info(
|
||||
"Subscription refreshed: tenant=%s subscription_id=%s result=%s",
|
||||
tenant_id,
|
||||
subscription.id,
|
||||
result.get("result"),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Subscription refresh failed: tenant=%s id=%s", tenant_id, subscription.id)
|
||||
|
||||
|
||||
@shared_task(queue="trigger_refresh_executor")
|
||||
def trigger_subscription_refresh(tenant_id: str, subscription_id: str) -> None:
|
||||
"""Refresh a trigger subscription if needed, guarded by a Redis in-flight lock."""
|
||||
lock_key: str = build_trigger_refresh_lock_key(tenant_id, subscription_id)
|
||||
if not redis_client.get(lock_key):
|
||||
logger.debug("Refresh lock missing, skip: %s", lock_key)
|
||||
return
|
||||
|
||||
logger.info("Begin subscription refresh: tenant=%s id=%s", tenant_id, subscription_id)
|
||||
try:
|
||||
now: int = _now_ts()
|
||||
with Session(db.engine) as session:
|
||||
subscription: TriggerSubscription | None = _load_subscription(session, tenant_id, subscription_id)
|
||||
|
||||
if not subscription:
|
||||
logger.warning("Subscription not found: tenant=%s id=%s", tenant_id, subscription_id)
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
"Loaded subscription: tenant=%s id=%s cred_exp=%s sub_exp=%s now=%s",
|
||||
tenant_id,
|
||||
subscription.id,
|
||||
subscription.credential_expires_at,
|
||||
subscription.expires_at,
|
||||
now,
|
||||
)
|
||||
|
||||
_refresh_oauth_if_expired(tenant_id=tenant_id, subscription=subscription, now=now)
|
||||
_refresh_subscription_if_expired(tenant_id=tenant_id, subscription=subscription, now=now)
|
||||
finally:
|
||||
try:
|
||||
redis_client.delete(lock_key)
|
||||
logger.debug("Lock released: %s", lock_key)
|
||||
except Exception:
|
||||
# Best-effort lock cleanup
|
||||
logger.warning("Failed to release lock: %s", lock_key, exc_info=True)
|
||||
Reference in New Issue
Block a user