shanxia2.0/backend/scripts/patch_pr3_quote_contract.py

#!/usr/bin/env python3
"""PR3 Task 2.6–2.7: quote/contract CRUD 数据权限补丁。"""
from __future__ import annotations

import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


def patch_quote_py():
    path = ROOT / "app/crud/crm/quote.py"
    c = path.read_text()
    if "apply_crm_owner_scope" in c:
        print("quote.py already patched")
        return

    import_block = """from app.crud.crm.data_scope_helpers import apply_crm_owner_scope, ensure_crm_owner_row_in_scope

"""
    c = c.replace(
        "from app.services.crm_quote_pdf import generate_quote_pdf\n\n",
        "from app.services.crm_quote_pdf import generate_quote_pdf\n"
        + import_block,
    )

    helper = '''
def _ensure_quote_in_scope(db: Session, row: CrmQuote, user_info, tenant_no: str) -> None:
    ensure_crm_owner_row_in_scope(db, row, user_info, tenant_no, message="无权访问该报价单")


def _get_quote_in_scope(
    db: Session,
    quote_id: str,
    *,
    tenant_no: str,
    user_info=None,
) -> CrmQuote:
    obj = get_quote(db, quote_id, tenant_no=tenant_no)
    if not obj:
        raise ValueError("报价单不存在")
    _ensure_quote_in_scope(db, obj, user_info, tenant_no)
    return obj

'''
    c = c.replace("_DATE_FMT = \"%Y-%m-%d\"\n\n\n", "_DATE_FMT = \"%Y-%m-%d\"\n" + helper + "\n")

    # list_quotes
    c = c.replace(
        "    current_user_id: str = None,\n) -> Tuple[List[Dict[str, Any]], int]:\n    \"\"\"Java: queryPageList\"\"\"\n    q = db.query(CrmQuote)\n    q = _build_quote_filters(db, q, filters or {}, tenant_no=tenant_no)",
        "    current_user_id: str = None,\n    user_info=None,\n) -> Tuple[List[Dict[str, Any]], int]:\n    \"\"\"Java: queryPageList\"\"\"\n    q = db.query(CrmQuote)\n    q = _build_quote_filters(db, q, filters or {}, tenant_no=tenant_no)\n    if user_info:\n        q = apply_crm_owner_scope(q, CrmQuote, user_info, db, tenant_no)",
    )

    # get_quote_detail
    c = c.replace(
        "def get_quote_detail(\n    db: Session,\n    quote_id: str,\n    *,\n    tenant_no: str = None,\n) -> Dict[str, Any]:",
        "def get_quote_detail(\n    db: Session,\n    quote_id: str,\n    *,\n    tenant_no: str = None,\n    user_info=None,\n) -> Dict[str, Any]:",
    )
    c = c.replace(
        '    obj = get_quote(db, quote_id, tenant_no=tenant_no)\n    if not obj:\n        raise ValueError("报价单不存在")\n\n    detail = _row_to_dict(obj)',
        '    obj = get_quote(db, quote_id, tenant_no=tenant_no)\n    if not obj:\n        raise ValueError("报价单不存在")\n    _ensure_quote_in_scope(db, obj, user_info, tenant_no or obj.tenant_no)\n\n    detail = _row_to_dict(obj)',
    )

    # Mutations: replace get_quote with _get_quote_in_scope where entity is used for mutation
    mutation_funcs = [
        "update_quote",
        "soft_delete_quotes",
        "submit_quote",
        "save_template",
        "get_template",
        "generate_pdf",
        "generate_share",
        "accept_quote",
        "reject_quote",
        "to_contract",
        "create_new_version",
        "query_versions",
        "withdraw_quote_approval",
        "invalid_quote",
    ]
    for fn in mutation_funcs:
        # add user_info to signature before closing paren of def
        pat = rf"(def {fn}\([^)]*)(\n\) ->)"
        if re.search(pat, c, re.DOTALL):
            block = re.search(rf"def {fn}\((.*?)\n\) ->", c, re.DOTALL)
            if block and "user_info" not in block.group(1):
                c = re.sub(
                    rf"(def {fn}\(.*?)((?:,\n)?\n\) ->)",
                    r"\1,\n    user_info=None,\2",
                    c,
                    count=1,
                    flags=re.DOTALL,
                )

    # Replace entity = get_quote patterns in mutation bodies
    c = re.sub(
        r"(\n    old = )get_quote\(db, quote_id, tenant_no=tenant_no\)",
        r"\1_get_quote_in_scope(db, quote_id, tenant_no=tenant_no, user_info=user_info)",
        c,
    )
    c = re.sub(
        r"(\n        entity = )get_quote\(db, qid, tenant_no=tenant_no\)",
        r"\1_get_quote_in_scope(db, qid, tenant_no=tenant_no, user_info=user_info)",
        c,
    )
    c = re.sub(
        r"(\n    entity = )get_quote\(db, quote_id, tenant_no=tenant_no\)",
        r"\1_get_quote_in_scope(db, quote_id, tenant_no=tenant_no, user_info=user_info)",
        c,
    )

    # statistics
    for fn in ("statistics_overview", "statistics_status", "statistics_trend"):
        pat = rf"def {fn}\("
        if f"def {fn}(" in c and f"def {fn}(" in c:
            block = re.search(rf"def {fn}\((.*?)\n\) ->", c, re.DOTALL)
            if block and "user_info" not in block.group(1):
                c = re.sub(
                    rf"(def {fn}\(.*?)((?:,\n)?\n\) ->)",
                    r"\1,\n    user_info=None,\2",
                    c,
                    count=1,
                    flags=re.DOTALL,
                )

    # Add scope to statistics queries - find q = db.query(CrmQuote) in statistics funcs
    c = c.replace(
        '    q = db.query(CrmQuote).filter(\n        CrmQuote.tenant_no == tenant_no,\n        CrmQuote.del_flag == 0,\n    )\n    if quote_type:',
        '    q = db.query(CrmQuote).filter(\n        CrmQuote.tenant_no == tenant_no,\n        CrmQuote.del_flag == 0,\n    )\n    if user_info:\n        q = apply_crm_owner_scope(q, CrmQuote, user_info, db, tenant_no)\n    if quote_type:',
    )

    path.write_text(c)
    print("patched quote.py")


