feat(pl): 运输计划模块接入数据权限过滤

新增 pl_data_scope 辅助，PL 列表/统计 apply_pl_scope，详情与改删 assert_pl_row_in_scope；运费结算待办审批 skip_scope。 Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-28 21:00:15 +08:00
1 parent bf73b6b2d8
commit eff99edeea
@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+"""为 PL API 批量接入数据权限（PR4 Phase 3）。"""
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+
+IMPORT_LINE = "from app.services.pl_data_scope import apply_pl_scope, assert_pl_row_in_scope\n"
+
+
+def ensure_import(content: str) -> str:
+    if "from app.services.pl_data_scope import" in content:
+        return content
+    marker = "from app.core.db_auth import DbUserContext, get_db_current_user\n"
+    if marker in content:
+        return content.replace(marker, marker + IMPORT_LINE)
+    return content
+
+
+def patch_get_helper(content: str, fn_name: str, model_name: str) -> str:
+    """_get_xxx(db, tenant_no, id) -> _get_xxx(db, tenant_no, id, user) + assert."""
+    old_sig = f"def {fn_name}(db: Session, tenant_no: str, "
+    if f"def {fn_name}(db: Session, tenant_no: str, " not in content:
+        return content
+    if f"def {fn_name}(db: Session, tenant_no: str, " in content and "user: DbUserContext" in content.split(f"def {fn_name}")[1].split("\n\n")[0]:
+        return content
+
+    pattern = rf"(def {fn_name}\(\n    db: Session,\n    tenant_no: str,\n    \w+_id: str,\n\) -> \w+:\n)(    \w+ = \([\s\S]*?\n    if not \w+:\n        raise HTTPException\(status_code=404[^\n]+\n    return \w+\n)"
+    repl = rf"\1\2".replace("\\2", "")  # noqa - use explicit replacement below
+
+    def replacer(m):
+        body = m.group(0)
+        if "assert_pl_row_in_scope" in body:
+            return body
+        var = re.search(r"if not (\w+):", body).group(1)
+        return (
+            body.replace(
+                f"def {fn_name}(db: Session, tenant_no: str, ",
+                f"def {fn_name}(db: Session, tenant_no: str, ",
+            )
+            .replace(
+                f") -> ",
+                f",\n    user: DbUserContext,\n) -> ",
+                1,
+            )
+            if ",\n    user: DbUserContext," not in body
+            else body
+        )
+
+    # simpler manual replace per function
+    old = f"""def {fn_name}(db: Session, tenant_no: str, """
+    # find entity id param name
+    m = re.search(rf"def {fn_name}\(db: Session, tenant_no: str, (\w+): str\)", content)
+    if not m:
+        return content
+    id_param = m.group(1)
+    old_block = f"""def {fn_name}(db: Session, tenant_no: str, {id_param}: str) ->"""
+    if f"user: DbUserContext" in content[content.find(old_block) : content.find(old_block) + 200]:
+        return content
+    new_block = f"""def {fn_name}(db: Session, tenant_no: str, {id_param}: str, user: DbUserContext) ->"""
+    content = content.replace(old_block, new_block, 1)
+    # add assert before return
+    needle = f"if not {id_param.split('_')[0] if False else ''}"
+    # find variable name from query
+    fn_body_start = content.find(new_block)
+    fn_chunk = content[fn_body_start : fn_body_start + 600]
+    var_m = re.search(r"(\w+) = \(\n        db\.query", fn_chunk)
+    if not var_m:
+        var_m = re.search(r"(\w+) = \(\n        db\.query", content[fn_body_start : fn_body_start + 800])
+    var = var_m.group(1) if var_m else id_param.replace("_id", "").replace("carrier", "carrier")
+    if var_m:
+        var = var_m.group(1)
+    else:
+        # guess from id_param
+        var = {"carrier_id": "carrier", "vehicle_id": "vehicle", "personnel_id": "personnel", "plan_id": "plan", "transport_id": "transport"}.get(id_param, "row")
+
+    insert = f"""    assert_pl_row_in_scope(db, {var}, {model_name}, user, tenant_no)
+    return {var}
+"""
+    old_return = f"    return {var}\n"
+    if f"assert_pl_row_in_scope(db, {var}," not in content:
+        content = content.replace(old_return, insert, 1)
+    # fix _get calls: _get_xxx(db, tenant_no, id) -> _get_xxx(db, tenant_no, id, user)
+    content = re.sub(
+        rf"{fn_name}\(db, tenant_no, ([^,\)]+)\)",
+        rf"{fn_name}(db, tenant_no, \1, user)",
+        content,
+    )
+    return content
+
+
+def add_scope_before_order(content: str, model: str, anchor: str) -> str:
+    """在 order_by 前插入 apply_pl_scope。"""
+    if f"apply_pl_scope(query, {model}," in content:
+        return content
+    idx = content.find(anchor)
+    if idx == -1:
+        return content
+    insert = f"    query = apply_pl_scope(query, {model}, user, db, tenant_no)\n\n    "
+    # insert before anchor's line
+    line_start = content.rfind("\n", 0, idx) + 1
+    return content[:line_start] + insert + content[line_start:]
+
+
+def patch_base_data():
+    path = ROOT / "app/api/pl/base_data.py"
+    c = ensure_import(path.read_text())
+    c = add_scope_before_order(c, "PlCarrier", "    query = query.order_by(\n        case((PlCarrier.order_num")
+    c = add_scope_before_order(c, "PlVehicle", "    vehicles = query.order_by(")
+    if "vehicles = query.order_by(" not in c:
+        c = add_scope_before_order(c, "PlVehicle", "    query = query.order_by(")
+    c = add_scope_before_order(c, "PlPersonnel", "    query = query.order_by(\n        case((PlPersonnel.order_num")
+    # driver list - second PlPersonnel order_by
+    if c.count("apply_pl_scope(query, PlPersonnel,") < 2:
+        c = c.replace(
+            "    query = query.order_by(\n        case((PlPersonnel.order_num.is_(None), 1), else_=0).asc(),\n        PlPersonnel.create_time.desc(),\n    )\n    return [_personnel_to_driver_read(p) for p in query.all()]",
+            "    query = apply_pl_scope(query, PlPersonnel, user, db, tenant_no)\n\n    query = query.order_by(\n        case((PlPersonnel.order_num.is_(None), 1), else_=0).asc(),\n        PlPersonnel.create_time.desc(),\n    )\n    return [_personnel_to_driver_read(p) for p in query.all()]",
+            1,
+        )
+    for fn, model in (
+        ("_get_carrier", "PlCarrier"),
+        ("_get_vehicle", "PlVehicle"),
+        ("_get_personnel", "PlPersonnel"),
+    ):
+        c = patch_get_helper(c, fn, model)
+    path.write_text(c)
+    print("patched base_data.py")
+
+
+def patch_plan():
+    path = ROOT / "app/api/pl/plan.py"
+    c = ensure_import(path.read_text())
+    if "apply_pl_scope(query, PlPlan," not in c:
+        c = c.replace(
+            "    total = query.count()\n    offset = (page - 1) * page_size\n    plans = query.order_by(PlPlan.create_time.desc())",
+            "    query = apply_pl_scope(query, PlPlan, user, db, tenant_no)\n\n    total = query.count()\n    offset = (page - 1) * page_size\n    plans = query.order_by(PlPlan.create_time.desc())",
+            1,
+        )
+    c = patch_get_helper(c, "_get_plan", "PlPlan")
+    path.write_text(c)
+    print("patched plan.py")
+
+
+def patch_transport():
+    path = ROOT / "app/api/pl/transport.py"
+    c = ensure_import(path.read_text())
+    if "apply_pl_scope(query, PlTransport," not in c:
+        c = c.replace(
+            "    total = query.count()\n    offset = (page - 1) * page_size\n    transports = query.order_by(PlTransport.create_time.desc())",
+            "    query = apply_pl_scope(query, PlTransport, user, db, tenant_no)\n\n    total = query.count()\n    offset = (page - 1) * page_size\n    transports = query.order_by(PlTransport.create_time.desc())",
+            1,
+        )
+    c = patch_get_helper(c, "_get_transport", "PlTransport")
+    path.write_text(c)
+    print("patched transport.py")
+
+
+def patch_inventory():
+    path = ROOT / "app/api/pl/inventory.py"
+    c = ensure_import(path.read_text())
+    replacements = [
+        (
+            "    total = query.count()\n    offset = (page - 1) * page_size\n    items = query.order_by(PlInventory.product_name)",
+            "    query = apply_pl_scope(query, PlInventory, user, db, tenant_no)\n\n    total = query.count()\n    offset = (page - 1) * page_size\n    items = query.order_by(PlInventory.product_name)",
+        ),
+        (
+            "    return query.order_by(PlInventoryWarning.create_time.desc()).all()",
+            "    query = apply_pl_scope(query, PlInventoryWarning, user, db, tenant_no)\n    return query.order_by(PlInventoryWarning.create_time.desc()).all()",
+        ),
+        (
+            "    total = query.count()\n    offset = (page - 1) * page_size\n    items = query.order_by(PlInventoryTransaction.create_time.desc())",
+            "    query = apply_pl_scope(query, PlInventoryTransaction, user, db, tenant_no)\n\n    total = query.count()\n    offset = (page - 1) * page_size\n    items = query.order_by(PlInventoryTransaction.create_time.desc())",
+        ),
+        (
+            "    query = db.query(PlInventoryTransaction).filter(\n        PlInventoryTransaction.tenant_no == tenant_no,\n        PlInventoryTransaction.del_flag == 0,\n        PlInventoryTransaction.create_time >= date_from,\n        PlInventoryTransaction.create_time <= date_to,\n    )\n\n    if product_id:",
+            "    query = db.query(PlInventoryTransaction).filter(\n        PlInventoryTransaction.tenant_no == tenant_no,\n        PlInventoryTransaction.del_flag == 0,\n        PlInventoryTransaction.create_time >= date_from,\n        PlInventoryTransaction.create_time <= date_to,\n    )\n    query = apply_pl_scope(query, PlInventoryTransaction, user, db, tenant_no)\n\n    if product_id:",
+        ),
+        (
+            "    return query.order_by(PlDocumentReminder.create_time.desc()).all()",
+            "    query = apply_pl_scope(query, PlDocumentReminder, user, db, tenant_no)\n    return query.order_by(PlDocumentReminder.create_time.desc()).all()",
+        ),
+    ]
+    for old, new in replacements:
+        if old in c and new not in c:
+            c = c.replace(old, new, 1)
+    # warning update assert
+    if "assert_pl_row_in_scope(db, warning," not in c:
+        c = c.replace(
+            "    if not warning:\n        raise HTTPException(status_code=404, detail=\"预警不存在\")\n\n    if payload.status == \"resolved\":",
+            "    if not warning:\n        raise HTTPException(status_code=404, detail=\"预警不存在\")\n    assert_pl_row_in_scope(db, warning, PlInventoryWarning, user, tenant_no)\n\n    if payload.status == \"resolved\":",
+            1,
+        )
+    if "assert_pl_row_in_scope(db, reminder," not in c:
+        c = c.replace(
+            "    if not reminder:\n        raise HTTPException(status_code=404, detail=\"提醒不存在\")",
+            "    if not reminder:\n        raise HTTPException(status_code=404, detail=\"提醒不存在\")\n    assert_pl_row_in_scope(db, reminder, PlDocumentReminder, user, tenant_no)",
+            1,
+        )
+    if "assert_pl_row_in_scope(db, inventory," not in c:
+        c = c.replace(
+            "    if inventory:\n        inventory.low_stock_warning = payload.low_stock_warning",
+            "    if inventory:\n        assert_pl_row_in_scope(db, inventory, PlInventory, user, tenant_no)\n        inventory.low_stock_warning = payload.low_stock_warning",
+            1,
+        )
+    path.write_text(c)
+    print("patched inventory.py")
+
+
+if __name__ == "__main__":
+    patch_base_data()
+    patch_plan()
+    patch_transport()
+    patch_inventory()