example.py

"""A minimal knowledge graph, end to end."""

import os
from neo4j import GraphDatabase

URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687")
AUTH = (
    os.environ.get("NEO4J_USER", "neo4j"),
    os.environ["NEO4J_PASSWORD"],
)

CONSTRAINTS = [
    "CREATE CONSTRAINT engineer_email IF NOT EXISTS FOR (e:Engineer) REQUIRE e.email IS UNIQUE",
    "CREATE CONSTRAINT service_name  IF NOT EXISTS FOR (s:Service)  REQUIRE s.name  IS UNIQUE",
    "CREATE CONSTRAINT incident_ref  IF NOT EXISTS FOR (i:Incident) REQUIRE i.ref   IS UNIQUE",
    "CREATE CONSTRAINT team_name     IF NOT EXISTS FOR (t:Team)     REQUIRE t.name  IS UNIQUE",
]

PEOPLE = [
    {"email": "ada@example.com",   "name": "Ada Okonjo",   "service": "payments", "team": "Platform"},
    {"email": "grace@example.com", "name": "Grace Lin",    "service": "payments", "team": "Platform"},
    {"email": "linus@example.com", "name": "Linus Berg",   "service": "checkout", "team": "Commerce"},
    {"email": "mira@example.com",  "name": "Mira Haddad",  "service": "auth",     "team": "Platform"},
    {"email": "tom@example.com",   "name": "Tom Ferreira", "service": "search",   "team": "Discovery"},
]

# One engineer who owns nothing, so the OPTIONAL MATCH example has something to
# show. Without her, that query looks identical to a plain MATCH.
UNASSIGNED = {"email": "nadia@example.com", "name": "Nadia Rossi"}

# Service dependencies, which the variable length path example walks.
DEPENDENCIES = [
    {"upstream": "auth",     "downstream": "payments"},
    {"upstream": "auth",     "downstream": "checkout"},
    {"upstream": "payments", "downstream": "checkout"},
    {"upstream": "search",   "downstream": "checkout"},
]

INCIDENT = {"ref": "INC-4471", "summary": "Elevated 5xx on card capture",
            "services": ["payments", "checkout"]}


def setup(driver):
    """Constraints first. They enforce correctness and create the indexes
    that stop MERGE from scanning every node."""
    for statement in CONSTRAINTS:
        driver.execute_query(statement, database_="neo4j")


def load(driver):
    """People and teams, then the unassigned engineer, then dependencies,
    then the incident. Four round trips for the whole dataset."""
    driver.execute_query(
        """
        UNWIND $rows AS row
        MERGE (e:Engineer {email: row.email})
          SET e.name = row.name
        MERGE (s:Service {name: row.service})
        MERGE (t:Team {name: row.team})
        MERGE (e)-[:OWNS]->(s)
        MERGE (e)-[:MEMBER_OF]->(t)
        """,
        rows=PEOPLE, database_="neo4j",
    )
    driver.execute_query(
        "MERGE (e:Engineer {email: $email}) SET e.name = $name",
        **UNASSIGNED, database_="neo4j",
    )
    driver.execute_query(
        """
        UNWIND $rows AS row
        MATCH (u:Service {name: row.upstream}), (d:Service {name: row.downstream})
        MERGE (d)-[:DEPENDS_ON]->(u)
        """,
        rows=DEPENDENCIES, database_="neo4j",
    )
    driver.execute_query(
        """
        MERGE (i:Incident {ref: $ref}) SET i.summary = $summary
        WITH i
        UNWIND $services AS svc
        MATCH (s:Service {name: svc})
        MERGE (i)-[:AFFECTS]->(s)
        """,
        **INCIDENT, database_="neo4j",
    )


def who_has_context(driver, ref):
    """The question from the introduction, in one pattern."""
    records, _, _ = driver.execute_query(
        """
        MATCH (i:Incident {ref: $ref})-[:AFFECTS]->(:Service)<-[:OWNS]-(e:Engineer)
        RETURN DISTINCT e.name AS name, e.email AS email
        ORDER BY name
        """,
        ref=ref, database_="neo4j",
    )
    return [dict(r) for r in records]


def teams_involved(driver, ref):
    """Split into two patterns on purpose. A single pattern would hit the
    relationship uniqueness rule and silently drop any team whose only
    member is also the owner."""
    records, _, _ = driver.execute_query(
        """
        MATCH (i:Incident {ref: $ref})-[:AFFECTS]->(:Service)<-[:OWNS]-(:Engineer)-[:MEMBER_OF]->(t:Team)
        WITH DISTINCT t
        MATCH (t)<-[:MEMBER_OF]-(e:Engineer)
        RETURN t.name AS team, collect(e.name) AS members
        ORDER BY team
        """,
        ref=ref, database_="neo4j",
    )
    return [dict(r) for r in records]


def main():
    with GraphDatabase.driver(URI, auth=AUTH) as driver:
        driver.verify_connectivity()
        setup(driver)
        load(driver)

        print("Engineers with context on INC-4471:")
        for row in who_has_context(driver, "INC-4471"):
            print(f"  {row['name']:<14} {row['email']}")

        print("\nTeams involved:")
        for row in teams_involved(driver, "INC-4471"):
            print(f"  {row['team']:<10} {', '.join(row['members'])}")


if __name__ == "__main__":
    main()