100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
步骤 3：批量更新 Neo4j 中所有节点 knowledge_source 字段里的旧服务 URL
|
||
在【迁移目标机】上执行本脚本
|
||
配置项请修改 migration_config.py
|
||
"""
|
||
|
||
import json
|
||
import sys
|
||
import os
|
||
|
||
sys.path.insert(0, os.path.dirname(__file__))
|
||
from migration_config import OLD_BASE_URL, NEW_BASE_URL, NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD
|
||
|
||
from neo4j import GraphDatabase
|
||
|
||
|
||
def update_urls(driver, old_url: str, new_url: str):
|
||
updated = 0
|
||
skipped = 0
|
||
errors = 0
|
||
|
||
with driver.session() as session:
|
||
result = session.run(
|
||
"MATCH (n) WHERE n.knowledge_source IS NOT NULL "
|
||
"RETURN id(n) AS node_id, n.名称 AS name, n.knowledge_source AS ks"
|
||
)
|
||
records = list(result)
|
||
|
||
print(f"共找到 {len(records)} 个含 knowledge_source 的节点")
|
||
|
||
for record in records:
|
||
node_id = record["node_id"]
|
||
name = record["name"]
|
||
ks_raw = record["ks"]
|
||
|
||
try:
|
||
if isinstance(ks_raw, str):
|
||
ks_list = json.loads(ks_raw)
|
||
elif isinstance(ks_raw, list):
|
||
ks_list = ks_raw
|
||
else:
|
||
skipped += 1
|
||
continue
|
||
except json.JSONDecodeError as e:
|
||
print(f" [WARN] 节点 {name}(id={node_id}) JSON 解析失败: {e}")
|
||
errors += 1
|
||
continue
|
||
|
||
has_old_url = any(
|
||
isinstance(item, dict) and old_url in item.get("url", "")
|
||
for item in ks_list
|
||
)
|
||
if not has_old_url:
|
||
skipped += 1
|
||
continue
|
||
|
||
for item in ks_list:
|
||
if isinstance(item, dict) and "url" in item:
|
||
item["url"] = item["url"].replace(old_url, new_url)
|
||
|
||
new_ks_str = json.dumps(ks_list, ensure_ascii=False)
|
||
|
||
with driver.session() as session:
|
||
session.run(
|
||
"MATCH (n) WHERE id(n) = $node_id SET n.knowledge_source = $ks",
|
||
node_id=node_id,
|
||
ks=new_ks_str,
|
||
)
|
||
|
||
print(f" [OK] 更新节点: {name}")
|
||
updated += 1
|
||
|
||
return updated, skipped, errors
|
||
|
||
|
||
def main():
|
||
print("==========================================")
|
||
print(" kgrag Neo4j URL 批量更新脚本 - 步骤 3")
|
||
print(f" 旧 URL 前缀: {OLD_BASE_URL}")
|
||
print(f" 新 URL 前缀: {NEW_BASE_URL}")
|
||
print(f" Neo4j: {NEO4J_URI}")
|
||
print("==========================================")
|
||
|
||
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
|
||
try:
|
||
updated, skipped, errors = update_urls(driver, OLD_BASE_URL, NEW_BASE_URL)
|
||
finally:
|
||
driver.close()
|
||
|
||
print("")
|
||
print("==========================================")
|
||
print(f" 完成！更新: {updated} 跳过: {skipped} 失败: {errors}")
|
||
print(" 下一步：执行 bash 04_update_config.sh")
|
||
print("==========================================")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|