From 54bddc7d9f72bcef741c5e8a1ccc66818116f6af Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sat, 5 Sep 2026 15:14:02 +0800 Subject: [PATCH 1/9] feat: add query_semantic_graph over the semantic graph views One tool with three views. `summary` reports the entity types, the relationship types, and the endpoint types each relationship connects, so a caller learns the graph's shape without paging through edges to find out -- reading an unfiltered dump is both a round trip and an invitation to read the result as a service call graph. Only `calls` edges carry request, error, and duration counts, so an unfiltered result is ordered by relationship type and endpoint. Ordering a mixed result by RED would sort every other relationship type below a NULL. Passing rel_type=calls opts into the RED ordering. The window is required and half-open over `observed_at`, and the result echoes the window it used. GreptimeDB rejects an unbounded lower bound and otherwise defaults to the last hour, which silently answers a different question than the caller asked. A zero result keeps `items: []` and adds `status`, `applied_filters` and guidance whose next query drops the identifier but keeps the type filters, since an unfiltered read of a large graph is the failure this tool avoids. The views only exist in GreptimeDB 1.3. Registration happens at import, before a connection exists, so lifespan probes -- both views present, carrying the columns read here, and readable by this account -- and withdraws the tool when the answer is conclusive. An unreachable database is not conclusive and leaves the tool in place. Integration coverage seeds declared edges rather than OTLP traces, which gives the graph one RED-bearing relationship type and one without. --- README.md | 3 + src/greptimedb_mcp_server/graph.py | 435 ++++++++++++++++++++++++++++ src/greptimedb_mcp_server/server.py | 139 ++++++++- tests/integration/conftest.py | 48 +++ tests/integration/test_e2e.py | 91 ++++++ tests/test_graph.py | 210 ++++++++++++++ 6 files changed, 925 insertions(+), 1 deletion(-) create mode 100644 src/greptimedb_mcp_server/graph.py create mode 100644 tests/test_graph.py diff --git a/README.md b/README.md index 1de6dc3..9695165 100644 --- a/README.md +++ b/README.md @@ -44,12 +44,15 @@ For Claude Desktop, add this to your config (`~/Library/Application Support/Clau | `execute_tql` | Execute TQL (PromQL-compatible) queries for time-series analysis | | `query_range` | Execute time-window aggregation queries with RANGE/ALIGN syntax | | `search_table_semantics` | Find tables by observability concept, ranked by matched terms; searches table names, semantic options, and entity declarations | +| `query_semantic_graph` | Query the semantic graph: `summary` (what it contains), `entities` (nodes), `relationships` (edges) over a required time window | | `describe_table` | Inspect a table profile: schema, semantic metadata, latest sample rows, and query guidance | | `explain_query` | Analyze SQL or TQL query execution plans (`analyze=true` for runtime stats; add `verbose=true` alongside `analyze=true` for per-partition scan metrics and index-pruning counters) | | `health_check` | Check database connection status and server version | `search_table_semantics` and the semantic metadata in `describe_table` read `information_schema.table_semantics`. A table appears there when it carries a `greptime.semantic.*` option or a built-in convention derives an entity declaration for it; other tables are absent. The server reads the view's column list once per process and selects only the columns it exposes. `entity_declarations` requires GreptimeDB 1.3; on earlier versions it is reported as a missing column rather than as an empty declaration set. +`query_semantic_graph` reads `greptime_private.semantic_entities` and `greptime_private.semantic_relationships`, which require GreptimeDB 1.3. At startup the server checks that both views exist, carry the columns it reads, and are readable by the connected account; when they are not, the tool is not offered and the reason is logged. Its time window is required and half-open, `[start_time, end_time)` over `observed_at`, and rows are aggregated across the 60-second observation buckets in that window. + ### Pipeline Management | Tool | Description | diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py new file mode 100644 index 0000000..e7e49b3 --- /dev/null +++ b/src/greptimedb_mcp_server/graph.py @@ -0,0 +1,435 @@ +"""Reads the semantic graph in `greptime_private`. + +Two computed views, derived at read time: `semantic_entities` holds the nodes +and `semantic_relationships` the witnessed edges. Both arrived in GreptimeDB +1.3, so `GraphView` decides whether the graph is usable at all before the +server offers a tool for it. + +Rows are observations in 60-second buckets, so one logical edge appears once +per bucket it was seen in. Reads here aggregate across the requested window, +which is why the window is a required argument rather than a default. +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime, timezone + +from mysql.connector import Error + +ENTITIES_VIEW = "greptime_private.semantic_entities" +RELATIONSHIPS_VIEW = "greptime_private.semantic_relationships" + +# Reading a column the view lacks fails the statement to plan, so a view +# missing any of these cannot be queried the way this module queries it. +ENTITY_REQUIRED_COLUMNS = frozenset( + {"observed_at", "entity_type", "entity_id", "entity_id_attrs", "source_tables"} +) +RELATIONSHIP_REQUIRED_COLUMNS = frozenset( + { + "observed_at", + "src_type", + "src_id", + "dst_type", + "dst_id", + "rel_type", + "provenance", + "confidence", + } +) + +# Summed across the window when present. A view without them still works; the +# fields are simply absent from the result. +RED_COLUMNS = ( + "request_count", + "unmatched_count", + "error_count", + "duration_sum", + "duration_count", +) + +VIEWS = ("summary", "entities", "relationships") + +ENTITY_FILTERS = ("entity_type", "entity_id", "scope") +RELATIONSHIP_FILTERS = ( + "rel_type", + "src_type", + "src_id", + "dst_type", + "dst_id", + "provenance", +) +# Filters that name a specific node. A zero-result query that used one is most +# often a wrong identifier rather than an absent relationship. +ID_FILTERS = ("entity_id", "src_id", "dst_id") + +MAX_LIMIT = 500 +DEFAULT_LIMIT = 100 + +OBSERVATION_BUCKET_SECONDS = 60 +TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S.%f" + +ERRNO_TABLE_NOT_FOUND = 1146 +ERRNO_PERMISSION_DENIED = frozenset({1044, 1045, 1142, 1143, 1227}) + + +@dataclass(frozen=True) +class GraphCapability: + """Whether the graph can be read, and why not when it cannot.""" + + status: str + entity_columns: frozenset[str] = frozenset() + relationship_columns: frozenset[str] = frozenset() + detail: str | None = None + + @property + def available(self) -> bool: + return self.status == "available" + + @property + def conclusive(self) -> bool: + """Whether the probe reached an answer about the server itself. + + An inconclusive probe -- the database was unreachable, or failed in a + way this module does not recognise -- says nothing about the graph, so + it is neither cached nor allowed to withdraw the tool. + """ + return self.status in ( + "available", + "unavailable", + "permission_denied", + "incompatible_schema", + ) + + def red_columns(self) -> list[str]: + return [c for c in RED_COLUMNS if c in self.relationship_columns] + + def has_duration_max(self) -> bool: + return "duration_max" in self.relationship_columns + + +@dataclass(frozen=True) +class TimeWindow: + """A half-open [start, end) range over `observed_at`.""" + + start: datetime + end: datetime + + @classmethod + def parse(cls, start_time: str, end_time: str) -> "TimeWindow": + start = _parse_timestamp(start_time, "start_time") + end = _parse_timestamp(end_time, "end_time") + if start >= end: + raise ValueError("start_time must be earlier than end_time") + return cls(start=start, end=end) + + @property + def params(self) -> list[str]: + return [ + self.start.strftime(TIMESTAMP_FORMAT), + self.end.strftime(TIMESTAMP_FORMAT), + ] + + def describe(self) -> dict: + return { + "start": self.start.isoformat(), + "end": self.end.isoformat(), + "bounds": "[start, end)", + "time_field": "observed_at", + "observation_bucket_seconds": OBSERVATION_BUCKET_SECONDS, + } + + +def _parse_timestamp(value: str, name: str) -> datetime: + """Read an RFC3339 timestamp, treating a naive one as UTC.""" + if not value or not str(value).strip(): + raise ValueError(f"{name} is required") + text = str(value).strip().replace("Z", "+00:00") + try: + parsed = datetime.fromisoformat(text) + except ValueError: + raise ValueError( + f"Invalid {name}: {value}. Use an RFC3339 timestamp such as " + "2026-09-05T07:00:00Z" + ) from None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +@dataclass(frozen=True) +class GraphRequest: + """A validated graph query.""" + + view: str + window: TimeWindow + filters: dict + limit: int + + @classmethod + def parse( + cls, + view: str, + start_time: str, + end_time: str, + limit: int = DEFAULT_LIMIT, + **filters, + ) -> "GraphRequest": + if view not in VIEWS: + raise ValueError( + f"Invalid view: {view}. Must be one of: {', '.join(VIEWS)}" + ) + allowed = ENTITY_FILTERS if view == "entities" else RELATIONSHIP_FILTERS + applied = {} + for name, value in filters.items(): + if value in (None, ""): + continue + if view == "summary": + raise ValueError(f"view=summary takes no filters, got {name}") + if name not in allowed: + raise ValueError( + f"Filter {name} does not apply to view={view}. " + f"Available: {', '.join(allowed)}" + ) + applied[name] = value + return cls( + view=view, + window=TimeWindow.parse(start_time, end_time), + filters=applied, + limit=max(1, min(limit, MAX_LIMIT)), + ) + + @property + def names_an_id(self) -> bool: + return any(name in ID_FILTERS for name in self.filters) + + +def _decode_json(value): + if not isinstance(value, str) or not value: + return value + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +def _row_dict(columns: list[str], row) -> dict: + return {name: _decode_json(value) for name, value in zip(columns, row)} + + +def _filter_sql(filters: dict) -> tuple[list[str], list]: + predicates = [f"{name} = %s" for name in filters] + return predicates, list(filters.values()) + + +def _no_match_guidance(request: GraphRequest) -> dict: + """Say what to try next, without proposing an unfiltered dump. + + An unfiltered read of a large graph is the failure this tool exists to + avoid, so the suggestion drops the identifier and keeps the type filters. + """ + if request.names_an_id: + kept = {k: v for k, v in request.filters.items() if k not in ID_FILTERS} + return { + "reason": ( + "The supplied identifier may not be a canonical graph entity ID. " + "IDs from alerts and telemetry are not interchangeable with them." + ), + "next_query": {"view": request.view, **kept}, + } + return { + "reason": ( + "No rows matched in this window. The relationship may not have been " + "witnessed, or the window may not cover it." + ), + "next_query": {"view": "summary"}, + } + + +@dataclass +class GraphView: + """A handle on the two graph views that remembers whether they work.""" + + capability: GraphCapability | None = field(default=None) + + def negotiate(self, cursor) -> GraphCapability: + """Decide once per process whether the graph is usable.""" + if self.capability is not None: + return self.capability + capability = _probe(cursor) + if capability.conclusive: + self.capability = capability + return capability + + def summary(self, cursor, window: TimeWindow) -> dict: + """Report what the graph contains, without returning the graph. + + This is what a caller should read first: the shape of the graph is an + interface fact, and learning it by paging through edges both costs a + round trip and invites reading the result as a service call graph. + """ + entity_types = self._entity_types(cursor, window) + relationship_types = self._relationship_types(cursor, window) + return { + "view": "summary", + "window": window.describe(), + "entity_types": entity_types, + "relationship_types": relationship_types, + "entity_count": sum(item["count"] for item in entity_types), + "relationship_count": sum(item["count"] for item in relationship_types), + "complete": True, + } + + def _entity_types(self, cursor, window: TimeWindow) -> list[dict]: + cursor.execute( + "SELECT entity_type, COUNT(*) AS entity_count FROM (" + " SELECT DISTINCT entity_type, entity_id" + f" FROM {ENTITIES_VIEW}" + " WHERE observed_at >= %s AND observed_at < %s" + ") t GROUP BY entity_type ORDER BY entity_type", + window.params, + ) + return [{"type": row[0], "count": int(row[1])} for row in cursor.fetchall()] + + def _relationship_types(self, cursor, window: TimeWindow) -> list[dict]: + cursor.execute( + "SELECT rel_type, src_type, dst_type, COUNT(*) AS edge_count FROM (" + " SELECT DISTINCT rel_type, src_type, dst_type, src_id, dst_id, provenance" + f" FROM {RELATIONSHIPS_VIEW}" + " WHERE observed_at >= %s AND observed_at < %s" + ") t GROUP BY rel_type, src_type, dst_type " + "ORDER BY rel_type, src_type, dst_type", + window.params, + ) + grouped: dict[str, dict] = {} + for rel_type, src_type, dst_type, count in cursor.fetchall(): + entry = grouped.setdefault( + rel_type, + { + "type": rel_type, + "source_types": [], + "destination_types": [], + "count": 0, + }, + ) + if src_type not in entry["source_types"]: + entry["source_types"].append(src_type) + if dst_type not in entry["destination_types"]: + entry["destination_types"].append(dst_type) + entry["count"] += int(count) + return list(grouped.values()) + + def entities(self, cursor, request: GraphRequest) -> dict: + """List distinct entities observed in the window.""" + predicates, params = _filter_sql(request.filters) + where = " AND ".join(["observed_at >= %s", "observed_at < %s", *predicates]) + identity = "entity_type, entity_id, entity_id_attrs, scope, source_tables" + cursor.execute( + f"SELECT {identity}," + " MIN(observed_at) AS first_seen," + " MAX(observed_at) AS last_seen" + f" FROM {ENTITIES_VIEW} WHERE {where}" + f" GROUP BY {identity}" + " ORDER BY entity_type, entity_id" + f" LIMIT {request.limit + 1}", + [*request.window.params, *params], + ) + return self._envelope(cursor, request) + + def relationships(self, cursor, request: GraphRequest) -> dict: + """List edges observed in the window, aggregated across buckets.""" + capability = self.capability + red = capability.red_columns() if capability else list(RED_COLUMNS) + aggregates = [f"SUM({column}) AS {column}" for column in red] + if capability is None or capability.has_duration_max(): + aggregates.append("MAX(duration_max) AS duration_max") + + predicates, params = _filter_sql(request.filters) + where = " AND ".join(["observed_at >= %s", "observed_at < %s", *predicates]) + identity = "src_type, src_id, dst_type, dst_id, rel_type, provenance" + cursor.execute( + f"SELECT {identity}, MAX(confidence) AS confidence, " + + ", ".join(aggregates) + + ", MIN(observed_at) AS first_seen, MAX(observed_at) AS last_seen" + f" FROM {RELATIONSHIPS_VIEW} WHERE {where}" + f" GROUP BY {identity}" + f" ORDER BY {_relationship_order(request)}" + f" LIMIT {request.limit + 1}", + [*request.window.params, *params], + ) + return self._envelope(cursor, request) + + def _envelope(self, cursor, request: GraphRequest) -> dict: + columns = [desc[0] for desc in cursor.description] + rows = cursor.fetchall() + complete = len(rows) <= request.limit + items = [_row_dict(columns, row) for row in rows[: request.limit]] + result = { + "view": request.view, + "status": "ok" if items else "no_match", + "window": request.window.describe(), + "applied_filters": dict(request.filters), + "items": items, + "item_count": len(items), + "complete": complete, + } + if not items: + result["guidance"] = _no_match_guidance(request) + return result + + +def _relationship_order(request: GraphRequest) -> str: + """Order edges by identity unless the caller asked for one RED-bearing type. + + Only `calls` edges carry request and error counts. Ordering a mixed result + by them would sort every other relationship type to the bottom on a NULL + and teach the caller that the graph is a service call graph. + """ + if request.filters.get("rel_type") == "calls": + return "error_count DESC, request_count DESC, src_id, dst_id" + return "rel_type, src_type, src_id, dst_type, dst_id" + + +def _probe(cursor) -> GraphCapability: + """Check both views exist, carry the columns used here, and can be read.""" + columns = {} + for view in (ENTITIES_VIEW, RELATIONSHIPS_VIEW): + try: + cursor.execute(f"DESC TABLE {view}") + columns[view] = frozenset(str(row[0]) for row in cursor.fetchall()) + except Error as e: + return _classify(e) + + missing = ENTITY_REQUIRED_COLUMNS - columns[ENTITIES_VIEW] + missing |= RELATIONSHIP_REQUIRED_COLUMNS - columns[RELATIONSHIPS_VIEW] + if missing: + return GraphCapability( + "incompatible_schema", + detail=f"missing columns: {', '.join(sorted(missing))}", + ) + + # DESC answers from the catalog, so it says nothing about whether this + # account may read the derivation. A bounded read does. + for view in (ENTITIES_VIEW, RELATIONSHIPS_VIEW): + try: + cursor.execute( + f"SELECT COUNT(*) FROM {view} " + "WHERE observed_at >= now() - INTERVAL '1' MINUTE" + ) + cursor.fetchall() + except Error as e: + return _classify(e) + + return GraphCapability( + "available", + entity_columns=columns[ENTITIES_VIEW], + relationship_columns=columns[RELATIONSHIPS_VIEW], + ) + + +def _classify(error: Error) -> GraphCapability: + errno = getattr(error, "errno", None) + if errno == ERRNO_TABLE_NOT_FOUND: + return GraphCapability("unavailable", detail=str(error)) + if errno in ERRNO_PERMISSION_DENIED: + return GraphCapability("permission_denied", detail=str(error)) + return GraphCapability("error", detail=str(error)) diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index 122b35f..9b91be9 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -8,7 +8,7 @@ if sys.platform == "win32" and any(t in sys.argv for t in ("sse", "streamable-http")): asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) -from greptimedb_mcp_server import semantics +from greptimedb_mcp_server import graph, semantics from greptimedb_mcp_server.config import Config from greptimedb_mcp_server.formatter import format_results, VALID_FORMATS from greptimedb_mcp_server.utils import ( @@ -70,6 +70,7 @@ class AppState: table_semantics: semantics.SemanticsView = field( default_factory=semantics.SemanticsView ) + semantic_graph: graph.GraphView = field(default_factory=graph.GraphView) def get_connection(self): """Get a connection from the pool, creating pool if needed.""" @@ -333,6 +334,7 @@ async def lifespan(mcp: MCPServer): "Do NOT use against production data." ) logger.info("Starting GreptimeDB MCP server...") + await asyncio.to_thread(_withdraw_graph_tool_if_unusable, _state) try: yield _state @@ -681,6 +683,141 @@ def _sync_search(): return json.dumps(result, ensure_ascii=False, indent=2, default=str) +GRAPH_TOOL_NAME = "query_semantic_graph" + + +def _withdraw_graph_tool_if_unusable(state: AppState) -> None: + """Stop advertising the graph tool when this server cannot serve it. + + Registration happens at import, before any connection exists, so the + decision is made here instead. An inconclusive probe leaves the tool in + place: a database that was briefly unreachable at startup is not evidence + about the graph. + """ + try: + with state.get_connection() as conn: + with conn.cursor() as cursor: + capability = state.semantic_graph.negotiate(cursor) + except Error as e: + logger.warning(f"Semantic graph probe failed, keeping the tool: {e}") + return + + if capability.available: + logger.info("Semantic graph: available") + return + if not capability.conclusive: + logger.warning(f"Semantic graph probe inconclusive: {capability.detail}") + return + + mcp.remove_tool(GRAPH_TOOL_NAME) + logger.info( + f"Semantic graph: {capability.status}, {GRAPH_TOOL_NAME} not offered " + f"({capability.detail})" + ) + + +@tool(name=GRAPH_TOOL_NAME) +async def query_semantic_graph( + view: Annotated[ + str, + "summary (what the graph contains), entities (nodes), or relationships " + "(edges)", + ], + start_time: Annotated[ + str, "Inclusive RFC3339 start of the window, e.g. 2026-09-05T07:00:00Z" + ], + end_time: Annotated[str, "Exclusive RFC3339 end of the window"], + entity_type: Annotated[str | None, "entities: filter by entity type"] = None, + entity_id: Annotated[str | None, "entities: filter by canonical entity id"] = None, + scope: Annotated[str | None, "entities: filter by namespace or environment"] = None, + rel_type: Annotated[ + str | None, "relationships: calls, runs_on, part_of, contains, uses, ..." + ] = None, + src_type: Annotated[str | None, "relationships: source endpoint type"] = None, + src_id: Annotated[str | None, "relationships: source endpoint id"] = None, + dst_type: Annotated[str | None, "relationships: destination endpoint type"] = None, + dst_id: Annotated[str | None, "relationships: destination endpoint id"] = None, + provenance: Annotated[ + str | None, "relationships: trace, attribute, declared, or agent" + ] = None, + limit: Annotated[ + int, f"Maximum rows to return (1-{graph.MAX_LIMIT}, default: 100)" + ] = graph.DEFAULT_LIMIT, +) -> str: + """Query the semantic graph: which entities exist and which are related. + + Start with view=summary. It returns the entity types, the relationship + types, and the endpoint types each relationship connects, so the shape of + the graph is known before any edge is read. + + The window is required and is half-open, [start_time, end_time), over + observed_at -- the 60-second bucket an observation was recorded in, taken + from the client side of a call. Rows are aggregated across the buckets in + the window, so one edge is one row and RED fields are summed over the + window. The result echoes the window it used. + + Not every relationship type carries request, error, and duration counts: + only `calls` does. An unfiltered result is ordered by relationship type and + endpoint, not by those counts, because ordering a mixed result by them + would rank every other kind of relationship last. Pass rel_type=calls to + order by error and request count. + + confidence is derivation certainty, not health: 1.0 when both sides were + observed, 0.5 when only the caller was and the callee is inferred from a + peer attribute. `calls` edges only cover calls that propagated trace + context, so a dependency invoked without it is absent. A missing edge is + not evidence that two entities are unrelated, and entities are not + deduplicated across identity schemes -- the same process can appear under + two ids if two sources named it differently. + + entity_id_attrs names the attributes an id was assembled from and + source_tables names the telemetry tables that witnessed it; go there for + the raw data. Identifiers from alerts and telemetry are not graph entity + ids unless a query returns that exact id. + """ + state = get_state() + request = graph.GraphRequest.parse( + view, + start_time, + end_time, + limit, + entity_type=entity_type, + entity_id=entity_id, + scope=scope, + rel_type=rel_type, + src_type=src_type, + src_id=src_id, + dst_type=dst_type, + dst_id=dst_id, + provenance=provenance, + ) + + def _sync_query(): + with state.get_connection() as conn: + with conn.cursor() as cursor: + capability = state.semantic_graph.negotiate(cursor) + if not capability.available: + return { + "view": request.view, + "status": "unavailable", + "reason": capability.status, + "error": capability.detail, + "items": [], + } + if request.view == "summary": + return state.semantic_graph.summary(cursor, request.window) + if request.view == "entities": + return state.semantic_graph.entities(cursor, request) + return state.semantic_graph.relationships(cursor, request) + + try: + result = await asyncio.to_thread(_sync_query) + except Error as e: + logger.error(f"Error querying the semantic graph: {e}") + return f"Error querying the semantic graph: {str(e)}" + return json.dumps(result, ensure_ascii=False, indent=2, default=str) + + @tool() async def health_check() -> str: """Check GreptimeDB connection status and server version.""" diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index b115e23..33f8fd7 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -28,6 +28,16 @@ METRICS_TABLE = "it_cpu_metrics" CREDENTIALS_TABLE = "it_credentials" +DECLARED_EDGES_TABLE = "greptime_private.semantic_relationships_declared" +# Declared edges are inserted directly, so the graph has both a RED-bearing +# relationship type and one without needing OTLP traces to pair up. +GRAPH_EDGES = ( + ("service", "it-frontend", "service", "it-checkout", "calls", 100, 7), + ("service", "it-checkout", "service", "it-payment", "calls", 100, 45), + ("service", "it-checkout", "k8s.pod", "it-pod-a", "runs_on", None, None), + ("k8s.pod", "it-pod-a", "k8s.node", "it-node-1", "runs_on", None, None), +) + METRIC_HOSTS = ("host-a", "host-b") METRIC_POINTS = 12 METRIC_INTERVAL_MS = 10_000 @@ -44,6 +54,7 @@ class SeedData: points_per_host: int # False before GreptimeDB 1.3, which is where entity options were added. entity_declared: bool = False + graph_seeded: bool = False @property def row_count(self) -> int: @@ -106,6 +117,8 @@ def seed(db): api_key STRING )""") + graph_seeded = _seed_declared_edges(cursor) + base_ms = int(time.time() * 1000) metric_rows = [ (base_ms - i * METRIC_INTERVAL_MS, host, 50.0 + i) @@ -128,13 +141,48 @@ def seed(db): hosts=METRIC_HOSTS, points_per_host=METRIC_POINTS, entity_declared=entity_declared, + graph_seeded=graph_seeded, ) for table in (METRICS_TABLE, CREDENTIALS_TABLE): cursor.execute(f"DROP TABLE IF EXISTS {table}") + if graph_seeded: + # Delete only the seeded edges: the table is shared and this suite did + # not create it. + for edge in GRAPH_EDGES: + cursor.execute( + f"DELETE FROM {DECLARED_EDGES_TABLE} WHERE src_id = %s AND dst_id = %s", + (edge[1], edge[3]), + ) db.commit() +def _seed_declared_edges(cursor) -> bool: + """Insert declared edges, reporting False when the graph does not exist.""" + columns = ( + "observed_at, src_type, src_id, dst_type, dst_id, rel_type, provenance, " + "scope, generation_id, confidence, request_count, error_count" + ) + try: + for ( + src_type, + src_id, + dst_type, + dst_id, + rel_type, + requests, + errors, + ) in GRAPH_EDGES: + cursor.execute( + f"INSERT INTO {DECLARED_EDGES_TABLE} ({columns}) VALUES " + "(now(), %s, %s, %s, %s, %s, 'declared', '', '', 1.0, %s, %s)", + (src_type, src_id, dst_type, dst_id, rel_type, requests, errors), + ) + except mysql.connector.Error: + return False + return True + + def server_argv(**overrides) -> list[str]: """Build the CLI argv for a server subprocess pointed at the test instance. diff --git a/tests/integration/test_e2e.py b/tests/integration/test_e2e.py index e2c32a4..27fdaf9 100644 --- a/tests/integration/test_e2e.py +++ b/tests/integration/test_e2e.py @@ -5,6 +5,7 @@ """ import json +from datetime import datetime, timedelta, timezone import pytest from mcp import MCPError @@ -330,3 +331,93 @@ async def test_describe_table_reports_entity_declarations(seed): "does not expose entity_declarations" in line for line in profile["guidance"] ) + + +def _graph_window(seed): + """A window wide enough to contain the seeded edges.""" + base = datetime.fromtimestamp(seed.base_ms / 1000, tz=timezone.utc) + return { + "start_time": (base - timedelta(minutes=30)).isoformat(), + "end_time": (base + timedelta(minutes=30)).isoformat(), + } + + +async def test_graph_tool_is_withdrawn_without_a_graph(seed): + """A server whose database has no graph must not advertise the tool.""" + async with stdio_session() as client: + names = {tool.name for tool in (await client.list_tools()).tools} + + assert ("query_semantic_graph" in names) is seed.graph_seeded + + +async def test_graph_summary_reports_shape_without_returning_edges(seed): + if not seed.graph_seeded: + pytest.skip("this GreptimeDB has no semantic graph") + + async with stdio_session() as client: + payload = json.loads( + await call_text( + client, + "query_semantic_graph", + {"view": "summary", **_graph_window(seed)}, + ) + ) + + by_type = {item["type"]: item for item in payload["relationship_types"]} + assert by_type["calls"]["source_types"] == ["service"] + assert by_type["calls"]["count"] == 2 + assert set(by_type["runs_on"]["destination_types"]) == {"k8s.pod", "k8s.node"} + assert "items" not in payload + + +async def test_graph_orders_mixed_relationships_by_type_not_red(seed): + """Only `calls` has RED; sorting a mixed result by it buries the rest.""" + if not seed.graph_seeded: + pytest.skip("this GreptimeDB has no semantic graph") + + window = _graph_window(seed) + async with stdio_session() as client: + mixed = json.loads( + await call_text( + client, "query_semantic_graph", {"view": "relationships", **window} + ) + ) + calls = json.loads( + await call_text( + client, + "query_semantic_graph", + {"view": "relationships", "rel_type": "calls", **window}, + ) + ) + + assert [item["rel_type"] for item in mixed["items"]] == sorted( + item["rel_type"] for item in mixed["items"] + ) + assert [item["error_count"] for item in calls["items"]] == [45, 7] + + +async def test_graph_zero_result_keeps_the_envelope(seed): + if not seed.graph_seeded: + pytest.skip("this GreptimeDB has no semantic graph") + + async with stdio_session() as client: + payload = json.loads( + await call_text( + client, + "query_semantic_graph", + { + "view": "relationships", + "rel_type": "calls", + "src_id": "not-a-graph-id", + **_graph_window(seed), + }, + ) + ) + + assert payload["status"] == "no_match" + assert payload["items"] == [] + assert payload["applied_filters"]["rel_type"] == "calls" + assert payload["guidance"]["next_query"] == { + "view": "relationships", + "rel_type": "calls", + } diff --git a/tests/test_graph.py b/tests/test_graph.py new file mode 100644 index 0000000..bbf094e --- /dev/null +++ b/tests/test_graph.py @@ -0,0 +1,210 @@ +"""Tests for the semantic graph window, ordering, and capability contracts.""" + +from datetime import datetime, timezone + +import pytest +from mysql.connector import Error + +from greptimedb_mcp_server import graph +from greptimedb_mcp_server.graph import ( + GraphCapability, + GraphRequest, + GraphView, + TimeWindow, + _no_match_guidance, + _relationship_order, +) + +START = "2026-09-05T07:00:00Z" +END = "2026-09-05T08:00:00Z" + +FULL_ENTITY_COLUMNS = frozenset(graph.ENTITY_REQUIRED_COLUMNS) +FULL_RELATIONSHIP_COLUMNS = frozenset( + graph.RELATIONSHIP_REQUIRED_COLUMNS | set(graph.RED_COLUMNS) | {"duration_max"} +) + + +class FakeCursor: + """Answers DESC from a column map and SELECT from a queued result.""" + + def __init__(self, columns=None, rows=None, errno=None, fail_on=None): + self.columns = columns or {} + self.rows = rows or [] + self.errno = errno + self.fail_on = fail_on + self.queries = [] + self.description = None + + def execute(self, query, args=None): + self.queries.append(query) + if self.errno is not None and (self.fail_on is None or self.fail_on in query): + error = Error("probe failed") + error.errno = self.errno + raise error + if query.startswith("DESC TABLE"): + view = query.split()[-1] + self._result = [(name,) for name in sorted(self.columns.get(view, ()))] + self.description = [("Column", None)] + else: + self._result = self.rows + self.description = [("count", None)] + + def fetchall(self): + return self._result + + +def request(view="relationships", limit=graph.DEFAULT_LIMIT, **filters): + return GraphRequest.parse(view, START, END, limit, **filters) + + +def test_window_treats_a_naive_timestamp_as_utc(): + window = TimeWindow.parse("2026-09-05T07:00:00", "2026-09-05T08:00:00") + + assert window.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) + assert window.describe()["bounds"] == "[start, end)" + + +def test_window_normalizes_an_offset_to_utc(): + window = TimeWindow.parse("2026-09-05T09:00:00+02:00", END) + + assert window.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) + + +def test_window_rejects_an_empty_or_inverted_range(): + with pytest.raises(ValueError, match="earlier than"): + TimeWindow.parse(END, START) + with pytest.raises(ValueError, match="earlier than"): + TimeWindow.parse(START, START) + + +def test_window_rejects_an_unparseable_timestamp(): + with pytest.raises(ValueError, match="RFC3339"): + TimeWindow.parse("last tuesday", END) + + +def test_request_rejects_a_filter_from_another_view(): + with pytest.raises(ValueError, match="does not apply"): + request(view="entities", rel_type="calls") + + +def test_request_rejects_filters_on_summary(): + with pytest.raises(ValueError, match="no filters"): + request(view="summary", rel_type="calls") + + +def test_request_rejects_an_unknown_view(): + with pytest.raises(ValueError, match="Invalid view"): + request(view="paths") + + +def test_mixed_relationships_are_not_ordered_by_red(): + """Only `calls` carries RED, so a mixed sort would rank the rest last.""" + order = _relationship_order(request()) + + assert "error_count" not in order + assert order.startswith("rel_type") + + +def test_calls_relationships_are_ordered_by_red(): + order = _relationship_order(request(rel_type="calls")) + + assert order.startswith("error_count DESC") + + +def test_no_match_guidance_drops_the_id_and_keeps_the_type(): + guidance = _no_match_guidance(request(rel_type="calls", src_id="unknown")) + + assert guidance["next_query"] == {"view": "relationships", "rel_type": "calls"} + assert "canonical graph entity ID" in guidance["reason"] + + +def test_no_match_guidance_without_an_id_points_at_the_summary(): + guidance = _no_match_guidance(request(rel_type="calls")) + + assert guidance["next_query"] == {"view": "summary"} + + +@pytest.mark.parametrize( + "errno,expected", + [(1146, "unavailable"), (1142, "permission_denied"), (1105, "error")], +) +def test_probe_classifies_failures(errno, expected): + capability = GraphView().negotiate(FakeCursor(errno=errno)) + + assert capability.status == expected + assert capability.available is False + + +def test_probe_reports_an_incompatible_schema(): + """A view that exists but lacks a column this module reads is not usable.""" + columns = { + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS - {"entity_id_attrs"}, + graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, + } + + capability = GraphView().negotiate(FakeCursor(columns=columns)) + + assert capability.status == "incompatible_schema" + assert "entity_id_attrs" in capability.detail + + +def test_probe_rejects_a_view_it_cannot_read(): + """DESC answers from the catalog, so it does not prove SELECT is allowed.""" + columns = { + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, + graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, + } + cursor = FakeCursor(columns=columns, errno=1142, fail_on="SELECT COUNT(*)") + + capability = GraphView().negotiate(cursor) + + assert capability.status == "permission_denied" + + +def test_probe_caches_a_conclusive_answer_only(): + conclusive = FakeCursor(errno=1146) + inconclusive = FakeCursor(errno=1105) + + view_a, view_b = GraphView(), GraphView() + view_a.negotiate(conclusive) + view_a.negotiate(conclusive) + view_b.negotiate(inconclusive) + view_b.negotiate(inconclusive) + + assert len(conclusive.queries) == 1 + assert len(inconclusive.queries) == 2 + + +def test_summary_groups_endpoint_types_by_relationship(): + """The caller learns the graph's shape without reading its edges.""" + cursor = FakeCursor( + rows=[ + ("runs_on", "service", "k8s.pod", 3), + ("runs_on", "k8s.pod", "k8s.node", 2), + ("calls", "service", "service", 14), + ] + ) + + types = GraphView()._relationship_types(cursor, TimeWindow.parse(START, END)) + + runs_on = next(t for t in types if t["type"] == "runs_on") + assert runs_on["source_types"] == ["service", "k8s.pod"] + assert runs_on["destination_types"] == ["k8s.pod", "k8s.node"] + assert runs_on["count"] == 5 + + +def test_relationships_omit_red_columns_the_view_lacks(): + """Selecting a column the view lacks fails the statement to plan.""" + view = GraphView( + capability=GraphCapability( + "available", + entity_columns=FULL_ENTITY_COLUMNS, + relationship_columns=graph.RELATIONSHIP_REQUIRED_COLUMNS, + ) + ) + cursor = FakeCursor(rows=[]) + + view.relationships(cursor, request()) + + assert "request_count" not in cursor.queries[0] + assert "duration_max" not in cursor.queries[0] From 58a353599136cd10d752d08a9b5f20d5fc088859 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sat, 5 Sep 2026 15:40:21 +0800 Subject: [PATCH 2/9] fix: correct identifier decoding, window offsets, and the graph contract Two defects that produced wrong results. Every returned string was run through `json.loads`, so an identifier that looked like a JSON scalar changed type: entity_id "123" came back as the number 123, "true" as a boolean, and "null" as nothing at all. A caller cannot feed those back in, and "null" is unrecoverable. Only the four JSON columns are decoded now. Window bounds were formatted without their offset, so the database read them in the session time zone. Measured on 1.3 with an edge at 07:05Z and a window of [07:00Z, 07:10Z): one row under a +00:00 session, none under +08:00. Since this server takes --timezone, that silently answered a different question. Bounds now carry +00:00, and an integration test runs the suite's window through an Asia/Shanghai session. Four contract defects: - `next_query` in a no-match result omitted the required window, so following it failed on a missing argument. It now carries the window, and the integration test runs the suggested query instead of pinning its shape. - Capability checks disagreed with the SQL: `scope` was selected and filterable without being required, an absent RED column produced `MAX(confidence), , MIN(...)`, and rel_type=calls ordered by columns the probe allowed to be missing. GreptimeDB 1.3 is the floor, so every column read is now required and the SQL is fixed rather than capability-shaped. - The input schema carried no argument semantics: the SDK drops a bare string in `Annotated`. view is a Literal, so it reaches the schema as an enum, and the filters and limit use pydantic Field with descriptions and bounds. - The description misstated the contract. observed_at is the observation bucket, not the client side of a call; confidence 1.0 also covers declared edges; an unpaired client span still produces a virtual-node edge, so the trace-context claim was too strong; unmatched_count was undocumented despite being what separates a callee that stopped answering from a caller that stopped asking. summary now reports endpoint pairs. Two sets lose which combinations exist: service->pod and pod->node read as service->node. Entities also return descriptive and freshness. --- src/greptimedb_mcp_server/graph.py | 132 +++++++++++++++---------- src/greptimedb_mcp_server/server.py | 147 ++++++++++++++++++++-------- tests/integration/test_e2e.py | 76 ++++++++++++-- tests/test_graph.py | 75 +++++++++----- 4 files changed, 306 insertions(+), 124 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index e7e49b3..650fcde 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -19,26 +19,28 @@ ENTITIES_VIEW = "greptime_private.semantic_entities" RELATIONSHIPS_VIEW = "greptime_private.semantic_relationships" -# Reading a column the view lacks fails the statement to plan, so a view -# missing any of these cannot be queried the way this module queries it. -ENTITY_REQUIRED_COLUMNS = frozenset( - {"observed_at", "entity_type", "entity_id", "entity_id_attrs", "source_tables"} -) -RELATIONSHIP_REQUIRED_COLUMNS = frozenset( - { - "observed_at", - "src_type", - "src_id", - "dst_type", - "dst_id", - "rel_type", - "provenance", - "confidence", - } +# Read verbatim, so a view missing any of them cannot be queried the way this +# module queries it. These are the GreptimeDB 1.3 columns; the graph does not +# exist before that, so there is no older shape to degrade to. +ENTITY_COLUMNS = ( + "entity_type", + "entity_id", + "entity_id_attrs", + "scope", + "descriptive", + "source_tables", ) +ENTITY_REQUIRED_COLUMNS = frozenset({"observed_at", "fresh_until", *ENTITY_COLUMNS}) -# Summed across the window when present. A view without them still works; the -# fields are simply absent from the result. +RELATIONSHIP_IDENTITY_COLUMNS = ( + "src_type", + "src_id", + "dst_type", + "dst_id", + "rel_type", + "provenance", +) +# Summed over the buckets in the window. RED_COLUMNS = ( "request_count", "unmatched_count", @@ -46,6 +48,23 @@ "duration_sum", "duration_count", ) +RELATIONSHIP_REQUIRED_COLUMNS = frozenset( + { + "observed_at", + "fresh_until", + "confidence", + "duration_max", + *RELATIONSHIP_IDENTITY_COLUMNS, + *RED_COLUMNS, + } +) + +# Columns holding JSON. Everything else is returned verbatim: decoding an +# identifier would turn "123" into a number and "null" into nothing, and the +# caller needs the exact string back to query with it. +JSON_COLUMNS = frozenset( + {"entity_id_attrs", "descriptive", "source_tables", "attributes"} +) VIEWS = ("summary", "entities", "relationships") @@ -66,7 +85,6 @@ DEFAULT_LIMIT = 100 OBSERVATION_BUCKET_SECONDS = 60 -TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S.%f" ERRNO_TABLE_NOT_FOUND = 1146 ERRNO_PERMISSION_DENIED = frozenset({1044, 1045, 1142, 1143, 1227}) @@ -100,12 +118,6 @@ def conclusive(self) -> bool: "incompatible_schema", ) - def red_columns(self) -> list[str]: - return [c for c in RED_COLUMNS if c in self.relationship_columns] - - def has_duration_max(self) -> bool: - return "duration_max" in self.relationship_columns - @dataclass(frozen=True) class TimeWindow: @@ -124,10 +136,12 @@ def parse(cls, start_time: str, end_time: str) -> "TimeWindow": @property def params(self) -> list[str]: - return [ - self.start.strftime(TIMESTAMP_FORMAT), - self.end.strftime(TIMESTAMP_FORMAT), - ] + """Bind values that carry their offset. + + A literal without one is read in the session time zone, so the same + window would select different rows depending on GREPTIMEDB_TIMEZONE. + """ + return [self.start.isoformat(), self.end.isoformat()] def describe(self) -> dict: return { @@ -203,8 +217,8 @@ def names_an_id(self) -> bool: return any(name in ID_FILTERS for name in self.filters) -def _decode_json(value): - if not isinstance(value, str) or not value: +def _decode_json(name: str, value): + if name not in JSON_COLUMNS or not isinstance(value, str) or not value: return value try: return json.loads(value) @@ -213,7 +227,7 @@ def _decode_json(value): def _row_dict(columns: list[str], row) -> dict: - return {name: _decode_json(value) for name, value in zip(columns, row)} + return {name: _decode_json(name, value) for name, value in zip(columns, row)} def _filter_sql(filters: dict) -> tuple[list[str], list]: @@ -227,6 +241,12 @@ def _no_match_guidance(request: GraphRequest) -> dict: An unfiltered read of a large graph is the failure this tool exists to avoid, so the suggestion drops the identifier and keeps the type filters. """ + # The window is a required argument, so a next_query without it would not + # run. + window = { + "start_time": request.window.start.isoformat(), + "end_time": request.window.end.isoformat(), + } if request.names_an_id: kept = {k: v for k, v in request.filters.items() if k not in ID_FILTERS} return { @@ -234,14 +254,14 @@ def _no_match_guidance(request: GraphRequest) -> dict: "The supplied identifier may not be a canonical graph entity ID. " "IDs from alerts and telemetry are not interchangeable with them." ), - "next_query": {"view": request.view, **kept}, + "next_query": {"view": request.view, **kept, **window}, } return { "reason": ( "No rows matched in this window. The relationship may not have been " "witnessed, or the window may not cover it." ), - "next_query": {"view": "summary"}, + "next_query": {"view": "summary", **window}, } @@ -300,21 +320,17 @@ def _relationship_types(self, cursor, window: TimeWindow) -> list[dict]: "ORDER BY rel_type, src_type, dst_type", window.params, ) + # Endpoint types are reported as pairs. Reducing them to a set of + # sources and a set of destinations would imply combinations that do + # not exist: service->pod and pod->node would read as service->node. grouped: dict[str, dict] = {} for rel_type, src_type, dst_type, count in cursor.fetchall(): entry = grouped.setdefault( - rel_type, - { - "type": rel_type, - "source_types": [], - "destination_types": [], - "count": 0, - }, + rel_type, {"type": rel_type, "endpoints": [], "count": 0} + ) + entry["endpoints"].append( + {"source": src_type, "destination": dst_type, "count": int(count)} ) - if src_type not in entry["source_types"]: - entry["source_types"].append(src_type) - if dst_type not in entry["destination_types"]: - entry["destination_types"].append(dst_type) entry["count"] += int(count) return list(grouped.values()) @@ -322,11 +338,15 @@ def entities(self, cursor, request: GraphRequest) -> dict: """List distinct entities observed in the window.""" predicates, params = _filter_sql(request.filters) where = " AND ".join(["observed_at >= %s", "observed_at < %s", *predicates]) - identity = "entity_type, entity_id, entity_id_attrs, scope, source_tables" + # descriptive is grouped rather than aggregated: MAX over a JSON + # column would silently return one bucket's value for an entity whose + # attributes changed inside the window. + identity = ", ".join(ENTITY_COLUMNS) cursor.execute( f"SELECT {identity}," " MIN(observed_at) AS first_seen," - " MAX(observed_at) AS last_seen" + " MAX(observed_at) AS last_seen," + " MAX(fresh_until) AS fresh_until" f" FROM {ENTITIES_VIEW} WHERE {where}" f" GROUP BY {identity}" " ORDER BY entity_type, entity_id" @@ -336,20 +356,24 @@ def entities(self, cursor, request: GraphRequest) -> dict: return self._envelope(cursor, request) def relationships(self, cursor, request: GraphRequest) -> dict: - """List edges observed in the window, aggregated across buckets.""" - capability = self.capability - red = capability.red_columns() if capability else list(RED_COLUMNS) - aggregates = [f"SUM({column}) AS {column}" for column in red] - if capability is None or capability.has_duration_max(): - aggregates.append("MAX(duration_max) AS duration_max") + """List edges observed in the window, aggregated across buckets. + + `attributes` is left out: it varies per observation, so grouping by it + would split one edge into several rows and break the RED totals, while + aggregating it would present one bucket's value as the edge's. Read it + per observation with execute_sql. + """ + aggregates = [f"SUM({column}) AS {column}" for column in RED_COLUMNS] + aggregates.append("MAX(duration_max) AS duration_max") predicates, params = _filter_sql(request.filters) where = " AND ".join(["observed_at >= %s", "observed_at < %s", *predicates]) - identity = "src_type, src_id, dst_type, dst_id, rel_type, provenance" + identity = ", ".join(RELATIONSHIP_IDENTITY_COLUMNS) cursor.execute( f"SELECT {identity}, MAX(confidence) AS confidence, " + ", ".join(aggregates) + ", MIN(observed_at) AS first_seen, MAX(observed_at) AS last_seen" + + ", MAX(fresh_until) AS fresh_until" f" FROM {RELATIONSHIPS_VIEW} WHERE {where}" f" GROUP BY {identity}" f" ORDER BY {_relationship_order(request)}" diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index 9b91be9..f71db21 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -33,10 +33,11 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Annotated +from typing import Annotated, Literal from urllib.parse import quote import aiohttp +from pydantic import Field from mcp.server.mcpserver import MCPServer from mcp.server.mcpserver.exceptions import ResourceError, ToolError from mcp.server.transport_security import TransportSecuritySettings @@ -719,61 +720,123 @@ def _withdraw_graph_tool_if_unusable(state: AppState) -> None: @tool(name=GRAPH_TOOL_NAME) async def query_semantic_graph( view: Annotated[ - str, - "summary (what the graph contains), entities (nodes), or relationships " - "(edges)", + Literal["summary", "entities", "relationships"], + Field( + description=( + "summary: which entity and relationship types exist and what " + "they connect. entities: the nodes. relationships: the edges." + ) + ), ], start_time: Annotated[ - str, "Inclusive RFC3339 start of the window, e.g. 2026-09-05T07:00:00Z" + str, + Field( + description=( + "Inclusive start of the window, RFC3339, e.g. " + "2026-09-05T07:00:00Z. Without an offset it is read as UTC." + ) + ), + ], + end_time: Annotated[ + str, Field(description="Exclusive end of the window, RFC3339.") ], - end_time: Annotated[str, "Exclusive RFC3339 end of the window"], - entity_type: Annotated[str | None, "entities: filter by entity type"] = None, - entity_id: Annotated[str | None, "entities: filter by canonical entity id"] = None, - scope: Annotated[str | None, "entities: filter by namespace or environment"] = None, + entity_type: Annotated[ + str | None, + Field(description="entities only: service, k8s.pod, host, ..."), + ] = None, + entity_id: Annotated[ + str | None, + Field(description="entities only: a canonical id this graph returned."), + ] = None, + scope: Annotated[ + str | None, + Field( + description="entities only: the namespace or environment an id is scoped to." + ), + ] = None, rel_type: Annotated[ - str | None, "relationships: calls, runs_on, part_of, contains, uses, ..." + str | None, + Field( + description=( + "relationships only: calls, runs_on, contains, part_of, uses, " + "invokes, depends_on, owns, or a custom declared value." + ) + ), + ] = None, + src_type: Annotated[ + str | None, Field(description="relationships only: source endpoint type.") + ] = None, + src_id: Annotated[ + str | None, + Field( + description="relationships only: a canonical source id this graph returned." + ), + ] = None, + dst_type: Annotated[ + str | None, Field(description="relationships only: destination endpoint type.") + ] = None, + dst_id: Annotated[ + str | None, + Field( + description="relationships only: a canonical destination id this graph returned." + ), ] = None, - src_type: Annotated[str | None, "relationships: source endpoint type"] = None, - src_id: Annotated[str | None, "relationships: source endpoint id"] = None, - dst_type: Annotated[str | None, "relationships: destination endpoint type"] = None, - dst_id: Annotated[str | None, "relationships: destination endpoint id"] = None, provenance: Annotated[ - str | None, "relationships: trace, attribute, declared, or agent" + str | None, + Field( + description=( + "relationships only: how the edge was obtained -- trace " + "(paired spans), attribute (identities on one row), declared, " + "or agent." + ) + ), ] = None, limit: Annotated[ - int, f"Maximum rows to return (1-{graph.MAX_LIMIT}, default: 100)" + int, + Field(description="Maximum rows to return.", ge=1, le=graph.MAX_LIMIT), ] = graph.DEFAULT_LIMIT, ) -> str: """Query the semantic graph: which entities exist and which are related. - Start with view=summary. It returns the entity types, the relationship - types, and the endpoint types each relationship connects, so the shape of - the graph is known before any edge is read. - - The window is required and is half-open, [start_time, end_time), over - observed_at -- the 60-second bucket an observation was recorded in, taken - from the client side of a call. Rows are aggregated across the buckets in - the window, so one edge is one row and RED fields are summed over the - window. The result echoes the window it used. - - Not every relationship type carries request, error, and duration counts: - only `calls` does. An unfiltered result is ordered by relationship type and - endpoint, not by those counts, because ordering a mixed result by them - would rank every other kind of relationship last. Pass rel_type=calls to - order by error and request count. - - confidence is derivation certainty, not health: 1.0 when both sides were - observed, 0.5 when only the caller was and the callee is inferred from a - peer attribute. `calls` edges only cover calls that propagated trace - context, so a dependency invoked without it is absent. A missing edge is - not evidence that two entities are unrelated, and entities are not - deduplicated across identity schemes -- the same process can appear under - two ids if two sources named it differently. + Start with view=summary. It returns each relationship type with the + endpoint type pairs it actually connects, so the shape of the graph is + known before any edge is read. + + The window is required and half-open, [start_time, end_time), over + observed_at -- the 60-second bucket an observation was recorded in. Rows + are aggregated across the buckets in the window, so one edge is one row and + request, error and duration fields are summed over it. The result echoes + the window it used. + + Ordering: only `calls` edges carry request, error and duration counts, so + an unfiltered result is ordered by relationship type and endpoint. Ordering + a mixed result by those counts would rank every other relationship type + below a null. Pass rel_type=calls to order by error and request count. + + confidence is derivation certainty, not health. A paired or declared edge + is 1.0; an edge whose callee was named by a client-side peer attribute + rather than observed is 0.5, and that endpoint is a virtual node. Read + provenance for how a row was obtained rather than inferring it from the + number. + + request_count counts calls whose client and server spans paired. + unmatched_count counts client spans with no matching server span, and is + not additive with it: a fall in request_count with unmatched_count present + means the callee stopped answering, while a fall in both means the caller + stopped asking. error_count aggregates span status verbatim, and some SDKs + mark a normal long-lived stream timeout as an error, so read the errors + before concluding from a rate. + + A missing edge is not evidence that two entities are unrelated: it can also + mean the call was not instrumented, was sampled out, or fell outside this + window. Entities are not deduplicated across identity schemes, so one + process can appear under two ids if two sources named it differently -- + compare their runs_on and part_of edges before treating them as two things. entity_id_attrs names the attributes an id was assembled from and - source_tables names the telemetry tables that witnessed it; go there for - the raw data. Identifiers from alerts and telemetry are not graph entity - ids unless a query returns that exact id. + source_tables names the telemetry tables that witnessed it; query those for + the underlying rows. Identifiers from alerts and other tools are not graph + ids unless a query here returned that exact string. """ state = get_state() request = graph.GraphRequest.parse( diff --git a/tests/integration/test_e2e.py b/tests/integration/test_e2e.py index 27fdaf9..92a3e3e 100644 --- a/tests/integration/test_e2e.py +++ b/tests/integration/test_e2e.py @@ -13,6 +13,7 @@ from .conftest import ( CREDENTIALS_TABLE, MASK_PLACEHOLDER, + GRAPH_EDGES, METRICS_TABLE, SECRET_API_KEY, SECRET_PASSWORD, @@ -364,12 +365,67 @@ async def test_graph_summary_reports_shape_without_returning_edges(seed): ) by_type = {item["type"]: item for item in payload["relationship_types"]} - assert by_type["calls"]["source_types"] == ["service"] assert by_type["calls"]["count"] == 2 - assert set(by_type["runs_on"]["destination_types"]) == {"k8s.pod", "k8s.node"} + assert by_type["calls"]["endpoints"] == [ + {"source": "service", "destination": "service", "count": 2} + ] + # Pairs, not two sets: service->pod and pod->node must not read as + # service->node. + assert by_type["runs_on"]["endpoints"] == [ + {"source": "k8s.pod", "destination": "k8s.node", "count": 1}, + {"source": "service", "destination": "k8s.pod", "count": 1}, + ] assert "items" not in payload +async def test_graph_window_is_read_in_utc_whatever_the_session(seed): + """A literal without an offset is read in the session time zone.""" + if not seed.graph_seeded: + pytest.skip("this GreptimeDB has no semantic graph") + + window = _graph_window(seed) + async with stdio_session(**{"--timezone": "Asia/Shanghai"}) as client: + payload = json.loads( + await call_text( + client, "query_semantic_graph", {"view": "summary", **window} + ) + ) + + assert payload["relationship_count"] == len(GRAPH_EDGES) + + +async def test_graph_returns_identifiers_verbatim(seed): + """A caller has to be able to feed a returned id straight back in.""" + if not seed.graph_seeded: + pytest.skip("this GreptimeDB has no semantic graph") + + window = _graph_window(seed) + async with stdio_session() as client: + payload = json.loads( + await call_text( + client, + "query_semantic_graph", + {"view": "relationships", "rel_type": "calls", **window}, + ) + ) + first = payload["items"][0] + refetched = json.loads( + await call_text( + client, + "query_semantic_graph", + { + "view": "relationships", + "src_id": first["src_id"], + "dst_id": first["dst_id"], + **window, + }, + ) + ) + + assert isinstance(first["src_id"], str) + assert refetched["item_count"] == 1 + + async def test_graph_orders_mixed_relationships_by_type_not_red(seed): """Only `calls` has RED; sorting a mixed result by it buries the rest.""" if not seed.graph_seeded: @@ -417,7 +473,15 @@ async def test_graph_zero_result_keeps_the_envelope(seed): assert payload["status"] == "no_match" assert payload["items"] == [] assert payload["applied_filters"]["rel_type"] == "calls" - assert payload["guidance"]["next_query"] == { - "view": "relationships", - "rel_type": "calls", - } + + # The suggested retry drops the identifier, keeps the type filter, and + # carries the window, so a caller can run it as given. + next_query = payload["guidance"]["next_query"] + assert "src_id" not in next_query + assert next_query["rel_type"] == "calls" + + async with stdio_session() as client: + retried = json.loads( + await call_text(client, "query_semantic_graph", next_query) + ) + assert retried["status"] == "ok" diff --git a/tests/test_graph.py b/tests/test_graph.py index bbf094e..9ca8e5c 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -19,9 +19,7 @@ END = "2026-09-05T08:00:00Z" FULL_ENTITY_COLUMNS = frozenset(graph.ENTITY_REQUIRED_COLUMNS) -FULL_RELATIONSHIP_COLUMNS = frozenset( - graph.RELATIONSHIP_REQUIRED_COLUMNS | set(graph.RED_COLUMNS) | {"duration_max"} -) +FULL_RELATIONSHIP_COLUMNS = frozenset(graph.RELATIONSHIP_REQUIRED_COLUMNS) class FakeCursor: @@ -114,14 +112,17 @@ def test_calls_relationships_are_ordered_by_red(): def test_no_match_guidance_drops_the_id_and_keeps_the_type(): guidance = _no_match_guidance(request(rel_type="calls", src_id="unknown")) - assert guidance["next_query"] == {"view": "relationships", "rel_type": "calls"} + next_query = guidance["next_query"] + assert next_query["view"] == "relationships" + assert next_query["rel_type"] == "calls" + assert "src_id" not in next_query assert "canonical graph entity ID" in guidance["reason"] def test_no_match_guidance_without_an_id_points_at_the_summary(): guidance = _no_match_guidance(request(rel_type="calls")) - assert guidance["next_query"] == {"view": "summary"} + assert guidance["next_query"]["view"] == "summary" @pytest.mark.parametrize( @@ -175,8 +176,8 @@ def test_probe_caches_a_conclusive_answer_only(): assert len(inconclusive.queries) == 2 -def test_summary_groups_endpoint_types_by_relationship(): - """The caller learns the graph's shape without reading its edges.""" +def test_summary_reports_endpoint_pairs_not_two_sets(): + """Two sets would imply service->node from service->pod and pod->node.""" cursor = FakeCursor( rows=[ ("runs_on", "service", "k8s.pod", 3), @@ -188,23 +189,53 @@ def test_summary_groups_endpoint_types_by_relationship(): types = GraphView()._relationship_types(cursor, TimeWindow.parse(START, END)) runs_on = next(t for t in types if t["type"] == "runs_on") - assert runs_on["source_types"] == ["service", "k8s.pod"] - assert runs_on["destination_types"] == ["k8s.pod", "k8s.node"] + assert runs_on["endpoints"] == [ + {"source": "service", "destination": "k8s.pod", "count": 3}, + {"source": "k8s.pod", "destination": "k8s.node", "count": 2}, + ] assert runs_on["count"] == 5 -def test_relationships_omit_red_columns_the_view_lacks(): - """Selecting a column the view lacks fails the statement to plan.""" - view = GraphView( - capability=GraphCapability( - "available", - entity_columns=FULL_ENTITY_COLUMNS, - relationship_columns=graph.RELATIONSHIP_REQUIRED_COLUMNS, - ) - ) - cursor = FakeCursor(rows=[]) +def test_a_view_missing_a_red_column_is_incompatible(): + """The query reads every RED column, so a view without one cannot serve it.""" + columns = { + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, + graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS - {"unmatched_count"}, + } + + capability = GraphView().negotiate(FakeCursor(columns=columns)) + + assert capability.status == "incompatible_schema" + assert "unmatched_count" in capability.detail + - view.relationships(cursor, request()) +def test_identifier_shaped_strings_survive_the_row_decode(): + """Decoding every string would make entity_id "123" a number.""" + columns = ["entity_id", "entity_id_attrs", "source_tables"] + row = ("123", '{"host":"123"}', '["public.t"]') + + decoded = graph._row_dict(columns, row) + + assert decoded["entity_id"] == "123" + assert decoded["entity_id_attrs"] == {"host": "123"} + assert decoded["source_tables"] == ["public.t"] + + +def test_identifiers_that_look_like_json_literals_survive(): + assert graph._row_dict(["src_id"], ("null",))["src_id"] == "null" + assert graph._row_dict(["dst_id"], ("true",))["dst_id"] == "true" + + +def test_window_params_carry_their_offset(): + """A literal without one is read in the session time zone.""" + params = TimeWindow.parse(START, END).params + + assert all(p.endswith("+00:00") for p in params) + + +def test_no_match_guidance_next_query_is_runnable(): + """start_time and end_time are required, so a retry without them fails.""" + guidance = _no_match_guidance(request(rel_type="calls", src_id="unknown")) - assert "request_count" not in cursor.queries[0] - assert "duration_max" not in cursor.queries[0] + assert set(guidance["next_query"]) >= {"view", "start_time", "end_time"} + GraphRequest.parse(**guidance["next_query"]) From 69cbe8c813d5270a96ab11d49c9ebf6daaa59f93 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 16:30:22 +0800 Subject: [PATCH 3/9] fix: mask graph attributes, keep confidence populations apart, and trim the description Three review findings and five API notes. The graph tool serialized rows directly, so `mask_enabled` and `mask_patterns` did not apply to it: an attribute map could publish what the same value would have had masked through execute_sql. Attribute values whose name matches a sensitive pattern are masked now, in `entity_id_attrs` and `descriptive`, and `entity_id` with them when it was assembled from one -- it is those values joined, so leaving it would undo the masking. `MAX(confidence)` flattened two different measurements. The database reports 1.0 for a bucket whose spans paired and 0.5 for one that saw only clients, and `real_wins` switches request_count, error_count and the durations to whichever population that bucket describes: a pair timed by the server span, an unmatched client by its own. Summing across both added those populations together and the maximum hid that it had happened. `confidence` is in the group key now, so such an edge returns one row per population. The description asserted that request_count counts paired calls, which is only true of a bucket that paired; where none did, the database reports the unmatched clients there. It also read a fall in those counts as the callee having stopped answering or the caller having stopped asking, which sampling and missing instrumentation produce equally. Both are gone, durations are stated in seconds, and the entities view now says it returns one row per set of attributes so item_count is not an entity count, and that first_seen and last_seen bound the observation inside the window rather than the entity's lifetime. `Start with view=summary` became a condition rather than an order, and the RCA guidance is out. A truncated result now carries the effective limit and which filters would narrow it, since there is no cursor and an unfiltered result ordered by type can drop later types entirely. A failed query raises instead of returning prose, so `status=no_match` only ever means the query ran and matched nothing. Not covered: an edge observed at both confidences end to end. Declared edges cannot express it -- the computed view keeps one revision per identity, so a second row for the same edge replaces the first whatever its scope or bucket. That needs trace-derived pairing, which the integration suite does not produce. --- src/greptimedb_mcp_server/graph.py | 87 ++++++++++++++++++++++++++++- src/greptimedb_mcp_server/server.py | 66 ++++++++++++---------- tests/integration/conftest.py | 23 +++----- tests/test_graph.py | 54 ++++++++++++++++++ 4 files changed, 182 insertions(+), 48 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index 650fcde..e425f8f 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -16,6 +16,12 @@ from mysql.connector import Error +from greptimedb_mcp_server.masking import ( + DEFAULT_SENSITIVE_PATTERNS, + MASK_PLACEHOLDER, + is_sensitive_column, +) + ENTITIES_VIEW = "greptime_private.semantic_entities" RELATIONSHIPS_VIEW = "greptime_private.semantic_relationships" @@ -32,6 +38,12 @@ ) ENTITY_REQUIRED_COLUMNS = frozenset({"observed_at", "fresh_until", *ENTITY_COLUMNS}) +# `confidence` is in the group key, not aggregated. The database reports 1.0 +# for a bucket that paired client and server spans and 0.5 for one that only +# saw clients, and it switches request_count, error_count and the durations to +# the matching population at the same time -- a pair is timed by the server +# span, an unmatched client by its own. Summing across both would add two +# different measurements and MAX(confidence) would hide that it happened. RELATIONSHIP_IDENTITY_COLUMNS = ( "src_type", "src_id", @@ -39,6 +51,7 @@ "dst_id", "rel_type", "provenance", + "confidence", ) # Summed over the buckets in the window. RED_COLUMNS = ( @@ -230,6 +243,45 @@ def _row_dict(columns: list[str], row) -> dict: return {name: _decode_json(name, value) for name, value in zip(columns, row)} +# Attribute maps carry telemetry values keyed by the column they came from, so +# the column-name rule that masks query results applies to them too. +MASKABLE_ATTRIBUTE_COLUMNS = ("entity_id_attrs", "descriptive") + + +def mask_patterns(mask_enabled: bool, extra: list[str] | None) -> list[str] | None: + """The pattern list to mask with, or None when masking is off.""" + if not mask_enabled: + return None + return [*DEFAULT_SENSITIVE_PATTERNS, *(extra or [])] + + +def _mask_entity(item: dict, patterns: list[str] | None) -> dict: + """Hide attribute values whose name matches a sensitive pattern. + + `entity_id` is masked as well when it was assembled from one of them: it is + those values joined, so leaving it would publish what the map just hid. + That does make the entity unqueryable by id, which is what masking a column + does everywhere else in this server. + """ + if not patterns: + return item + masked = dict(item) + for column in MASKABLE_ATTRIBUTE_COLUMNS: + value = masked.get(column) + if not isinstance(value, dict): + continue + masked[column] = { + name: (MASK_PLACEHOLDER if is_sensitive_column(name, patterns) else attr) + for name, attr in value.items() + } + identifying = masked.get("entity_id_attrs") + if isinstance(identifying, dict) and any( + is_sensitive_column(name, patterns) for name in identifying + ): + masked["entity_id"] = MASK_PLACEHOLDER + return masked + + def _filter_sql(filters: dict) -> tuple[list[str], list]: predicates = [f"{name} = %s" for name in filters] return predicates, list(filters.values()) @@ -265,11 +317,32 @@ def _no_match_guidance(request: GraphRequest) -> dict: } +def _truncation_guidance(request: GraphRequest) -> dict: + """Say how to reach the rows the limit cut off. + + There is no cursor. An unfiltered result is ordered by relationship type, + so a type with many edges can push later types out of the window entirely; + narrowing by type is what brings them back. + """ + narrow = [name for name in RELATIONSHIP_FILTERS if name not in request.filters] + if request.view == "entities": + narrow = [name for name in ENTITY_FILTERS if name not in request.filters] + return { + "reason": ( + f"More rows matched than the limit of {request.limit}. Ordering is " + "by type and endpoint, so kinds sorting later may be missing " + "entirely rather than merely cut short." + ), + "narrow_by": narrow, + } + + @dataclass class GraphView: """A handle on the two graph views that remembers whether they work.""" capability: GraphCapability | None = field(default=None) + mask: list[str] | None = field(default=None) def negotiate(self, cursor) -> GraphCapability: """Decide once per process whether the graph is usable.""" @@ -358,6 +431,10 @@ def entities(self, cursor, request: GraphRequest) -> dict: def relationships(self, cursor, request: GraphRequest) -> dict: """List edges observed in the window, aggregated across buckets. + One edge yields one row per confidence, so an edge whose buckets were + partly paired and partly client-only comes back as two rows rather than + one sum over both populations. + `attributes` is left out: it varies per observation, so grouping by it would split one edge into several rows and break the RED totals, while aggregating it would present one bucket's value as the edge's. Read it @@ -370,7 +447,7 @@ def relationships(self, cursor, request: GraphRequest) -> dict: where = " AND ".join(["observed_at >= %s", "observed_at < %s", *predicates]) identity = ", ".join(RELATIONSHIP_IDENTITY_COLUMNS) cursor.execute( - f"SELECT {identity}, MAX(confidence) AS confidence, " + f"SELECT {identity}, " + ", ".join(aggregates) + ", MIN(observed_at) AS first_seen, MAX(observed_at) AS last_seen" + ", MAX(fresh_until) AS fresh_until" @@ -386,18 +463,24 @@ def _envelope(self, cursor, request: GraphRequest) -> dict: columns = [desc[0] for desc in cursor.description] rows = cursor.fetchall() complete = len(rows) <= request.limit - items = [_row_dict(columns, row) for row in rows[: request.limit]] + items = [ + _mask_entity(_row_dict(columns, row), self.mask) + for row in rows[: request.limit] + ] result = { "view": request.view, "status": "ok" if items else "no_match", "window": request.window.describe(), "applied_filters": dict(request.filters), + "limit": request.limit, "items": items, "item_count": len(items), "complete": complete, } if not items: result["guidance"] = _no_match_guidance(request) + elif not complete: + result["guidance"] = _truncation_guidance(request) return result diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index f71db21..c382f63 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -324,6 +324,9 @@ async def lifespan(mcp: MCPServer): mask_patterns=mask_patterns, allow_write=config.allow_write, http_session=aiohttp.ClientSession(), + semantic_graph=graph.GraphView( + mask=graph.mask_patterns(config.mask_enabled, mask_patterns) + ), ) safe_db_config = {**db_config, "password": "***" if config.password else ""} @@ -798,45 +801,44 @@ async def query_semantic_graph( ) -> str: """Query the semantic graph: which entities exist and which are related. - Start with view=summary. It returns each relationship type with the - endpoint type pairs it actually connects, so the shape of the graph is - known before any edge is read. + Use view=summary when the entity and relationship types in this graph are + not known yet; it reports them with the endpoint type pairs each + relationship connects. With a type or an id already in hand, query + entities or relationships directly. The window is required and half-open, [start_time, end_time), over observed_at -- the 60-second bucket an observation was recorded in. Rows - are aggregated across the buckets in the window, so one edge is one row and - request, error and duration fields are summed over it. The result echoes - the window it used. - - Ordering: only `calls` edges carry request, error and duration counts, so - an unfiltered result is ordered by relationship type and endpoint. Ordering - a mixed result by those counts would rank every other relationship type - below a null. Pass rel_type=calls to order by error and request count. - - confidence is derivation certainty, not health. A paired or declared edge - is 1.0; an edge whose callee was named by a client-side peer attribute - rather than observed is 0.5, and that endpoint is a virtual node. Read - provenance for how a row was obtained rather than inferring it from the - number. - - request_count counts calls whose client and server spans paired. - unmatched_count counts client spans with no matching server span, and is - not additive with it: a fall in request_count with unmatched_count present - means the callee stopped answering, while a fall in both means the caller - stopped asking. error_count aggregates span status verbatim, and some SDKs - mark a normal long-lived stream timeout as an error, so read the errors - before concluding from a rate. + are aggregated across the buckets in the window, and the result echoes the + window and the limit it used. + + relationships returns one row per edge per confidence. The database reports + confidence 1.0 for a bucket whose client and server spans paired and 0.5 + for one where only the client was seen, and it switches request_count, + error_count and the durations to whichever population that bucket + describes: paired requests timed by the server span, or unmatched clients + timed by their own. An edge observed both ways therefore comes back as two + rows. unmatched_count reports client spans with no paired server span. + Durations are in seconds. + + entities returns one row per distinct set of attributes, so an entity whose + descriptive attributes changed inside the window appears more than once; + item_count counts rows, not entities. first_seen and last_seen bound where + the row was observed inside this window, not when the entity first existed. + + Ordering is by type and endpoint. Only `calls` edges carry request, error + and duration counts, so pass rel_type=calls to order by error and request + count instead. When complete is false, more rows matched than the limit and + later types may be absent entirely rather than merely cut short. A missing edge is not evidence that two entities are unrelated: it can also mean the call was not instrumented, was sampled out, or fell outside this window. Entities are not deduplicated across identity schemes, so one - process can appear under two ids if two sources named it differently -- - compare their runs_on and part_of edges before treating them as two things. + process can appear under two ids if two sources named it differently. entity_id_attrs names the attributes an id was assembled from and - source_tables names the telemetry tables that witnessed it; query those for - the underlying rows. Identifiers from alerts and other tools are not graph - ids unless a query here returned that exact string. + source_tables names the telemetry tables that witnessed it. Identifiers + from alerts and other tools are not graph ids unless a query here returned + that exact string. """ state = get_state() request = graph.GraphRequest.parse( @@ -876,8 +878,10 @@ def _sync_query(): try: result = await asyncio.to_thread(_sync_query) except Error as e: + # A failed read is not an empty graph: raising keeps `status=no_match` + # meaning "the query ran and matched nothing". logger.error(f"Error querying the semantic graph: {e}") - return f"Error querying the semantic graph: {str(e)}" + raise ToolError(f"Error querying the semantic graph: {str(e)}") from e return json.dumps(result, ensure_ascii=False, indent=2, default=str) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 33f8fd7..1cbd7c8 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -32,10 +32,10 @@ # Declared edges are inserted directly, so the graph has both a RED-bearing # relationship type and one without needing OTLP traces to pair up. GRAPH_EDGES = ( - ("service", "it-frontend", "service", "it-checkout", "calls", 100, 7), - ("service", "it-checkout", "service", "it-payment", "calls", 100, 45), - ("service", "it-checkout", "k8s.pod", "it-pod-a", "runs_on", None, None), - ("k8s.pod", "it-pod-a", "k8s.node", "it-node-1", "runs_on", None, None), + ("service", "it-frontend", "service", "it-checkout", "calls", 100, 7, 1.0, ""), + ("service", "it-checkout", "service", "it-payment", "calls", 100, 45, 1.0, ""), + ("service", "it-checkout", "k8s.pod", "it-pod-a", "runs_on", None, None, 1.0, ""), + ("k8s.pod", "it-pod-a", "k8s.node", "it-node-1", "runs_on", None, None, 1.0, ""), ) METRIC_HOSTS = ("host-a", "host-b") @@ -164,19 +164,12 @@ def _seed_declared_edges(cursor) -> bool: "scope, generation_id, confidence, request_count, error_count" ) try: - for ( - src_type, - src_id, - dst_type, - dst_id, - rel_type, - requests, - errors, - ) in GRAPH_EDGES: + for edge in GRAPH_EDGES: + src, sid, dst, did, rel, requests, errors, confidence, scope = edge cursor.execute( f"INSERT INTO {DECLARED_EDGES_TABLE} ({columns}) VALUES " - "(now(), %s, %s, %s, %s, %s, 'declared', '', '', 1.0, %s, %s)", - (src_type, src_id, dst_type, dst_id, rel_type, requests, errors), + "(now(), %s, %s, %s, %s, %s, 'declared', %s, '', %s, %s, %s)", + (src, sid, dst, did, rel, scope, confidence, requests, errors), ) except mysql.connector.Error: return False diff --git a/tests/test_graph.py b/tests/test_graph.py index 9ca8e5c..d8ed364 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -239,3 +239,57 @@ def test_no_match_guidance_next_query_is_runnable(): assert set(guidance["next_query"]) >= {"view", "start_time", "end_time"} GraphRequest.parse(**guidance["next_query"]) + + +def test_confidence_is_grouped_not_aggregated(): + """A paired bucket and a client-only bucket measure different populations.""" + cursor = FakeCursor(rows=[]) + GraphView().relationships(cursor, request()) + + assert "MAX(confidence)" not in cursor.queries[0] + assert ( + "GROUP BY src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence" + in (cursor.queries[0]) + ) + + +def test_sensitive_attributes_are_masked_by_name(): + """The column-name rule that masks query results covers attribute maps.""" + item = { + "entity_type": "service", + "entity_id": "checkout,hunter2", + "entity_id_attrs": {"service_name": "checkout", "api_key": "sk-live"}, + "descriptive": {"team": "payments", "access_token": "t-123"}, + } + + masked = graph._mask_entity(item, graph.mask_patterns(True, None)) + + assert masked["entity_id_attrs"] == { + "service_name": "checkout", + "api_key": "******", + } + assert masked["descriptive"] == {"team": "payments", "access_token": "******"} + # the id is those values joined, so publishing it would undo the masking + assert masked["entity_id"] == "******" + + +def test_masking_off_returns_attributes_untouched(): + item = {"entity_id": "checkout", "entity_id_attrs": {"api_key": "sk-live"}} + + assert graph._mask_entity(item, graph.mask_patterns(False, ["api_key"])) == item + + +def test_custom_patterns_extend_the_defaults(): + item = {"entity_id": "checkout", "entity_id_attrs": {"internal_ref": "r-1"}} + + masked = graph._mask_entity(item, graph.mask_patterns(True, ["internal_ref"])) + + assert masked["entity_id_attrs"]["internal_ref"] == "******" + + +def test_truncated_result_says_how_to_narrow(): + """There is no cursor, so the caller needs to know what to filter by.""" + guidance = graph._truncation_guidance(request(rel_type="calls")) + + assert "rel_type" not in guidance["narrow_by"] + assert "src_id" in guidance["narrow_by"] From 729f1c9612b80f6b815f53c7c655a72774d9cbc0 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 17:32:56 +0800 Subject: [PATCH 4/9] fix: apply the column-name masking rule to graph fields, not only inside them The previous change masked the names inside attribute maps and stopped there, so a returned field was never checked against the patterns itself: `mask_patterns=descriptive` left `descriptive` fully readable, and a pattern naming `entity_id`, `src_id` or `dst_id` did nothing. A field whose own name matches is hidden first now, exactly as the column would be through execute_sql; maps that survive that are then masked by the names inside them. Hiding `entity_id` when a masked attribute helped build it does not extend to `relationships`: its view carries no attribute names, so the same value can still appear there as `src_id` or `dst_id` unless a pattern names those columns. The previous message implied those values were hidden everywhere. They are not, and the tool description now says so. The truncation hint claimed ordering by type and endpoint, which is wrong under rel_type=calls; it states that later rows are missing and lists what to narrow by. --- src/greptimedb_mcp_server/graph.py | 34 +++++++++++++++++++---------- src/greptimedb_mcp_server/server.py | 7 ++++++ tests/test_graph.py | 28 ++++++++++++++++++++---- 3 files changed, 53 insertions(+), 16 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index e425f8f..3fa2595 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -255,17 +255,27 @@ def mask_patterns(mask_enabled: bool, extra: list[str] | None) -> list[str] | No return [*DEFAULT_SENSITIVE_PATTERNS, *(extra or [])] -def _mask_entity(item: dict, patterns: list[str] | None) -> dict: - """Hide attribute values whose name matches a sensitive pattern. - - `entity_id` is masked as well when it was assembled from one of them: it is - those values joined, so leaving it would publish what the map just hid. - That does make the entity unqueryable by id, which is what masking a column - does everywhere else in this server. +def _mask_item(item: dict, patterns: list[str] | None) -> dict: + """Apply the column-name masking rule to one returned row. + + A returned field whose own name matches a pattern is hidden outright, as + the column would be through execute_sql. Attribute maps that survive that + are then masked by the names inside them, since those names are the columns + the values came from. + + `entity_id` is hidden as well when it was assembled from a masked + attribute, because it is those values joined. That does not reach + `relationships`: its view carries no attribute names, so the same value can + still surface there as `src_id` or `dst_id` unless a pattern matches those + column names. """ if not patterns: return item - masked = dict(item) + + masked = { + name: (MASK_PLACEHOLDER if is_sensitive_column(name, patterns) else value) + for name, value in item.items() + } for column in MASKABLE_ATTRIBUTE_COLUMNS: value = masked.get(column) if not isinstance(value, dict): @@ -329,9 +339,9 @@ def _truncation_guidance(request: GraphRequest) -> dict: narrow = [name for name in ENTITY_FILTERS if name not in request.filters] return { "reason": ( - f"More rows matched than the limit of {request.limit}. Ordering is " - "by type and endpoint, so kinds sorting later may be missing " - "entirely rather than merely cut short." + f"More rows matched than the limit of {request.limit}, so rows " + "sorting later are missing entirely rather than merely cut short. " + "Narrow with the filters below." ), "narrow_by": narrow, } @@ -464,7 +474,7 @@ def _envelope(self, cursor, request: GraphRequest) -> dict: rows = cursor.fetchall() complete = len(rows) <= request.limit items = [ - _mask_entity(_row_dict(columns, row), self.mask) + _mask_item(_row_dict(columns, row), self.mask) for row in rows[: request.limit] ] result = { diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index c382f63..e53e41d 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -839,6 +839,13 @@ async def query_semantic_graph( source_tables names the telemetry tables that witnessed it. Identifiers from alerts and other tools are not graph ids unless a query here returned that exact string. + + When masking is on, a returned field is hidden if its own name matches a + sensitive pattern, and an attribute map is also masked by the names inside + it. entities additionally hides entity_id when a masked attribute helped + build it; relationships cannot do the same, because its view does not carry + attribute names, so such a value can still appear there as src_id or + dst_id. """ state = get_state() request = graph.GraphRequest.parse( diff --git a/tests/test_graph.py b/tests/test_graph.py index d8ed364..1ebc31e 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -253,8 +253,28 @@ def test_confidence_is_grouped_not_aggregated(): ) +def test_a_matching_column_name_hides_the_whole_field(): + """The same rule execute_sql applies to a column applies to a field.""" + item = {"entity_type": "service", "descriptive": {"team": "payments"}} + + masked = graph._mask_item(item, graph.mask_patterns(True, ["descriptive"])) + + assert masked["descriptive"] == "******" + assert masked["entity_type"] == "service" + + +def test_endpoint_ids_are_hidden_when_the_pattern_names_them(): + """relationships has no attribute names, so the column rule is all it has.""" + item = {"src_id": "checkout", "dst_id": "payment", "rel_type": "calls"} + + masked = graph._mask_item(item, graph.mask_patterns(True, ["src_id"])) + + assert masked["src_id"] == "******" + assert masked["dst_id"] == "payment" + + def test_sensitive_attributes_are_masked_by_name(): - """The column-name rule that masks query results covers attribute maps.""" + """The column-name rule reaches inside attribute maps too.""" item = { "entity_type": "service", "entity_id": "checkout,hunter2", @@ -262,7 +282,7 @@ def test_sensitive_attributes_are_masked_by_name(): "descriptive": {"team": "payments", "access_token": "t-123"}, } - masked = graph._mask_entity(item, graph.mask_patterns(True, None)) + masked = graph._mask_item(item, graph.mask_patterns(True, None)) assert masked["entity_id_attrs"] == { "service_name": "checkout", @@ -276,13 +296,13 @@ def test_sensitive_attributes_are_masked_by_name(): def test_masking_off_returns_attributes_untouched(): item = {"entity_id": "checkout", "entity_id_attrs": {"api_key": "sk-live"}} - assert graph._mask_entity(item, graph.mask_patterns(False, ["api_key"])) == item + assert graph._mask_item(item, graph.mask_patterns(False, ["api_key"])) == item def test_custom_patterns_extend_the_defaults(): item = {"entity_id": "checkout", "entity_id_attrs": {"internal_ref": "r-1"}} - masked = graph._mask_entity(item, graph.mask_patterns(True, ["internal_ref"])) + masked = graph._mask_item(item, graph.mask_patterns(True, ["internal_ref"])) assert masked["entity_id_attrs"]["internal_ref"] == "******" From b98a6cb6d83bcf2ac50befa5c084b7639f545b02 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 18:22:22 +0800 Subject: [PATCH 5/9] fix: decide the masked identity from the original row Whether to hide `entity_id` was read back off the already-masked row, so a pattern that hid the whole attribute map left nothing that looked like one and the id went out in the clear. Configuring `mask_patterns=entity_id_attrs` therefore re-exposed a value that the default rules had hidden: adding a rule made masking weaker. The decision is taken from the original row now, and a map hidden in full also hides the id built from it. A regression test asserts the property directly -- across pattern sets, the id stays hidden and the sensitive value never appears in the output. --- src/greptimedb_mcp_server/graph.py | 20 ++++++++++++++++---- tests/test_graph.py | 21 +++++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index 3fa2595..e1888cc 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -272,6 +272,11 @@ def _mask_item(item: dict, patterns: list[str] | None) -> dict: if not patterns: return item + # Decided from the original row: reading it back after masking would miss + # the case where the whole map was hidden, and adding a pattern would then + # expose an id that a narrower rule had hidden. + hide_id = _identity_is_sensitive(item, patterns) + masked = { name: (MASK_PLACEHOLDER if is_sensitive_column(name, patterns) else value) for name, value in item.items() @@ -284,14 +289,21 @@ def _mask_item(item: dict, patterns: list[str] | None) -> dict: name: (MASK_PLACEHOLDER if is_sensitive_column(name, patterns) else attr) for name, attr in value.items() } - identifying = masked.get("entity_id_attrs") - if isinstance(identifying, dict) and any( - is_sensitive_column(name, patterns) for name in identifying - ): + if hide_id: masked["entity_id"] = MASK_PLACEHOLDER return masked +def _identity_is_sensitive(item: dict, patterns: list[str]) -> bool: + """Whether `entity_id` was assembled from anything the patterns hide.""" + attrs = item.get("entity_id_attrs") + if not isinstance(attrs, dict): + return False + if is_sensitive_column("entity_id_attrs", patterns): + return True + return any(is_sensitive_column(name, patterns) for name in attrs) + + def _filter_sql(filters: dict) -> tuple[list[str], list]: predicates = [f"{name} = %s" for name in filters] return predicates, list(filters.values()) diff --git a/tests/test_graph.py b/tests/test_graph.py index 1ebc31e..3b1634c 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -1,5 +1,6 @@ """Tests for the semantic graph window, ordering, and capability contracts.""" +import json from datetime import datetime, timezone import pytest @@ -293,6 +294,26 @@ def test_sensitive_attributes_are_masked_by_name(): assert masked["entity_id"] == "******" +def test_adding_a_pattern_never_exposes_what_a_narrower_one_hid(): + """Masking must be monotonic: more patterns can only hide more. + + The id is decided from the original row, so hiding the whole attribute map + does not stop it from being recognised as the source of the id. + """ + item = { + "entity_id": "checkout,sk-live", + "entity_id_attrs": {"service_name": "checkout", "api_key": "sk-live"}, + } + + exposed = set() + for extra in (None, ["entity_id_attrs"], ["entity_id"], ["service_name"]): + masked = graph._mask_item(item, graph.mask_patterns(True, extra)) + assert masked["entity_id"] == "******", extra + exposed.add(json.dumps(masked, sort_keys=True).count("sk-live")) + + assert exposed == {0} + + def test_masking_off_returns_attributes_untouched(): item = {"entity_id": "checkout", "entity_id_attrs": {"api_key": "sk-live"}} From 9b72174d5e8509cb3d6fd14eb1086fdeafd968f7 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 18:28:16 +0800 Subject: [PATCH 6/9] refactor: drop state and response fields the graph tool does not need `GraphCapability` carried the column set of each view, which nothing read once the queries stopped being capability-shaped; the probe validates the columns and reports availability. `summary` returned entity_count and relationship_count, both sums of the arrays beside them, and a `complete` that was always true. Every response repeated `bounds`, `time_field` and `observation_bucket_seconds`, which are facts about the tool rather than about the result and are already in its description. Tests lost four duplicates: two that drove the same incompatible-schema branch with a different column, two inputs of one decoding rule, a separate assertion that the no-match retry runs, and two spellings of "the window parses to UTC". --- src/greptimedb_mcp_server/graph.py | 60 ++++++++---------------------- tests/integration/test_e2e.py | 3 +- tests/test_graph.py | 54 +++++++-------------------- 3 files changed, 31 insertions(+), 86 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index e1888cc..cc09b40 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -38,12 +38,10 @@ ) ENTITY_REQUIRED_COLUMNS = frozenset({"observed_at", "fresh_until", *ENTITY_COLUMNS}) -# `confidence` is in the group key, not aggregated. The database reports 1.0 -# for a bucket that paired client and server spans and 0.5 for one that only -# saw clients, and it switches request_count, error_count and the durations to -# the matching population at the same time -- a pair is timed by the server -# span, an unmatched client by its own. Summing across both would add two -# different measurements and MAX(confidence) would hide that it happened. +# `confidence` is in the group key, not aggregated: the database reports 1.0 +# for a bucket whose spans paired and 0.5 for one that saw only clients, and it +# switches the counts and durations to that population at the same time. +# Summing across both would add two different measurements. RELATIONSHIP_IDENTITY_COLUMNS = ( "src_type", "src_id", @@ -53,7 +51,6 @@ "provenance", "confidence", ) -# Summed over the buckets in the window. RED_COLUMNS = ( "request_count", "unmatched_count", @@ -97,7 +94,6 @@ MAX_LIMIT = 500 DEFAULT_LIMIT = 100 -OBSERVATION_BUCKET_SECONDS = 60 ERRNO_TABLE_NOT_FOUND = 1146 ERRNO_PERMISSION_DENIED = frozenset({1044, 1045, 1142, 1143, 1227}) @@ -108,8 +104,6 @@ class GraphCapability: """Whether the graph can be read, and why not when it cannot.""" status: str - entity_columns: frozenset[str] = frozenset() - relationship_columns: frozenset[str] = frozenset() detail: str | None = None @property @@ -157,13 +151,7 @@ def params(self) -> list[str]: return [self.start.isoformat(), self.end.isoformat()] def describe(self) -> dict: - return { - "start": self.start.isoformat(), - "end": self.end.isoformat(), - "bounds": "[start, end)", - "time_field": "observed_at", - "observation_bucket_seconds": OBSERVATION_BUCKET_SECONDS, - } + return {"start": self.start.isoformat(), "end": self.end.isoformat()} def _parse_timestamp(value: str, name: str) -> datetime: @@ -225,10 +213,6 @@ def parse( limit=max(1, min(limit, MAX_LIMIT)), ) - @property - def names_an_id(self) -> bool: - return any(name in ID_FILTERS for name in self.filters) - def _decode_json(name: str, value): if name not in JSON_COLUMNS or not isinstance(value, str) or not value: @@ -258,23 +242,17 @@ def mask_patterns(mask_enabled: bool, extra: list[str] | None) -> list[str] | No def _mask_item(item: dict, patterns: list[str] | None) -> dict: """Apply the column-name masking rule to one returned row. - A returned field whose own name matches a pattern is hidden outright, as - the column would be through execute_sql. Attribute maps that survive that - are then masked by the names inside them, since those names are the columns - the values came from. - - `entity_id` is hidden as well when it was assembled from a masked - attribute, because it is those values joined. That does not reach - `relationships`: its view carries no attribute names, so the same value can - still surface there as `src_id` or `dst_id` unless a pattern matches those - column names. + A field whose own name matches is hidden, as the column would be through + execute_sql; surviving attribute maps are then masked by the names inside + them. `entity_id` goes too when a masked attribute built it. That last part + cannot reach `relationships`, whose view carries no attribute names, so the + value can still surface there as `src_id` or `dst_id`. """ if not patterns: return item - # Decided from the original row: reading it back after masking would miss - # the case where the whole map was hidden, and adding a pattern would then - # expose an id that a narrower rule had hidden. + # From the original row: read back after masking, a wholly hidden map no + # longer looks like one, and adding a pattern would expose the id. hide_id = _identity_is_sensitive(item, patterns) masked = { @@ -315,13 +293,12 @@ def _no_match_guidance(request: GraphRequest) -> dict: An unfiltered read of a large graph is the failure this tool exists to avoid, so the suggestion drops the identifier and keeps the type filters. """ - # The window is a required argument, so a next_query without it would not - # run. + # next_query has to carry the window: it is a required argument. window = { "start_time": request.window.start.isoformat(), "end_time": request.window.end.isoformat(), } - if request.names_an_id: + if any(name in ID_FILTERS for name in request.filters): kept = {k: v for k, v in request.filters.items() if k not in ID_FILTERS} return { "reason": ( @@ -389,9 +366,6 @@ def summary(self, cursor, window: TimeWindow) -> dict: "window": window.describe(), "entity_types": entity_types, "relationship_types": relationship_types, - "entity_count": sum(item["count"] for item in entity_types), - "relationship_count": sum(item["count"] for item in relationship_types), - "complete": True, } def _entity_types(self, cursor, window: TimeWindow) -> list[dict]: @@ -548,11 +522,7 @@ def _probe(cursor) -> GraphCapability: except Error as e: return _classify(e) - return GraphCapability( - "available", - entity_columns=columns[ENTITIES_VIEW], - relationship_columns=columns[RELATIONSHIPS_VIEW], - ) + return GraphCapability("available") def _classify(error: Error) -> GraphCapability: diff --git a/tests/integration/test_e2e.py b/tests/integration/test_e2e.py index 92a3e3e..78abf43 100644 --- a/tests/integration/test_e2e.py +++ b/tests/integration/test_e2e.py @@ -391,7 +391,8 @@ async def test_graph_window_is_read_in_utc_whatever_the_session(seed): ) ) - assert payload["relationship_count"] == len(GRAPH_EDGES) + edges = sum(item["count"] for item in payload["relationship_types"]) + assert edges == len(GRAPH_EDGES) async def test_graph_returns_identifiers_verbatim(seed): diff --git a/tests/test_graph.py b/tests/test_graph.py index 3b1634c..b4ca7cc 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -21,6 +21,10 @@ FULL_ENTITY_COLUMNS = frozenset(graph.ENTITY_REQUIRED_COLUMNS) FULL_RELATIONSHIP_COLUMNS = frozenset(graph.RELATIONSHIP_REQUIRED_COLUMNS) +FULL_VIEWS = { + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, + graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, +} class FakeCursor: @@ -56,17 +60,13 @@ def request(view="relationships", limit=graph.DEFAULT_LIMIT, **filters): return GraphRequest.parse(view, START, END, limit, **filters) -def test_window_treats_a_naive_timestamp_as_utc(): - window = TimeWindow.parse("2026-09-05T07:00:00", "2026-09-05T08:00:00") +def test_window_normalizes_to_utc(): + """A naive timestamp is read as UTC; an offset one is converted.""" + naive = TimeWindow.parse("2026-09-05T07:00:00", "2026-09-05T08:00:00") + offset = TimeWindow.parse("2026-09-05T09:00:00+02:00", END) - assert window.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) - assert window.describe()["bounds"] == "[start, end)" - - -def test_window_normalizes_an_offset_to_utc(): - window = TimeWindow.parse("2026-09-05T09:00:00+02:00", END) - - assert window.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) + assert naive.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) + assert offset.start == datetime(2026, 9, 5, 7, tzinfo=timezone.utc) def test_window_rejects_an_empty_or_inverted_range(): @@ -118,6 +118,8 @@ def test_no_match_guidance_drops_the_id_and_keeps_the_type(): assert next_query["rel_type"] == "calls" assert "src_id" not in next_query assert "canonical graph entity ID" in guidance["reason"] + # the window is required, so a retry without it would not run + GraphRequest.parse(**next_query) def test_no_match_guidance_without_an_id_points_at_the_summary(): @@ -140,8 +142,8 @@ def test_probe_classifies_failures(errno, expected): def test_probe_reports_an_incompatible_schema(): """A view that exists but lacks a column this module reads is not usable.""" columns = { + **FULL_VIEWS, graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS - {"entity_id_attrs"}, - graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, } capability = GraphView().negotiate(FakeCursor(columns=columns)) @@ -152,11 +154,7 @@ def test_probe_reports_an_incompatible_schema(): def test_probe_rejects_a_view_it_cannot_read(): """DESC answers from the catalog, so it does not prove SELECT is allowed.""" - columns = { - graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, - graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, - } - cursor = FakeCursor(columns=columns, errno=1142, fail_on="SELECT COUNT(*)") + cursor = FakeCursor(columns=FULL_VIEWS, errno=1142, fail_on="SELECT COUNT(*)") capability = GraphView().negotiate(cursor) @@ -197,19 +195,6 @@ def test_summary_reports_endpoint_pairs_not_two_sets(): assert runs_on["count"] == 5 -def test_a_view_missing_a_red_column_is_incompatible(): - """The query reads every RED column, so a view without one cannot serve it.""" - columns = { - graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, - graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS - {"unmatched_count"}, - } - - capability = GraphView().negotiate(FakeCursor(columns=columns)) - - assert capability.status == "incompatible_schema" - assert "unmatched_count" in capability.detail - - def test_identifier_shaped_strings_survive_the_row_decode(): """Decoding every string would make entity_id "123" a number.""" columns = ["entity_id", "entity_id_attrs", "source_tables"] @@ -220,9 +205,6 @@ def test_identifier_shaped_strings_survive_the_row_decode(): assert decoded["entity_id"] == "123" assert decoded["entity_id_attrs"] == {"host": "123"} assert decoded["source_tables"] == ["public.t"] - - -def test_identifiers_that_look_like_json_literals_survive(): assert graph._row_dict(["src_id"], ("null",))["src_id"] == "null" assert graph._row_dict(["dst_id"], ("true",))["dst_id"] == "true" @@ -234,14 +216,6 @@ def test_window_params_carry_their_offset(): assert all(p.endswith("+00:00") for p in params) -def test_no_match_guidance_next_query_is_runnable(): - """start_time and end_time are required, so a retry without them fails.""" - guidance = _no_match_guidance(request(rel_type="calls", src_id="unknown")) - - assert set(guidance["next_query"]) >= {"view", "start_time", "end_time"} - GraphRequest.parse(**guidance["next_query"]) - - def test_confidence_is_grouped_not_aggregated(): """A paired bucket and a client-only bucket measure different populations.""" cursor = FakeCursor(rows=[]) From 06ad5ab73b413c97cb1bdcc4ba87852a82086b78 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 18:46:32 +0800 Subject: [PATCH 7/9] fix: bound the startup probe and fail on an inconclusive one at call time A probe that could not run was reported as a result. If startup left the capability undecided and the retry inside a call then failed to read, the tool returned status=unavailable with reason=error and the audit recorded a success, while the identical connection failure one step later, in the query itself, raised. It raises now; only a conclusive answer -- the view being absent, rejected, or the wrong shape -- is still an answer. The startup probe also had no time bound. It runs before the server can serve anything, over a connection whose driver defaults leave both timeouts unset, so a database that accepted the connection and then went quiet would hold initialisation open and take every other tool with it. Wrapping the thread would not have helped, since the blocking read would continue behind it; the probe opens its own connection with connection_timeout and read_timeout instead. Measured against a socket that accepts and never answers: startup completes in 5.0s and keeps the tool, the timeout being inconclusive. Both are driven from the tool entry point in the tests, not from _classify. --- src/greptimedb_mcp_server/graph.py | 6 +++ src/greptimedb_mcp_server/server.py | 14 +++++- tests/test_graph.py | 72 ++++++++++++++++++++++++++++- 3 files changed, 90 insertions(+), 2 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index cc09b40..6e95915 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -95,6 +95,12 @@ DEFAULT_LIMIT = 100 +# The startup probe opens its own bounded connection: it runs before the +# server can serve anything, and a database that accepts the connection but +# never answers would otherwise hold initialisation open indefinitely. A +# timeout surfaces as an error, which is inconclusive, so the tool stays. +PROBE_TIMEOUT_SECONDS = 5 + ERRNO_TABLE_NOT_FOUND = 1146 ERRNO_PERMISSION_DENIED = frozenset({1044, 1045, 1142, 1143, 1227}) diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index e53e41d..409b663 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -698,8 +698,13 @@ def _withdraw_graph_tool_if_unusable(state: AppState) -> None: place: a database that was briefly unreachable at startup is not evidence about the graph. """ + probe_config = { + **state.db_config, + "connection_timeout": graph.PROBE_TIMEOUT_SECONDS, + "read_timeout": graph.PROBE_TIMEOUT_SECONDS, + } try: - with state.get_connection() as conn: + with connect(**probe_config) as conn: with conn.cursor() as cursor: capability = state.semantic_graph.negotiate(cursor) except Error as e: @@ -868,6 +873,13 @@ def _sync_query(): with state.get_connection() as conn: with conn.cursor() as cursor: capability = state.semantic_graph.negotiate(cursor) + if capability.status == "error": + # A probe that could not run is a failure, the same as a + # failed query; only a conclusive answer is a result. + raise ToolError( + f"Could not determine whether the semantic graph is " + f"readable: {capability.detail}" + ) if not capability.available: return { "view": request.view, diff --git a/tests/test_graph.py b/tests/test_graph.py index b4ca7cc..0645c2b 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -6,7 +6,9 @@ import pytest from mysql.connector import Error -from greptimedb_mcp_server import graph +from mcp.server.mcpserver.exceptions import ToolError + +from greptimedb_mcp_server import graph, server from greptimedb_mcp_server.graph import ( GraphCapability, GraphRequest, @@ -308,3 +310,71 @@ def test_truncated_result_says_how_to_narrow(): assert "rel_type" not in guidance["narrow_by"] assert "src_id" in guidance["narrow_by"] + + +@pytest.fixture +def app_state(): + """Application state backed by the mocked MySQL connection.""" + server._state = server.AppState( + db_config={ + "host": "localhost", + "port": 4002, + "user": "", + "password": "", + "database": "testdb", + "time_zone": "", + }, + pool_config={"pool_name": "greptimedb_pool", "pool_size": 5}, + templates={}, + http_base_url="http://localhost:4000", + ) + yield server._state + server._state = None + + +@pytest.mark.asyncio +async def test_a_probe_that_could_not_run_raises(app_state): + """The same connection failure must not read as a result on one path and + a failure on another.""" + app_state.semantic_graph = GraphView( + capability=GraphCapability("error", detail="2013: Lost connection") + ) + + with pytest.raises(ToolError) as excinfo: + await server.query_semantic_graph( + view="summary", start_time=START, end_time=END + ) + + assert "2013: Lost connection" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_a_conclusive_probe_still_answers(app_state): + """An absent graph is an answer, not a failure.""" + app_state.semantic_graph = GraphView( + capability=GraphCapability("unavailable", detail="Table not found") + ) + + payload = json.loads( + await server.query_semantic_graph( + view="summary", start_time=START, end_time=END + ) + ) + + assert payload["status"] == "unavailable" + assert payload["reason"] == "unavailable" + + +def test_the_startup_probe_is_time_bounded(app_state, monkeypatch): + """It runs before the server can serve, so it cannot wait indefinitely.""" + captured = {} + + def fake_connect(**kwargs): + captured.update(kwargs) + raise Error("refused") + + monkeypatch.setattr(server, "connect", fake_connect) + server._withdraw_graph_tool_if_unusable(app_state) + + assert captured["connection_timeout"] == graph.PROBE_TIMEOUT_SECONDS + assert captured["read_timeout"] == graph.PROBE_TIMEOUT_SECONDS From f2a1d256fd77fb73c4f20011c976e3e6672eb852 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 18:53:10 +0800 Subject: [PATCH 8/9] fix: give the startup probe 10s rather than 5 Five seconds is tight for a cold graph derivation on a busy instance, and the cost of waiting longer is only paid when the database has stopped answering. --- src/greptimedb_mcp_server/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index 6e95915..0ecbdfa 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -99,7 +99,7 @@ # server can serve anything, and a database that accepts the connection but # never answers would otherwise hold initialisation open indefinitely. A # timeout surfaces as an error, which is inconclusive, so the tool stays. -PROBE_TIMEOUT_SECONDS = 5 +PROBE_TIMEOUT_SECONDS = 10 ERRNO_TABLE_NOT_FOUND = 1146 ERRNO_PERMISSION_DENIED = frozenset({1044, 1045, 1142, 1143, 1227}) From 3cc9fd834fe158fab6fc0c4ca9ed4a329db1eed8 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 9 Sep 2026 19:01:10 +0800 Subject: [PATCH 9/9] fix: give every graph outcome the shape of its view, and scope the teardown `summary` carried no `status` when it succeeded but was given one when the graph could not be read, and that unavailable response dropped `window`, `applied_filters` and `limit` from the item views. Each view now answers with its own shape whether the read happened or not: `view`, `status` and `window` always, plus the empty type lists for summary or the empty item envelope for the others. The integration teardown deleted seeded edges by src_id and dst_id alone. `semantic_relationships_declared` is shared and this suite does not create it, so the delete now matches the endpoint types, the relationship type and the provenance as well. --- src/greptimedb_mcp_server/graph.py | 36 +++++++++++++++++++++++++---- src/greptimedb_mcp_server/server.py | 8 +------ tests/integration/conftest.py | 7 ++++-- tests/test_graph.py | 22 ++++++++++++------ 4 files changed, 53 insertions(+), 20 deletions(-) diff --git a/src/greptimedb_mcp_server/graph.py b/src/greptimedb_mcp_server/graph.py index 0ecbdfa..15dd901 100644 --- a/src/greptimedb_mcp_server/graph.py +++ b/src/greptimedb_mcp_server/graph.py @@ -322,6 +322,35 @@ def _no_match_guidance(request: GraphRequest) -> dict: } +def unavailable_result(request: GraphRequest, capability: GraphCapability) -> dict: + """What a view returns when the graph cannot be read. + + Shaped like that view's successful answer with no data in it, so a caller + parses one shape per view rather than two. + """ + result = { + "view": request.view, + "status": "unavailable", + "reason": capability.status, + "error": capability.detail, + "window": request.window.describe(), + } + if request.view == "summary": + result["entity_types"] = [] + result["relationship_types"] = [] + return result + result.update( + { + "applied_filters": dict(request.filters), + "limit": request.limit, + "items": [], + "item_count": 0, + "complete": True, + } + ) + return result + + def _truncation_guidance(request: GraphRequest) -> dict: """Say how to reach the rows the limit cut off. @@ -365,13 +394,12 @@ def summary(self, cursor, window: TimeWindow) -> dict: interface fact, and learning it by paging through edges both costs a round trip and invites reading the result as a service call graph. """ - entity_types = self._entity_types(cursor, window) - relationship_types = self._relationship_types(cursor, window) return { "view": "summary", + "status": "ok", "window": window.describe(), - "entity_types": entity_types, - "relationship_types": relationship_types, + "entity_types": self._entity_types(cursor, window), + "relationship_types": self._relationship_types(cursor, window), } def _entity_types(self, cursor, window: TimeWindow) -> list[dict]: diff --git a/src/greptimedb_mcp_server/server.py b/src/greptimedb_mcp_server/server.py index 409b663..b5d4f4c 100644 --- a/src/greptimedb_mcp_server/server.py +++ b/src/greptimedb_mcp_server/server.py @@ -881,13 +881,7 @@ def _sync_query(): f"readable: {capability.detail}" ) if not capability.available: - return { - "view": request.view, - "status": "unavailable", - "reason": capability.status, - "error": capability.detail, - "items": [], - } + return graph.unavailable_result(request, capability) if request.view == "summary": return state.semantic_graph.summary(cursor, request.window) if request.view == "entities": diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1cbd7c8..a6810b2 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -150,9 +150,12 @@ def seed(db): # Delete only the seeded edges: the table is shared and this suite did # not create it. for edge in GRAPH_EDGES: + src, sid, dst, did, rel = edge[:5] cursor.execute( - f"DELETE FROM {DECLARED_EDGES_TABLE} WHERE src_id = %s AND dst_id = %s", - (edge[1], edge[3]), + f"DELETE FROM {DECLARED_EDGES_TABLE} WHERE src_type = %s AND " + "src_id = %s AND dst_type = %s AND dst_id = %s AND rel_type = %s " + "AND provenance = 'declared'", + (src, sid, dst, did, rel), ) db.commit() diff --git a/tests/test_graph.py b/tests/test_graph.py index 0645c2b..67c3528 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -348,21 +348,29 @@ async def test_a_probe_that_could_not_run_raises(app_state): assert "2013: Lost connection" in str(excinfo.value) +@pytest.mark.parametrize( + "view,expected", + [ + ("summary", {"entity_types", "relationship_types"}), + ("relationships", {"applied_filters", "limit", "items", "item_count"}), + ], +) @pytest.mark.asyncio -async def test_a_conclusive_probe_still_answers(app_state): - """An absent graph is an answer, not a failure.""" +async def test_a_conclusive_probe_answers_in_the_views_own_shape( + app_state, view, expected +): + """An absent graph is an answer, and it looks like that view's answer.""" app_state.semantic_graph = GraphView( - capability=GraphCapability("unavailable", detail="Table not found") + capability=GraphCapability("permission_denied", detail="denied") ) payload = json.loads( - await server.query_semantic_graph( - view="summary", start_time=START, end_time=END - ) + await server.query_semantic_graph(view=view, start_time=START, end_time=END) ) assert payload["status"] == "unavailable" - assert payload["reason"] == "unavailable" + assert payload["reason"] == "permission_denied" + assert {"view", "status", "window"} | expected <= set(payload) def test_the_startup_probe_is_time_bounded(app_state, monkeypatch):