399 lines
17 KiB
Python
399 lines
17 KiB
Python
# 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]])
|