def patch_contract_py():
    path = ROOT / "app/crud/crm/contract.py"
    c = path.read_text()
    if "apply_crm_owner_scope" in c:
        print("contract.py already patched")
        return

    # find first import from app.
    c = c.replace(
        "from app.core.utils import generate_uuid\n",
        "from app.core.utils import generate_uuid\n"
        "from app.crud.crm.data_scope_helpers import apply_crm_owner_scope, ensure_crm_owner_row_in_scope\n",
    )

    helper = '''
def _ensure_contract_in_scope(db: Session, row: CrmContract, user_info, tenant_no: str) -> None:
    ensure_crm_owner_row_in_scope(db, row, user_info, tenant_no, message="无权访问该合同")


def _get_contract_row_in_scope(
    db: Session,
    contract_id: str,
    *,
    tenant_no: str,
    user_info=None,
) -> CrmContract:
    row = get_contract_row(db, contract_id, tenant_no=tenant_no)
    if not row:
        raise ValueError("合同不存在")
    _ensure_contract_in_scope(db, row, user_info, tenant_no)
    return row

'''
    # Insert after CrmContract import line area - find _FMT
    c = c.replace("_FMT = \"%Y-%m-%d %H:%M:%S\"\n", "_FMT = \"%Y-%m-%d %H:%M:%S\"\n" + helper + "\n")

    c = c.replace(
        "    current_user_id: str = None,\n) -> Tuple[List[Dict[str, Any]], int]:\n    \"\"\"Java: queryPageList\"\"\"\n    q = db.query(CrmContract)\n    q = _build_contract_filters(db, q, filters or {}, tenant_no=tenant_no)",
        "    current_user_id: str = None,\n    user_info=None,\n) -> Tuple[List[Dict[str, Any]], int]:\n    \"\"\"Java: queryPageList\"\"\"\n    q = db.query(CrmContract)\n    q = _build_contract_filters(db, q, filters or {}, tenant_no=tenant_no)\n    if user_info:\n        q = apply_crm_owner_scope(q, CrmContract, user_info, db, tenant_no)",
    )

    c = c.replace(
        "def get_contract(\n    db: Session,\n    contract_id: str,\n    *,\n    tenant_no: str,\n    current_user_id: str = None,\n) -> Optional[Dict[str, Any]]:",
        "def get_contract(\n    db: Session,\n    contract_id: str,\n    *,\n    tenant_no: str,\n    current_user_id: str = None,\n    user_info=None,\n) -> Optional[Dict[str, Any]]:",
    )
    c = c.replace(
        "    if not row:\n        return None\n\n    detail = _row_to_dict(row)",
        "    if not row:\n        return None\n    _ensure_contract_in_scope(db, row, user_info, tenant_no)\n\n    detail = _row_to_dict(row)",
    )

    c = c.replace(
        "def get_contract_detail(\n    db: Session,\n    contract_id: str,\n    *,\n    tenant_no: str,\n    current_user_id: str = None,\n) -> Dict[str, Any]:",
        "def get_contract_detail(\n    db: Session,\n    contract_id: str,\n    *,\n    tenant_no: str,\n    current_user_id: str = None,\n    user_info=None,\n) -> Dict[str, Any]:",
    )
    c = c.replace(
        "    detail = get_contract(\n        db,\n        contract_id,\n        tenant_no=tenant_no,\n        current_user_id=current_user_id,\n    )",
        "    detail = get_contract(\n        db,\n        contract_id,\n        tenant_no=tenant_no,\n        current_user_id=current_user_id,\n        user_info=user_info,\n    )",
    )

    mutation_funcs = [
        "update_contract",
        "soft_delete_contracts",
        "submit_contract_approval",
        "withdraw_contract_approval",
        "approve_contract_by_contract_id",
        "sign_contract",
        "start_execute",
        "finish_contract",
        "invalid_contract",
        "renew_contract",
        "batch_assign_owner",
        "batch_submit_contract_approval",
        "upload_attachments",
        "upload_attachments_with_category",
        "upload_file",
        "save_template",
        "get_template",
        "download_contract_pdf",
        "create_from_quote",
        "create_from_opportunity",
        "export_contracts",
        "query_expiring",
        "statistics_overview",
        "statistics_by_status",
        "statistics_by_type",
        "statistics_trend",
    ]
    for fn in mutation_funcs:
        block = re.search(rf"def {fn}\((.*?)\n\) ->", c, re.DOTALL)
        if block and "user_info" not in block.group(1):
            c = re.sub(
                rf"(def {fn}\(.*?)((?:,\n)?\n\) ->)",
                r"\1,\n    user_info=None,\2",
                c,
                count=1,
                flags=re.DOTALL,
            )

    # Replace get_contract_row with scoped version in mutations
    c = re.sub(
        r"(\n    row = )get_contract_row\(db, contract_id, tenant_no=tenant_no\)",
        r"\1_get_contract_row_in_scope(db, contract_id, tenant_no=tenant_no, user_info=user_info)",
        c,
    )
    c = re.sub(
        r"(\n        entity = )get_contract_row\(db, cid, tenant_no=tenant_no\)",
        r"\1_get_contract_row_in_scope(db, cid, tenant_no=tenant_no, user_info=user_info)",
        c,
    )
    c = re.sub(
        r"(\n    entity = )get_contract_row\(db, contract_id, tenant_no=tenant_no\)",
        r"\1_get_contract_row_in_scope(db, contract_id, tenant_no=tenant_no, user_info=user_info)",
        c,
    )
    c = re.sub(
        r"(\n    old = )get_contract_row\(db, contract_id, tenant_no=tenant_no\)",
        r"\1_get_contract_row_in_scope(db, contract_id, tenant_no=tenant_no, user_info=user_info)",
        c,
    )

    path.write_text(c)
    print("patched contract.py")


