From 36fd5f346137c7cacb74c4d57613c8c86f43c7b3 Mon Sep 17 00:00:00 2001 From: Jayway Liu <14022498+bluejayway@user.noreply.gitee.com> Date: Sat, 26 Sep 2026 01:18:18 +0800 Subject: [PATCH 1/4] =?UTF-8?q?fix(exporting):=20F01=20=E8=B7=AF=E5=BE=84?= =?UTF-8?q?=E5=AE=8C=E6=95=B4=E9=AA=8C=E8=AF=81=E5=89=8D=E4=B8=8D=E7=99=BB?= =?UTF-8?q?=E8=AE=B0=E6=B8=85=E7=90=86=E7=9B=AE=E6=A0=87=EF=BC=8C=E6=B8=85?= =?UTF-8?q?=E7=90=86=E4=B8=8D=E8=B7=9F=E9=9A=8F=E7=AC=A6=E5=8F=B7=E9=93=BE?= =?UTF-8?q?=E6=8E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 根目录/用户目录/目标文件先通过 lstat 完整校验（拒绝符号链接与非目录）， 校验失败时不删除任何磁盘文件（此前父目录为链接时会沿链接 unlink 同名文件） - 失败清理只处理本次登记的受控文件，并通过 O_NOFOLLOW 目录描述符删除； 目录被换成符号链接时拒绝删除，保留任务编号与路径的对应关系 - 测试覆盖：用户目录链接、目标链接、运行中目录被换成链接（受保护文件字节不变）、 失败后正常重试成功 --- finance/exporting.py | 89 ++++++++++++++----- tests/test_formula_export_reliability.py | 108 +++++++++++++++++++++++ 2 files changed, 176 insertions(+), 21 deletions(-) diff --git a/finance/exporting.py b/finance/exporting.py index 28e5e18..b08ecf2 100644 --- a/finance/exporting.py +++ b/finance/exporting.py @@ -1,7 +1,9 @@ import hashlib import hmac +import logging import os import secrets +import stat import warnings from datetime import timedelta from decimal import Decimal, InvalidOperation @@ -28,6 +30,58 @@ FORMULA_VERSION = "3" MAX_ROWS = 100_000 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): if value is None: @@ -594,41 +648,34 @@ def build_formula_workbook(ds, member, export_context, output): def process_formula_export(export_id): job = FormulaExport.objects.select_related("dataset", "user").get(pk=export_id) - temporary = None target = None + created = [] 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) if set(member.departments) != set(job.departments): 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) + created.append(temporary) with os.fdopen(descriptor, "wb") as 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.finished_at = timezone.now() job.expires_at = timezone.now() + timedelta(hours=24) job.save(update_fields=["path", "status", "error", "finished_at", "expires_at"]) return stats except Exception as exc: - for path in (temporary, target): - try: - if path is not None and not path.is_symlink(): - path.unlink(missing_ok=True) - except OSError: - pass - job.status, job.error, job.path = "failed", str(exc)[:500], "" - job.finished_at, job.expires_at = timezone.now(), None + job.status, job.error, job.finished_at, job.expires_at = "failed", str(exc)[:500], timezone.now(), None + # 只清理本次登记且确认位于受控目录内的文件；目录被换成符号链接时拒绝删除。 + leftover = [path for path in created if _unlink_created_file(path) == "blocked"] + job.path = str(target) if target in leftover else "" job.save(update_fields=["status", "error", "path", "finished_at", "expires_at"]) raise diff --git a/tests/test_formula_export_reliability.py b/tests/test_formula_export_reliability.py index 0530117..9e15b8d 100644 --- a/tests/test_formula_export_reliability.py +++ b/tests/test_formula_export_reliability.py @@ -198,6 +198,114 @@ def test_formula_export_worker_records_path_preparation_failure(workspace, setti 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_formula_workbook_rejects_unclassified_label_collision(workspace): from decimal import Decimal -- 2.54.0 From 54dc58d2a593aeb91f277f477a93be496b26462a Mon Sep 17 00:00:00 2001 From: Jayway Liu <14022498+bluejayway@user.noreply.gitee.com> Date: Sat, 26 Sep 2026 01:20:50 +0800 Subject: [PATCH 2/4] =?UTF-8?q?fix(cleanup):=20F02=20=E6=B8=85=E7=90=86?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E8=AE=B0=E5=BD=95=E5=90=AB=E4=BB=BB=E5=8A=A1?= =?UTF-8?q?=E7=BC=96=E5=8F=B7=E7=9A=84=E5=91=8A=E8=AD=A6=EF=BC=8C=E5=AD=A4?= =?UTF-8?q?=E5=84=BF=20xlsx=20=E6=8C=89=E4=BF=9D=E7=95=99=E6=9C=9F?= =?UTF-8?q?=E5=9B=9E=E6=94=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - process_formula_export 失败清理受阻时记录不含敏感数据的 WARNING（任务编号+文件名）， 并保留 job.path 使残留文件仍可追踪，而不是清空后失联 - cleanup_formula_exports 新增孤儿 .xlsx 回收：超过 24h 保留期且不属于任何 未过期成功任务的文件统一回收（覆盖失败任务残留、任务行被级联删除两种来源）， 不触碰仍处于活动任务的文件 - 测试：注入 unlink 权限错误→任务失败状态仍保存+告警可见；恢复权限后回收成功 且其他任务文件不受影响；失败任务残留/任务行已删/未过期成功/近期文件/非 UUID 命名 --- finance/exporting.py | 7 ++ .../commands/cleanup_formula_exports.py | 51 ++++++++++-- tests/test_formula_cleanup.py | 78 ++++++++++++++++++- tests/test_formula_export_reliability.py | 57 ++++++++++++++ 4 files changed, 186 insertions(+), 7 deletions(-) diff --git a/finance/exporting.py b/finance/exporting.py index b08ecf2..3051a5e 100644 --- a/finance/exporting.py +++ b/finance/exporting.py @@ -675,6 +675,13 @@ def process_formula_export(export_id): job.status, job.error, job.finished_at, job.expires_at = "failed", str(exc)[:500], timezone.now(), None # 只清理本次登记且确认位于受控目录内的文件；目录被换成符号链接时拒绝删除。 leftover = [path for path in created if _unlink_created_file(path) == "blocked"] + if leftover: + # 日志只含任务编号与文件名，不包含敏感数据；残留文件由清理命令按保留期回收。 + logger.warning( + "公式导出任务 %s 失败后清理文件受阻，待清理命令回收：%s", + 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"]) raise diff --git a/finance/management/commands/cleanup_formula_exports.py b/finance/management/commands/cleanup_formula_exports.py index 692e1aa..20cd91c 100644 --- a/finance/management/commands/cleanup_formula_exports.py +++ b/finance/management/commands/cleanup_formula_exports.py @@ -1,4 +1,5 @@ import time +import uuid from datetime import timedelta from pathlib import Path @@ -9,9 +10,11 @@ from django.utils import timezone from finance.models import FormulaExport +ORPHAN_RETENTION_HOURS = 24 + 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): parser.add_argument("--loop", action="store_true") @@ -19,8 +22,10 @@ class Command(BaseCommand): def handle(self, *args, **options): while True: - jobs, files = self.clean_expired() - self.stdout.write(f"Expired {jobs} jobs; removed {files} files") + jobs, files, orphans = self.clean_expired() + self.stdout.write( + f"Expired {jobs} jobs; removed {files} files; reclaimed {orphans} orphan files" + ) if not options["loop"]: return time.sleep(max(60, options["interval"])) @@ -31,7 +36,7 @@ class Command(BaseCommand): raise ValueError("公式导出目录不能是符号链接") root = export_root.resolve() now = timezone.now() - jobs = files = 0 + jobs = files = orphans = 0 for job in FormulaExport.objects.filter(status="succeeded", expires_at__lte=now).iterator(): path = Path(job.path) if job.path else None @@ -61,5 +66,41 @@ class Command(BaseCommand): files += 1 except OSError as 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 diff --git a/tests/test_formula_cleanup.py b/tests/test_formula_cleanup.py index 19d5475..33acc9d 100644 --- a/tests/test_formula_cleanup.py +++ b/tests/test_formula_cleanup.py @@ -1,4 +1,5 @@ import os +import uuid from datetime import timedelta import pytest @@ -27,6 +28,12 @@ def cleanup_root(settings, tmp_path): 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): root, dataset, user = cleanup_root 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.touch() recent.touch() - old = timezone.now().timestamp() - 25 * 3600 - os.utime(stale, (old, old)) + _aged(stale) 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() assert job.status == "expired" 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() diff --git a/tests/test_formula_export_reliability.py b/tests/test_formula_export_reliability.py index 9e15b8d..9376041 100644 --- a/tests/test_formula_export_reliability.py +++ b/tests/test_formula_export_reliability.py @@ -1,3 +1,4 @@ +import os import stat from decimal import Decimal from io import BytesIO @@ -306,6 +307,62 @@ def test_formula_export_retry_after_failure_succeeds(workspace, monkeypatch): 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): from decimal import Decimal -- 2.54.0 From 5ebe83afbbae9575339137365c9302b642c421ba Mon Sep 17 00:00:00 2001 From: Jayway Liu <14022498+bluejayway@user.noreply.gitee.com> Date: Sat, 26 Sep 2026 01:22:37 +0800 Subject: [PATCH 3/4] =?UTF-8?q?fix(worker):=20F08=20=E6=89=A7=E8=A1=8C/?= =?UTF-8?q?=E4=B8=8B=E8=BD=BD=E7=BB=9F=E4=B8=80=E6=A0=B8=E9=AA=8C=E8=A7=92?= =?UTF-8?q?=E8=89=B2=E4=B8=8E=E8=B4=A6=E5=8F=B7=E7=8A=B6=E6=80=81=EF=BC=8C?= =?UTF-8?q?=E6=81=A2=E5=A4=8D=E5=8D=A1=E6=AD=BB=E7=9A=84=20processing=20?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - worker 执行前重新核验账号 is_active 与 admin/analyst 角色：撤权/停用后排队任务 直接失败，不发布敏感工作簿 - formula-exports 状态与下载接口补齐与创建接口一致的角色白名单，被降为 viewer 的账号不能再查询/下载旧导出 - import_worker 对超过 15 分钟仍处于 processing 的任务重新认领（PostgreSQL 依赖 咨询锁立即恢复），崩溃重启不再无限停留 processing；新近 processing 不被抢占 - 测试：撤权/停用执行、viewer 下载与状态 403、半成品/缺文件下载 409/410、 超时任务恢复且无重复临时文件、新近 processing 不被重复认领 --- finance/exporting.py | 5 + finance/formula_views.py | 4 +- finance/management/commands/import_worker.py | 36 ++++-- tests/test_worker_paths.py | 118 +++++++++++++++++++ 4 files changed, 148 insertions(+), 15 deletions(-) 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" -- 2.54.0 From 05168f671ead68dac003f1a628dc422d06ae3d85 Mon Sep 17 00:00:00 2001 From: Jayway Liu <14022498+bluejayway@user.noreply.gitee.com> Date: Sat, 26 Sep 2026 16:19:02 +0800 Subject: [PATCH 4/4] =?UTF-8?q?test(worker):=20=E9=80=82=E9=85=8D=20Postgr?= =?UTF-8?q?eSQL=20=E5=92=A8=E8=AF=A2=E9=94=81=E4=B8=8B=E7=9A=84=E7=AB=8B?= =?UTF-8?q?=E5=8D=B3=E9=87=8D=E8=AE=A4=E9=A2=86=E8=AF=AD=E4=B9=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PG 上 worker 持咨询锁会立即重认领 processing 任务，第二次 --once 被导入任务 抢占导致导出断言失败（该测试此前从未在 PG 上真正运行过）。第一次调用后把导入 任务置为 ready 模拟完成，两种后端行为一致 --- tests/test_worker_paths.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_worker_paths.py b/tests/test_worker_paths.py index a08a0b1..a8a00a3 100644 --- a/tests/test_worker_paths.py +++ b/tests/test_worker_paths.py @@ -6,7 +6,7 @@ import pytest from django.core.management import call_command from django.utils import timezone -from finance.models import FormulaExport +from finance.models import Dataset, FormulaExport pytestmark = pytest.mark.django_db @@ -20,6 +20,10 @@ def test_worker_processes_queued_import_and_reports_export_failure(workspace): ) as process_import: call_command("import_worker", once=True) process_import.assert_called_once_with(ds.pk) + # PostgreSQL 上 worker 持咨询锁会立即重认领 processing 任务：模拟导入完成， + # 否则第二次 --once 又被它抢占，轮不到下面的导出任务（SQLite 无咨询锁， + # 不走立即恢复路径，两种后端行为一致）。 + Dataset.objects.filter(pk=ds.pk).update(status="ready") export = FormulaExport.objects.create(dataset=ds, user=ds.created_by, context={}) with patch( -- 2.54.0