company-ai-platform/app/modules/subscriptions/services/scanner.py

from datetime import UTC, datetime, time, timedelta
from hashlib import sha256
from uuid import NAMESPACE_URL, uuid4, uuid5

from sqlalchemy import func, or_, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core.http.pagination import bounded_limit
from app.core.utils.time import utc_now
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.feishu_users.constants import FeishuUserRole, FeishuUserStatus
from app.modules.feishu_users.models import FeishuUser
from app.modules.subscriptions.constants import (
    DAILY_DELIVERY_LIMIT_REACHED,
    MAX_DAILY_DELIVERIES,
    PushDeliveryStatus,
    PushSubscriptionStatus,
    SUBSCRIPTION_LEASE_SECONDS,
    SubscriptionAuditAction,
    SubscriptionScheduleType,
    SubscriptionTargetType,
)
from app.modules.subscriptions.models import PushDelivery, PushSubscription
from app.modules.subscriptions.services.schedule import (
    is_in_quiet_hours,
    next_occurrence,
    next_quiet_end,
    validate_timezone,
)


class SubscriptionScanner:
    """Claim due plans and materialize one durable delivery per schedule window."""

    def __init__(self, db: Session, *, lease_seconds: int = SUBSCRIPTION_LEASE_SECONDS):
        self.db = db
        self.lease_seconds = lease_seconds

    def scan_due(
        self,
        *,
        now: datetime | None = None,
        limit: int = 100,
        worker_id: str = "subscription-scanner",
    ) -> list[PushDelivery]:
        current = _naive_utc(now or utc_now())
        claims = self._claim_due_subscriptions(
            current=current,
            limit=limit,
            worker_id=worker_id,
        )
        deliveries: list[PushDelivery] = []
        for subscription_id, lock_owner in claims:
            try:
                delivery = self._materialize_delivery(
                    subscription_id=subscription_id,
                    lock_owner=lock_owner,
                    current=current,
                )
            except Exception:
                self.db.rollback()
                self._release_claim(subscription_id, lock_owner)
                raise
            if delivery is not None:
                deliveries.append(delivery)
        return deliveries

    def _claim_due_subscriptions(
        self,
        *,
        current: datetime,
        limit: int,
        worker_id: str,
    ) -> list[tuple[int, str]]:
        stmt = (
            select(PushSubscription.id)
            .where(
                PushSubscription.status == PushSubscriptionStatus.ACTIVE,
                PushSubscription.next_run_at.is_not(None),
                PushSubscription.next_run_at <= current,
                or_(
                    PushSubscription.locked_until.is_(None),
                    PushSubscription.locked_until <= current,
                ),
            )
            .order_by(PushSubscription.next_run_at.asc(), PushSubscription.id.asc())
            .limit(bounded_limit(limit))
        )
        if self.db.get_bind().dialect.name == "postgresql":
            stmt = stmt.with_for_update(skip_locked=True)
        candidate_ids = list(self.db.execute(stmt).scalars())
        claims: list[tuple[int, str]] = []
        locked_until = current + timedelta(seconds=self.lease_seconds)
        for subscription_id in candidate_ids:
            lock_owner = f"{worker_id}:{uuid4().hex}"
            result = self.db.execute(
                update(PushSubscription)
                .where(
                    PushSubscription.id == subscription_id,
                    PushSubscription.status == PushSubscriptionStatus.ACTIVE,
                    PushSubscription.next_run_at.is_not(None),
                    PushSubscription.next_run_at <= current,
                    or_(
                        PushSubscription.locked_until.is_(None),
                        PushSubscription.locked_until <= current,
                    ),
                )
                .values(locked_by=lock_owner, locked_until=locked_until)
                .execution_options(synchronize_session=False)
            )
            if result.rowcount == 1:
                claims.append((subscription_id, lock_owner))
        self.db.commit()
        return claims

    def _materialize_delivery(
        self,
        *,
        subscription_id: int,
        lock_owner: str,
        current: datetime,
    ) -> PushDelivery | None:
        subscription = self.db.execute(
            select(PushSubscription)
            .where(
                PushSubscription.id == subscription_id,
                PushSubscription.locked_by == lock_owner,
            )
            .with_for_update()
        ).scalar_one_or_none()
        if subscription is None:
            self.db.rollback()
            return None
        if (
            subscription.status != PushSubscriptionStatus.ACTIVE
            or subscription.next_run_at is None
        ):
            subscription.locked_by = None
            subscription.locked_until = None
            self.db.commit()
            return None

        owner = self.db.execute(
            select(FeishuUser)
            .where(FeishuUser.id == subscription.owner_id)
            .with_for_update()
        ).scalar_one_or_none()
        scheduled_for = subscription.next_run_at
        skip_reason = self._delivery_skip_reason(subscription, owner, current)
        next_attempt_at = current
        if (
            skip_reason is None
            and owner is not None
            and is_in_quiet_hours(
                current,
                owner.timezone,
                owner.quiet_hours_start,
                owner.quiet_hours_end,
            )
        ):
            next_attempt_at = next_quiet_end(
                current,
                owner.timezone,
                owner.quiet_hours_start,
                owner.quiet_hours_end,
            )

        idempotency_key = _delivery_key(subscription.id, scheduled_for)
        message_uuid = str(uuid5(NAMESPACE_URL, f"company-ai-platform:{idempotency_key}"))
        delivery = PushDelivery(
            code=f"DEL-{uuid4().hex.upper()}",
            subscription_id=subscription.id,
            scheduled_for=scheduled_for,
            idempotency_key=idempotency_key,
            message_uuid=message_uuid,
            status=(
                PushDeliveryStatus.SKIPPED
                if skip_reason is not None
                else PushDeliveryStatus.PENDING
            ),
            next_attempt_at=None if skip_reason is not None else next_attempt_at,
            last_error=skip_reason,
            created_at=current,
            updated_at=current,
        )
        try:
            with self.db.begin_nested():
                self.db.add(delivery)
                self.db.flush()
        except IntegrityError:
            delivery = self.db.execute(
                select(PushDelivery).where(
                    PushDelivery.idempotency_key == idempotency_key
                )
            ).scalar_one()

        subscription.last_run_at = scheduled_for
        if subscription.schedule_type == SubscriptionScheduleType.ONCE:
            subscription.status = PushSubscriptionStatus.COMPLETED
            subscription.next_run_at = None
        else:
            subscription.next_run_at = next_occurrence(
                subscription.schedule_type,
                subscription.schedule_config,
                subscription.timezone,
                after=current,
            )
        subscription.locked_by = None
        subscription.locked_until = None
        if skip_reason is not None:
            self._audit_skipped(owner, delivery, skip_reason)
        self.db.commit()
        self.db.refresh(delivery)
        return delivery

    def _delivery_skip_reason(
        self,
        subscription: PushSubscription,
        owner: FeishuUser | None,
        current: datetime,
    ) -> str | None:
        if owner is None or owner.status != FeishuUserStatus.ACTIVE:
            return "Feishu user is disabled"
        if (
            subscription.target_type == SubscriptionTargetType.USER
            and subscription.target_id != owner.open_id
        ):
            return "Private subscription target no longer matches its owner"
        if subscription.target_type == SubscriptionTargetType.CHAT and (
            owner.role != FeishuUserRole.ADMIN or not subscription.target_id
        ):
            return "Group subscription owner is no longer an administrator"
        if subscription.target_type not in {
            SubscriptionTargetType.USER,
            SubscriptionTargetType.CHAT,
        }:
            return "Unsupported subscription target"
        if self._daily_delivery_count(owner, current) >= MAX_DAILY_DELIVERIES:
            return DAILY_DELIVERY_LIMIT_REACHED
        return None

    def _daily_delivery_count(self, owner: FeishuUser, current: datetime) -> int:
        zone = validate_timezone(owner.timezone)
        local_now = current.replace(tzinfo=UTC).astimezone(zone)
        local_start = datetime.combine(local_now.date(), time.min, tzinfo=zone)
        local_end = local_start + timedelta(days=1)
        start_utc = local_start.astimezone(UTC).replace(tzinfo=None)
        end_utc = local_end.astimezone(UTC).replace(tzinfo=None)
        return int(
            self.db.scalar(
                select(func.count())
                .select_from(PushDelivery)
                .join(PushSubscription)
                .where(
                    PushSubscription.owner_id == owner.id,
                    PushDelivery.created_at >= start_utc,
                    PushDelivery.created_at < end_utc,
                    PushDelivery.status.not_in(
                        [
                            PushDeliveryStatus.FAILED,
                            PushDeliveryStatus.SKIPPED,
                        ]
                    ),
                )
            )
            or 0
        )

    def _audit_skipped(
        self,
        owner: FeishuUser | None,
        delivery: PushDelivery,
        reason: str,
    ) -> None:
        AuditService(self.db).record(
            AuditLogCreate(
                actor=owner.code if owner is not None else "subscription-system",
                source="subscriptions",
                action=SubscriptionAuditAction.DELIVERY_SKIPPED,
                target_type="push-delivery",
                target_id=delivery.code,
                response_payload={"status": PushDeliveryStatus.SKIPPED, "reason": reason},
            )
        )

    def _release_claim(self, subscription_id: int, lock_owner: str) -> None:
        self.db.execute(
            update(PushSubscription)
            .where(
                PushSubscription.id == subscription_id,
                PushSubscription.locked_by == lock_owner,
            )
            .values(locked_by=None, locked_until=None)
            .execution_options(synchronize_session=False)
        )
        self.db.commit()


def _delivery_key(subscription_id: int, scheduled_for: datetime) -> str:
    material = f"{subscription_id}:{_naive_utc(scheduled_for).isoformat(timespec='microseconds')}"
    return sha256(material.encode("utf-8")).hexdigest()


def _naive_utc(value: datetime) -> datetime:
    if value.tzinfo is None:
        return value
    return value.astimezone(UTC).replace(tzinfo=None)