"""Optional graph mirror — BetterDB (managed Neo4j) or direct Aura. The ABM graph (accounts, contacts, enrichment events) mirrors into Neo4j when a connection is available; the engine runs fully without it (OpenSearch is the source of truth). Two transports, auto-detected at boot: 1. BETTERDB_URL + BETTERDB_TOKEN — managed Neo4j behind betterdb.com. The HTTP query API is probed at a few known paths; BetterDB answers "waiting" until the owner connects their instance in the console, in which case the mirror stays disabled and /health reports it honestly. 2. NEO4J_URI (+user/password) — a direct Neo4j/Aura endpoint via the official driver (bolt+ssc for Aura). """ import logging from datetime import datetime, timezone from typing import Any, Dict, List, Optional import httpx from ..config import settings logger = logging.getLogger("abm.graph") class GraphStore: def __init__(self) -> None: self._mode: Optional[str] = None # "betterdb" | "neo4j" | None self._driver = None self._http: Optional[httpx.AsyncClient] = None self._query_path: Optional[str] = None @property def mode(self) -> Optional[str]: return self._mode async def connect(self) -> bool: if settings.neo4j_uri: try: from neo4j import GraphDatabase # optional dependency scheme = "neo4j+ssc" if "+ssc" not in settings.neo4j_uri else None uri = settings.neo4j_uri if uri.startswith("neo4j://"): uri = uri.replace("neo4j://", "neo4j+ssc://", 1) elif scheme and not uri.startswith("neo4j+"): uri = f"{scheme}://{uri}" self._driver = GraphDatabase.driver( uri, auth=(settings.neo4j_user or "neo4j", settings.neo4j_password or "") ) await asyncio_wrap(lambda: self._driver.verify_connectivity()) self._mode = "neo4j" logger.info("graph: direct Neo4j connected (%s)", uri) return True except Exception as exc: logger.warning("graph: direct Neo4j connect failed: %s", exc) self._driver = None if settings.betterdb_url and settings.betterdb_token: self._http = httpx.AsyncClient( base_url=settings.betterdb_url.rstrip("/"), headers={"Authorization": f"Bearer {settings.betterdb_token}"}, timeout=httpx.Timeout(20.0), ) for path in ( "/api/neo4j/query/v2", "/neo4j/query/v2", "/api/db/neo4j/query/v2", "/db/neo4j/query/v2", "/api/query", ): try: res = await self._http.post( path, json={"statement": "RETURN 1 AS ok"} ) if res.status_code == 200: self._query_path = path self._mode = "betterdb" logger.info("graph: BetterDB query API live at %s", path) return True if res.status_code in (401, 402, 403): logger.warning("graph: BetterDB answered %s at %s — token %s", res.status_code, path, "rejected") except Exception: continue logger.info("graph: BetterDB reachable but no query API answered yet " "(owner must connect the instance in the console)") return False async def close(self) -> None: if self._http is not None: await self._http.aclose() if self._driver is not None: self._driver.close() async def run(self, cypher: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: """Run a write/query statement; returns rows (best-effort, logged).""" if self._mode == "betterdb" and self._query_path: try: res = await self._http.post( self._query_path, json={"statement": cypher, "parameters": params or {}} ) if res.status_code != 200: logger.warning("graph query %s: %s", res.status_code, res.text[:150]) return [] body = res.json() return body.get("records") or body.get("data") or [] except Exception as exc: logger.warning("graph query failed: %s", exc) return [] if self._mode == "neo4j" and self._driver is not None: def _run(): with self._driver.session() as session: result = session.run(cypher, params or {}) return [r.data() for r in result] try: return await asyncio_wrap(_run) except Exception as exc: logger.warning("graph query failed: %s", exc) return [] return [] # ── mirror API (idempotent MERGEs) ─────────────────────────────────── async def upsert_account(self, account: Dict[str, Any]) -> None: if self._mode is None: return await self.run( """ MERGE (a:Account {account_id: $account_id}) SET a.name = $name, a.domain = $domain, a.industry = $industry, a.hq_country = $hq_country, a.icp_score = $icp_score, a.updated_at = $updated_at """, { "account_id": account.get("account_id"), "name": account.get("name"), "domain": account.get("domain"), "industry": account.get("industry"), "hq_country": account.get("hq_country"), "icp_score": account.get("icp_score"), "updated_at": datetime.now(timezone.utc).isoformat(), }, ) async def upsert_contact(self, account_id: str, contact: Dict[str, Any]) -> None: if self._mode is None: return await self.run( """ MATCH (a:Account {account_id: $account_id}) MERGE (c:Contact {contact_id: $contact_id}) SET c.full_name = $full_name, c.title = $title, c.email = $email, c.linkedin_url = $linkedin_url MERGE (c)-[:WORKS_AT]->(a) """, { "account_id": account_id, "contact_id": contact.get("contact_id"), "full_name": contact.get("full_name"), "title": contact.get("title"), "email": contact.get("email"), "linkedin_url": contact.get("linkedin_url"), }, ) async def overview(self) -> Dict[str, Any]: if self._mode is None: return {"mode": None, "status": "not connected"} rows = await self.run( "MATCH (a:Account) WITH count(a) AS accounts " "OPTIONAL MATCH (c:Contact) RETURN accounts, count(c) AS contacts" ) return {"mode": self._mode, "status": "connected", "counts": rows or []} import asyncio as _asyncio async def asyncio_wrap(fn): return await _asyncio.to_thread(fn) graph = GraphStore()