安全收尾:F01/F02/F08 导出清理边界、孤儿回收、撤权与任务恢复 #1

Open
jianwel wants to merge 5 commits from fix/security-f01-f02-f08 into main
7 changed files with 510 additions and 43 deletions
+80 -21
View File
@@ -1,7 +1,9 @@
import hashlib import hashlib
import hmac import hmac
import logging
import os import os
import secrets import secrets
import stat
import warnings import warnings
from datetime import timedelta from datetime import timedelta
from decimal import Decimal, InvalidOperation from decimal import Decimal, InvalidOperation
@@ -28,6 +30,58 @@ FORMULA_VERSION = "3"
MAX_ROWS = 100_000 MAX_ROWS = 100_000
MAX_GROUPS = 200 MAX_GROUPS = 200
logger = logging.getLogger(__name__)
def _validated_export_directory(path, description):
"""导出目录必须不存在或为真实目录;符号链接一律拒绝并保持文件原样。"""
try:
info = os.lstat(path)
except FileNotFoundError:
path.mkdir(parents=True, mode=0o700)
info = os.lstat(path)
if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode):
raise ValueError(f"{description}必须是真实目录,不能是符号链接")
os.chmod(path, 0o700)
return path
def _validated_export_target(user_dir, job):
"""目标文件若已存在必须是普通文件;符号链接或其他类型一律拒绝覆盖。"""
target = user_dir / f"{job.pk}.xlsx"
try:
info = os.lstat(target)
except FileNotFoundError:
return target
if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
raise ValueError("公式导出目标路径已存在且不是普通文件,已拒绝生成")
return target
def _unlink_created_file(path):
"""只通过不跟随链接的目录句柄删除受控目录内的普通文件。
返回 "removed"(已删除)、"missing"(本就不存在)或 "blocked"(无法安全删除,
例如目录已被换成符号链接或权限不足),调用方据此保留失败任务与磁盘的一致性。
"""
try:
dir_fd = os.open(path.parent, os.O_RDONLY | os.O_NOFOLLOW)
except OSError:
return "blocked"
try:
try:
info = os.lstat(path.name, dir_fd=dir_fd)
except FileNotFoundError:
return "missing"
if not stat.S_ISREG(info.st_mode):
return "blocked"
os.unlink(path.name, dir_fd=dir_fd)
return "removed"
except OSError:
return "blocked"
finally:
os.close(dir_fd)
def _excel_number(value): def _excel_number(value):
if value is None: if value is None:
@@ -594,41 +648,46 @@ def build_formula_workbook(ds, member, export_context, output):
def process_formula_export(export_id): def process_formula_export(export_id):
job = FormulaExport.objects.select_related("dataset", "user").get(pk=export_id) job = FormulaExport.objects.select_related("dataset", "user").get(pk=export_id)
temporary = None
target = None target = None
created = []
try: try:
root = Path(settings.MEDIA_ROOT) / "formula-exports"
if root.is_symlink():
raise ValueError("公式导出目录不能是符号链接")
root.mkdir(parents=True, exist_ok=True, mode=0o700)
root.chmod(0o700)
target = root / str(job.user_id) / f"{job.pk}.xlsx"
target.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
if target.parent.is_symlink() or target.is_symlink():
raise ValueError("公式导出路径不能是符号链接")
target.parent.chmod(0o700)
temporary = target.with_name(f"{job.pk}.{secrets.token_hex(8)}.tmp")
member = Membership.objects.get(user=job.user, organization=job.dataset.organization) 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): if set(member.departments) != set(job.departments):
raise PermissionError("科室授权已变更,请重新提交导出") raise PermissionError("科室授权已变更,请重新提交导出")
root = _validated_export_directory(Path(settings.MEDIA_ROOT) / "formula-exports", "公式导出目录")
user_dir = _validated_export_directory(root / str(job.user_id), "公式导出用户目录")
target = _validated_export_target(user_dir, job)
# 路径完整验证通过后才把目标登记为可清理文件;此前任何失败都不会删除磁盘文件。
created.append(target)
temporary = target.with_name(f"{job.pk}.{secrets.token_hex(8)}.tmp")
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
created.append(temporary)
with os.fdopen(descriptor, "wb") as file: with os.fdopen(descriptor, "wb") as file:
stats = build_formula_workbook(job.dataset, member, job.context, file) stats = build_formula_workbook(job.dataset, member, job.context, file)
temporary.replace(target) os.replace(temporary, target)
created = [target]
job.path, job.status, job.error = str(target), "succeeded", "" job.path, job.status, job.error = str(target), "succeeded", ""
job.finished_at = timezone.now() job.finished_at = timezone.now()
job.expires_at = timezone.now() + timedelta(hours=24) job.expires_at = timezone.now() + timedelta(hours=24)
job.save(update_fields=["path", "status", "error", "finished_at", "expires_at"]) job.save(update_fields=["path", "status", "error", "finished_at", "expires_at"])
return stats return stats
except Exception as exc: except Exception as exc:
for path in (temporary, target): job.status, job.error, job.finished_at, job.expires_at = "failed", str(exc)[:500], timezone.now(), None
try: # 只清理本次登记且确认位于受控目录内的文件;目录被换成符号链接时拒绝删除。
if path is not None and not path.is_symlink(): leftover = [path for path in created if _unlink_created_file(path) == "blocked"]
path.unlink(missing_ok=True) if leftover:
except OSError: # 日志只含任务编号与文件名,不包含敏感数据;残留文件由清理命令按保留期回收。
pass logger.warning(
job.status, job.error, job.path = "failed", str(exc)[:500], "" "公式导出任务 %s 失败后清理文件受阻,待清理命令回收:%s",
job.finished_at, job.expires_at = timezone.now(), None job.pk,
"、".join(path.name for path in leftover),
)
job.path = str(target) if target in leftover else ""
job.save(update_fields=["status", "error", "path", "finished_at", "expires_at"]) job.save(update_fields=["status", "error", "path", "finished_at", "expires_at"])
raise raise
+2 -2
View File
@@ -292,7 +292,7 @@ def _owned_job(request, dataset_id, export_id):
return job return job
@endpoint(("GET",)) @endpoint(("GET",), roles=("admin", "analyst"))
def formula_export_status(request, dataset_id, export_id): def formula_export_status(request, dataset_id, export_id):
job = _owned_job(request, dataset_id, export_id) job = _owned_job(request, dataset_id, export_id)
return JsonResponse( 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): def formula_export_download(request, dataset_id, export_id):
job = _owned_job(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(): if job.status != "succeeded" or not job.path or not job.expires_at or job.expires_at <= timezone.now():
@@ -1,4 +1,5 @@
import time import time
import uuid
from datetime import timedelta from datetime import timedelta
from pathlib import Path from pathlib import Path
@@ -9,9 +10,11 @@ from django.utils import timezone
from finance.models import FormulaExport from finance.models import FormulaExport
ORPHAN_RETENTION_HOURS = 24
class Command(BaseCommand): class Command(BaseCommand):
help = "Remove expired formula workbooks and stale temporary files" help = "Remove expired formula workbooks and reclaim orphan files"
def add_arguments(self, parser): def add_arguments(self, parser):
parser.add_argument("--loop", action="store_true") parser.add_argument("--loop", action="store_true")
@@ -19,8 +22,10 @@ class Command(BaseCommand):
def handle(self, *args, **options): def handle(self, *args, **options):
while True: while True:
jobs, files = self.clean_expired() jobs, files, orphans = self.clean_expired()
self.stdout.write(f"Expired {jobs} jobs; removed {files} files") self.stdout.write(
f"Expired {jobs} jobs; removed {files} files; reclaimed {orphans} orphan files"
)
if not options["loop"]: if not options["loop"]:
return return
time.sleep(max(60, options["interval"])) time.sleep(max(60, options["interval"]))
@@ -31,7 +36,7 @@ class Command(BaseCommand):
raise ValueError("公式导出目录不能是符号链接") raise ValueError("公式导出目录不能是符号链接")
root = export_root.resolve() root = export_root.resolve()
now = timezone.now() now = timezone.now()
jobs = files = 0 jobs = files = orphans = 0
for job in FormulaExport.objects.filter(status="succeeded", expires_at__lte=now).iterator(): for job in FormulaExport.objects.filter(status="succeeded", expires_at__lte=now).iterator():
path = Path(job.path) if job.path else None path = Path(job.path) if job.path else None
@@ -61,5 +66,41 @@ class Command(BaseCommand):
files += 1 files += 1
except OSError as exc: except OSError as exc:
self.stderr.write(f"Could not remove stale temporary export {path.name}: {exc}") self.stderr.write(f"Could not remove stale temporary export {path.name}: {exc}")
orphans += self._reclaim_orphan_workbooks(user_dir, now)
return jobs, files return jobs, files, orphans
def _reclaim_orphan_workbooks(self, user_dir, now):
"""回收超过保留期且不再属于任何有效导出任务的 .xlsx 孤儿文件。
有效任务指状态 succeeded、路径一致且尚未过期的任务;失败任务残留的文件和任务行
已被删除(如数据集级联删除)的文件在保留期后回收,绝不触碰活动任务的文件。
"""
retention = (now - timedelta(hours=ORPHAN_RETENTION_HOURS)).timestamp()
removed = 0
for path in user_dir.glob("*.xlsx"):
try:
if path.is_symlink() or not path.is_file() or path.stat().st_mtime > retention:
continue
job_id = uuid.UUID(path.stem)
except (OSError, ValueError):
continue
job = FormulaExport.objects.filter(pk=job_id).first()
if (
job is not None
and job.status == "succeeded"
and job.path == str(path)
and job.expires_at is not None
and job.expires_at > now
):
continue
try:
path.unlink()
removed += 1
if job is not None and job.path == str(path):
job.path = ""
job.save(update_fields=["path"])
self.stdout.write(f"Reclaimed leftover workbook of job {job.pk}")
except OSError as exc:
self.stderr.write(f"Could not reclaim orphan export {path.stem}: {exc}")
return removed
+23 -13
View File
@@ -1,13 +1,19 @@
import time import time
from datetime import timedelta
from django.core.management.base import BaseCommand from django.core.management.base import BaseCommand
from django.db import connection, transaction from django.db import connection, transaction
from django.db.models import Q
from django.utils import timezone from django.utils import timezone
from finance.exporting import process_formula_export from finance.exporting import process_formula_export
from finance.importing import process_dataset from finance.importing import process_dataset
from finance.models import Dataset, FormulaExport from finance.models import Dataset, FormulaExport
# 无咨询锁的数据库(如 SQLite)上等待该时限后才重新认领 processing 任务,
# 避免与仍在运行的 worker 抢任务;PostgreSQL 由咨询锁保证单 worker,可立即恢复。
STALE_PROCESSING_SECONDS = 15 * 60
class Command(BaseCommand): class Command(BaseCommand):
help = "Process durable import and formula-export jobs" help = "Process durable import and formula-export jobs"
@@ -30,30 +36,35 @@ class Command(BaseCommand):
with connection.cursor() as cursor: with connection.cursor() as cursor:
cursor.execute("SELECT pg_advisory_unlock(782146291)") cursor.execute("SELECT pg_advisory_unlock(782146291)")
def run_jobs(self, options, recover): def run_jobs(self, options, recover_immediately):
while True: while True:
processed = self._next_import(recover) processed = self._next_import(recover_immediately)
if not processed: if not processed:
processed = self._next_export(recover) processed = self._next_export(recover_immediately)
if options["once"]: if options["once"]:
return return
if not processed: if not processed:
time.sleep(2) 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(): 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: if connection.features.has_select_for_update_skip_locked:
qs = qs.select_for_update(skip_locked=True) qs = qs.select_for_update(skip_locked=True)
item = qs.first() item = qs.first()
if not item: if not item:
return None return None
updated = model.objects.filter(pk=item.pk, status__in=states).update(status="processing") # CAS 认领 queued 任务互斥;processing 仅在超时或咨询锁保护下重认领,
return item if updated else None # 重复发布对同一目标路径写入相同内容,不会产生重复文件。
claimed = model.objects.filter(pk=item.pk, status=item.status).update(status="processing")
return item if claimed else None
def _next_import(self, recover): def _next_import(self, recover_immediately):
states = ["queued", "processing"] if recover else ["queued"] item = self._claim(Dataset, recover_immediately)
item = self._claim(Dataset, states)
if not item: if not item:
return False return False
Dataset.objects.filter(pk=item.pk).update(started_at=timezone.now()) 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") self.stdout.write(f"{result.pk}: {result.status}, {result.row_count} rows")
return True return True
def _next_export(self, recover): def _next_export(self, recover_immediately):
states = ["queued", "processing"] if recover else ["queued"] item = self._claim(FormulaExport, recover_immediately)
item = self._claim(FormulaExport, states)
if not item: if not item:
return False return False
try: try:
+76 -2
View File
@@ -1,4 +1,5 @@
import os import os
import uuid
from datetime import timedelta from datetime import timedelta
import pytest import pytest
@@ -27,6 +28,12 @@ def cleanup_root(settings, tmp_path):
return tmp_path / "formula-exports", dataset, user return tmp_path / "formula-exports", dataset, user
def _aged(path, hours=25):
stamp = (timezone.now() - timedelta(hours=hours)).timestamp()
os.utime(path, (stamp, stamp))
return path
def test_cleanup_expires_jobs_and_removes_files(cleanup_root): def test_cleanup_expires_jobs_and_removes_files(cleanup_root):
root, dataset, user = cleanup_root root, dataset, user = cleanup_root
path = root / str(user.pk) / "expired.xlsx" path = root / str(user.pk) / "expired.xlsx"
@@ -55,8 +62,7 @@ def test_cleanup_removes_only_stale_temporary_files(cleanup_root):
stale, recent = user_dir / "stale.tmp", user_dir / "recent.tmp" stale, recent = user_dir / "stale.tmp", user_dir / "recent.tmp"
stale.touch() stale.touch()
recent.touch() recent.touch()
old = timezone.now().timestamp() - 25 * 3600 _aged(stale)
os.utime(stale, (old, old))
call_command("cleanup_formula_exports") call_command("cleanup_formula_exports")
@@ -81,3 +87,71 @@ def test_cleanup_does_not_delete_files_outside_export_root(cleanup_root, tmp_pat
job.refresh_from_db() job.refresh_from_db()
assert job.status == "expired" assert job.status == "expired"
assert outside.read_bytes() == b"keep" assert outside.read_bytes() == b"keep"
def test_cleanup_reclaims_failed_job_leftover_after_retention(cleanup_root):
root, dataset, user = cleanup_root
job = FormulaExport.objects.create(dataset=dataset, user=user, status="failed")
path = root / str(user.pk) / f"{job.pk}.xlsx"
path.parent.mkdir(parents=True)
path.write_bytes(b"leftover")
job.path = str(path)
job.save(update_fields=["path"])
_aged(path)
call_command("cleanup_formula_exports")
job.refresh_from_db()
assert not path.exists()
assert job.path == ""
def test_cleanup_keeps_workbook_of_unexpired_succeeded_job(cleanup_root):
root, dataset, user = cleanup_root
job = FormulaExport.objects.create(
dataset=dataset,
user=user,
status="succeeded",
expires_at=timezone.now() + timedelta(hours=1),
)
path = root / str(user.pk) / f"{job.pk}.xlsx"
path.parent.mkdir(parents=True)
path.write_bytes(b"active")
job.path = str(path)
job.save(update_fields=["path"])
_aged(path)
call_command("cleanup_formula_exports")
assert path.read_bytes() == b"active"
job.refresh_from_db()
assert job.path == str(path)
def test_cleanup_reclaims_orphan_whose_job_row_was_deleted(cleanup_root):
root, _, user = cleanup_root
user_dir = root / str(user.pk)
user_dir.mkdir(parents=True)
orphan = user_dir / f"{uuid.uuid4()}.xlsx"
orphan.write_bytes(b"orphan")
recent = user_dir / f"{uuid.uuid4()}.xlsx"
recent.write_bytes(b"recent orphan")
_aged(orphan)
call_command("cleanup_formula_exports")
assert not orphan.exists()
assert recent.exists()
def test_cleanup_ignores_non_uuid_workbook_names(cleanup_root):
root, _, user = cleanup_root
user_dir = root / str(user.pk)
user_dir.mkdir(parents=True)
odd = user_dir / "notes.xlsx"
odd.write_bytes(b"keep")
_aged(odd)
call_command("cleanup_formula_exports")
assert odd.exists()
+165
View File
@@ -1,3 +1,4 @@
import os
import stat import stat
from decimal import Decimal from decimal import Decimal
from io import BytesIO from io import BytesIO
@@ -198,6 +199,170 @@ def test_formula_export_worker_records_path_preparation_failure(workspace, setti
assert job.status == "failed" assert job.status == "failed"
def _queued_job(ds):
return FormulaExport.objects.create(
dataset=ds,
user=ds.created_by,
context={"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"},
)
def test_failed_export_rejects_symlinked_user_directory_without_touching_files(workspace, settings):
from finance.exporting import process_formula_export
ds, *_ = workspace
job = _queued_job(ds)
root = Path(settings.MEDIA_ROOT) / "formula-exports"
root.mkdir(parents=True)
protected = Path(settings.MEDIA_ROOT) / "protected"
protected.mkdir()
victim = protected / f"{job.pk}.xlsx"
victim.write_bytes(b"protected workbook")
(root / str(job.user_id)).symlink_to(protected, target_is_directory=True)
with pytest.raises(ValueError, match="符号链接"):
process_formula_export(job.pk)
job.refresh_from_db()
assert job.status == "failed"
assert not job.path
assert victim.read_bytes() == b"protected workbook"
assert list(protected.glob("*.tmp")) == []
def test_failed_export_rejects_symlinked_target_and_keeps_link(workspace, settings):
from finance.exporting import process_formula_export
ds, *_ = workspace
job = _queued_job(ds)
user_dir = Path(settings.MEDIA_ROOT) / "formula-exports" / str(job.user_id)
user_dir.mkdir(parents=True)
protected_file = Path(settings.MEDIA_ROOT) / "protected.xlsx"
protected_file.write_bytes(b"protected workbook")
(user_dir / f"{job.pk}.xlsx").symlink_to(protected_file)
with pytest.raises(ValueError, match="普通文件"):
process_formula_export(job.pk)
job.refresh_from_db()
assert job.status == "failed"
assert protected_file.read_bytes() == b"protected workbook"
assert (user_dir / f"{job.pk}.xlsx").is_symlink()
assert list(user_dir.glob("*.tmp")) == []
def test_failed_export_refuses_cleanup_when_parent_directory_swapped_to_symlink(
workspace, settings, monkeypatch
):
from finance.exporting import process_formula_export
ds, *_ = workspace
job = _queued_job(ds)
root = Path(settings.MEDIA_ROOT) / "formula-exports"
user_dir = root / str(job.user_id)
user_dir.mkdir(parents=True)
protected = Path(settings.MEDIA_ROOT) / "protected"
protected.mkdir()
victim = protected / f"{job.pk}.xlsx"
victim.write_bytes(b"protected workbook")
def swap_and_fail(*args, **kwargs):
args[-1].write(b"partial")
user_dir.rename(user_dir.with_name("swapped"))
user_dir.symlink_to(protected, target_is_directory=True)
raise ValueError("controlled failure after path validation")
monkeypatch.setattr("finance.exporting.build_formula_workbook", swap_and_fail)
with pytest.raises(ValueError, match="controlled failure"):
process_formula_export(job.pk)
job.refresh_from_db()
assert job.status == "failed"
assert victim.read_bytes() == b"protected workbook"
assert user_dir.is_symlink()
# 本次创建的临时文件残留在被换名的原目录中,由清理命令按保留期回收。
assert list((root / "swapped").glob(f"{job.pk}.*.tmp"))
def test_formula_export_retry_after_failure_succeeds(workspace, monkeypatch):
from finance.exporting import process_formula_export
ds, *_ = workspace
job = _queued_job(ds)
def fail(*args, **kwargs):
args[-1].write(b"partial")
raise ValueError("controlled export failure")
monkeypatch.setattr("finance.exporting.build_formula_workbook", fail)
with pytest.raises(ValueError, match="controlled export failure"):
process_formula_export(job.pk)
monkeypatch.undo()
stats = process_formula_export(job.pk)
job.refresh_from_db()
assert stats["rows"] == 5
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_cleanup_failure_is_logged_and_reclaimed_after_permissions_restored(
workspace, settings, monkeypatch, caplog
):
import logging as logging_module
from datetime import timedelta
from django.core.management import call_command
from django.utils import timezone
from finance.exporting import process_formula_export
ds, *_ = workspace
job = _queued_job(ds)
user_dir = Path(settings.MEDIA_ROOT) / "formula-exports" / str(job.user_id)
user_dir.mkdir(parents=True)
target = user_dir / f"{job.pk}.xlsx"
target.write_bytes(b"stale workbook")
other = FormulaExport.objects.create(
dataset=ds,
user=ds.created_by,
context={"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"},
status="succeeded",
expires_at=timezone.now() + timedelta(hours=1),
)
kept = user_dir / f"{other.pk}.xlsx"
other.path = str(kept)
other.save(update_fields=["path"])
kept.write_bytes(b"active workbook")
def fail(*args, **kwargs):
args[-1].write(b"partial")
os.chmod(user_dir, 0o500)
raise ValueError("controlled export failure")
monkeypatch.setattr("finance.exporting.build_formula_workbook", fail)
with pytest.raises(ValueError, match="controlled export failure"):
process_formula_export(job.pk)
job.refresh_from_db()
assert job.status == "failed"
assert job.path == str(target)
warnings = [r for r in caplog.records if r.levelno == logging_module.WARNING]
assert warnings and all(str(job.pk) in record.getMessage() for record in warnings)
os.chmod(user_dir, 0o700)
old = (timezone.now() - timedelta(hours=25)).timestamp()
for leftover in [target, *user_dir.glob(f"{job.pk}.*.tmp"), kept]:
os.utime(leftover, (old, old))
call_command("cleanup_formula_exports")
job.refresh_from_db()
assert not target.exists()
assert not list(user_dir.glob(f"{job.pk}.*.tmp"))
assert job.path == ""
assert kept.read_bytes() == b"active workbook"
def test_formula_workbook_rejects_unclassified_label_collision(workspace): def test_formula_workbook_rejects_unclassified_label_collision(workspace):
from decimal import Decimal from decimal import Decimal
+118
View File
@@ -1,7 +1,10 @@
from datetime import timedelta
from pathlib import Path
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
from django.core.management import call_command from django.core.management import call_command
from django.utils import timezone
from finance.models import Dataset, FormulaExport from finance.models import Dataset, FormulaExport
@@ -30,3 +33,118 @@ def test_worker_processes_queued_import_and_reports_export_failure(workspace):
call_command("import_worker", once=True) call_command("import_worker", once=True)
export.refresh_from_db() export.refresh_from_db()
assert export.status == "processing" 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"