diff --git a/finance/exporting.py b/finance/exporting.py index 3051a5e..29cc4d1 100644 --- a/finance/exporting.py +++ b/finance/exporting.py @@ -652,6 +652,11 @@ def process_formula_export(export_id): created = [] try: member = Membership.objects.get(user=job.user, organization=job.dataset.organization) + # 执行阶段重新核验角色与账号状态：撤权或停用后排队中的任务不得继续发布敏感工作簿。 + if not job.user.is_active: + raise PermissionError("账号已停用，公式导出已拒绝执行") + if member.role not in ("admin", "analyst"): + raise PermissionError("账号角色已变更，公式导出已拒绝执行") if set(member.departments) != set(job.departments): raise PermissionError("科室授权已变更，请重新提交导出") root = _validated_export_directory(Path(settings.MEDIA_ROOT) / "formula-exports", "公式导出目录") diff --git a/finance/formula_views.py b/finance/formula_views.py index 898b741..075e22e 100644 --- a/finance/formula_views.py +++ b/finance/formula_views.py @@ -292,7 +292,7 @@ def _owned_job(request, dataset_id, export_id): return job -@endpoint(("GET",)) +@endpoint(("GET",), roles=("admin", "analyst")) def formula_export_status(request, dataset_id, export_id): job = _owned_job(request, dataset_id, export_id) return JsonResponse( @@ -305,7 +305,7 @@ def formula_export_status(request, dataset_id, export_id): ) -@endpoint(("GET",)) +@endpoint(("GET",), roles=("admin", "analyst")) def formula_export_download(request, dataset_id, export_id): job = _owned_job(request, dataset_id, export_id) if job.status != "succeeded" or not job.path or not job.expires_at or job.expires_at <= timezone.now(): diff --git a/finance/management/commands/import_worker.py b/finance/management/commands/import_worker.py index 469b76b..f11ada7 100644 --- a/finance/management/commands/import_worker.py +++ b/finance/management/commands/import_worker.py @@ -1,13 +1,19 @@ 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" @@ -30,30 +36,35 @@ class Command(BaseCommand): with connection.cursor() as cursor: cursor.execute("SELECT pg_advisory_unlock(782146291)") - def run_jobs(self, options, recover): + def run_jobs(self, options, recover_immediately): while True: - processed = self._next_import(recover) + processed = self._next_import(recover_immediately) if not processed: - processed = self._next_export(recover) + processed = self._next_export(recover_immediately) if options["once"]: return if not processed: time.sleep(2) - def _claim(self, model, states): + 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(status__in=states).order_by("created_at") + 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 - updated = model.objects.filter(pk=item.pk, status__in=states).update(status="processing") - return item if updated else 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): - states = ["queued", "processing"] if recover else ["queued"] - item = self._claim(Dataset, states) + 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()) @@ -61,9 +72,8 @@ class Command(BaseCommand): self.stdout.write(f"{result.pk}: {result.status}, {result.row_count} rows") return True - def _next_export(self, recover): - states = ["queued", "processing"] if recover else ["queued"] - item = self._claim(FormulaExport, states) + def _next_export(self, recover_immediately): + item = self._claim(FormulaExport, recover_immediately) if not item: return False try: diff --git a/tests/test_worker_paths.py b/tests/test_worker_paths.py index 61e1402..a08a0b1 100644 --- a/tests/test_worker_paths.py +++ b/tests/test_worker_paths.py @@ -1,7 +1,10 @@ +from datetime import timedelta +from pathlib import Path from unittest.mock import patch import pytest from django.core.management import call_command +from django.utils import timezone from finance.models import FormulaExport @@ -26,3 +29,118 @@ def test_worker_processes_queued_import_and_reports_export_failure(workspace): call_command("import_worker", once=True) export.refresh_from_db() assert export.status == "processing" + + +def _export_job(ds, **kwargs): + return FormulaExport.objects.create( + dataset=ds, + user=ds.created_by, + context={"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"}, + **kwargs, + ) + + +def test_worker_refuses_export_after_role_revocation(workspace): + from finance.exporting import process_formula_export + + ds, member, *_ = workspace + job = _export_job(ds, status="processing") + member.role = "viewer" + member.save(update_fields=["role"]) + + with pytest.raises(PermissionError, match="角色"): + process_formula_export(job.pk) + + job.refresh_from_db() + assert job.status == "failed" + assert not job.path + assert not list(Path(ds.file.storage.location).rglob(f"{job.pk}*")) + + +def test_worker_refuses_export_after_account_deactivation(workspace): + from finance.exporting import process_formula_export + + ds, _, _, _, _ = workspace + user = ds.created_by + user.is_active = False + user.save(update_fields=["is_active"]) + job = _export_job(ds, status="processing") + + with pytest.raises(PermissionError, match="停用"): + process_formula_export(job.pk) + + job.refresh_from_db() + assert job.status == "failed" + + +def test_status_and_download_deny_revoked_viewer(workspace, settings): + ds, member, client, *_ = workspace + job = _export_job( + ds, + status="succeeded", + finished_at=timezone.now(), + expires_at=timezone.now() + timedelta(hours=1), + ) + target = Path(settings.MEDIA_ROOT) / "formula-exports" / str(job.user_id) / f"{job.pk}.xlsx" + target.parent.mkdir(parents=True) + target.write_bytes(b"PK workbook") + job.path = str(target) + job.save(update_fields=["path"]) + + assert client.get(f"/api/datasets/{ds.pk}/formula-exports/{job.pk}").status_code == 200 + assert ( + client.get(f"/api/datasets/{ds.pk}/formula-exports/{job.pk}/download").status_code == 200 + ) + + member.role = "viewer" + member.save(update_fields=["role"]) + assert client.get(f"/api/datasets/{ds.pk}/formula-exports/{job.pk}").status_code == 403 + assert ( + client.get(f"/api/datasets/{ds.pk}/formula-exports/{job.pk}/download").status_code == 403 + ) + + +def test_half_finished_export_cannot_be_downloaded(workspace, settings): + ds, _, client, *_ = workspace + for status, code in (("processing", 409), ("failed", 409), ("queued", 409)): + job = _export_job(ds, status=status) + assert ( + client.get(f"/api/datasets/{ds.pk}/formula-exports/{job.pk}/download").status_code + == code + ), status + succeeded = _export_job( + ds, + status="succeeded", + path=str(Path(settings.MEDIA_ROOT) / "nowhere.xlsx"), + finished_at=timezone.now(), + expires_at=timezone.now() + timedelta(hours=1), + ) + assert ( + client.get(f"/api/datasets/{ds.pk}/formula-exports/{succeeded.pk}/download").status_code + == 410 + ) + + +def test_stale_processing_export_is_reclaimed_after_timeout(workspace): + ds, *_ = workspace + job = _export_job(ds, status="processing") + FormulaExport.objects.filter(pk=job.pk).update( + created_at=timezone.now() - timedelta(minutes=20) + ) + + call_command("import_worker", once=True) + + job.refresh_from_db() + assert job.status == "succeeded" + assert Path(job.path).read_bytes()[:2] == b"PK" + assert not list(Path(job.path).parent.glob(f"{job.pk}.*.tmp")) + + +def test_recent_processing_export_is_not_reclaimed(workspace): + ds, *_ = workspace + job = _export_job(ds, status="processing") + + call_command("import_worker", once=True) + + job.refresh_from_db() + assert job.status == "processing"