Jayway Liu
5ebe83afbb
- worker 执行前重新核验账号 is_active 与 admin/analyst 角色：撤权/停用后排队任务 直接失败，不发布敏感工作簿 - formula-exports 状态与下载接口补齐与创建接口一致的角色白名单，被降为 viewer 的账号不能再查询/下载旧导出 - import_worker 对超过 15 分钟仍处于 processing 的任务重新认领（PostgreSQL 依赖 咨询锁立即恢复），崩溃重启不再无限停留 processing；新近 processing 不被抢占 - 测试：撤权/停用执行、viewer 下载与状态 403、半成品/缺文件下载 409/410、 超时任务恢复且无重复临时文件、新近 processing 不被重复认领
85 lines
3.5 KiB
Python
85 lines
3.5 KiB
Python
import time
|
|
from datetime import timedelta
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.db import connection, transaction
|
|
from django.db.models import Q
|
|
from django.utils import timezone
|
|
|
|
from finance.exporting import process_formula_export
|
|
from finance.importing import process_dataset
|
|
from finance.models import Dataset, FormulaExport
|
|
|
|
# 无咨询锁的数据库（如 SQLite）上等待该时限后才重新认领 processing 任务，
|
|
# 避免与仍在运行的 worker 抢任务；PostgreSQL 由咨询锁保证单 worker，可立即恢复。
|
|
STALE_PROCESSING_SECONDS = 15 * 60
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Process durable import and formula-export jobs"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument("--once", action="store_true")
|
|
|
|
def handle(self, *args, **options):
|
|
postgres = connection.vendor == "postgresql"
|
|
if postgres:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute("SELECT pg_try_advisory_lock(782146291)")
|
|
if not cursor.fetchone()[0]:
|
|
self.stdout.write("Another worker owns the database lock")
|
|
return
|
|
try:
|
|
self.run_jobs(options, postgres)
|
|
finally:
|
|
if postgres:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute("SELECT pg_advisory_unlock(782146291)")
|
|
|
|
def run_jobs(self, options, recover_immediately):
|
|
while True:
|
|
processed = self._next_import(recover_immediately)
|
|
if not processed:
|
|
processed = self._next_export(recover_immediately)
|
|
if options["once"]:
|
|
return
|
|
if not processed:
|
|
time.sleep(2)
|
|
|
|
def _claim(self, model, immediate_recovery):
|
|
stale_before = timezone.now() - timedelta(seconds=STALE_PROCESSING_SECONDS)
|
|
claimable = Q(status="queued") | Q(status="processing", created_at__lt=stale_before)
|
|
if immediate_recovery:
|
|
claimable = Q(status="queued") | Q(status="processing")
|
|
with transaction.atomic():
|
|
qs = model.objects.filter(claimable).order_by("created_at")
|
|
if connection.features.has_select_for_update_skip_locked:
|
|
qs = qs.select_for_update(skip_locked=True)
|
|
item = qs.first()
|
|
if not item:
|
|
return None
|
|
# CAS 认领 queued 任务互斥；processing 仅在超时或咨询锁保护下重认领，
|
|
# 重复发布对同一目标路径写入相同内容，不会产生重复文件。
|
|
claimed = model.objects.filter(pk=item.pk, status=item.status).update(status="processing")
|
|
return item if claimed else None
|
|
|
|
def _next_import(self, recover_immediately):
|
|
item = self._claim(Dataset, recover_immediately)
|
|
if not item:
|
|
return False
|
|
Dataset.objects.filter(pk=item.pk).update(started_at=timezone.now())
|
|
result = process_dataset(item.pk)
|
|
self.stdout.write(f"{result.pk}: {result.status}, {result.row_count} rows")
|
|
return True
|
|
|
|
def _next_export(self, recover_immediately):
|
|
item = self._claim(FormulaExport, recover_immediately)
|
|
if not item:
|
|
return False
|
|
try:
|
|
stats = process_formula_export(item.pk)
|
|
self.stdout.write(f"{item.pk}: formula export ready, {stats['rows']} rows")
|
|
except Exception as exc:
|
|
self.stderr.write(f"Formula export {item.pk} failed: {exc}")
|
|
return True
|