def patch_quote_product_py():
    path = ROOT / "app/crud/crm/quote_product.py"
    c = path.read_text()
    if "_ensure_quote_product_scope" in c:
        print("quote_product.py already patched")
        return
    c = c.replace(
        "from app.core.utils import generate_uuid\n",
        "from app.core.utils import generate_uuid\n"
        "from app.crud.crm.quote import _get_quote_in_scope\n",
    )

    helper = '''
def _ensure_quote_parent_scope(db: Session, quote_id: str, tenant_no: str, user_info) -> None:
    if user_info:
        _get_quote_in_scope(db, quote_id, tenant_no=tenant_no, user_info=user_info)

'''
    c = c.replace("_FMT = \"%Y-%m-%d %H:%M:%S\"\n\n", "_FMT = \"%Y-%m-%d %H:%M:%S\"\n" + helper + "\n")

    for fn in ("list_quote_products", "insert_quote_product", "update_quote_product", "soft_delete_quote_products"):
        block = re.search(rf"def {fn}\((.*?)\n\) ->", c, re.DOTALL)
        if block and "user_info" not in block.group(1):
            c = re.sub(
                rf"(def {fn}\(.*?)((?:,\n)?\n\) ->)",
                r"\1,\n    user_info=None,\2",
                c,
                count=1,
                flags=re.DOTALL,
            )

    c = c.replace(
        '    """Java: queryListByQuoteId"""\n    q = db.query(CrmQuoteProduct)',
        '    """Java: queryListByQuoteId"""\n    if tenant_no and user_info:\n        _ensure_quote_parent_scope(db, quote_id, tenant_no, user_info)\n    q = db.query(CrmQuoteProduct)',
    )

    path.write_text(c)
    print("patched quote_product.py (partial)")


if __name__ == "__main__":
    patch_quote_py()
    patch_contract_py()
    patch_quote_product_py()