605 lines
21 KiB
Python
605 lines
21 KiB
Python
#!/usr/bin/env python3
|
||
"""Export any live FalkorDB graph as a lossless project import bundle.
|
||
|
||
The selected live graph is the authority for this export. Earlier project
|
||
fixtures rebuilt a different graph from PostgreSQL rows: duplicate AMap POIs
|
||
were split, media and links were modeled twice, and unrelated custom tables
|
||
were imported automatically. That made the JSON totals differ from the graph
|
||
shown by the project.
|
||
|
||
This exporter snapshots every live node, label, property, relationship and
|
||
relationship property, then compares the JSON with independent FalkorDB count
|
||
queries before writing a passed manifest. Multi-label nodes are preserved so
|
||
shared hotel/food POIs and BusLine/BusRoute nodes do not change shape on import.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import math
|
||
import sys
|
||
from collections import Counter, defaultdict
|
||
from datetime import date, datetime, timezone
|
||
from decimal import Decimal
|
||
from pathlib import Path
|
||
from typing import Any, Iterable, Mapping
|
||
from uuid import UUID
|
||
|
||
from falkordb import FalkorDB
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
if str(ROOT) not in sys.path:
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
from app.config import settings # noqa: E402
|
||
from app.project_lifecycle import normalize_provision_payload # noqa: E402
|
||
|
||
|
||
PROJECT_ID = "yunyou_libo"
|
||
GRAPH_NAME = "yunyou_libo"
|
||
DISPLAY_NAME = "云游荔波"
|
||
SCHEMA_VERSION = "3.0.0"
|
||
LABEL_PRIORITY = (
|
||
"Hotel",
|
||
"FoodPlace",
|
||
"ScenicSpot",
|
||
"TransitFacility",
|
||
"BusStop",
|
||
"BusRoute",
|
||
"BusLine",
|
||
"GeoCell",
|
||
"ScenicArea",
|
||
"Area",
|
||
"Place",
|
||
)
|
||
|
||
NODE_LABELS = {
|
||
"Place": "地点",
|
||
"Hotel": "酒店",
|
||
"FoodPlace": "美食店铺",
|
||
"ScenicSpot": "景点",
|
||
"TransitFacility": "交通设施",
|
||
"BusStop": "公交站",
|
||
"BusLine": "公交线路",
|
||
"BusRoute": "公交运行方向",
|
||
"GeoCell": "H3 空间网格",
|
||
"ScenicArea": "景区片区",
|
||
"Area": "行政区域",
|
||
}
|
||
|
||
RELATION_LABELS = {
|
||
"LOCATED_IN": "位于行政区域",
|
||
"IN_H3_R9": "位于 H3 网格",
|
||
"STOPS_AT": "途经站点",
|
||
"NEXT_STOP": "下一站",
|
||
"PART_OF": "行政隶属",
|
||
"PART_OF_SCENIC_AREA": "属于景区片区",
|
||
}
|
||
|
||
MAP_POI_LABELS = {"Hotel", "FoodPlace", "ScenicSpot", "TransitFacility"}
|
||
|
||
|
||
def jsonable(value: Any) -> Any:
|
||
if value is None or isinstance(value, (str, bool, int)):
|
||
return value
|
||
if isinstance(value, float):
|
||
return value if math.isfinite(value) else str(value)
|
||
if isinstance(value, Decimal):
|
||
return int(value) if value == value.to_integral_value() else float(value)
|
||
if isinstance(value, (datetime, date)):
|
||
return value.isoformat()
|
||
if isinstance(value, UUID):
|
||
return str(value)
|
||
if isinstance(value, bytes):
|
||
return value.decode("utf-8", errors="replace")
|
||
if isinstance(value, Mapping):
|
||
return {str(key): jsonable(item) for key, item in value.items()}
|
||
if isinstance(value, (list, tuple, set)):
|
||
return [jsonable(item) for item in value]
|
||
return str(value)
|
||
|
||
|
||
def properties_of(value: Any) -> dict[str, Any]:
|
||
properties = getattr(value, "properties", None) or {}
|
||
return {str(key): jsonable(item) for key, item in dict(properties).items()}
|
||
|
||
|
||
def internal_id(value: Any) -> str:
|
||
identifier = getattr(value, "id", None)
|
||
if identifier is None:
|
||
raise RuntimeError("FalkorDB 返回了缺少内部 ID 的节点")
|
||
return str(identifier)
|
||
|
||
|
||
def labels_of(value: Any) -> list[str]:
|
||
labels = [str(item) for item in (getattr(value, "labels", None) or [])]
|
||
if not labels:
|
||
raise RuntimeError(f"FalkorDB 节点 {internal_id(value)} 没有标签")
|
||
order = {label: index for index, label in enumerate(LABEL_PRIORITY)}
|
||
return sorted(set(labels), key=lambda label: (order.get(label, len(order)), label))
|
||
|
||
|
||
def primary_label(labels: Iterable[str]) -> str:
|
||
label_set = set(labels)
|
||
for label in LABEL_PRIORITY:
|
||
if label in label_set and label not in {"Place", "BusLine"}:
|
||
return label
|
||
for label in LABEL_PRIORITY:
|
||
if label in label_set:
|
||
return label
|
||
return sorted(label_set)[0]
|
||
|
||
|
||
def relation_type_of(edge: Any) -> str:
|
||
relation_type = str(getattr(edge, "relation", "") or "").strip()
|
||
if not relation_type:
|
||
raise RuntimeError("FalkorDB 返回了缺少类型的关系")
|
||
return relation_type
|
||
|
||
|
||
def inferred_value_type(key: str, values: list[Any]) -> str:
|
||
populated = [value for value in values if value is not None]
|
||
if not populated:
|
||
return "any"
|
||
kinds: set[str] = set()
|
||
for value in populated:
|
||
if isinstance(value, bool):
|
||
kinds.add("boolean")
|
||
elif isinstance(value, int):
|
||
kinds.add("integer")
|
||
elif isinstance(value, float):
|
||
kinds.add("number")
|
||
elif isinstance(value, list):
|
||
kinds.add("array")
|
||
elif isinstance(value, dict):
|
||
kinds.add("object")
|
||
else:
|
||
kinds.add("string")
|
||
if kinds <= {"integer", "number"}:
|
||
return "number" if "number" in kinds else "integer"
|
||
if len(kinds) == 1:
|
||
only = next(iter(kinds))
|
||
if only == "string" and (key.endswith("_at") or key.endswith("_time")):
|
||
return "datetime"
|
||
return only
|
||
return "any"
|
||
|
||
|
||
def infer_fields(property_rows: Iterable[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||
values: dict[str, list[Any]] = defaultdict(list)
|
||
for properties in property_rows:
|
||
for key, value in properties.items():
|
||
values[key].append(value)
|
||
return {
|
||
key: {"type": inferred_value_type(key, items), "required": False}
|
||
for key, items in sorted(values.items())
|
||
}
|
||
|
||
|
||
def make_client() -> FalkorDB:
|
||
options: dict[str, Any] = {
|
||
"host": settings.falkordb_host,
|
||
"port": settings.falkordb_port,
|
||
"socket_timeout": 60,
|
||
"socket_connect_timeout": 5,
|
||
}
|
||
if settings.falkordb_password:
|
||
options["password"] = settings.falkordb_password
|
||
return FalkorDB(**options)
|
||
|
||
|
||
def export_live_graph(graph: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
nodes: list[dict[str, Any]] = []
|
||
exported_id_by_internal: dict[str, str] = {}
|
||
exported_ids: set[str] = set()
|
||
|
||
page_size = 5000
|
||
offset = 0
|
||
while True:
|
||
rows = graph.query(
|
||
"MATCH (n) RETURN n ORDER BY id(n) "
|
||
f"SKIP {offset} LIMIT {page_size}"
|
||
).result_set
|
||
for row in rows:
|
||
raw_node = row[0]
|
||
source_internal_id = internal_id(raw_node)
|
||
properties = properties_of(raw_node)
|
||
node_id = str(properties.get("__kg_node_id") or f"falkor:{source_internal_id}")
|
||
if node_id in exported_ids:
|
||
raise RuntimeError(f"图谱存在重复导出节点 ID：{node_id}")
|
||
labels = labels_of(raw_node)
|
||
exported_ids.add(node_id)
|
||
exported_id_by_internal[source_internal_id] = node_id
|
||
nodes.append(
|
||
{
|
||
"id": node_id,
|
||
"type": primary_label(labels),
|
||
"labels": labels,
|
||
"properties": properties,
|
||
}
|
||
)
|
||
if len(rows) < page_size:
|
||
break
|
||
offset += page_size
|
||
|
||
relations: list[dict[str, Any]] = []
|
||
offset = 0
|
||
while True:
|
||
rows = graph.query(
|
||
"MATCH (source)-[relation]->(target) "
|
||
"RETURN source, relation, target ORDER BY id(relation) "
|
||
f"SKIP {offset} LIMIT {page_size}"
|
||
).result_set
|
||
for source, edge, target in rows:
|
||
source_id = exported_id_by_internal.get(internal_id(source))
|
||
target_id = exported_id_by_internal.get(internal_id(target))
|
||
if not source_id or not target_id:
|
||
raise RuntimeError("关系引用了未导出的节点")
|
||
relations.append(
|
||
{
|
||
"type": relation_type_of(edge),
|
||
"source": source_id,
|
||
"target": target_id,
|
||
"properties": properties_of(edge),
|
||
}
|
||
)
|
||
if len(rows) < page_size:
|
||
break
|
||
offset += page_size
|
||
return nodes, relations
|
||
|
||
|
||
def build_schema(
|
||
nodes: list[dict[str, Any]],
|
||
relations: list[dict[str, Any]],
|
||
*,
|
||
project_id: str,
|
||
graph_name: str,
|
||
display_name: str,
|
||
) -> dict[str, Any]:
|
||
rows_by_label: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
primary_type_by_id: dict[str, str] = {}
|
||
for item in nodes:
|
||
primary_type_by_id[item["id"]] = item["type"]
|
||
for label in item["labels"]:
|
||
rows_by_label[label].append(item["properties"])
|
||
|
||
relation_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
relation_endpoints: dict[str, tuple[set[str], set[str]]] = {}
|
||
for item in relations:
|
||
relation_rows[item["type"]].append(item["properties"])
|
||
sources, targets = relation_endpoints.setdefault(item["type"], (set(), set()))
|
||
sources.add(primary_type_by_id[item["source"]])
|
||
targets.add(primary_type_by_id[item["target"]])
|
||
|
||
entity_types = {}
|
||
for label, rows in sorted(rows_by_label.items()):
|
||
entity_types[label] = {
|
||
"label": NODE_LABELS.get(label, label),
|
||
"primary_key": "element_id" if any("element_id" in row for row in rows) else "id",
|
||
"fields": infer_fields(rows),
|
||
}
|
||
|
||
relation_types = {}
|
||
for relation_type, rows in sorted(relation_rows.items()):
|
||
sources, targets = relation_endpoints[relation_type]
|
||
relation_types[relation_type] = {
|
||
"label": RELATION_LABELS.get(relation_type, relation_type),
|
||
"from": "|".join(sorted(sources)),
|
||
"to": "|".join(sorted(targets)),
|
||
"properties": infer_fields(rows),
|
||
}
|
||
|
||
return {
|
||
"namespace": project_id,
|
||
"version": SCHEMA_VERSION,
|
||
"display_name": f"{display_name}当前图谱快照 Schema v3",
|
||
"description": f"从当前 {graph_name} FalkorDB 图谱无损导出，保留多标签节点。",
|
||
"entity_types": entity_types,
|
||
"relation_types": relation_types,
|
||
}
|
||
|
||
|
||
def label_set_key(labels: Iterable[str]) -> str:
|
||
return "+".join(sorted(labels))
|
||
|
||
|
||
def exported_counts(
|
||
nodes: list[dict[str, Any]],
|
||
relations: list[dict[str, Any]],
|
||
) -> dict[str, Any]:
|
||
label_counts: Counter[str] = Counter()
|
||
label_set_counts: Counter[str] = Counter()
|
||
primary_type_counts: Counter[str] = Counter()
|
||
relation_type_counts: Counter[str] = Counter()
|
||
coordinate_nodes = 0
|
||
map_poi_nodes = 0
|
||
bus_stop_nodes = 0
|
||
for item in nodes:
|
||
labels = set(item["labels"])
|
||
label_counts.update(labels)
|
||
label_set_counts[label_set_key(labels)] += 1
|
||
primary_type_counts[item["type"]] += 1
|
||
properties = item["properties"]
|
||
has_coordinates = properties.get("lng") is not None and properties.get("lat") is not None
|
||
if not has_coordinates:
|
||
has_coordinates = (
|
||
properties.get("longitude") is not None
|
||
and properties.get("latitude") is not None
|
||
)
|
||
if has_coordinates:
|
||
coordinate_nodes += 1
|
||
if labels & MAP_POI_LABELS or properties.get("place_type") or properties.get("type_label"):
|
||
map_poi_nodes += 1
|
||
if "BusStop" in labels:
|
||
bus_stop_nodes += 1
|
||
for item in relations:
|
||
relation_type_counts[item["type"]] += 1
|
||
return {
|
||
"nodes": len(nodes),
|
||
"relations": len(relations),
|
||
"primary_type_counts": dict(sorted(primary_type_counts.items())),
|
||
"label_counts": dict(sorted(label_counts.items())),
|
||
"label_set_counts": dict(sorted(label_set_counts.items())),
|
||
"relation_type_counts": dict(sorted(relation_type_counts.items())),
|
||
"coordinate_nodes": coordinate_nodes,
|
||
"map_poi_nodes": map_poi_nodes,
|
||
"bus_stop_nodes": bus_stop_nodes,
|
||
}
|
||
|
||
|
||
def source_counts(graph: Any) -> dict[str, Any]:
|
||
node_count = int(graph.query("MATCH (n) RETURN count(n)").result_set[0][0])
|
||
relation_count = int(graph.query("MATCH ()-[r]->() RETURN count(r)").result_set[0][0])
|
||
coordinate_count = int(
|
||
graph.query(
|
||
"MATCH (n) WHERE coalesce(n.lng,n.longitude) IS NOT NULL "
|
||
"AND coalesce(n.lat,n.latitude) IS NOT NULL RETURN count(n)"
|
||
).result_set[0][0]
|
||
)
|
||
label_counts: Counter[str] = Counter()
|
||
label_set_counts: Counter[str] = Counter()
|
||
for labels, count in graph.query(
|
||
"MATCH (n) RETURN labels(n), count(n) ORDER BY count(n) DESC"
|
||
).result_set:
|
||
normalized = [str(label) for label in labels]
|
||
amount = int(count)
|
||
for label in normalized:
|
||
label_counts[label] += amount
|
||
label_set_counts[label_set_key(normalized)] += amount
|
||
relation_type_counts = {
|
||
str(relation_type): int(count)
|
||
for relation_type, count in graph.query(
|
||
"MATCH ()-[r]->() RETURN type(r), count(r) ORDER BY type(r)"
|
||
).result_set
|
||
}
|
||
return {
|
||
"nodes": node_count,
|
||
"relations": relation_count,
|
||
"label_counts": dict(sorted(label_counts.items())),
|
||
"label_set_counts": dict(sorted(label_set_counts.items())),
|
||
"relation_type_counts": dict(sorted(relation_type_counts.items())),
|
||
"coordinate_nodes": coordinate_count,
|
||
}
|
||
|
||
|
||
def validate_snapshot(
|
||
source: dict[str, Any],
|
||
exported: dict[str, Any],
|
||
nodes: list[dict[str, Any]],
|
||
relations: list[dict[str, Any]],
|
||
) -> dict[str, bool]:
|
||
node_ids = [item["id"] for item in nodes]
|
||
node_id_set = set(node_ids)
|
||
checks = {
|
||
"node_count_matches_live_graph": exported["nodes"] == source["nodes"],
|
||
"relation_count_matches_live_graph": exported["relations"] == source["relations"],
|
||
"label_counts_match_live_graph": exported["label_counts"] == source["label_counts"],
|
||
"label_sets_match_live_graph": exported["label_set_counts"] == source["label_set_counts"],
|
||
"relation_types_match_live_graph": (
|
||
exported["relation_type_counts"] == source["relation_type_counts"]
|
||
),
|
||
"coordinate_count_matches_live_graph": (
|
||
exported["coordinate_nodes"] == source["coordinate_nodes"]
|
||
),
|
||
"node_ids_are_unique": len(node_ids) == len(node_id_set),
|
||
"all_relation_endpoints_exist": all(
|
||
item["source"] in node_id_set and item["target"] in node_id_set
|
||
for item in relations
|
||
),
|
||
}
|
||
failed = [name for name, passed in checks.items() if not passed]
|
||
if failed:
|
||
raise RuntimeError("图谱快照校验失败：" + "、".join(failed))
|
||
return checks
|
||
|
||
|
||
def sha256_file(path: Path) -> str:
|
||
digest = hashlib.sha256()
|
||
with path.open("rb") as handle:
|
||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||
digest.update(chunk)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def dump_json(path: Path, value: Any, *, pretty: bool = False) -> None:
|
||
with path.open("w", encoding="utf-8") as handle:
|
||
if pretty:
|
||
json.dump(value, handle, ensure_ascii=False, indent=2)
|
||
else:
|
||
json.dump(value, handle, ensure_ascii=False, separators=(",", ":"))
|
||
handle.write("\n")
|
||
|
||
|
||
def safe_prefix(value: str) -> str:
|
||
normalized = "".join(
|
||
char.lower() if char.isalnum() else "_"
|
||
for char in value.strip()
|
||
)
|
||
return "_".join(part for part in normalized.split("_") if part) or "graph"
|
||
|
||
|
||
def export(
|
||
output_dir: Path,
|
||
*,
|
||
project_id: str = PROJECT_ID,
|
||
graph_name: str = GRAPH_NAME,
|
||
display_name: str = DISPLAY_NAME,
|
||
map_enabled: bool = True,
|
||
spatial_scope: str = "libo",
|
||
region_name: str | None = None,
|
||
region_adcode: str | None = None,
|
||
region_level: str | None = None,
|
||
output_prefix: str | None = None,
|
||
) -> dict[str, Any]:
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
generated_at = datetime.now(timezone.utc).isoformat()
|
||
default_region = (
|
||
("荔波县", "522722", "district")
|
||
if spatial_scope == "libo"
|
||
else ("贵州省", "520000", "province")
|
||
)
|
||
spatial_map = {
|
||
"enabled": map_enabled,
|
||
"scope": spatial_scope,
|
||
"region_name": (region_name or default_region[0]) if map_enabled else "",
|
||
"region_adcode": (region_adcode or default_region[1]) if map_enabled else "",
|
||
"region_level": (region_level or default_region[2]) if map_enabled else "province",
|
||
}
|
||
prefix = safe_prefix(output_prefix or project_id)
|
||
client = make_client()
|
||
try:
|
||
graph_names = {
|
||
item.decode("utf-8") if isinstance(item, bytes) else str(item)
|
||
for item in client.list_graphs()
|
||
}
|
||
if graph_name not in graph_names:
|
||
raise RuntimeError(f"FalkorDB 中不存在图谱 {graph_name!r}")
|
||
graph = client.select_graph(graph_name)
|
||
live_counts = source_counts(graph)
|
||
nodes, relations = export_live_graph(graph)
|
||
finally:
|
||
client.close()
|
||
|
||
counts = exported_counts(nodes, relations)
|
||
checks = validate_snapshot(live_counts, counts, nodes, relations)
|
||
schema = build_schema(
|
||
nodes,
|
||
relations,
|
||
project_id=project_id,
|
||
graph_name=graph_name,
|
||
display_name=display_name,
|
||
)
|
||
graph_data = {
|
||
"_bundle": {
|
||
"format": "znkg-falkordb-snapshot-v3",
|
||
"source": "falkordb",
|
||
"project_id": project_id,
|
||
"graph_name": graph_name,
|
||
"generated_at": generated_at,
|
||
"spatial_map": spatial_map,
|
||
"source_snapshot_counts": live_counts,
|
||
},
|
||
"nodes": nodes,
|
||
"relations": relations,
|
||
}
|
||
provision = {
|
||
"format": "znkg-project-bundle-v3",
|
||
"project_id": project_id,
|
||
"display_name": display_name,
|
||
"spatial_map": spatial_map,
|
||
"schema": schema,
|
||
"graph_data": graph_data,
|
||
}
|
||
normalized = normalize_provision_payload(provision)
|
||
if normalized["counts"]["nodes"] != counts["nodes"]:
|
||
raise RuntimeError("后端校验后的节点数量不一致")
|
||
if normalized["counts"]["relations"] != counts["relations"]:
|
||
raise RuntimeError("后端校验后的关系数量不一致")
|
||
|
||
schema_path = output_dir / f"{prefix}_full_schema.v3.json"
|
||
graph_path = output_dir / f"{prefix}_full_graph_data.v3.json"
|
||
bundle_path = output_dir / f"{prefix}_full_bundle.v3.json"
|
||
manifest_path = output_dir / f"{prefix}_full_manifest.v3.json"
|
||
dump_json(schema_path, schema, pretty=True)
|
||
dump_json(graph_path, graph_data)
|
||
dump_json(bundle_path, provision)
|
||
|
||
files = {
|
||
path.name: {"bytes": path.stat().st_size, "sha256": sha256_file(path)}
|
||
for path in (schema_path, graph_path, bundle_path)
|
||
}
|
||
manifest = {
|
||
"format": "znkg-full-graph-manifest-v3",
|
||
"project_id": project_id,
|
||
"graph_name": graph_name,
|
||
"schema_version": SCHEMA_VERSION,
|
||
"generated_at": generated_at,
|
||
"source": "falkordb",
|
||
"spatial_map": spatial_map,
|
||
"source_snapshot_counts": live_counts,
|
||
"export_counts": counts,
|
||
"validation": "passed",
|
||
"validation_checks": checks,
|
||
"files": files,
|
||
}
|
||
dump_json(manifest_path, manifest, pretty=True)
|
||
manifest["manifest_file"] = str(manifest_path)
|
||
return manifest
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument(
|
||
"--output-dir",
|
||
type=Path,
|
||
required=True,
|
||
help="Directory for schema, graph-data, bundle and manifest JSON files.",
|
||
)
|
||
parser.add_argument("--project-id", default=PROJECT_ID)
|
||
parser.add_argument("--graph-name", default=GRAPH_NAME)
|
||
parser.add_argument("--display-name", default=DISPLAY_NAME)
|
||
parser.add_argument(
|
||
"--map-enabled",
|
||
action=argparse.BooleanOptionalAction,
|
||
default=True,
|
||
help="Whether the exported project should enable the knowledge-map template.",
|
||
)
|
||
parser.add_argument(
|
||
"--map-scope",
|
||
choices=("libo", "guizhou"),
|
||
default="libo",
|
||
help="Boundary scope used by the single knowledge-map template.",
|
||
)
|
||
parser.add_argument(
|
||
"--output-prefix",
|
||
default=None,
|
||
help="Optional filename prefix; defaults to project-id.",
|
||
)
|
||
parser.add_argument("--region-name", default=None)
|
||
parser.add_argument("--region-adcode", default=None)
|
||
parser.add_argument(
|
||
"--region-level",
|
||
choices=("province", "city", "district"),
|
||
default=None,
|
||
)
|
||
args = parser.parse_args()
|
||
manifest = export(
|
||
args.output_dir.expanduser().resolve(),
|
||
project_id=args.project_id,
|
||
graph_name=args.graph_name,
|
||
display_name=args.display_name,
|
||
map_enabled=args.map_enabled,
|
||
spatial_scope=args.map_scope,
|
||
region_name=args.region_name,
|
||
region_adcode=args.region_adcode,
|
||
region_level=args.region_level,
|
||
output_prefix=args.output_prefix,
|
||
)
|
||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|