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..15dd901 --- /dev/null +++ b/src/greptimedb_mcp_server/graph.py @@ -0,0 +1,568 @@ +"""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 + +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" + +# 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}) + +# `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", + "dst_type", + "dst_id", + "rel_type", + "provenance", + "confidence", +) +RED_COLUMNS = ( + "request_count", + "unmatched_count", + "error_count", + "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") + +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 + + +# 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 = 10 + +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 + 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", + ) + + +@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]: + """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 {"start": self.start.isoformat(), "end": self.end.isoformat()} + + +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)), + ) + + +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) + except json.JSONDecodeError: + return value + + +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_item(item: dict, patterns: list[str] | None) -> dict: + """Apply the column-name masking rule to one returned row. + + 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 + + # 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 = { + 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): + continue + masked[column] = { + name: (MASK_PLACEHOLDER if is_sensitive_column(name, patterns) else attr) + for name, attr in value.items() + } + 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()) + + +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. + """ + # 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 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": ( + "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, **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", **window}, + } + + +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. + + 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}, so rows " + "sorting later are missing entirely rather than merely cut short. " + "Narrow with the filters below." + ), + "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.""" + 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. + """ + return { + "view": "summary", + "status": "ok", + "window": window.describe(), + "entity_types": self._entity_types(cursor, window), + "relationship_types": self._relationship_types(cursor, window), + } + + 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, + ) + # 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, "endpoints": [], "count": 0} + ) + entry["endpoints"].append( + {"source": src_type, "destination": dst_type, "count": int(count)} + ) + 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]) + # 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(fresh_until) AS fresh_until" + 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. + + 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 + 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 = ", ".join(RELATIONSHIP_IDENTITY_COLUMNS) + cursor.execute( + f"SELECT {identity}, " + + ", ".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)}" + 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 = [ + _mask_item(_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 + + +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") + + +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..b5d4f4c 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 ( @@ -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 @@ -70,6 +71,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.""" @@ -322,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 ""} @@ -333,6 +338,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 +687,217 @@ 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. + """ + probe_config = { + **state.db_config, + "connection_timeout": graph.PROBE_TIMEOUT_SECONDS, + "read_timeout": graph.PROBE_TIMEOUT_SECONDS, + } + try: + with connect(**probe_config) 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[ + 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, + 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.") + ], + 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, + 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, + provenance: Annotated[ + 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, + 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. + + 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, 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. + + entity_id_attrs names the attributes an id was assembled from and + 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( + 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 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 graph.unavailable_result(request, capability) + 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: + # 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}") + raise ToolError(f"Error querying the semantic graph: {str(e)}") from 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..a6810b2 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, 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") 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,44 @@ 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: + src, sid, dst, did, rel = edge[:5] + cursor.execute( + 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() +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 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', %s, '', %s, %s, %s)", + (src, sid, dst, did, rel, scope, confidence, 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..78abf43 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 @@ -12,6 +13,7 @@ from .conftest import ( CREDENTIALS_TABLE, MASK_PLACEHOLDER, + GRAPH_EDGES, METRICS_TABLE, SECRET_API_KEY, SECRET_PASSWORD, @@ -330,3 +332,157 @@ 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"]["count"] == 2 + 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} + ) + ) + + edges = sum(item["count"] for item in payload["relationship_types"]) + assert edges == 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: + 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" + + # 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 new file mode 100644 index 0000000..67c3528 --- /dev/null +++ b/tests/test_graph.py @@ -0,0 +1,388 @@ +"""Tests for the semantic graph window, ordering, and capability contracts.""" + +import json +from datetime import datetime, timezone + +import pytest +from mysql.connector import Error + +from mcp.server.mcpserver.exceptions import ToolError + +from greptimedb_mcp_server import graph, server +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) +FULL_VIEWS = { + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS, + graph.RELATIONSHIPS_VIEW: FULL_RELATIONSHIP_COLUMNS, +} + + +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_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 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(): + 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")) + + 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"] + # 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(): + 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 = { + **FULL_VIEWS, + graph.ENTITIES_VIEW: FULL_ENTITY_COLUMNS - {"entity_id_attrs"}, + } + + 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.""" + cursor = FakeCursor(columns=FULL_VIEWS, 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_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), + ("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["endpoints"] == [ + {"source": "service", "destination": "k8s.pod", "count": 3}, + {"source": "k8s.pod", "destination": "k8s.node", "count": 2}, + ] + assert runs_on["count"] == 5 + + +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"] + 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_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_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 reaches inside attribute maps too.""" + 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_item(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_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"}} + + 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_item(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"] + + +@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.parametrize( + "view,expected", + [ + ("summary", {"entity_types", "relationship_types"}), + ("relationships", {"applied_filters", "limit", "items", "item_count"}), + ], +) +@pytest.mark.asyncio +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("permission_denied", detail="denied") + ) + + payload = json.loads( + await server.query_semantic_graph(view=view, start_time=START, end_time=END) + ) + + assert payload["status"] == "unavailable" + assert payload["reason"] == "permission_denied" + assert {"view", "status", "window"} | expected <= set(payload) + + +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