Skip to content
Open
Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
d213f54
feat(nodes,ai): RocketRide cloud DB nodes phase 1 — rocketride_sql + …
dylan-savage Jul 22, 2026
398068c
docs: fix stale db_neo4j references after graph_neo4j rename (#1611)
dylan-savage Jul 22, 2026
9b62fcd
feat(nodes,ai): RocketRide cloud DB nodes phase 2 — Cypher→AGE layer …
dylan-savage Jul 22, 2026
3b1a0fb
feat(ai): cloud DSN delivery for RocketRide DB nodes — broker resolve…
dylan-savage Jul 23, 2026
8cf951b
feat(nodes,ai): OSS cloud sign-in via per-node RocketRide API key
dylan-savage Jul 23, 2026
688744e
test(nodes): derive expected db name from RR_TEST_PG_DSN, not hardcoded
dylan-savage Jul 27, 2026
5855625
docs(ai): align broker/door contract notes and test fakes with shippe…
dylan-savage Jul 28, 2026
7b6977c
feat(ai): promote all four TBD AGE capability cells to REJECT (verifi…
dylan-savage Jul 28, 2026
5e732d2
revert(nodes,ai): defer the OSS API-key sign-in path (revert 8cf951ba)
dylan-savage Jul 28, 2026
ac68fcb
chore: mark age layer generated/vendored paths for collapsed diffs
dylan-savage Jul 28, 2026
06ab8c5
fix(nodes,ai): apply review fixes across DB nodes and AGE layer
dylan-savage Jul 28, 2026
61ae447
Merge remote-tracking branch 'origin/develop' into feat/rocketride-db…
dylan-savage Jul 28, 2026
6d0737c
fix(nodes,ai,ci): address human review — firewall ordering, env hygie…
dylan-savage Jul 30, 2026
d9b108d
Merge remote-tracking branch 'origin/develop' into feat/rocketride-db…
dylan-savage Jul 30, 2026
935bf3e
fix(nodes,ci): pin pgvector CI image to 0.8.0; shared autocommit guard
dylan-savage Jul 30, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion .gitattributes
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,9 @@
*.whl binary


.github/workflows/*.lock.yml linguist-generated=true merge=ours
.github/workflows/*.lock.yml linguist-generated=true merge=ours
# ANTLR-generated openCypher/agtype parsers (committed so devs don't need Java);
# vendored apache/age agtype decoder. Collapsed in diffs — see the age layer README.
packages/ai/src/ai/common/graph/age/_cypher/gen/** linguist-generated=true
packages/ai/src/ai/common/graph/age/_agtype/gen/** linguist-generated=true
packages/ai/src/ai/common/graph/age/_agtype/** linguist-vendored=true
335 changes: 335 additions & 0 deletions nodes/src/nodes/rocketride_graph/IGlobal.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,335 @@
# =============================================================================
# MIT License
# Copyright (c) 2026 Aparavi Software AG
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# =============================================================================

"""RocketRide cloud graph global state.

Mirrors ``graph_neo4j`` on :class:`GraphGlobalBase`, with two differences:

1. **Connection** — no host/user/password config; the per-tenant DSN comes
from the account layer (``ai.common.rocketride_db``), same seam as
``rocketride_sql`` / ``rocketride_vector``. The tenant database is
Postgres + Apache AGE.
2. **Query paths** — Apache AGE cannot run bare Cypher, so ``_run_query``
(safe), ``_run_query_raw`` (EXECUTE) and ``_validate_query`` all route
through the Cypher->AGE translation layer (``ai.common.graph.age``).

Session discipline: the cloud endpoint is a transaction-mode pooler, so every
query runs inside one transaction whose settings are ``SET LOCAL`` only (the
translated plan carries them). The safe path additionally makes its
transaction ``READ ONLY`` — server-side write protection, with the base's
``is_cypher_safe`` regex and the layer's semantic firewall as defence-in-depth.
"""

from typing import Any, Dict, List, Optional, Tuple

import psycopg2
from psycopg2 import sql

from rocketlib import warning
from ai.common.graph import GraphGlobalBase
from ai.common.graph.age import (
AgeTranslationError,
FirewallConfig,
TranslateMode,
TranslatedQuery,
decode_agtype,
decode_row,
translate,
)
from ai.common.rocketride_db import resolve_rocketride_dsn

DEFAULT_GRAPH_NAME = 'rocketride'
DEFAULT_MAX_ROWS = 1000
DEFAULT_QUERY_TIMEOUT_MS = 30000

# Reflection fan-out bounds (mirrors graph_falkordb's sampling posture).
SCHEMA_LABEL_LIMIT = 200
SCHEMA_REL_LIMIT = 100


class IGlobal(GraphGlobalBase):
"""Postgres + Apache AGE global for the RocketRide cloud graph node."""

client = None
graph_name: str = DEFAULT_GRAPH_NAME
max_rows: int = DEFAULT_MAX_ROWS
query_timeout_ms: int = DEFAULT_QUERY_TIMEOUT_MS
age_version: str = ''

# ------------------------------------------------------------------
# Driver lifecycle
# ------------------------------------------------------------------

def _open_driver(self, config: Dict[str, Any]) -> None:
"""Resolve the tenant DSN, connect, and fail fast on a missing graph."""
self.graph_name = str(config.get('graph') or DEFAULT_GRAPH_NAME).strip()
self.max_rows = self._config_int(config, 'max_rows', DEFAULT_MAX_ROWS)
self.query_timeout_ms = self._config_int(config, 'query_timeout_ms', DEFAULT_QUERY_TIMEOUT_MS)

# The one RocketRide difference: DSN from the account seam, no config.
dsn = resolve_rocketride_dsn()
# Bounded connect so an unreachable pooler fails fast (matches the
# probe's 3s timeout) instead of waiting out the OS TCP timeout.
self.client = psycopg2.connect(dsn, connect_timeout=3)

# Fail-fast round-trip: AGE present + the tenant graph exists.
with self.client.cursor() as cur:
cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'age'")
row = cur.fetchone()
if row is None:
raise RuntimeError('The tenant database does not have the Apache AGE extension installed')
self.age_version = row[0]
cur.execute('SELECT 1 FROM ag_catalog.ag_graph WHERE name = %s', (self.graph_name,))
graph_exists = cur.fetchone() is not None
self.client.commit()

if not graph_exists:
# OPEN QUESTION (create_graph ownership, see the node README and
# the design's open-questions): whether the node or the cloud
# provisioner creates the per-tenant graph is undecided. Until it
# is, the node refuses to run against a missing graph rather than
# silently creating one.
raise RuntimeError(
f'AGE graph {self.graph_name!r} does not exist in the tenant database. '
'Graph provisioning ownership is pending; create the graph '
f"(SELECT create_graph('{self.graph_name}')) or configure an existing one."
)

def _close_driver(self) -> None:
if self.client is not None:
self.client.close()
self.client = None

def _probe_connection(self, config: Dict[str, Any]) -> None:
"""Save-time probe: warn-only (the user is still editing the node)."""
conn = None
try:
dsn = resolve_rocketride_dsn()
conn = psycopg2.connect(dsn, connect_timeout=3)
with conn.cursor() as cur:
cur.execute('SELECT 1')
cur.execute("SELECT 1 FROM pg_extension WHERE extname = 'age'")
if cur.fetchone() is None:
warning('The tenant database does not have the Apache AGE extension installed')
return
graph = str(config.get('graph') or DEFAULT_GRAPH_NAME).strip()
cur.execute('SELECT 1 FROM ag_catalog.ag_graph WHERE name = %s', (graph,))
if cur.fetchone() is None:
warning(f'AGE graph {graph!r} does not exist yet in the tenant database')
except NotImplementedError as e:
# The OSS account stub: no cloud sign-in available.
warning(str(e))
except Exception as e:
warning(str(e).strip().splitlines()[0] if str(e).strip() else repr(e))
finally:
try:
if conn is not None:
conn.close()
except Exception:
pass

# ------------------------------------------------------------------
# Query paths — everything routes through the translation layer
# ------------------------------------------------------------------

@property
def read_row_cap(self) -> int:
"""Node-owner row ceiling for the safe read path."""
return self.max_rows

def _firewall_config(self) -> FirewallConfig:
return FirewallConfig(statement_timeout_ms=self.query_timeout_ms)

def _translate(
self,
query: str,
params: Optional[Dict[str, Any]] = None,
limit: Optional[int] = None,
mode: TranslateMode = TranslateMode.SAFE,
) -> TranslatedQuery:
return translate(
query,
params=params,
limit=limit,
mode=mode,
graph_name=self.graph_name,
age_version=self.age_version or '1.5.0',
firewall=self._firewall_config(),
)

def _execute_plan(self, plan: TranslatedQuery, fetch_cap: Optional[int] = None) -> List[Tuple]:
"""Run one translated plan in its own transaction; return raw rows.

The plan's READ ONLY flag is applied with ``SET TRANSACTION READ
ONLY`` as the first in-transaction statement (psycopg2 has already
opened the transaction), giving server-side write protection on the
safe path.
"""
if self.client is None:
raise RuntimeError('Database connection is not initialized')
rows: List[Tuple] = []
try:
with self.client.cursor() as cur:
if plan.read_only:
cur.execute('SET TRANSACTION READ ONLY')

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both server-side controls on this path rest on an assumption nothing checks. SET TRANSACTION READ ONLY only applies when it is the first statement of an open transaction, and the SET LOCAL statement_timeout that emit() prepends only applies inside one. Today both hold because the connection is left with autocommit = False, so psycopg2 opens the transaction implicitly on this very execute.

The failure mode is what worries me: PostgreSQL accepts both statements outside a transaction without an errorSET TRANSACTION READ ONLY becomes a no-op and SET LOCAL only emits a warning. So if that connection ever becomes autocommit (a later change here, or a pooler handing it back that way), the write protection and the timeout cap both vanish silently, and the suite wouldn't notice: test_write_rejected_before_reaching_db asserts the rejection in the translation layer, not at the database. Two of the firewall's three caps hang off an unverified precondition.

Cheap ways to close it, whichever you prefer: set readonly at the session level (self.client.set_session(readonly=True)) so it doesn't depend on statement ordering, assert self.client.autocommit is False on entry here, or add a test that reads back SHOW transaction_read_only inside a plan so a regression fails loudly.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 6d0737c: _execute_plan now asserts autocommit is False up front (loud failure instead of the silent no-op), and a new integration test runs SHOW transaction_read_only through a safe plan and asserts on — executing in CI as of this PR.

for i, (stmt, binds) in enumerate(zip(plan.statements, plan.binds)):
cur.execute(stmt, binds or None)
if i == plan.result_index and cur.description is not None:
rows = cur.fetchmany(fetch_cap) if fetch_cap else cur.fetchall()
self.client.commit()
except Exception:
self.client.rollback()
raise
return rows

def _run_query(self, query: str, params: Optional[Dict] = None, limit: Optional[int] = None) -> List[Dict]:
"""Read-only path: translate (safe mode) and run in a READ ONLY txn."""
cap = self.read_row_cap if limit is None else min(int(limit), self.read_row_cap) + 1
plan = self._translate(query, params=params, limit=cap, mode=TranslateMode.SAFE)
raw_rows = self._execute_plan(plan, fetch_cap=cap)
return [decode_row(plan, row) for row in raw_rows]

def _run_query_raw(self, query: str) -> Dict[str, Any]:
"""EXECUTE path: translation still applies (AGE can't run bare Cypher);
only the semantic firewall is skipped, never the resource caps.
"""
plan = self._translate(query, mode=TranslateMode.RAW)
max_rows = self.max_execute_rows
raw_rows = self._execute_plan(plan, fetch_cap=max_rows + 1)
if len(raw_rows) > max_rows:
raise ValueError(f'EXECUTE query exceeded max_execute_rows={max_rows}')
rows = [decode_row(plan, row) for row in raw_rows]
# AGE exposes no write counters (unlike Bolt); affected_rows reports
# the result-row count — 0 for writes without a RETURN clause.
return {'rows': rows, 'affected_rows': len(rows)}

def _validate_query(self, query: str) -> Tuple[bool, str]:
"""EXPLAIN the translated statement without executing it.

Translation errors (syntax, firewall, capability) surface here too, so
the LLM repair loop receives the layer's actionable messages.
"""
try:
plan = self._translate(query, mode=TranslateMode.SAFE)
except AgeTranslationError as e:
return False, str(e)
if self.client is None:
return False, 'Database connection is not initialized'
try:
with self.client.cursor() as cur:
cur.execute('SET TRANSACTION READ ONLY')

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — _validate_query did not get the autocommit guard that _execute_plan gained.

The new assert at line 201 is the right fix, and the comment above it explains exactly why:

Both server-side controls below (SET TRANSACTION READ ONLY and the plan's SET LOCAL statement_timeout) are silent no-ops outside a transaction.

That reasoning applies word-for-word here. This method runs the same SET TRANSACTION READ ONLY and then the plan's SET LOCAL statement_timeout, with no precondition check.

The risk today is low — EXPLAIN without ANALYZE does not execute the query, so nothing can write. The part that would actually degrade is the timeout: outside a transaction SET LOCAL only emits a warning, so a pathological EXPLAIN would run uncapped.

My concern is less the behaviour than the shape. Someone reading this file later sees one of two identical transaction-opening paths guarded and the other not, and has to work out whether that was deliberate. Either hoist the check into a small helper both call, or add the same three lines here.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 935bf3e — took the hoist option: _require_transactional_client() now carries the precondition (client present + autocommit off) and both transaction-opening paths call it — _execute_plan raising, _validate_query folding the RuntimeError into its return-a-message contract. The two identical paths are now guarded identically, so no future reader has to wonder whether the asymmetry was deliberate.

for i, (stmt, binds) in enumerate(zip(plan.statements, plan.binds)):
if i == plan.result_index:
cur.execute('EXPLAIN ' + stmt, binds or None)
break
cur.execute(stmt, binds or None)
self.client.rollback()
return True, ''
except Exception as e:
self.client.rollback()
return False, str(e)

# ------------------------------------------------------------------
# Schema reflection — AGE catalog + bounded sampling
# ------------------------------------------------------------------

def _reflect_schema(self) -> Dict[str, Any]:
"""Labels from ``ag_catalog.ag_label``; property types by sampling.

Best-effort by contract: any failure warns and degrades to a partial
(or empty) schema rather than blocking startup.
"""
schema: Dict[str, Any] = {'nodes': {}, 'relationships': []}
if self.client is None:
return schema

try:
with self.client.cursor() as cur:
cur.execute(
'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 NOT LIKE '\\_ag\\_%%'",
(self.graph_name,),
)
labels = cur.fetchall()
self.client.commit()
except Exception as e:
self.client.rollback()
warning(f'rocketride_graph: unable to list graph labels: {e}')
return schema

vertex_labels = [name for name, kind in labels if kind == 'v'][:SCHEMA_LABEL_LIMIT]
edge_labels = [name for name, kind in labels if kind == 'e'][:SCHEMA_REL_LIMIT]

for label in vertex_labels:
schema['nodes'][label] = self._sample_properties(label)

for etype in edge_labels:
start, end = self._sample_edge_endpoints(etype)
schema['relationships'].append({'type': etype, 'start': start, 'end': end})

return schema

def _sample_properties(self, label: str) -> List[Tuple[str, str]]:
"""One-row property sample from the label's backing table."""
try:
with self.client.cursor() as cur:
cur.execute(sql.SQL('SELECT properties FROM {} LIMIT 1').format(sql.Identifier(self.graph_name, label)))
row = cur.fetchone()
self.client.commit()
except Exception as e:
self.client.rollback()
warning(f'rocketride_graph: unable to sample label {label!r}: {e}')
return []
if row is None:
return []
properties = decode_agtype(row[0]) or {}
return [(key, type(value).__name__) for key, value in properties.items()]
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def _sample_edge_endpoints(self, etype: str) -> Tuple[str, str]:
"""Start/end labels for one sampled edge of the given type."""
try:
plan = self._translate(
f'MATCH (a)-[r:{etype}]->(b) RETURN label(a) AS s, label(b) AS e',
Comment thread
asclearuc marked this conversation as resolved.
limit=1,
mode=TranslateMode.SAFE,
)
raw = self._execute_plan(plan, fetch_cap=1)
except Exception as e:
warning(f'rocketride_graph: unable to sample relationship {etype!r}: {e}')
return '', ''
if not raw:
return '', ''
row = decode_row(plan, raw[0])
return str(row.get('s') or ''), str(row.get('e') or '')

# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

@staticmethod
def _config_int(config: Dict[str, Any], key: str, default: int) -> int:
try:
return max(1, int(config.get(key, default)))
except (TypeError, ValueError):
return default
25 changes: 25 additions & 0 deletions nodes/src/nodes/rocketride_graph/IInstance.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
# =============================================================================
# MIT License
# Copyright (c) 2026 Aparavi Software AG
# =============================================================================

from ai.common.graph import GraphInstanceBase
from .IGlobal import IGlobal


class IInstance(GraphInstanceBase):
"""RocketRide cloud graph instance.

All tool methods (get_data, get_schema, get_query, execute, dialect) and
lane handlers are inherited from GraphInstanceBase; the Cypher->AGE
translation lives entirely behind IGlobal's query hooks. The tool surface
speaks Cypher to callers, so _query_language stays the base default.
"""

IGlobal: IGlobal

def _db_display_name(self) -> str:
return 'RocketRide Graph'

def _db_dialect(self) -> str:
return 'age'
Loading
Loading