group_xinghuo_jinrong/scripts/sync/sync_neo4j.py

#!/usr/bin/env python3
"""Core → Neo4j 关系图同步（P0 节点与关系）。"""

from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from neo4j import GraphDatabase
from sqlalchemy import create_engine, text

from app.config.settings import settings


def db_url(database: str) -> str:
    pwd = settings.mysql_password
    auth = f"{settings.mysql_user}:{pwd}" if pwd else settings.mysql_user
    return (
        f"mysql+pymysql://{auth}@{settings.mysql_host}:{settings.mysql_port}"
        f"/{database}?charset=utf8mb4"
    )


def main() -> None:
    if not settings.neo4j_password:
        print("WARN: NEO4J_PASSWORD empty, skip neo4j sync")
        return

    engine = create_engine(db_url(settings.mysql_core_database), pool_pre_ping=True)
    driver = GraphDatabase.driver(
        settings.neo4j_uri,
        auth=(settings.neo4j_user, settings.neo4j_password),
    )

    with engine.connect() as conn:
        customers = conn.execute(
            text("SELECT customer_id, display_name FROM core_customer WHERE is_active=1")
        ).mappings().all()
        advisors = conn.execute(
            text(
                "SELECT staff_id, display_name FROM core_staff WHERE staff_type='advisor' AND is_active=1"
            )
        ).mappings().all()
        products = conn.execute(
            text("SELECT product_id, product_name, min_risk_code, industry_code FROM core_product")
        ).mappings().all()
        grades = conn.execute(
            text("SELECT code, display_name, grade_type FROM core_risk_grade")
        ).mappings().all()
        industries = conn.execute(
            text("SELECT industry_code, industry_name FROM core_industry")
        ).mappings().all()
        assignments = conn.execute(
            text(
                """
                SELECT customer_id, advisor_id, effective_from, rel_status
                FROM core_customer_advisor WHERE rel_status='active'
                """
            )
        ).mappings().all()
        cust_risks = conn.execute(
            text("SELECT customer_id, risk_code, evaluated_at FROM core_customer_risk")
        ).mappings().all()
        holdings = conn.execute(
            text(
                """
                SELECT customer_id, product_id, qty, market_value, cost_amount, pnl_pct, as_of
                FROM core_holding
                """
            )
        ).mappings().all()

    with driver.session() as session:
        session.run("CREATE CONSTRAINT customer_id IF NOT EXISTS FOR (c:Customer) REQUIRE c.customer_id IS UNIQUE")
        session.run("CREATE CONSTRAINT advisor_id IF NOT EXISTS FOR (a:Advisor) REQUIRE a.advisor_id IS UNIQUE")
        session.run("CREATE CONSTRAINT product_id IF NOT EXISTS FOR (p:Product) REQUIRE p.product_id IS UNIQUE")
        session.run("CREATE CONSTRAINT risk_code IF NOT EXISTS FOR (r:RiskGrade) REQUIRE r.code IS UNIQUE")
        session.run("CREATE CONSTRAINT industry_code IF NOT EXISTS FOR (i:Industry) REQUIRE i.industry_code IS UNIQUE")

        for g in grades:
            session.run(
                "MERGE (r:RiskGrade {code: $code}) SET r.display_name=$display_name, r.grade_type=$grade_type",
                **g,
            )
        for ind in industries:
            session.run(
                "MERGE (i:Industry {industry_code: $industry_code}) SET i.industry_name=$industry_name",
                **ind,
            )
        for c in customers:
            session.run(
                "MERGE (c:Customer {customer_id: $customer_id}) SET c.display_name=$display_name",
                **c,
            )
        for a in advisors:
            session.run(
                "MERGE (a:Advisor {advisor_id: $staff_id}) SET a.display_name=$display_name",
                staff_id=a["staff_id"],
                display_name=a["display_name"],
            )
        for p in products:
            session.run(
                """
                MERGE (p:Product {product_id: $product_id})
                SET p.product_name=$product_name, p.min_risk_code=$min_risk_code
                WITH p
                OPTIONAL MATCH (i:Industry {industry_code: $industry_code})
                FOREACH (_ IN CASE WHEN i IS NOT NULL THEN [1] ELSE [] END |
                    MERGE (p)-[:BELONGS_TO]->(i)
                )
                """,
                **p,
            )
            session.run(
                """
                MATCH (p:Product {product_id: $product_id}), (r:RiskGrade {code: $min_risk_code})
                MERGE (p)-[:REQUIRES_MIN_RISK {rule_id: 'R02-PROD-MIN'}]->(r)
                """,
                product_id=p["product_id"],
                min_risk_code=p["min_risk_code"],
            )
        for row in assignments:
            session.run(
                """
                MATCH (c:Customer {customer_id: $customer_id}), (a:Advisor {advisor_id: $advisor_id})
                MERGE (c)-[r:ASSIGNED_TO]->(a)
                SET r.since=$effective_from, r.status=$rel_status
                """,
                **row,
            )
        for row in cust_risks:
            session.run(
                """
                MATCH (c:Customer {customer_id: $customer_id}), (r:RiskGrade {code: $risk_code})
                MERGE (c)-[hr:HAS_RISK_LEVEL]->(r)
                SET hr.source='l0', hr.evaluated_at=$evaluated_at
                """,
                customer_id=row["customer_id"],
                risk_code=row["risk_code"],
                evaluated_at=str(row["evaluated_at"]),
            )
        session.run("MATCH ()-[h:HOLDS]->() DELETE h")
        for h in holdings:
            session.run(
                """
                MATCH (c:Customer {customer_id: $customer_id}), (p:Product {product_id: $product_id})
                MERGE (c)-[h:HOLDS]->(p)
                SET h.qty=$qty, h.market_value=$market_value, h.cost=$cost_amount,
                    h.pnl_pct=$pnl_pct, h.as_of=$as_of
                """,
                customer_id=h["customer_id"],
                product_id=h["product_id"],
                qty=float(h["qty"]),
                market_value=float(h["market_value"]),
                cost_amount=float(h["cost_amount"]),
                pnl_pct=float(h["pnl_pct"]),
                as_of=str(h["as_of"]),
            )

    driver.close()
    print(
        f"sync_neo4j: customers={len(customers)} advisors={len(advisors)} "
        f"products={len(products)} holdings={len(holdings)}"
    )


if __name__ == "__main__":
    main()