Spaces:
Sleeping
Sleeping
| from neo4j import GraphDatabase | |
| import os | |
| import time | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| class Neo4jConnection: | |
| def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"): | |
| self._uri = uri | |
| self._auth = (user, password) | |
| self.database = database | |
| self.driver = None | |
| self._connect_with_retry() | |
| def _connect_with_retry(self, retries: int = 10, delay: int = 6): | |
| auth = None if (self._auth == ("neo4j", "") or self._auth[1] == "") else self._auth | |
| for attempt in range(retries): | |
| try: | |
| self.driver = GraphDatabase.driver(self._uri, auth=auth) | |
| self.driver.verify_connectivity() | |
| print(f"[neo4j] Connected on attempt {attempt + 1}") | |
| return | |
| except Exception as e: | |
| print(f"[neo4j] Attempt {attempt + 1}/{retries} failed: {e}") | |
| if attempt < retries - 1: | |
| time.sleep(delay) | |
| print("[neo4j] WARNING: Could not connect — queries will fail until Neo4j is ready") | |
| def close(self): | |
| if self.driver: | |
| self.driver.close() | |
| def run_query(self, query: str, parameters: dict | None = None) -> list: | |
| if not self.driver: | |
| self._connect_with_retry() | |
| try: | |
| with self.driver.session(database=self.database) as session: | |
| result = session.run(query, parameters or {}) | |
| return [record.data() for record in result] | |
| except Exception: | |
| # Driver may have gone stale — reconnect once and retry | |
| self._connect_with_retry(retries=3, delay=3) | |
| with self.driver.session(database=self.database) as session: | |
| result = session.run(query, parameters or {}) | |
| return [record.data() for record in result] | |
| _auth_disabled = (os.getenv("NEO4J_AUTH", "") == "none") | |
| neo4j_conn = Neo4jConnection( | |
| uri=os.getenv("NEO4J_URI") or "bolt://127.0.0.1:7687", | |
| user=os.getenv("NEO4J_USERNAME") or "neo4j", | |
| password="" if _auth_disabled else (os.getenv("NEO4J_PASSWORD") or "clinicalmatch2024"), | |
| database=os.getenv("NEO4J_DATABASE") or "neo4j", | |
| ) | |
| def setup_schema(): | |
| constraints = [ | |
| "CREATE CONSTRAINT patient_id IF NOT EXISTS FOR (p:Patient) REQUIRE p.id IS UNIQUE", | |
| "CREATE CONSTRAINT trial_id IF NOT EXISTS FOR (t:Trial) REQUIRE t.id IS UNIQUE", | |
| "CREATE CONSTRAINT diagnosis_code IF NOT EXISTS FOR (d:Diagnosis) REQUIRE d.code IS UNIQUE", | |
| "CREATE CONSTRAINT site_id IF NOT EXISTS FOR (s:StudySite) REQUIRE s.id IS UNIQUE", | |
| ] | |
| indexes = [ | |
| "CREATE INDEX patient_age IF NOT EXISTS FOR (p:Patient) ON (p.age)", | |
| "CREATE INDEX trial_phase IF NOT EXISTS FOR (t:Trial) ON (t.phase)", | |
| "CREATE INDEX trial_condition IF NOT EXISTS FOR (t:Trial) ON (t.condition)", | |
| "CREATE INDEX trial_status IF NOT EXISTS FOR (t:Trial) ON (t.status)", | |
| ] | |
| for query in constraints + indexes: | |
| try: | |
| neo4j_conn.run_query(query) | |
| except Exception as e: | |
| print(f"Schema warning: {e}") | |
| print("Schema setup complete.") | |
| if __name__ == "__main__": | |
| setup_schema() | |
| neo4j_conn.close() | |