All code that runs outside a request context (APScheduler jobs, lifespan startup) has no app.current_account_id set, so the app-role session returns 0 rows from every RLS-protected table. Changed to _admin_session_factory (BYPASSRLS) in: - knowledge_flywheel_scheduler.py — queries ai_sessions - psa_retry_scheduler.py — queries psa_post_log - retention_cleanup.py — queries assistant_chats - scheduler.py (_fire_maintenance_schedule, _cleanup_expired_ai_conversations) - main.py (archive_stale_ai_sessions, _process_notification_retries, load_all_schedules at startup) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""Chat retention cleanup job.
|
|
|
|
Runs daily via APScheduler to enforce account-level retention settings:
|
|
- Delete non-pinned chats older than chat_retention_days
|
|
- Delete oldest non-pinned chats when count exceeds chat_retention_max_count
|
|
"""
|
|
import logging
|
|
from datetime import datetime, timezone, timedelta
|
|
|
|
from sqlalchemy import select, delete, func
|
|
|
|
from app.core.admin_database import _admin_session_factory as async_session_maker
|
|
from app.models.account import Account
|
|
from app.models.assistant_chat import AssistantChat
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def cleanup_expired_chats() -> None:
|
|
"""Enforce chat retention policies for all accounts."""
|
|
async with async_session_maker() as db:
|
|
try:
|
|
result = await db.execute(select(Account))
|
|
accounts = result.scalars().all()
|
|
|
|
total_deleted = 0
|
|
for account in accounts:
|
|
deleted = await _cleanup_account_chats(account, db)
|
|
total_deleted += deleted
|
|
|
|
await db.commit()
|
|
if total_deleted > 0:
|
|
logger.info("[retention] Cleaned up %d expired chats", total_deleted)
|
|
except Exception as e:
|
|
logger.error("[retention] Chat cleanup failed: %s", e)
|
|
await db.rollback()
|
|
|
|
|
|
async def _cleanup_account_chats(account: Account, db) -> int:
|
|
"""Enforce retention for a single account. Returns count deleted."""
|
|
deleted = 0
|
|
|
|
# Age-based retention
|
|
if account.chat_retention_days:
|
|
cutoff = datetime.now(timezone.utc) - timedelta(days=account.chat_retention_days)
|
|
result = await db.execute(
|
|
delete(AssistantChat)
|
|
.where(
|
|
AssistantChat.account_id == account.id,
|
|
AssistantChat.pinned == False, # noqa: E712
|
|
AssistantChat.updated_at < cutoff,
|
|
)
|
|
.returning(AssistantChat.id)
|
|
)
|
|
deleted += len(result.all())
|
|
|
|
# Count-based retention
|
|
if account.chat_retention_max_count:
|
|
total = await db.scalar(
|
|
select(func.count(AssistantChat.id)).where(
|
|
AssistantChat.account_id == account.id,
|
|
)
|
|
) or 0
|
|
|
|
if total > account.chat_retention_max_count:
|
|
excess = total - account.chat_retention_max_count
|
|
# Get oldest non-pinned chat IDs
|
|
oldest = await db.execute(
|
|
select(AssistantChat.id)
|
|
.where(
|
|
AssistantChat.account_id == account.id,
|
|
AssistantChat.pinned == False, # noqa: E712
|
|
)
|
|
.order_by(AssistantChat.updated_at.asc())
|
|
.limit(excess)
|
|
)
|
|
ids_to_delete = [row[0] for row in oldest.all()]
|
|
if ids_to_delete:
|
|
await db.execute(
|
|
delete(AssistantChat).where(AssistantChat.id.in_(ids_to_delete))
|
|
)
|
|
deleted += len(ids_to_delete)
|
|
|
|
return deleted
|