shanxia2.0/backend/scripts/apply_pl_data_scope.py

#!/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()