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)