Move peo package to a different spot to prepare it for PyPI

peo/backend.py

@@ -0,0 +1,398 @@
+# SPDX-FileCopyrightText: 2026 Bart Broere
+# SPDX-License-Identifier: GPL-3.0-or-later
+"""Executes Bolt queries against PostgreSQL + Apache AGE.
+
+Every Neo4j "database" maps to an AGE graph. The special ``system`` database
+only answers the administrative queries that neo4j-browser issues.
+"""
+
+import logging
+import re
+import secrets
+import time
+from dataclasses import dataclass, field
+
+import psycopg
+from psycopg import errors as pg_errors
+from psycopg import sql
+
+from . import agtype
+from .cypher import (CypherError, expand_id_in_lists, first_keyword, inline_parameters,
+                     return_columns, rewrite_for_age, strip_statement)
+from .packstream import Node, Relationship
+
+log = logging.getLogger("peo.backend")
+
+NEO4J_VERSION = "5.26.0"
+SYSTEM_DB = "system"
+SYSTEM_LABELS = ("_ag_label_vertex", "_ag_label_edge")
+GRAPH_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.\-]{2,62}$")
+
+# Procedures emulated by peo, reported through SHOW PROCEDURES.
+PROCEDURES = [
+    ("db.labels", "List all node labels in the graph.", "db.labels() :: (label :: STRING)"),
+    ("db.relationshipTypes", "List all relationship types in the graph.",
+     "db.relationshipTypes() :: (relationshipType :: STRING)"),
+    ("db.propertyKeys", "List property keys in the graph.", "db.propertyKeys() :: (propertyKey :: STRING)"),
+    ("db.schema.visualization", "Visualize the schema of the graph.",
+     "db.schema.visualization() :: (nodes :: LIST<NODE>, relationships :: LIST<RELATIONSHIP>)"),
+    ("dbms.components", "List DBMS components and their versions.",
+     "dbms.components() :: (name :: STRING, versions :: LIST<STRING>, edition :: STRING)"),
+    ("dbms.clientConfig", "Return config settings interesting to clients.",
+     "dbms.clientConfig() :: (name :: STRING, value :: STRING)"),
+    ("dbms.showCurrentUser", "Show the current user.",
+     "dbms.showCurrentUser() :: (username :: STRING, roles :: LIST<STRING>, flags :: LIST<STRING>)"),
+]
+
+CLIENT_CONFIG = {
+    "browser.allow_outgoing_connections": "true",
+    "browser.credential_timeout": "0s",
+    "browser.retain_connection_credentials": "true",
+    "browser.retain_editor_history": "true",
+    "client.allow_telemetry": "false",
+    "dbms.security.auth_enabled": "true",
+}
+
+
+@dataclass
+class Config:
+    pg_host: str
+    pg_port: int
+    pg_database: str
+    default_graph: str
+    advertised_bolt: str
+
+
+@dataclass
+class Result:
+    fields: list
+    rows: list = field(default_factory=list)
+    db: str = ""
+    query_type: str = "r"
+    t_first: int = 0
+
+
+class BoltFailure(Exception):
+    def __init__(self, code, message):
+        super().__init__(message)
+        self.code = code
+        self.message = message
+
+
+def failure_from_pg(e: Exception) -> BoltFailure:
+    """Map a PostgreSQL/AGE error onto a Neo4j status code."""
+    sqlstate = getattr(e, "sqlstate", None) or ""
+    diag = getattr(e, "diag", None)
+    message = ((diag.message_primary if diag else None) or str(e)).strip()
+    if sqlstate == "42601":
+        code = "Neo.ClientError.Statement.SyntaxError"
+    elif sqlstate.startswith("42"):
+        code = "Neo.ClientError.Statement.SemanticError"
+    elif sqlstate.startswith("22"):
+        code = "Neo.ClientError.Statement.ArgumentError"
+    elif sqlstate.startswith("23"):
+        code = "Neo.ClientError.Schema.ConstraintValidationFailed"
+    elif sqlstate.startswith("28"):
+        code = "Neo.ClientError.Security.Unauthorized"
+    elif sqlstate.startswith("40"):
+        code = "Neo.TransientError.Transaction.DeadlockDetected"
+    elif sqlstate == "57014":
+        code = "Neo.ClientError.Transaction.TransactionTimedOut"
+    elif sqlstate.startswith("08") or isinstance(e, psycopg.OperationalError):
+        code = "Neo.TransientError.General.DatabaseUnavailable"
+    else:
+        code = "Neo.DatabaseError.General.UnknownError"
+    return BoltFailure(code, message)
+
+
+def _norm(query: str) -> str:
+    return re.sub(r"\s+", " ", strip_statement(query)).strip()
+
+
+class Session:
+    """One PostgreSQL connection, owned by one Bolt connection."""
+
+    def __init__(self, config: Config, user: str):
+        self.config = config
+        self.user = user
+        self.conn: psycopg.AsyncConnection | None = None
+        self.tx_db: str | None = None
+
+    async def connect(self, password: str):
+        try:
+            self.conn = await psycopg.AsyncConnection.connect(
+                host=self.config.pg_host, port=self.config.pg_port, dbname=self.config.pg_database,
+                user=self.user, password=password, autocommit=True,
+                application_name="peo")
+        except psycopg.OperationalError as e:
+            msg = str(e)
+            if "password authentication failed" in msg or "role" in msg and "does not exist" in msg:
+                raise BoltFailure("Neo.ClientError.Security.Unauthorized",
+                                  "The client is unauthorized due to authentication failure.") from e
+            raise BoltFailure("Neo.TransientError.General.DatabaseUnavailable",
+                              f"Could not connect to PostgreSQL: {msg.strip()}") from e
+        try:
+            await self.conn.execute("LOAD 'age'")
+        except psycopg.Error:
+            pass  # Already preloaded, or not permitted for this role.
+        await self.conn.execute('SET search_path = ag_catalog, "$user", public')
+
+    async def close(self):
+        if self.conn is not None:
+            await self.conn.close()
+            self.conn = None
+
+    # -- transactions ---------------------------------------------------------
+
+    @property
+    def in_transaction(self):
+        return self.tx_db is not None
+
+    async def begin(self, db):
+        db = await self.resolve_db(db)
+        await self._sql("BEGIN")
+        self.tx_db = db
+
+    async def commit(self):
+        self.tx_db = None
+        await self._sql("COMMIT")
+
+    async def rollback(self):
+        self.tx_db = None
+        if self.conn is not None and not self.conn.closed:
+            try:
+                await self._sql("ROLLBACK")
+            except BoltFailure:
+                pass
+
+    async def reset(self):
+        await self.rollback()
+
+    # -- query execution ------------------------------------------------------
+
+    async def _sql(self, query, params=None):
+        try:
+            cur = await self.conn.execute(query, params)
+        except psycopg.Error as e:
+            raise failure_from_pg(e) from e
+        return await cur.fetchall() if cur.description else []
+
+    async def graphs(self):
+        rows = await self._sql("SELECT name FROM ag_catalog.ag_graph ORDER BY name")
+        return [r[0] for r in rows]
+
+    async def resolve_db(self, db):
+        if not db:
+            db = self.config.default_graph
+        if db.lower() == SYSTEM_DB:
+            return SYSTEM_DB
+        for name in await self.graphs():
+            if name == db or name.lower() == db.lower():
+                return name
+        raise BoltFailure("Neo.ClientError.Database.DatabaseNotFound",
+                          f"Database does not exist. Database name: '{db}'.")
+
+    async def run(self, query: str, params: dict, db: str | None) -> Result:
+        started = time.monotonic()
+        db = self.tx_db if self.in_transaction else await self.resolve_db(db)
+        try:
+            result = await self._run_system(query, db)
+            if result is None:
+                if db == SYSTEM_DB:
+                    raise BoltFailure("Neo.ClientError.Statement.NotSystemDatabaseCommand",
+                                      "This is an administration command and it should be executed "
+                                      "against the system database: peo only supports a few of them.")
+                result = await self._run_cypher(query, params or {}, db)
+        except CypherError as e:
+            raise BoltFailure(e.code, e.message) from e
+        result.db = db
+        result.t_first = int((time.monotonic() - started) * 1000)
+        return result
+
+    async def _run_cypher(self, query, params, graph) -> Result:
+        text = strip_statement(query)
+        keyword = first_keyword(text)
+        if keyword in ("EXPLAIN", "PROFILE"):
+            raise CypherError("Neo.ClientError.Statement.FeatureNotSupported",
+                              f"{keyword} is not supported by Apache AGE.")
+        text = expand_id_in_lists(inline_parameters(rewrite_for_age(text), params))
+        columns = return_columns(text)
+        tag = "peo"
+        while f"${tag}$" in text:
+            tag = "peo_" + secrets.token_hex(4)
+        col_defs = ", ".join(f"c{i} ag_catalog.agtype" for i in range(max(len(columns), 1)))
+        stmt = (f"SELECT * FROM ag_catalog.cypher({sql.Literal(graph).as_string(self.conn)}, "
+                f"${tag}$ {text} ${tag}$) AS ({col_defs})")
+        log.debug("SQL: %s", stmt)
+        rows = await self._sql(stmt)
+        if columns:
+            rows = [[agtype.parse(v) for v in row] for row in rows]
+        else:
+            rows = []
+        writes = re.search(r"(?<![\w.])(CREATE|MERGE|SET|DELETE|REMOVE)(?!\w)", text, re.IGNORECASE)
+        return Result(columns, rows, query_type="rw" if writes else "r")
+
+    # -- emulated Neo4j procedures ---------------------------------------------
+
+    async def _run_system(self, query, db) -> Result | None:
+        q = _norm(query)
+        ql = q.lower()
+
+        if ql.startswith("call dbms.components()"):
+            return Result(["name", "versions", "edition"], [["Neo4j Kernel", [NEO4J_VERSION], "community"]])
+
+        if ql.startswith("call dbms.clientconfig()") or ql.startswith("call dbms.listconfig()"):
+            return Result(["name", "value"], [[k, v] for k, v in CLIENT_CONFIG.items()])
+
+        if ql.startswith("call dbms.showcurrentuser()") or ql == "show current user":
+            return Result(["username", "roles", "flags"], [[self.user, ["admin"], []]])
+
+        if ql.startswith("show procedures"):
+            return Result(["name", "description", "signature"], [list(p) for p in PROCEDURES])
+
+        if ql.startswith("show functions"):
+            return Result(["name", "description", "signature"], [])
+
+        if re.match(r"^show (range |lookup |fulltext |text |point |vector |all )?indexes", ql) \
+                or ql.startswith("call db.indexes()"):
+            return Result(["id", "name", "state", "populationPercent", "type", "entityType",
+                           "labelsOrTypes", "properties", "indexProvider", "owningConstraint"], [])
+
+        if re.match(r"^show (all |unique |exist(ence)? |node key )?constraints", ql) \
+                or ql.startswith("call db.constraints()"):
+            return Result(["id", "name", "type", "entityType", "labelsOrTypes", "properties",
+                           "ownedIndex", "propertyType"], [])
+
+        m = re.match(r"^show (?:(home|default) database|databases|database (`[^`]+`|[\w.\-]+))(?!\w)",
+                     q, re.IGNORECASE)
+        if m:
+            return await self._show_databases(m.group(1), m.group(2))
+
+        m = re.match(r"^create (?:database|graph) (`[^`]+`|[\w.\-]+)( if not exists)?$", q, re.IGNORECASE)
+        if m:
+            return await self._create_graph(m.group(1).strip("`"), bool(m.group(2)))
+
+        m = re.match(r"^drop (?:database|graph) (`[^`]+`|[\w.\-]+)( if exists)?", q, re.IGNORECASE)
+        if m:
+            return await self._drop_graph(m.group(1).strip("`"), bool(m.group(2)))
+
+        if db == SYSTEM_DB:
+            return None
+
+        # neo4j-browser's combined metadata queries (see dbMetaDuck.ts upstream).
+        if "call db.labels() yield label" in ql and "db.relationshiptypes()" in ql and "union all" in ql:
+            labels, types = await self._labels(db)
+            return Result(["result"], [[{"name": "labels", "data": labels}],
+                                       [{"name": "relationshipTypes", "data": types}],
+                                       [{"name": "propertyKeys", "data": await self._property_keys(db)}]])
+
+        if "return { name:'nodes', data:count(*) } as result" in ql:
+            nodes, rels = await self._counts(db)
+            return Result(["result"], [[{"name": "nodes", "data": nodes}],
+                                       [{"name": "relationships", "data": rels}]])
+
+        if re.match(r"^call db\.labels\(\)", ql):
+            labels, _ = await self._labels(db)
+            return Result(["label"], [[x] for x in labels])
+
+        if re.match(r"^call db\.relationshiptypes\(\)", ql):
+            _, types = await self._labels(db)
+            return Result(["relationshipType"], [[x] for x in types])
+
+        if re.match(r"^call db\.propertykeys\(\)", ql):
+            return Result(["propertyKey"], [[x] for x in await self._property_keys(db)])
+
+        if re.match(r"^call db\.schema\.visualization\(\)", ql):
+            return await self._schema_visualization(db)
+
+        m = re.match(r"^call (db|dbms|apoc|gds)\.[\w.]+", ql)
+        if m:
+            raise BoltFailure("Neo.ClientError.Procedure.ProcedureNotFound",
+                              f"There is no procedure with the name `{m.group(0)[5:]}` registered "
+                              "for this database instance (peo).")
+        return None
+
+    async def _show_databases(self, special, name):
+        graphs = await self.graphs()
+        default = self.config.default_graph
+        if default not in graphs and graphs:
+            default = graphs[0]
+        rows = []
+        for g in [SYSTEM_DB] + graphs:
+            is_default = g == default
+            rows.append([g, "system" if g == SYSTEM_DB else "standard", [], "read-write",
+                         self.config.advertised_bolt, "primary", True, "online", "online", "",
+                         is_default, is_default, []])
+        if special:
+            rows = [r for r in rows if r[10]]
+        elif name:
+            name = name.strip("`").lower()
+            rows = [r for r in rows if r[0].lower() == name]
+        return Result(["name", "type", "aliases", "access", "address", "role", "writer",
+                       "requestedStatus", "currentStatus", "statusMessage", "default", "home",
+                       "constituents"], rows)
+
+    async def _create_graph(self, name, if_not_exists):
+        if not GRAPH_NAME.match(name):
+            raise BoltFailure("Neo.ClientError.Statement.ArgumentError",
+                              f"Invalid graph name '{name}': AGE graph names must be 3-63 characters "
+                              "of letters, digits, '_', '.' or '-', starting with a letter or '_'.")
+        if name in await self.graphs():
+            if if_not_exists:
+                return Result([], query_type="s")
+            raise BoltFailure("Neo.ClientError.Database.ExistingDatabaseFound",
+                              f"Failed to create the specified database '{name}': Database already exists.")
+        await self._sql("SELECT ag_catalog.create_graph(%s)", [name])
+        return Result([], query_type="s")
+
+    async def _drop_graph(self, name, if_exists):
+        graphs = {g.lower(): g for g in await self.graphs()}
+        if name.lower() not in graphs:
+            if if_exists:
+                return Result([], query_type="s")
+            raise BoltFailure("Neo.ClientError.Database.DatabaseNotFound",
+                              f"Failed to delete the specified database '{name}': Database does not exist.")
+        await self._sql("SELECT ag_catalog.drop_graph(%s, true)", [graphs[name.lower()]])
+        return Result([], query_type="s")
+
+    async def _labels(self, graph):
+        rows = await self._sql(
+            "SELECT l.name, l.kind FROM ag_catalog.ag_label l "
+            "JOIN ag_catalog.ag_graph g ON l.graph = g.graphid "
+            "WHERE g.name = %s AND l.name <> ALL(%s) ORDER BY l.name",
+            [graph, list(SYSTEM_LABELS)])
+        return [r[0] for r in rows if r[1] == "v"], [r[0] for r in rows if r[1] == "e"]
+
+    async def _cypher_values(self, graph, query):
+        rows = await self._sql(
+            f"SELECT * FROM ag_catalog.cypher({sql.Literal(graph).as_string(self.conn)}, "
+            f"$peo$ {query} $peo$) AS (c0 ag_catalog.agtype)")
+        return [agtype.parse(r[0]) for r in rows]
+
+    async def _property_keys(self, graph):
+        # Sample rather than scan everything: this runs every 20s from the browser.
+        keys = set()
+        for q in ("MATCH (n) WITH n LIMIT 1000 UNWIND keys(n) AS k RETURN DISTINCT k",
+                  "MATCH ()-[r]->() WITH r LIMIT 1000 UNWIND keys(r) AS k RETURN DISTINCT k"):
+            keys.update(await self._cypher_values(graph, q))
+        return sorted(k for k in keys if isinstance(k, str))
+
+    async def _counts(self, graph):
+        nodes = await self._sql(sql.SQL("SELECT count(*) FROM {}._ag_label_vertex").format(sql.Identifier(graph)))
+        rels = await self._sql(sql.SQL("SELECT count(*) FROM {}._ag_label_edge").format(sql.Identifier(graph)))
+        return nodes[0][0], rels[0][0]
+
+    async def _schema_visualization(self, graph):
+        labels, _ = await self._labels(graph)
+        triples = await self._cypher_values(
+            graph, "MATCH (a)-[r]->(b) RETURN DISTINCT [label(a), label(r), label(b)]")
+        node_ids = {}
+        nodes = []
+        for label in labels + [t[0] for t in triples] + [t[2] for t in triples]:
+            if label not in node_ids:
+                node_ids[label] = -(len(node_ids) + 1)
+                nodes.append(Node(node_ids[label], [label or "(no label)"],
+                                  {"name": label or "(no label)", "indexes": [], "constraints": []}))
+        rels = [Relationship(-(i + 1), node_ids[a], node_ids[b], t, {"name": t})
+                for i, (a, t, b) in enumerate(triples)]
+        return Result(["nodes", "relationships"], [[nodes, rels]])