diff --git a/robosystems_client/clients/__init__.py b/robosystems_client/clients/__init__.py index 5fc9f907..fc038a20 100644 --- a/robosystems_client/clients/__init__.py +++ b/robosystems_client/clients/__init__.py @@ -21,6 +21,7 @@ OperatorQueryRequest, OperatorOptions, QueuedOperatorError, + OperatorRunError, ) from .operation_client import ( OperationClient, @@ -144,6 +145,7 @@ "OperatorQueryRequest", "OperatorOptions", "QueuedOperatorError", + "OperatorRunError", # Operation Client "OperationClient", "OperationStatus", diff --git a/robosystems_client/clients/auth_integration.py b/robosystems_client/clients/auth_integration.py index 1fffed7b..e1b42cc3 100644 --- a/robosystems_client/clients/auth_integration.py +++ b/robosystems_client/clients/auth_integration.py @@ -82,9 +82,14 @@ def authenticated_client(self) -> AuthenticatedClient: def execute_cypher_query( self, graph_id: str, query: str, parameters: Dict[str, Any] = None ): - """Execute Cypher query using authenticated SDK client""" + """Execute Cypher query using authenticated SDK client. + + A query the API queues (202) is followed over its SSE stream to the + result, as :meth:`QueryClient.execute_query` does. + """ from ..api.query.execute_cypher import sync_detailed from ..models.cypher_statement_request import CypherStatementRequest + from .query_client import QueryOptions, QueryResult, _interpret_response request = CypherStatementRequest(query=query, parameters=parameters or {}) @@ -95,15 +100,17 @@ def execute_cypher_query( body=request, ) - if response.parsed: - return { - "data": getattr(response.parsed, "data", []), - "columns": getattr(response.parsed, "columns", []), - "row_count": getattr(response.parsed, "row_count", 0), - "execution_time_ms": getattr(response.parsed, "execution_time_ms", 0), - } - else: - raise Exception(f"Query failed: {response.status_code}") + outcome = _interpret_response(response, graph_id) + if not isinstance(outcome, QueryResult): + outcome = self.query._wait_for_query_completion( + outcome.operation_id, QueryOptions() + ) + return { + "data": outcome.data, + "columns": outcome.columns, + "row_count": outcome.row_count, + "execution_time_ms": outcome.execution_time_ms, + } class CookieAuthClients(RoboSystemsClients): diff --git a/robosystems_client/clients/document_client.py b/robosystems_client/clients/document_client.py index 479d642a..ced72395 100644 --- a/robosystems_client/clients/document_client.py +++ b/robosystems_client/clients/document_client.py @@ -30,6 +30,7 @@ from ..models.search_request import SearchRequest from ..models.search_response import SearchResponse from ..types import UNSET +from .token_utils import current_token class DocumentClient: @@ -43,11 +44,12 @@ def __init__(self, config: Dict[str, Any]): self.timeout = config.get("timeout", 60) def _get_client(self) -> AuthenticatedClient: - if not self.token: + token = current_token(self.config, self.token) + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") return retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) diff --git a/robosystems_client/clients/facade.py b/robosystems_client/clients/facade.py index 7d9713d2..598b1254 100644 --- a/robosystems_client/clients/facade.py +++ b/robosystems_client/clients/facade.py @@ -161,6 +161,7 @@ def __init__(self, config: RoboSystemsClientConfig = None): "max_retries": config.max_retries, "retry_delay": config.retry_delay, "timeout": config.timeout, + "token_provider": config.token_provider, } # Initialize async clients diff --git a/robosystems_client/clients/file_client.py b/robosystems_client/clients/file_client.py index 8dca59fc..5c2f3e37 100644 --- a/robosystems_client/clients/file_client.py +++ b/robosystems_client/clients/file_client.py @@ -29,10 +29,17 @@ from ..models.file_upload_request import FileUploadRequest from ..models.ingest_file_op import IngestFileOp from ..models.delete_file_op import DeleteFileOp +from ..types import Unset +from .token_utils import current_token logger = logging.getLogger(__name__) +def _unset_to_none(value: Any) -> Any: + """A generated-model field as a plain value: ``UNSET`` becomes ``None``.""" + return None if isinstance(value, Unset) else value + + @dataclass class FileUploadOptions: """Options for file upload operations""" @@ -139,12 +146,14 @@ def upload( from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) @@ -295,12 +304,14 @@ def list( try: from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) @@ -330,11 +341,11 @@ def list( file_name=f.file_name, file_format=f.file_format, size_bytes=f.size_bytes or 0, - row_count=f.row_count, + row_count=_unset_to_none(f.row_count), upload_status=f.upload_status, - table_name=getattr(f, "table_name", ""), - created_at=f.created_at, - uploaded_at=f.uploaded_at, + table_name=_unset_to_none(getattr(f, "table_name", None)) or "", + created_at=_unset_to_none(f.created_at), + uploaded_at=_unset_to_none(f.uploaded_at), ) for f in files ] @@ -357,12 +368,14 @@ def get(self, graph_id: str, file_id: str) -> Optional[FileInfo]: try: from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) @@ -386,12 +399,12 @@ def get(self, graph_id: str, file_id: str) -> Optional[FileInfo]: file_name=file_data.file_name, file_format=file_data.file_format, size_bytes=file_data.size_bytes or 0, - row_count=file_data.row_count, + row_count=_unset_to_none(file_data.row_count), upload_status=file_data.upload_status, - table_name=file_data.table_name or "", - created_at=file_data.created_at, - uploaded_at=file_data.uploaded_at, - layers=getattr(file_data, "layers", None), + table_name=_unset_to_none(file_data.table_name) or "", + created_at=_unset_to_none(file_data.created_at), + uploaded_at=_unset_to_none(file_data.uploaded_at), + layers=_unset_to_none(getattr(file_data, "layers", None)), ) except Exception as e: @@ -413,12 +426,14 @@ def delete(self, graph_id: str, file_id: str, cascade: bool = False) -> bool: try: from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) diff --git a/robosystems_client/clients/graph_client.py b/robosystems_client/clients/graph_client.py index 6c2bc640..4d85dd85 100644 --- a/robosystems_client/clients/graph_client.py +++ b/robosystems_client/clients/graph_client.py @@ -4,10 +4,11 @@ automatic operation monitoring. Supports both SSE (Server-Sent Events) for real-time updates and polling fallback. -Graph lifecycle operations (create-subgraph, delete-subgraph, create-backup, -change-tier, materialize) all go through the operations surface at +The lifecycle operations this client wraps (materialize, change-tier) go +through the operations surface at ``POST /v1/graphs/{graph_id}/operations/{op_name}`` and return an -``OperationEnvelope``. +``OperationEnvelope``. Subgraph and backup operations live on the same +surface but have no method here — call the generated ``api`` functions. Backups are a download capability: there is no customer-facing restore. Graphs with an upstream rebuild from it — entity graphs re-materialize, @@ -23,7 +24,9 @@ import httpx +from ..types import Unset from .operation_client import OperationClient, OperationProgress, MonitorOptions +from .token_utils import apply_auth_header, current_token logger = logging.getLogger(__name__) @@ -96,8 +99,8 @@ class MaterializationResult: class GraphClient: """Client for graph management and lifecycle operations. - Covers graph creation, info retrieval, and all graph lifecycle - operations (materialize, subgraphs, backups, tier changes). + Covers graph creation and deletion, info retrieval, materialization + and tier changes. """ def __init__(self, config: Dict[str, Any]): @@ -118,12 +121,14 @@ def _get_authenticated_client(self): """Build an AuthenticatedClient for API calls.""" from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise ValueError("No API key provided. Set X-API-Key in headers.") return retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) @@ -250,18 +255,18 @@ def get_graph_info(self, graph_id: str) -> GraphInfo: client = self._get_authenticated_client() response = get_graphs(client=client) - if not response.parsed: + if response.status_code != 200 or not response.parsed: raise RuntimeError(f"Failed to get graphs: {response.status_code}") + # Every generated model carries `additional_properties` (the fields the + # schema does not name), so the declared `graphs` field is read first. data = response.parsed - graphs = None - if isinstance(data, dict): graphs = data.get("graphs", []) - elif hasattr(data, "additional_properties"): - graphs = data.additional_properties.get("graphs", []) elif hasattr(data, "graphs"): graphs = data.graphs + elif hasattr(data, "additional_properties"): + graphs = data.additional_properties.get("graphs", []) else: raise RuntimeError("Unexpected response format from get_graphs") @@ -290,16 +295,19 @@ def get_graph_info(self, graph_id: str) -> GraphInfo: status=graph_data.get("status"), ) else: + + def field(name: str) -> Any: + value = getattr(graph_data, name, None) + return None if isinstance(value, Unset) else value + return GraphInfo( - graph_id=getattr(graph_data, "graph_id", None) - or getattr(graph_data, "id", graph_id), - graph_name=getattr(graph_data, "graph_name", None) - or getattr(graph_data, "name", ""), - description=getattr(graph_data, "description", None), - schema_extensions=getattr(graph_data, "schema_extensions", None), - tags=getattr(graph_data, "tags", None), - created_at=getattr(graph_data, "created_at", None), - status=getattr(graph_data, "status", None), + graph_id=field("graph_id") or field("id") or graph_id, + graph_name=field("graph_name") or field("name") or "", + description=field("description"), + schema_extensions=field("schema_extensions"), + tags=field("tags"), + created_at=field("created_at"), + status=field("status"), ) # --------------------------------------------------------------------------- @@ -535,7 +543,10 @@ def _wait_with_sse( ) -> str: """Wait for operation completion using SSE stream.""" stream_url = f"{self.base_url}/v1/operations/{operation_id}/stream" - headers = {"X-API-Key": self.token, "Accept": "text/event-stream"} + headers = {"Accept": "text/event-stream"} + token = current_token(self.config, self.token) + if token: + apply_auth_header(headers, token) with httpx.Client(timeout=httpx.Timeout(timeout + 5.0)) as http_client: with http_client.stream("GET", stream_url, headers=headers) as response: diff --git a/robosystems_client/clients/ledger_client.py b/robosystems_client/clients/ledger_client.py index 0e0238d4..525cf1aa 100644 --- a/robosystems_client/clients/ledger_client.py +++ b/robosystems_client/clients/ledger_client.py @@ -547,6 +547,7 @@ from ..models.initialize_ledger_request import InitializeLedgerRequest from ..models.create_publish_list_request import CreatePublishListRequest from ..models.create_report_request import CreateReportRequest +from ..models.period_spec import PeriodSpec from ..models.delete_publish_list_operation import DeletePublishListOperation from ..models.delete_report_operation import DeleteReportOperation from ..models.file_report_request import FileReportRequest @@ -617,6 +618,13 @@ } +def _as_date(value: str | datetime.date) -> datetime.date: + """An ISO ``YYYY-MM-DD`` string or a date, as the date the request models need.""" + if isinstance(value, datetime.date): + return value + return datetime.date.fromisoformat(value) + + def _parse_filename(content_disposition: str) -> str | None: """Extract the ``filename`` value from a Content-Disposition header. @@ -730,12 +738,6 @@ def _query( def _is_envelope(self, value: Any) -> bool: return all(hasattr(value, f) for f in self._ENVELOPE_FIELDS) - def _unwrap(self, label: str, envelope: Any) -> Any: - """Unwrap an operation envelope and return `result` (None on failure).""" - if not self._is_envelope(envelope): - raise RuntimeError(f"{label} failed: {envelope!r}") - return envelope.result - def _typed_result( self, label: str, @@ -784,8 +786,8 @@ def _call_op(self, label: str, response: Any) -> Any: Returns the parsed envelope unchanged. Typed-envelope ops surface ``envelope.result`` as the SDK's typed attrs class (e.g. - ``ReportResponse``); untyped ops surface it as a plain dict via - ``OperationEnvelopeResultType0``. Facade methods are responsible + ``ReportResponse``); untyped ops (the plain ``OperationEnvelope``) + surface it as the raw JSON value, usually a dict. Facade methods are responsible for asserting / casting the result to the type they advertise. """ if response.status_code not in (HTTPStatus.OK, HTTPStatus.ACCEPTED): @@ -1396,11 +1398,14 @@ def list_information_blocks( category: str | None = None, limit: int | None = None, offset: int | None = None, + scenario_id: str | None = None, ) -> list[ListInformationBlocksInformationBlocks]: """List Information Block envelopes, optionally filtered. Replaces the old ``list_schedules`` method — use ``block_type='schedule'`` for the same set of blocks. + ``scenario_id`` (a forecast block's structure id) threads into each + envelope's FactSet binding. """ data = self._query( graph_id, @@ -1410,6 +1415,7 @@ def list_information_blocks( "category": category, "limit": limit, "offset": offset, + "scenarioId": scenario_id, }, ) return ListInformationBlocks.model_validate(data).information_blocks @@ -1552,8 +1558,8 @@ def dispose_schedule( graph_id: str, structure_id: str, disposal_date: str, - memo: str, - reason: str, + memo: str | None = None, + reason: str | None = None, sale_proceeds: int | None = None, proceeds_element_id: str | None = None, gain_loss_element_id: str | None = None, @@ -1564,13 +1570,14 @@ def dispose_schedule( Routes through ``create-event-block`` with ``event_type='asset_disposed'``. ``source`` defaults to ``"manual"`` - (user-initiated disposal); sync adapters override. + (user-initiated disposal); sync adapters override. ``memo`` and + ``reason`` are optional; the server fills in its own defaults. """ - metadata: dict[str, Any] = { - "schedule_id": structure_id, - "memo": memo, - "reason": reason, - } + metadata: dict[str, Any] = {"schedule_id": structure_id} + if memo is not None: + metadata["memo"] = memo + if reason is not None: + metadata["reason"] = reason if sale_proceeds is not None: metadata["proceeds"] = sale_proceeds if proceeds_element_id is not None: @@ -2489,22 +2496,32 @@ def create_report( taxonomy_id: str = "rs-gaap", period_type: str = "quarterly", comparative: bool = True, + periods: list[PeriodSpec | dict[str, Any]] | None = None, ) -> ReportResponse: """Generate report facts from the ledger and publish a Report definition. Synchronous — returns the published report header. ``taxonomy_id`` defaults to ``"rs-gaap"``, the canonical reporting vocabulary (matches the backend's ``CreateReportRequest`` default). + + ``periods`` sets the columns explicitly — ``PeriodSpec`` or + ``{"start", "end", "label"}`` dicts — for multi-period layouts + (YTD-by-quarter, multi-year). When given, the server ignores + ``period_start`` / ``period_end`` / ``comparative``. """ body = CreateReportRequest( name=name, mapping_id=mapping_id, - period_start=period_start, - period_end=period_end, + period_start=_as_date(period_start), + period_end=_as_date(period_end), taxonomy_id=taxonomy_id, period_type=period_type, comparative=comparative, ) + if periods: + body.periods = [ + p if isinstance(p, PeriodSpec) else PeriodSpec.from_dict(p) for p in periods + ] response = op_create_report(graph_id=graph_id, body=body, client=self._get_client()) envelope = self._call_op("Create report", response) return self._typed_result("Create report", envelope, ReportResponse) @@ -2569,8 +2586,8 @@ def regenerate_report( """ body = RegenerateReportOperation( report_id=report_id, - period_start=period_start if period_start is not None else UNSET, - period_end=period_end if period_end is not None else UNSET, + period_start=_as_date(period_start) if period_start is not None else UNSET, + period_end=_as_date(period_end) if period_end is not None else UNSET, ) response = op_regenerate_report( graph_id=graph_id, body=body, client=self._get_client() diff --git a/robosystems_client/clients/operation_client.py b/robosystems_client/clients/operation_client.py index ef0704b1..d4432b8b 100644 --- a/robosystems_client/clients/operation_client.py +++ b/robosystems_client/clients/operation_client.py @@ -3,12 +3,14 @@ Provides comprehensive operation monitoring with SSE support. """ +import asyncio import logging from dataclasses import dataclass from typing import Dict, Any, Optional, Callable, List, cast from datetime import datetime from enum import Enum +from ..client import Client from .sse_client import SSEClient, AsyncSSEClient, SSEConfig, EventType from .token_utils import resolve_auth_headers @@ -87,6 +89,33 @@ def _parsed_dict(parsed: Any) -> Dict[str, Any] | None: return None +def _progress_percent(data: Dict[str, Any]) -> Optional[float]: + """Progress of an `operation_progress` event; the API emits `progress_percent`.""" + percent = data.get("progress_percent") + return percent if percent is not None else data.get("percentage") + + +def _status_result(operation_id: str, payload: Dict[str, Any]) -> Dict[str, Any]: + return { + "operation_id": operation_id, + "status": payload.get("status", "unknown"), + "progress": payload.get("progress"), + "result": payload.get("result"), + "error": payload.get("error"), + } + + +def _is_cancelled(response: Any, payload: Dict[str, Any] | None) -> bool: + """Whether a `DELETE /v1/operations/{id}` answer reports the cancel landed. + + The API answers 200 with ``{operation_id, status: "cancelled", message}``. + """ + code = getattr(response, "status_code", None) + if isinstance(code, int) and code != 200: + return False + return payload is not None and payload.get("status") == "cancelled" + + class OperationClient: """Client for monitoring operations via SSE""" @@ -138,7 +167,7 @@ def on_operation_started(data): def on_operation_progress(data): progress = OperationProgress( message=data.get("message", "Processing..."), - percentage=data.get("percentage"), + percentage=_progress_percent(data), current_step=data.get("current_step"), total_steps=data.get("total_steps"), ) @@ -170,7 +199,7 @@ def on_operation_error(err): error = Exception(result.error) completed = True - def on_operation_cancelled(): + def on_operation_cancelled(_data=None): nonlocal completed result.status = OperationStatus.CANCELLED result.completed_at = datetime.now() @@ -196,30 +225,19 @@ def on_connection_error(err): sse_client.on("error", on_connection_error) sse_client.on("max_retries_exceeded", on_connection_error) - # Connect and monitor + # Connect and monitor. Registered first so cancel_operation() can close + # the stream; connect() blocks until the stream ends (or never opens). try: - sse_client.connect(operation_id) with self._lock: self.active_operations[operation_id] = sse_client + sse_client.connect(operation_id) - # Wait for completion - import time - - start_time = datetime.now() - while not completed: - if error: - raise error - - # Check timeout - if options.timeout: - elapsed = (datetime.now() - start_time).total_seconds() - if elapsed > options.timeout: - sse_client.close() - raise TimeoutError( - f"Operation {operation_id} timed out after {options.timeout}s" - ) - - time.sleep(options.poll_interval or 0.1) + if not completed: + # The stream ended without a terminal event: no verdict to report, + # and nothing left to wait on. + raise Exception( + f"Operation {operation_id} stream ended before a terminal event" + ) finally: # Clean up with thread safety @@ -257,13 +275,7 @@ def get_operation_status(self, operation_id: str) -> Dict[str, Any]: response = get_operation_status(operation_id=operation_id, client=client) payload = _parsed_dict(response.parsed) if payload is not None: - return { - "operation_id": operation_id, - "status": payload.get("status", "unknown"), - "progress": payload.get("progress"), - "result": payload.get("result"), - "error": payload.get("error"), - } + return _status_result(operation_id, payload) except Exception as e: # Logged rather than silently shaped into a result: swallowing here is # what hid the TypeError above for as long as it lived. @@ -289,7 +301,7 @@ def cancel_operation(self, operation_id: str) -> bool: # through `to_dict()` because attribute access always raises. response = cancel_operation(operation_id=operation_id, client=client) payload = _parsed_dict(response.parsed) - cancelled = bool(payload.get("cancelled")) if payload is not None else False + cancelled = _is_cancelled(response, payload) except Exception as e: logger.warning("Failed to cancel operation %s: %s", operation_id, e) return False @@ -362,7 +374,7 @@ def on_operation_started(data): def on_operation_progress(data): progress = OperationProgress( message=data.get("message", "Processing..."), - percentage=data.get("percentage"), + percentage=_progress_percent(data), current_step=data.get("current_step"), total_steps=data.get("total_steps"), ) @@ -394,7 +406,7 @@ def on_operation_error(err): error = Exception(result.error) completed = True - def on_operation_cancelled(): + def on_operation_cancelled(_data=None): nonlocal completed result.status = OperationStatus.CANCELLED result.completed_at = datetime.now() @@ -420,29 +432,20 @@ def on_connection_error(err): sse_client.on("error", on_connection_error) sse_client.on("max_retries_exceeded", on_connection_error) - # Connect and monitor + # Connect and monitor; connect() returns once the stream has ended. try: - await sse_client.connect(operation_id) self.active_operations[operation_id] = sse_client - - # Wait for completion - import asyncio - - start_time = datetime.now() - while not completed: - if error: - raise error - - # Check timeout - if options.timeout: - elapsed = (datetime.now() - start_time).total_seconds() - if elapsed > options.timeout: - await sse_client.close() - raise TimeoutError( - f"Operation {operation_id} timed out after {options.timeout}s" - ) - - await asyncio.sleep(options.poll_interval or 0.1) + try: + await asyncio.wait_for(sse_client.connect(operation_id), options.timeout) + except asyncio.TimeoutError: + raise TimeoutError( + f"Operation {operation_id} timed out after {options.timeout}s" + ) from None + + if not completed: + raise Exception( + f"Operation {operation_id} stream ended before a terminal event" + ) finally: # Clean up @@ -452,15 +455,43 @@ def on_connection_error(err): return result + def _rest_client(self) -> Client: + """A REST client for one call, carrying the credential current now.""" + return Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + async def get_operation_status(self, operation_id: str) -> Dict[str, Any]: """Get current status of an operation (async API call)""" - # Would use async version of the generated client - pass + from ..api.operations.get_operation_status import asyncio_detailed + + try: + async with self._rest_client() as client: + response = await asyncio_detailed(operation_id=operation_id, client=client) + payload = _parsed_dict(response.parsed) + if payload is not None: + return _status_result(operation_id, payload) + except Exception as e: + logger.warning("Failed to get status for operation %s: %s", operation_id, e) + return {"operation_id": operation_id, "status": "error", "error": str(e)} + + return {"operation_id": operation_id, "status": "unknown"} async def cancel_operation(self, operation_id: str) -> bool: """Cancel an operation (async)""" - # Would use async version of the generated client - pass + from ..api.operations.cancel_operation import asyncio_detailed + + try: + async with self._rest_client() as client: + response = await asyncio_detailed(operation_id=operation_id, client=client) + cancelled = _is_cancelled(response, _parsed_dict(response.parsed)) + except Exception as e: + logger.warning("Failed to cancel operation %s: %s", operation_id, e) + return False + + if operation_id in self.active_operations: + await self.active_operations[operation_id].close() + del self.active_operations[operation_id] + + return cancelled async def close_all(self): """Close all active operation monitors (async)""" diff --git a/robosystems_client/clients/operator_client.py b/robosystems_client/clients/operator_client.py index f7e790b6..441c8f13 100644 --- a/robosystems_client/clients/operator_client.py +++ b/robosystems_client/clients/operator_client.py @@ -3,9 +3,10 @@ Provides intelligent operator execution with automatic strategy selection. """ +import json import time from dataclasses import dataclass -from typing import Dict, Any, Optional, Callable, cast +from typing import Dict, Any, List, Optional, Callable, cast from datetime import datetime from ..api.operator.auto_select_operator import sync_detailed as auto_select_operator @@ -19,7 +20,8 @@ from .retry import retrying_client from ..models.operator_request import OperatorRequest from ..models.operator_message import OperatorMessage -from ..types import UNSET +from ..models.response_mode import ResponseMode +from ..types import UNSET, Unset from .sse_client import SSEClient, SSEConfig, EventType, event_error_message from .token_utils import resolve_auth_headers, resolve_config_token @@ -95,16 +97,6 @@ def _operator_result(data: Dict[str, Any]) -> OperatorResult: ) -def _attr_or_none(data: Any, name: str) -> Any: - """An attrs-model field as a plain value: UNSET → None, models → dicts.""" - value = getattr(data, name, None) - if value is UNSET: - return None - if hasattr(value, "to_dict"): - return value.to_dict() - return value - - @dataclass class QueuedOperatorResponse: """Response when operator execution is queued""" @@ -133,6 +125,81 @@ def _response_detail(response: Any) -> str: return text[:200] or "empty response" +def _status_code(response: Any) -> Optional[int]: + """The HTTP status of a generated-client response, when it carries one.""" + code = getattr(response, "status_code", None) + return int(code) if isinstance(code, int) else None + + +def _response_body(response: Any) -> Optional[Dict[str, Any]]: + """The JSON body of an operator response as a dict. + + The generated parser yields ``None`` for a 202, so its body (the queued + operation and its links) is read off the raw content instead. + """ + parsed = getattr(response, "parsed", None) + if isinstance(parsed, dict): + return parsed + if parsed is not None and hasattr(parsed, "to_dict"): + body = parsed.to_dict() + if isinstance(body, dict): + return body + content = getattr(response, "content", None) + if isinstance(content, (bytes, str)) and content: + try: + body = json.loads(content) + except ValueError: + return None + return body if isinstance(body, dict) else None + return None + + +def _response_mode(mode: Optional[str]) -> ResponseMode | Unset: + """`OperatorOptions.mode` as the endpoint's `?mode=` value, when it is one.""" + try: + return ResponseMode(mode) if mode else UNSET + except ValueError: + return UNSET + + +def _wrap_call_error(error: Exception) -> Exception: + """A transport-level failure of the operator call, labelled for the caller.""" + error_msg = str(error) + if "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower(): + return Exception(f"Authentication failed during operator execution: {error_msg}") + return Exception(f"Operator execution failed: {error_msg}") + + +def _read_writes(source: Any) -> List[Dict[str, Any]]: + """The `writes` a stopped run had already landed, from whichever payload has them.""" + writes = source.get("writes") if isinstance(source, dict) else None + if not isinstance(writes, list): + return [] + return [w for w in writes if isinstance(w, dict)] + + +class OperatorRunError(Exception): + """A run that failed or was cancelled. + + ``writes`` lists what the run had already changed on the graph before it + stopped, so a caller can show a receipt rather than invite a retry that + repeats them. ``status`` is ``"failed"`` or ``"cancelled"``. + """ + + def __init__( + self, + message: str, + status: str, + writes: Optional[List[Dict[str, Any]]] = None, + operation_id: Optional[str] = None, + ): + super().__init__(message) + self.message = message + self.status = status + self.writes: List[Dict[str, Any]] = list(writes or []) + self.operation_id = operation_id + + class QueuedOperatorError(Exception): """Exception thrown when operator execution is queued and maxWait is 0""" @@ -201,118 +268,12 @@ def execute_query( graph_id=graph_id, client=client, body=operator_request, + mode=_response_mode(options.mode), ) - - # Check response type and handle accordingly - if hasattr(response, "parsed") and response.parsed: - response_data = response.parsed - - # Handle both dict and attrs object responses - if isinstance(response_data, dict): - data = response_data - else: - # Response is an attrs object - data = response_data - - # Check if this is an immediate response (sync or SSE execution) - has_content = False - if isinstance(data, dict): - has_content = "content" in data and "operator_used" in data - else: - has_content = hasattr(data, "content") and hasattr(data, "operator_used") - - if has_content: - # Extract data from either dict or attrs object - if isinstance(data, dict): - return OperatorResult( - content=data["content"], - operator_used=data["operator_used"], - mode_used=data["mode_used"], - metadata=data.get("metadata"), - tokens_used=data.get("tokens_used"), - confidence_score=data.get("confidence_score"), - execution_time=data.get("execution_time"), - timestamp=data.get("timestamp", datetime.now().isoformat()), - error_details=data.get("error_details"), - ) - else: - # attrs object - access attributes directly - return OperatorResult( - content=data.content if data.content is not UNSET else "", - operator_used=data.operator_used - if data.operator_used is not UNSET - else "unknown", - mode_used=data.mode_used.value - if hasattr(data.mode_used, "value") - else data.mode_used - if data.mode_used is not UNSET - else "standard", - metadata=data.metadata if data.metadata is not UNSET else None, - tokens_used=data.tokens_used if data.tokens_used is not UNSET else None, - confidence_score=data.confidence_score - if data.confidence_score is not UNSET - else None, - execution_time=data.execution_time - if data.execution_time is not UNSET - else None, - timestamp=data.timestamp - if hasattr(data, "timestamp") and data.timestamp is not UNSET - else datetime.now().isoformat(), - error_details=_attr_or_none(data, "error_details"), - ) - - # Check if this is a queued response (async background task execution) - is_queued = False - queued_response = None - - if isinstance(data, dict): - is_queued = "operation_id" in data - if is_queued: - queued_response = QueuedOperatorResponse( - status=data.get("status", "queued"), - operation_id=data["operation_id"], - message=data.get("message", "Operator execution queued"), - sse_endpoint=data.get("sse_endpoint"), - ) - else: - is_queued = hasattr(data, "operation_id") - if is_queued: - queued_response = QueuedOperatorResponse( - status=data.status if hasattr(data, "status") else "queued", - operation_id=data.operation_id, - message=data.message - if hasattr(data, "message") and data.message is not UNSET - else "Operator execution queued", - sse_endpoint=data.sse_endpoint - if hasattr(data, "sse_endpoint") and data.sse_endpoint is not UNSET - else None, - ) - - if is_queued and queued_response: - # If user doesn't want to wait, raise with queue info - if options.max_wait == 0: - raise QueuedOperatorError(queued_response) - - # Use SSE to monitor the operation - return self._wait_for_operator_completion( - queued_response.operation_id, options - ) - except Exception as e: - if isinstance(e, QueuedOperatorError): - raise - - error_msg = str(e) - # Check for authentication errors - if ( - "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower() - ): - raise Exception(f"Authentication failed during operator execution: {error_msg}") - else: - raise Exception(f"Operator execution failed: {error_msg}") - - # Unexpected response format - raise Exception("Unexpected response format from operator endpoint") + raise _wrap_call_error(e) from e + + return self._settle(response, options) def execute_operator( self, @@ -347,115 +308,48 @@ def execute_operator( operator_type=operator_type, client=client, body=operator_request, + mode=_response_mode(options.mode), ) + except Exception as e: + raise _wrap_call_error(e) from e - # Check response type and handle accordingly - if hasattr(response, "parsed") and response.parsed: - response_data = response.parsed - - # Handle both dict and attrs object responses - if isinstance(response_data, dict): - data = response_data - else: - data = response_data - - # Check if this is an immediate response - has_content = False - if isinstance(data, dict): - has_content = "content" in data and "operator_used" in data - else: - has_content = hasattr(data, "content") and hasattr(data, "operator_used") - - if has_content: - # Extract data from either dict or attrs object - if isinstance(data, dict): - return OperatorResult( - content=data["content"], - operator_used=data["operator_used"], - mode_used=data["mode_used"], - metadata=data.get("metadata"), - tokens_used=data.get("tokens_used"), - confidence_score=data.get("confidence_score"), - execution_time=data.get("execution_time"), - timestamp=data.get("timestamp", datetime.now().isoformat()), - error_details=data.get("error_details"), - ) - else: - # attrs object - return OperatorResult( - content=data.content if data.content is not UNSET else "", - operator_used=data.operator_used - if data.operator_used is not UNSET - else "unknown", - mode_used=data.mode_used.value - if hasattr(data.mode_used, "value") - else data.mode_used - if data.mode_used is not UNSET - else "standard", - metadata=data.metadata if data.metadata is not UNSET else None, - tokens_used=data.tokens_used if data.tokens_used is not UNSET else None, - confidence_score=data.confidence_score - if data.confidence_score is not UNSET - else None, - execution_time=data.execution_time - if data.execution_time is not UNSET - else None, - timestamp=data.timestamp - if hasattr(data, "timestamp") and data.timestamp is not UNSET - else datetime.now().isoformat(), - error_details=_attr_or_none(data, "error_details"), - ) - - # Check if this is a queued response - is_queued = False - queued_response = None - - if isinstance(data, dict): - is_queued = "operation_id" in data - if is_queued: - queued_response = QueuedOperatorResponse( - status=data.get("status", "queued"), - operation_id=data["operation_id"], - message=data.get("message", "Operator execution queued"), - sse_endpoint=data.get("sse_endpoint"), - ) - else: - is_queued = hasattr(data, "operation_id") - if is_queued: - queued_response = QueuedOperatorResponse( - status=data.status if hasattr(data, "status") else "queued", - operation_id=data.operation_id, - message=data.message - if hasattr(data, "message") and data.message is not UNSET - else "Operator execution queued", - sse_endpoint=data.sse_endpoint - if hasattr(data, "sse_endpoint") and data.sse_endpoint is not UNSET - else None, - ) + return self._settle(response, options) - if is_queued and queued_response: - # If user doesn't want to wait, raise with queue info - if options.max_wait == 0: - raise QueuedOperatorError(queued_response) + def _settle(self, response: Any, options: OperatorOptions) -> OperatorResult: + """Resolve an operator endpoint response. - # Use SSE to monitor the operation - return self._wait_for_operator_completion( - queued_response.operation_id, options - ) + A 200 body is the result. A 202 carries the queued operation's links + — the generated parser yields ``None`` for it, so the body is read off + the raw response — and is followed to completion. + """ + code = _status_code(response) + if code is not None and code >= 400: + detail = _response_detail(response) + if code in (401, 403): + raise Exception( + f"Authentication failed during operator execution: {code}: {detail}" + ) + raise Exception(f"Operator execution failed: {code}: {detail}") + + data = _response_body(response) + if data is None: + raise Exception("Unexpected response format from operator endpoint") + + if "content" in data and "operator_used" in data: + return _operator_result(data) + + if data.get("operation_id"): + links = data.get("_links") or {} + queued_response = QueuedOperatorResponse( + status=data.get("status") or "queued", + operation_id=data["operation_id"], + message=data.get("message") or "Operator execution queued", + sse_endpoint=data.get("sse_endpoint") or links.get("stream"), + ) + if options.max_wait == 0: + raise QueuedOperatorError(queued_response) + return self._wait_for_operator_completion(queued_response.operation_id, options) - except Exception as e: - if isinstance(e, QueuedOperatorError): - raise - - error_msg = str(e) - if ( - "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower() - ): - raise Exception(f"Authentication failed during operator execution: {error_msg}") - else: - raise Exception(f"Operator execution failed: {error_msg}") - - # Unexpected response format raise Exception("Unexpected response format from operator endpoint") def _wait_for_operator_completion( @@ -473,9 +367,10 @@ def _wait_for_operator_completion( def on_progress(data): if options.on_progress: - options.on_progress( - data.get("message", "Processing..."), data.get("percentage") - ) + percent = data.get("progress_percent") + if percent is None: + percent = data.get("percentage") + options.on_progress(data.get("message", "Processing..."), percent) def on_operator_started(data): if options.on_progress: @@ -500,12 +395,17 @@ def on_completed(data): def on_error(err): # The run itself failed — a verdict, not a transport problem. nonlocal error, completed - error = Exception(event_error_message(err)) + details = err.get("error_details") if isinstance(err, dict) else None + error = OperatorRunError( + event_error_message(err), "failed", _read_writes(details), operation_id + ) completed = True def on_cancelled(_data=None): nonlocal error, completed - error = Exception("Operator execution cancelled") + error = OperatorRunError( + "Operator execution cancelled", "cancelled", [], operation_id + ) completed = True def on_transport_error(err): @@ -610,11 +510,19 @@ def _poll_for_completion( if state == "completed": return _operator_result(status.get("result") or {}) if state == "failed": - raise Exception( - status.get("error") or status.get("message") or "Operator run failed" + raise OperatorRunError( + status.get("error") or status.get("message") or "Operator run failed", + "failed", + _read_writes(status), + operation_id, ) if state == "cancelled": - raise Exception("Operator execution cancelled") + raise OperatorRunError( + "Operator execution cancelled", + "cancelled", + _read_writes(status.get("result")), + operation_id, + ) if options.on_progress and status.get("message"): options.on_progress(status["message"], None) time.sleep(interval) diff --git a/robosystems_client/clients/query_client.py b/robosystems_client/clients/query_client.py index faa52e63..8a9c061d 100644 --- a/robosystems_client/clients/query_client.py +++ b/robosystems_client/clients/query_client.py @@ -3,6 +3,7 @@ Provides intelligent query execution with automatic strategy selection. """ +import json from dataclasses import dataclass from typing import ( Dict, @@ -18,6 +19,7 @@ from datetime import datetime from ..api.query.execute_cypher import sync_detailed as execute_cypher_query +from ..api.query.execute_cypher import asyncio_detailed as execute_cypher_query_async from ..models.cypher_statement_request import CypherStatementRequest from ..client import Client from .retry import retrying_client @@ -28,7 +30,26 @@ EventType, event_error_message, ) -from .token_utils import resolve_auth_headers, resolve_config_token +from .token_utils import has_auth_header, resolve_auth_headers, resolve_config_token +from ..types import UNSET + + +def _status_code(response: Any) -> Optional[int]: + """The HTTP status of a generated-client response, when it carries one.""" + code = getattr(response, "status_code", None) + return int(code) if isinstance(code, int) else None + + +def _json_body(response: Any) -> Optional[Dict[str, Any]]: + """The raw JSON body of a response as a dict, or ``None``.""" + content = getattr(response, "content", None) + if not isinstance(content, (bytes, str)) or not content: + return None + try: + body = json.loads(content) + except ValueError: + return None + return body if isinstance(body, dict) else None @dataclass @@ -83,6 +104,199 @@ def __init__(self, queue_info: QueuedQueryResponse): self.queue_info = queue_info +class _QueryFailed(Exception): + """A definitive answer from the endpoint that is not a result: raised as is.""" + + +def _cypher_kwargs( + graph_id: str, client: Any, body: CypherStatementRequest, options: QueryOptions +) -> Dict[str, Any]: + return { + "graph_id": graph_id, + "client": client, + "body": body, + "mode": options.mode if options.mode else None, + "chunk_size": options.chunk_size if options.chunk_size else 1000, + "test_mode": options.test_mode if options.test_mode else False, + } + + +def _wrap_query_error(error: Exception) -> Exception: + error_msg = str(error) + if "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower(): + return Exception(f"Authentication failed during query execution: {error_msg}") + return Exception(f"Query execution failed: {error_msg}") + + +def _notify_queued(queued: QueuedQueryResponse, options: QueryOptions) -> None: + """Report the queue position, and stop here when the caller will not wait.""" + if options.on_queue_update: + options.on_queue_update(queued.queue_position, queued.estimated_wait_seconds) + if options.max_wait == 0: + raise QueuedQueryError(queued) + + +def _interpret_response( + response: Any, graph_id: str +) -> Union[QueryResult, QueuedQueryResponse]: + """Classify a `/query/cypher` response: a result, or a queued query. + + Raises :class:`_QueryFailed` for an error status or an unreadable body. + """ + # Check if this is an NDJSON streaming response (parsed will be None for NDJSON) + if ( + hasattr(response, "headers") + and ( + "application/x-ndjson" in response.headers.get("content-type", "") + or response.headers.get("x-stream-format") == "ndjson" + ) + ) or ( + hasattr(response, "parsed") + and response.parsed is None + and response.status_code == 200 + ): + return _parse_ndjson_response(response, graph_id) + + # A queued query answers 202, which the generated parser yields as None; + # its body (operation id, queue position) is on the raw response. + data = getattr(response, "parsed", None) + if data is None and _status_code(response) == 202: + data = _json_body(response) + + if data: + if isinstance(data, dict): + if "data" in data and "columns" in data: + return QueryResult( + data=data["data"], + columns=data["columns"], + row_count=data.get("row_count", len(data["data"])), + execution_time_ms=data.get("execution_time_ms", 0), + graph_id=graph_id, + timestamp=data.get("timestamp", datetime.now().isoformat()), + ) + if data.get("status") == "queued" and "operation_id" in data: + return QueuedQueryResponse( + status=data["status"], + operation_id=data["operation_id"], + queue_position=data.get("queue_position", 0), + estimated_wait_seconds=data.get("estimated_wait_seconds", 0), + message=data.get("message", "Query queued"), + ) + elif hasattr(data, "data") and hasattr(data, "columns"): + # attrs object - access attributes directly + raw_data = data.data if data.data is not UNSET else [] + result_data = [] + for item in raw_data: + if hasattr(item, "to_dict"): + result_data.append(item.to_dict()) + elif hasattr(item, "additional_properties"): + result_data.append(item.additional_properties) + else: + result_data.append(item) + return QueryResult( + data=result_data, + columns=data.columns if data.columns is not UNSET else [], + row_count=data.row_count if data.row_count is not UNSET else len(result_data), + execution_time_ms=( + data.execution_time_ms if data.execution_time_ms is not UNSET else 0 + ), + graph_id=graph_id, + timestamp=( + data.timestamp if data.timestamp is not UNSET else datetime.now().isoformat() + ), + ) + elif ( + hasattr(data, "status") + and hasattr(data, "operation_id") + and getattr(data, "status", None) == "queued" + ): + return QueuedQueryResponse( + status=data.status, + operation_id=data.operation_id, + queue_position=data.queue_position + if hasattr(data, "queue_position") and data.queue_position is not UNSET + else 0, + estimated_wait_seconds=data.estimated_wait_seconds + if hasattr(data, "estimated_wait_seconds") + and data.estimated_wait_seconds is not UNSET + else 0, + message=data.message + if hasattr(data, "message") and data.message is not UNSET + else "Query queued", + ) + + # Error responses (4xx/5xx) + code = _status_code(response) + if code is not None and code >= 400: + detail: Any = f"HTTP {code}" + try: + body = ( + response.content.decode("utf-8") + if isinstance(response.content, bytes) + else str(response.content) + ) + error_data = json.loads(body) + detail = error_data.get("detail", error_data.get("message", body)) + except Exception: + pass + raise _QueryFailed(f"Query failed ({code}): {detail}") + + raise _QueryFailed("Unexpected response format from query endpoint") + + +def _parse_ndjson_response(response: Any, graph_id: str) -> QueryResult: + """Parse NDJSON streaming response and aggregate into QueryResult""" + import json + + all_data = [] + columns = None + total_rows = 0 + execution_time_ms = 0 + + # Parse NDJSON line by line + content = ( + response.content.decode("utf-8") + if isinstance(response.content, bytes) + else response.content + ) + + for line in content.strip().split("\n"): + if not line.strip(): + continue + + try: + chunk = json.loads(line) + + # Extract columns from first chunk + if columns is None and "columns" in chunk: + columns = chunk["columns"] + + # Aggregate data rows (NDJSON uses "rows", regular JSON uses "data") + if "rows" in chunk: + all_data.extend(chunk["rows"]) + total_rows += len(chunk["rows"]) + elif "data" in chunk: + all_data.extend(chunk["data"]) + total_rows += len(chunk["data"]) + + # Track execution time (use max from all chunks) + if "execution_time_ms" in chunk: + execution_time_ms = max(execution_time_ms, chunk["execution_time_ms"]) + + except json.JSONDecodeError as e: + raise Exception(f"Failed to parse NDJSON line: {e}") + + # Return aggregated result + return QueryResult( + data=all_data, + columns=columns or [], + row_count=total_rows, + execution_time_ms=execution_time_ms, + graph_id=graph_id, + timestamp=datetime.now().isoformat(), + ) + + class QueryClient: """Enhanced query client with SSE streaming support""" @@ -125,230 +339,26 @@ def execute_query( client = self._rest_client() try: - kwargs = { - "graph_id": graph_id, - "client": client, - "body": query_request, - "mode": options.mode if options.mode else None, - "chunk_size": options.chunk_size if options.chunk_size else 1000, - "test_mode": options.test_mode if options.test_mode else False, - } - response = execute_cypher_query(**kwargs) - - # Check if this is an NDJSON streaming response (parsed will be None for NDJSON) - if ( - hasattr(response, "headers") - and ( - "application/x-ndjson" in response.headers.get("content-type", "") - or response.headers.get("x-stream-format") == "ndjson" - ) - ) or ( - hasattr(response, "parsed") - and response.parsed is None - and response.status_code == 200 - ): - return self._parse_ndjson_response(response, graph_id) - - # Check response type and handle accordingly - if hasattr(response, "parsed") and response.parsed: - response_data = response.parsed - - # Handle both dict and attrs object responses - if isinstance(response_data, dict): - # Response is a plain dict - data = response_data - has_data = "data" in data and "columns" in data - else: - # Response is an attrs object - check for attributes directly - data = response_data - has_data = hasattr(data, "data") and hasattr(data, "columns") - - # Check if this is an immediate response - if has_data: - # Extract data from either dict or attrs object - if isinstance(data, dict): - result_data = data["data"] - result_columns = data["columns"] - result_row_count = data.get("row_count", len(data["data"])) - result_execution_time = data.get("execution_time_ms", 0) - result_timestamp = data.get("timestamp", datetime.now().isoformat()) - else: - # attrs object - access attributes directly - from ..types import UNSET - - raw_data = data.data if data.data is not UNSET else [] - # Convert data items to dicts if they're objects - result_data = [] - for item in raw_data: - if hasattr(item, "to_dict"): - result_data.append(item.to_dict()) - elif hasattr(item, "additional_properties"): - result_data.append(item.additional_properties) - else: - result_data.append(item) - result_columns = data.columns if data.columns is not UNSET else [] - result_row_count = ( - data.row_count if data.row_count is not UNSET else len(result_data) - ) - result_execution_time = ( - data.execution_time_ms if data.execution_time_ms is not UNSET else 0 - ) - result_timestamp = ( - data.timestamp - if data.timestamp is not UNSET - else datetime.now().isoformat() - ) - - return QueryResult( - data=result_data, - columns=result_columns, - row_count=result_row_count, - execution_time_ms=result_execution_time, - graph_id=graph_id, - timestamp=result_timestamp, - ) - - # Check if this is a queued response - is_queued = False - queued_response = None - - if isinstance(data, dict): - is_queued = data.get("status") == "queued" and "operation_id" in data - if is_queued: - queued_response = QueuedQueryResponse( - status=data["status"], - operation_id=data["operation_id"], - queue_position=data.get("queue_position", 0), - estimated_wait_seconds=data.get("estimated_wait_seconds", 0), - message=data.get("message", "Query queued"), - ) - else: - is_queued = ( - hasattr(data, "status") - and hasattr(data, "operation_id") - and getattr(data, "status", None) == "queued" - ) - if is_queued: - from ..types import UNSET - - queued_response = QueuedQueryResponse( - status=data.status, - operation_id=data.operation_id, - queue_position=data.queue_position - if hasattr(data, "queue_position") and data.queue_position is not UNSET - else 0, - estimated_wait_seconds=data.estimated_wait_seconds - if hasattr(data, "estimated_wait_seconds") - and data.estimated_wait_seconds is not UNSET - else 0, - message=data.message - if hasattr(data, "message") and data.message is not UNSET - else "Query queued", - ) - - if is_queued and queued_response: - # Notify about queue status - if options.on_queue_update: - options.on_queue_update( - queued_response.queue_position, queued_response.estimated_wait_seconds - ) - - # If user doesn't want to wait, raise with queue info - if options.max_wait == 0: - raise QueuedQueryError(queued_response) - - # Use SSE to monitor the operation - if options.mode == "stream": - return self._stream_query_results(queued_response.operation_id, options) - else: - return self._wait_for_query_completion( - queued_response.operation_id, options - ) + response = execute_cypher_query( + **_cypher_kwargs(graph_id, client, query_request, options) + ) + outcome = _interpret_response(response, graph_id) + if isinstance(outcome, QueryResult): + return outcome - except Exception as e: - if isinstance(e, QueuedQueryError): - raise - - error_msg = str(e) - # Check for authentication errors - if ( - "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower() - ): - raise Exception(f"Authentication failed during query execution: {error_msg}") - else: - raise Exception(f"Query execution failed: {error_msg}") - - # Handle error responses (4xx/5xx) where parsed is None - if hasattr(response, "status_code") and response.status_code >= 400: - import json as _json - - detail = f"HTTP {response.status_code}" - try: - body = ( - response.content.decode("utf-8") - if isinstance(response.content, bytes) - else str(response.content) - ) - error_data = _json.loads(body) - detail = error_data.get("detail", error_data.get("message", body)) - except Exception: - pass - raise Exception(f"Query failed ({response.status_code}): {detail}") + _notify_queued(outcome, options) + if options.mode == "stream": + return self._stream_query_results(outcome.operation_id, options) + return self._wait_for_query_completion(outcome.operation_id, options) - # Unexpected response format - raise Exception("Unexpected response format from query endpoint") + except (QueuedQueryError, _QueryFailed): + raise + except Exception as e: + raise _wrap_query_error(e) from e def _parse_ndjson_response(self, response, graph_id: str) -> QueryResult: """Parse NDJSON streaming response and aggregate into QueryResult""" - import json - - all_data = [] - columns = None - total_rows = 0 - execution_time_ms = 0 - - # Parse NDJSON line by line - content = ( - response.content.decode("utf-8") - if isinstance(response.content, bytes) - else response.content - ) - - for line in content.strip().split("\n"): - if not line.strip(): - continue - - try: - chunk = json.loads(line) - - # Extract columns from first chunk - if columns is None and "columns" in chunk: - columns = chunk["columns"] - - # Aggregate data rows (NDJSON uses "rows", regular JSON uses "data") - if "rows" in chunk: - all_data.extend(chunk["rows"]) - total_rows += len(chunk["rows"]) - elif "data" in chunk: - all_data.extend(chunk["data"]) - total_rows += len(chunk["data"]) - - # Track execution time (use max from all chunks) - if "execution_time_ms" in chunk: - execution_time_ms = max(execution_time_ms, chunk["execution_time_ms"]) - - except json.JSONDecodeError as e: - raise Exception(f"Failed to parse NDJSON line: {e}") - - # Return aggregated result - return QueryResult( - data=all_data, - columns=columns or [], - row_count=total_rows, - execution_time_ms=execution_time_ms, - graph_id=graph_id, - timestamp=datetime.now().isoformat(), - ) + return _parse_ndjson_response(response, graph_id) def _stream_query_results( self, operation_id: str, options: QueryOptions @@ -641,21 +651,156 @@ class AsyncQueryClient: def __init__(self, config: Dict[str, Any]): self.config = config self.base_url = config["base_url"] + self.headers = config.get("headers", {}) + self.token = config.get("token") self.sse_client: Optional[AsyncSSEClient] = None + def _rest_client(self) -> Client: + """A REST client for one call, carrying the credential current now.""" + headers = resolve_auth_headers(self.config) + if not has_auth_header(headers): + raise Exception("No API key provided. Set X-API-Key in headers.") + return Client(base_url=self.base_url, headers=headers) + + def _sse_config(self) -> SSEConfig: + """Stream config for one connect; headers carry the credential current now.""" + return SSEConfig(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + async def execute_query( self, graph_id: str, request: QueryRequest, options: QueryOptions = None ) -> Union[QueryResult, AsyncIterator[Any]]: - """Execute a query asynchronously""" - # Similar implementation to sync version but with async/await - # Would need async version of the generated client - pass + """Execute a query asynchronously. + + Same contract as :meth:`QueryClient.execute_query`: a result, or for a + queued query (202) the SSE-followed result — an async iterator of rows + under ``mode="stream"``. + """ + if options is None: + options = QueryOptions() + + query_request = CypherStatementRequest( + query=request.query, parameters=request.parameters or {} + ) + + try: + async with self._rest_client() as client: + response = await execute_cypher_query_async( + **_cypher_kwargs(graph_id, client, query_request, options) + ) + outcome = _interpret_response(response, graph_id) + if isinstance(outcome, QueryResult): + return outcome + + _notify_queued(outcome, options) + if options.mode == "stream": + return self._stream_query_results(outcome.operation_id, options) + return await self._wait_for_query_completion(outcome.operation_id, options) + + except (QueuedQueryError, _QueryFailed): + raise + except Exception as e: + raise _wrap_query_error(e) from e + + async def _follow(self, operation_id: str, options: QueryOptions) -> Dict[str, Any]: + """Follow a queued query's stream to its end; the verdict and any rows.""" + state: Dict[str, Any] = {"rows": [], "result": None, "error": None} + + def on_data_chunk(data): + if isinstance(data.get("rows"), list): + state["rows"].extend(data["rows"]) + elif isinstance(data.get("data"), list): + state["rows"].extend(data["data"]) + + def on_queue_update(data): + if options.on_queue_update: + options.on_queue_update( + data.get("position", 0), data.get("estimated_wait_seconds", 0) + ) + + def on_progress(data): + if options.on_progress: + options.on_progress(data.get("message", "Processing...")) + + def on_completed(data): + state["result"] = data.get("result", data) or {} + + def on_error(err): + state["error"] = ( + err if isinstance(err, Exception) else Exception(event_error_message(err)) + ) + + def on_cancelled(_data=None): + state["error"] = Exception("Query cancelled") + + sse_client = AsyncSSEClient(self._sse_config()) + self.sse_client = sse_client + sse_client.on(EventType.DATA_CHUNK.value, on_data_chunk) + sse_client.on(EventType.QUEUE_UPDATE.value, on_queue_update) + sse_client.on(EventType.OPERATION_PROGRESS.value, on_progress) + sse_client.on(EventType.OPERATION_COMPLETED.value, on_completed) + sse_client.on(EventType.OPERATION_ERROR.value, on_error) + sse_client.on(EventType.OPERATION_CANCELLED.value, on_cancelled) + sse_client.on("error", on_error) + sse_client.on("max_retries_exceeded", on_error) + + # connect() returns once the stream has ended (or never opened). + try: + await sse_client.connect(operation_id) + finally: + await sse_client.close() + if self.sse_client is sse_client: + self.sse_client = None + + if state["error"] is not None: + raise state["error"] + if state["result"] is None: + raise Exception( + f"Query stream for operation {operation_id} ended before a result" + ) + return state + + async def _wait_for_query_completion( + self, operation_id: str, options: QueryOptions + ) -> QueryResult: + """Wait for a queued query's completion and return its result.""" + state = await self._follow(operation_id, options) + query_result = state["result"] + return QueryResult( + data=query_result.get("data", []), + columns=query_result.get("columns", []), + row_count=query_result.get("row_count", 0), + execution_time_ms=query_result.get("execution_time_ms", 0), + graph_id=query_result.get("graph_id"), + timestamp=query_result.get("timestamp", datetime.now().isoformat()), + ) + + async def _stream_query_results( + self, operation_id: str, options: QueryOptions + ) -> AsyncIterator[Any]: + """Yield a queued query's rows once its stream has ended.""" + state = await self._follow(operation_id, options) + for item in state["rows"]: + yield item + for item in state["result"].get("data") or []: + yield item async def query( self, graph_id: str, cypher: str, parameters: Dict[str, Any] = None ) -> QueryResult: """Async convenience method for simple queries""" - pass + request = QueryRequest(query=cypher, parameters=parameters) + result = await self.execute_query(graph_id, request, QueryOptions(mode="auto")) + if isinstance(result, QueryResult): + return result + data = [item async for item in result] + return QueryResult( + data=data, + columns=[], + row_count=len(data), + execution_time_ms=0, + graph_id=graph_id, + timestamp=datetime.now().isoformat(), + ) async def stream_query( self, @@ -664,8 +809,22 @@ async def stream_query( parameters: Dict[str, Any] = None, chunk_size: int = 1000, ) -> AsyncIterator[Any]: - """Async streaming query for large results""" - pass + """Async streaming query for large results. + + Rows are yielded once the query's stream has ended — the same + buffering :meth:`QueryClient.stream_query` does — or straight from the + result when the query answered inline. + """ + request = QueryRequest(query=cypher, parameters=parameters) + result = await self.execute_query( + graph_id, request, QueryOptions(mode="stream", chunk_size=chunk_size) + ) + if isinstance(result, QueryResult): + for item in result.data: + yield item + else: + async for item in result: + yield item async def close(self): """Cancel any active SSE connections""" diff --git a/robosystems_client/clients/retry.py b/robosystems_client/clients/retry.py index 99125caf..bf865ff3 100644 --- a/robosystems_client/clients/retry.py +++ b/robosystems_client/clients/retry.py @@ -3,10 +3,15 @@ The API rate-limits per user per endpoint category and answers an exhausted budget with ``429`` plus ``Retry-After`` / ``X-RateLimit-*`` headers. That rejection is raised by a request dependency *before* the -endpoint handler runs, so the request had no effect and is always safe -to replay — including a ``POST`` carrying no idempotency key. Nothing -other than ``429`` is retried here, precisely because nothing else -carries that guarantee. +endpoint handler runs, so the request had no effect and is safe to +replay — including a ``POST`` carrying no idempotency key. Nothing other +than ``429`` is retried here, because nothing else carries that +guarantee. + +A few handlers also answer ``429`` for a quota of their own (a backup +quota, the concurrent-query cap). Those refuse the request before doing +its work, so a replay is still harmless, though it can only succeed once +the quota frees up. The motivating case is a bulk backfill: an integrator loading a year of history through per-event write calls runs at the category budget for @@ -145,29 +150,48 @@ def retrying_authenticated_client( base_url: str, token: str, headers: dict[str, str] | None = None, - auth_header_name: str = "X-API-Key", - prefix: str = "", + auth_header_name: str | None = None, + prefix: str | None = None, config: dict[str, Any] | None = None, ) -> AuthenticatedClient: """An :class:`AuthenticatedClient` whose transport replays 429s. Mirrors what ``AuthenticatedClient.get_httpx_client()`` would build — - including stamping the credential onto the header the caller named — - and installs it through the public ``set_httpx_client`` hook, so the - generated ``api/`` layer is untouched and survives ``just - generate-sdk``. + including stamping the credential onto its header — and installs it + through the public ``set_httpx_client`` hook, so the generated ``api/`` + layer is untouched and survives ``just generate-sdk``. + + Without an explicit ``auth_header_name`` the credential is routed by + shape, the same rule the GraphQL reads use (``apply_auth_header``): + ``rfs…`` API keys in ``X-API-Key``, anything else (a JWT) as + ``Authorization: Bearer …``. Any credential header already in + ``headers`` is dropped first, so exactly one is sent. ``timeout`` is left unset to match the generated client's own default; callers that need one set it on the facade. """ - request_headers = dict(headers or {}) + base_headers = dict(headers or {}) + if auth_header_name is None: + base_headers = { + k: v + for k, v in base_headers.items() + if k.lower() not in ("x-api-key", "authorization") + } + if token.startswith("rfs"): + auth_header_name = "X-API-Key" + prefix = prefix or "" + else: + auth_header_name = "Authorization" + prefix = "Bearer" if prefix is None else prefix + prefix = prefix or "" + request_headers = dict(base_headers) request_headers[auth_header_name] = f"{prefix} {token}" if prefix else token client = AuthenticatedClient( base_url=base_url, token=token, prefix=prefix, auth_header_name=auth_header_name, - headers=dict(headers or {}), + headers=base_headers, ) return client.set_httpx_client( build_httpx_client( diff --git a/robosystems_client/clients/table_client.py b/robosystems_client/clients/table_client.py index d41fe062..442d84b9 100644 --- a/robosystems_client/clients/table_client.py +++ b/robosystems_client/clients/table_client.py @@ -15,6 +15,7 @@ sync_detailed as query_tables, ) from ..models.sql_statement_request import SqlStatementRequest +from .token_utils import current_token logger = logging.getLogger(__name__) @@ -24,10 +25,12 @@ class TableInfo: """Information about a DuckDB staging table""" table_name: str - table_type: str + # The API no longer reports a table type; kept for compatibility, always None. + table_type: Optional[str] row_count: int file_count: int total_size_bytes: int + s3_location: Optional[str] = None @dataclass @@ -64,12 +67,14 @@ def list(self, graph_id: str) -> list[TableInfo]: try: from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) @@ -91,10 +96,11 @@ def list(self, graph_id: str) -> list[TableInfo]: return [ TableInfo( table_name=t.table_name, - table_type=t.table_type, + table_type=None, row_count=t.row_count, file_count=t.file_count or 0, total_size_bytes=t.total_size_bytes or 0, + s3_location=t.s3_location if isinstance(t.s3_location, str) else None, ) for t in tables ] @@ -134,12 +140,14 @@ def query( from .retry import retrying_authenticated_client - if not self.token: + token = current_token(self.config, self.token) + + if not token: raise Exception("No API key provided. Set X-API-Key in headers.") client = retrying_authenticated_client( base_url=self.base_url, - token=self.token, + token=token, headers=self.headers, config=self.config, ) diff --git a/robosystems_client/clients/token_utils.py b/robosystems_client/clients/token_utils.py index 443bc45d..6aeb0696 100644 --- a/robosystems_client/clients/token_utils.py +++ b/robosystems_client/clients/token_utils.py @@ -446,6 +446,18 @@ def resolve_config_token(config: Dict[str, Any]) -> Optional[str]: return config.get("token") +def current_token(config: Dict[str, Any], static_token: Optional[str]) -> Optional[str]: + """The credential for one call: ``token_provider`` wins over ``static_token``. + + For the facades that keep the static token on ``self.token``; the same + precedence as :func:`resolve_config_token`. + """ + provider = config.get("token_provider") + if provider is not None: + return provider() + return static_token + + _AUTH_HEADER_NAMES = ("x-api-key", "authorization") @@ -469,6 +481,11 @@ def apply_auth_header(headers: Dict[str, str], credential: str) -> None: headers["Authorization"] = f"Bearer {credential}" +def has_auth_header(headers: Dict[str, str]) -> bool: + """Whether ``headers`` carry a credential (``X-API-Key`` or ``Authorization``).""" + return any(k.lower() in _AUTH_HEADER_NAMES and v for k, v in headers.items()) + + def resolve_auth_headers(config: Dict[str, Any]) -> Dict[str, str]: """Headers for one request or stream, carrying the credential current *now*. diff --git a/robosystems_client/graphql/generated/client.py b/robosystems_client/graphql/generated/client.py index 29583d04..8194210c 100644 --- a/robosystems_client/graphql/generated/client.py +++ b/robosystems_client/graphql/generated/client.py @@ -568,6 +568,7 @@ def list_information_blocks( category: Union[Optional[str], UnsetType] = UNSET, limit: Union[Optional[int], UnsetType] = UNSET, offset: Union[Optional[int], UnsetType] = UNSET, + scenario_id: Union[Optional[str], UnsetType] = UNSET, **kwargs: Any, ) -> ListInformationBlocks: variables: dict[str, object] = { @@ -575,6 +576,7 @@ def list_information_blocks( "category": category, "limit": limit, "offset": offset, + "scenarioId": scenario_id, } response = self.execute( query=LIST_INFORMATION_BLOCKS_GQL, diff --git a/robosystems_client/graphql/generated/operations.py b/robosystems_client/graphql/generated/operations.py index 33ffdb57..a5c00cd8 100644 --- a/robosystems_client/graphql/generated/operations.py +++ b/robosystems_client/graphql/generated/operations.py @@ -1177,12 +1177,13 @@ """ LIST_INFORMATION_BLOCKS_GQL = """ -query ListInformationBlocks($blockType: String, $category: String, $limit: Int, $offset: Int) { +query ListInformationBlocks($blockType: String, $category: String, $limit: Int, $offset: Int, $scenarioId: String) { informationBlocks( blockType: $blockType category: $category limit: $limit offset: $offset + scenarioId: $scenarioId ) { id blockType diff --git a/robosystems_client/graphql/operations/ledger/ListInformationBlocks.graphql b/robosystems_client/graphql/operations/ledger/ListInformationBlocks.graphql index 6001fc20..76617536 100644 --- a/robosystems_client/graphql/operations/ledger/ListInformationBlocks.graphql +++ b/robosystems_client/graphql/operations/ledger/ListInformationBlocks.graphql @@ -3,12 +3,14 @@ query ListInformationBlocks( $category: String $limit: Int $offset: Int + $scenarioId: String ) { informationBlocks( blockType: $blockType category: $category limit: $limit offset: $offset + scenarioId: $scenarioId ) { id blockType name displayName category taxonomyId taxonomyName diff --git a/tests/test_extensions_integration.py b/tests/test_extensions_integration.py index 7f9a7374..64a41ca4 100644 --- a/tests/test_extensions_integration.py +++ b/tests/test_extensions_integration.py @@ -86,13 +86,18 @@ def test_create_clients_factory(self, mock_api_key): @patch("robosystems_client.api.query.execute_cypher.sync_detailed") def test_cypher_query_execution(self, mock_sync_detailed, extensions): """Test executing Cypher queries through authenticated client""" - # Mock the response + # The generated parser hands a 200 body back as a plain dict. mock_response = Mock() - mock_response.parsed = Mock() - mock_response.parsed.data = [{"name": "Company A", "revenue": 1000000}] - mock_response.parsed.columns = ["name", "revenue"] - mock_response.parsed.row_count = 1 - mock_response.parsed.execution_time_ms = 150 + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.parsed = { + "success": True, + "data": [{"name": "Company A", "revenue": 1000000}], + "columns": ["name", "revenue"], + "row_count": 1, + "execution_time_ms": 150, + "graph_id": "test_graph", + } mock_sync_detailed.return_value = mock_response result = extensions.execute_cypher_query( diff --git a/tests/test_facade_wire_shapes.py b/tests/test_facade_wire_shapes.py new file mode 100644 index 00000000..b162acc9 --- /dev/null +++ b/tests/test_facade_wire_shapes.py @@ -0,0 +1,701 @@ +"""Facade behaviour against the API's real responses. + +Each test answers the facade with the status code and JSON body the server +actually sends (`robosystems` routers, cited per test), through the real +generated `api/` parsers — mocked 200s are how a 202 the generated layer +parses as ``None``, or a body shape the facade misread, went unnoticed. +""" + +import json +import re + +import pytest + +from robosystems_client.clients import OperatorRunError +from robosystems_client.clients.auth_integration import AuthenticatedClients +from robosystems_client.clients.facade import ( + AsyncRoboSystemsClients, + RoboSystemsClientConfig, +) +from robosystems_client.clients.file_client import FileClient +from robosystems_client.clients.graph_client import GraphClient +from robosystems_client.clients.ledger_client import LedgerClient +from robosystems_client.clients.operation_client import ( + AsyncOperationClient, + MonitorOptions, + OperationClient, + OperationStatus, +) +from robosystems_client.clients.operator_client import ( + OperatorClient, + OperatorOptions, + OperatorQueryRequest, + QueuedOperatorError, +) +from robosystems_client.clients.query_client import ( + AsyncQueryClient, + QueryClient, + QueryOptions, + QueryRequest, + QueryResult, + QueuedQueryError, +) +from robosystems_client.clients.table_client import TableClient + +BASE = "http://localhost:8000" +GRAPH = "kg_test" +OP = "op_01J0000000000000000000000A" +CONFIG = {"base_url": BASE, "token": "rfs_test", "headers": {}, "max_retries": 1} + + +def _sse(*events: tuple[str, dict]) -> bytes: + return "".join( + f"id: {i}\nevent: {name}\ndata: {json.dumps(data)}\n\n" + for i, (name, data) in enumerate(events) + ).encode() + + +def _stream_url(op_id: str = OP) -> re.Pattern: + return re.compile(rf"{BASE}/v1/operations/{op_id}/stream\?from_sequence=\d+") + + +# Body of `enqueue_task` / `create_operation_response` — what the operator +# endpoint returns with its 202 (routers/graphs/operator/execute.py `_dispatch`). +OPERATOR_QUEUED = { + "operation_id": OP, + "status": "pending", + "operation_type": "operator", + "created_at": "2026-10-05T00:00:00+00:00", + "graph_id": GRAPH, + "_links": { + "stream": f"/v1/operations/{OP}/stream", + "status": f"/v1/operations/{OP}/status", + "cancel": f"/v1/operations/{OP}", + }, + "message": "Operation operator queued. Connect to stream endpoint for real-time updates.", +} + +OPERATOR_RESULT = { + "content": "Burn is ~$1,500/month.", + "operator_used": "analyst", + "mode_used": "standard", + "metadata": {}, + "execution_time": 2.1, +} + +# routers/graphs/query/execute.py — the queued 202. +QUERY_QUEUED = { + "status": "queued", + "query_id": "q_1", + "operation_id": OP, + "queue_position": 2, + "estimated_wait_seconds": 10, + "message": "Query has been queued for execution", + "_links": {"self": "...", "monitor": f"/v1/operations/{OP}/stream"}, +} + +QUERY_ROWS = { + "data": [{"n": 1}], + "columns": ["n"], + "row_count": 1, + "execution_time_ms": 5, + "graph_id": GRAPH, +} + +OPERATOR_URL = re.compile(rf"{BASE}/v1/graphs/{GRAPH}/operator(\?.*)?$") +QUERY_URL = re.compile(rf"{BASE}/v1/graphs/{GRAPH}/query/cypher(\?.*)?$") + + +# ── Operator runs ──────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestOperatorQueuedRun: + def test_202_is_followed_over_the_stream(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=OPERATOR_URL, status_code=202, json=OPERATOR_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + headers={"content-type": "text/event-stream"}, + content=_sse( + ("operation_progress", {"message": "Thinking", "progress_percent": 40}), + ("operation_completed", {"message": "done", "result": OPERATOR_RESULT}), + ), + ) + progress = [] + + result = OperatorClient(CONFIG).execute_query( + GRAPH, + OperatorQueryRequest(message="burn?"), + OperatorOptions(on_progress=lambda msg, pct: progress.append((msg, pct))), + ) + + assert result.content == "Burn is ~$1,500/month." + assert result.operator_used == "analyst" + assert ("Thinking", 40) in progress + + def test_specific_operator_202_is_followed(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/v1/graphs/{GRAPH}/operator/analyst?mode=auto", + status_code=202, + json=OPERATOR_QUEUED, + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_completed", {"result": OPERATOR_RESULT})), + ) + + result = OperatorClient(CONFIG).execute_operator( + GRAPH, "analyst", OperatorQueryRequest(message="burn?") + ) + + assert result.operator_used == "analyst" + + def test_202_with_max_wait_zero_hands_back_the_links(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=OPERATOR_URL, status_code=202, json=OPERATOR_QUEUED + ) + + with pytest.raises(QueuedOperatorError) as exc_info: + OperatorClient(CONFIG).execute_query( + GRAPH, OperatorQueryRequest(message="burn?"), OperatorOptions(max_wait=0) + ) + + info = exc_info.value.queue_info + assert info.operation_id == OP + assert info.status == "pending" + assert info.sse_endpoint == f"/v1/operations/{OP}/stream" + + def test_failed_run_raises_operator_run_error_with_writes(self, httpx_mock): + writes = [{"op": "create-event-block", "id": "evt_1"}] + httpx_mock.add_response( + method="POST", url=OPERATOR_URL, status_code=202, json=OPERATOR_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse( + ( + "operation_error", + {"error": "Tool budget exhausted", "error_details": {"writes": writes}}, + ) + ), + ) + + with pytest.raises(OperatorRunError) as exc_info: + OperatorClient(CONFIG).execute_query(GRAPH, OperatorQueryRequest(message="x")) + + err = exc_info.value + assert str(err) == "Tool budget exhausted" + assert err.status == "failed" + assert err.writes == writes + assert err.operation_id == OP + + def test_cancelled_run_raises_operator_run_error(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=OPERATOR_URL, status_code=202, json=OPERATOR_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_cancelled", {"reason": "Cancelled by user request"})), + ) + + with pytest.raises(OperatorRunError) as exc_info: + OperatorClient(CONFIG).execute_query(GRAPH, OperatorQueryRequest(message="x")) + + assert exc_info.value.status == "cancelled" + + def test_polled_failure_carries_writes(self, httpx_mock): + writes = [{"op": "update-entity"}] + httpx_mock.add_response( + method="POST", url=OPERATOR_URL, status_code=202, json=OPERATOR_QUEUED + ) + # The stream ends with no terminal event, so the run is followed over + # `/status` — whose failed body carries the landed writes. + httpx_mock.add_response(url=_stream_url(), content=b": keepalive\n\n") + httpx_mock.add_response( + url=f"{BASE}/v1/operations/{OP}/status", + json={ + "operation_id": OP, + "operation_type": "operator", + "status": "failed", + "error": "Operator crashed", + "writes": writes, + "_links": {"stream": f"/v1/operations/{OP}/stream"}, + }, + ) + + with pytest.raises(OperatorRunError) as exc_info: + OperatorClient(CONFIG).execute_query( + GRAPH, OperatorQueryRequest(message="x"), OperatorOptions(poll_interval=0) + ) + + assert exc_info.value.writes == writes + + +@pytest.mark.unit +class TestOperatorSyncResponse: + def test_sync_mode_is_sent_and_a_200_returns_the_result(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/v1/graphs/{GRAPH}/operator?mode=sync", + json={**OPERATOR_RESULT, "operation_id": OP, "is_partial": False}, + ) + + result = OperatorClient(CONFIG).execute_query( + GRAPH, OperatorQueryRequest(message="x"), OperatorOptions(mode="sync") + ) + + assert result.content == "Burn is ~$1,500/month." + assert result.metadata == {} + + def test_cancelled_sync_run_returns_error_details(self, httpx_mock): + # `_response_from_operation`: a cancelled run is a 200 whose + # `error_details` describe it and whose `content` is the explanation. + httpx_mock.add_response( + method="POST", + url=f"{BASE}/v1/graphs/{GRAPH}/operator?mode=sync", + json={ + "content": "Operator run was cancelled", + "operator_used": "analyst", + "mode_used": "standard", + "metadata": {}, + "error_details": { + "code": "OPERATOR_CANCELLED", + "message": "Operator run was cancelled", + }, + "operation_id": OP, + "is_partial": False, + }, + ) + + result = OperatorClient(CONFIG).execute_query( + GRAPH, OperatorQueryRequest(message="x"), OperatorOptions(mode="sync") + ) + + assert result.content == "Operator run was cancelled" + assert result.error_details == { + "code": "OPERATOR_CANCELLED", + "message": "Operator run was cancelled", + } + + def test_402_reports_the_status(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=OPERATOR_URL, + status_code=402, + json={"detail": {"code": "INSUFFICIENT_CREDITS"}}, + ) + + with pytest.raises(Exception, match="Operator execution failed: 402"): + OperatorClient(CONFIG).execute_query(GRAPH, OperatorQueryRequest(message="x")) + + +# ── Cypher queries ─────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestQueuedQuery: + def test_202_is_followed_over_the_stream(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, status_code=202, json=QUERY_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_completed", {"result": QUERY_ROWS})), + ) + + result = QueryClient(CONFIG).execute_query(GRAPH, QueryRequest(query="RETURN 1")) + + assert isinstance(result, QueryResult) + assert result.data == [{"n": 1}] + + def test_202_with_max_wait_zero_raises_queued(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, status_code=202, json=QUERY_QUEUED + ) + positions = [] + + with pytest.raises(QueuedQueryError) as exc_info: + QueryClient(CONFIG).execute_query( + GRAPH, + QueryRequest(query="RETURN 1"), + QueryOptions(max_wait=0, on_queue_update=lambda p, w: positions.append(p)), + ) + + assert exc_info.value.queue_info.operation_id == OP + assert exc_info.value.queue_info.queue_position == 2 + assert positions == [2] + + def test_authenticated_clients_cypher_reads_the_dict_body(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, json={"success": True, **QUERY_ROWS} + ) + + clients = AuthenticatedClients("rfs_test", base_url=BASE) + result = clients.execute_cypher_query(GRAPH, "RETURN 1") + + assert result == { + "data": [{"n": 1}], + "columns": ["n"], + "row_count": 1, + "execution_time_ms": 5, + } + + def test_authenticated_clients_cypher_follows_a_202(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, status_code=202, json=QUERY_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_completed", {"result": QUERY_ROWS})), + ) + + clients = AuthenticatedClients("rfs_test", base_url=BASE) + + assert clients.execute_cypher_query(GRAPH, "RETURN 1")["data"] == [{"n": 1}] + + +# ── Async facade ───────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestAsyncClients: + @pytest.mark.asyncio + async def test_async_query_200(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, json={"success": True, **QUERY_ROWS} + ) + + clients = AsyncRoboSystemsClients( + RoboSystemsClientConfig(base_url=BASE, headers={"X-API-Key": "rfs_test"}) + ) + result = await clients.execute_query(GRAPH, "RETURN 1") + + assert result.data == [{"n": 1}] + assert httpx_mock.get_requests()[0].headers["X-API-Key"] == "rfs_test" + + @pytest.mark.asyncio + async def test_async_query_202_is_followed(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, status_code=202, json=QUERY_QUEUED + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_completed", {"result": QUERY_ROWS})), + ) + + client = AsyncQueryClient({**CONFIG}) + result = await client.query(GRAPH, "RETURN 1") + + assert result.data == [{"n": 1}] + + @pytest.mark.asyncio + async def test_async_stream_query_yields_rows(self, httpx_mock): + httpx_mock.add_response( + method="POST", url=QUERY_URL, json={"success": True, **QUERY_ROWS} + ) + + clients = AsyncRoboSystemsClients( + RoboSystemsClientConfig(base_url=BASE, headers={"X-API-Key": "rfs_test"}) + ) + rows = [row async for row in clients.stream_query(GRAPH, "RETURN 1")] + + assert rows == [{"n": 1}] + + @pytest.mark.asyncio + async def test_async_operation_status_and_cancel(self, httpx_mock): + httpx_mock.add_response( + url=f"{BASE}/v1/operations/{OP}/status", + json={"operation_id": OP, "operation_type": "operator", "status": "running"}, + ) + httpx_mock.add_response( + method="DELETE", + url=f"{BASE}/v1/operations/{OP}", + json={ + "operation_id": OP, + "status": "cancelled", + "message": "Operation has been cancelled", + }, + ) + + client = AsyncOperationClient({**CONFIG}) + + status = await client.get_operation_status(OP) + assert status["status"] == "running" + assert await client.cancel_operation(OP) is True + + +# ── Operations monitoring ──────────────────────────────────────────── + + +@pytest.mark.unit +class TestOperationMonitoring: + def test_cancelled_event_settles_the_monitor(self, httpx_mock): + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_cancelled", {"reason": "Cancelled by user request"})), + ) + + result = OperationClient(CONFIG).monitor_operation(OP) + + assert result.status == OperationStatus.CANCELLED + + def test_stream_ending_without_a_verdict_raises(self, httpx_mock): + httpx_mock.add_response( + url=_stream_url(), + content=_sse( + ("operation_progress", {"message": "Working", "progress_percent": 5}) + ), + ) + + with pytest.raises(Exception, match="ended before a terminal event"): + OperationClient(CONFIG).monitor_operation(OP) + + def test_progress_reads_progress_percent(self, httpx_mock): + httpx_mock.add_response( + url=_stream_url(), + content=_sse( + ("operation_progress", {"message": "Loading", "progress_percent": 55.0}), + ("operation_completed", {"message": "done", "result": {"ok": True}}), + ), + ) + seen = [] + + result = OperationClient(CONFIG).monitor_operation( + OP, MonitorOptions(on_progress=seen.append) + ) + + assert result.status == OperationStatus.COMPLETED + assert seen[0].percentage == 55.0 + + def test_cancel_reads_the_status_field(self, httpx_mock): + httpx_mock.add_response( + method="DELETE", + url=f"{BASE}/v1/operations/{OP}", + json={ + "operation_id": OP, + "status": "cancelled", + "message": "Operation has been cancelled", + }, + ) + + assert OperationClient(CONFIG).cancel_operation(OP) is True + + +# ── Credential routing on writes ───────────────────────────────────── + +EVENT_BLOCK_ENVELOPE = { + "operation": "create-event-block", + "operationId": "op_1", + "status": "completed", + "at": "2026-02-12T00:00:00+00:00", + "result": { + "id": "evt_1", + "event_type": "asset_disposed", + "event_category": "adjustment", + "status": "posted", + "occurred_at": "2026-02-12T00:00:00+00:00", + "source": "manual", + "currency": "USD", + "metadata": {}, + "dimension_ids": [], + "event_class": "economic", + "created_at": "2026-02-12T00:00:00+00:00", + "created_by": "u_1", + }, +} + + +@pytest.mark.unit +class TestWriteCredentialRouting: + def test_jwt_from_token_provider_rides_as_bearer(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/extensions/roboledger/{GRAPH}/operations/create-event-block", + json=EVENT_BLOCK_ENVELOPE, + ) + ledger = LedgerClient( + {"base_url": BASE, "headers": {}, "token_provider": lambda: "eyJ.jwt.sig"} + ) + + ledger.dispose_schedule(GRAPH, "str_1", "2026-02-12") + + request = httpx_mock.get_requests()[0] + assert request.headers["Authorization"] == "Bearer eyJ.jwt.sig" + assert "X-API-Key" not in request.headers + # memo / reason left out for the server's defaults + assert json.loads(request.content)["metadata"] == {"schedule_id": "str_1"} + + def test_api_key_rides_as_x_api_key(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/extensions/roboledger/{GRAPH}/operations/create-event-block", + json=EVENT_BLOCK_ENVELOPE, + ) + + LedgerClient(CONFIG).dispose_schedule( + GRAPH, "str_1", "2026-02-12", "Sold the truck", "sale" + ) + + request = httpx_mock.get_requests()[0] + assert request.headers["X-API-Key"] == "rfs_test" + assert "Authorization" not in request.headers + assert json.loads(request.content)["metadata"] == { + "schedule_id": "str_1", + "memo": "Sold the truck", + "reason": "sale", + } + + def test_graph_client_honours_token_provider(self, httpx_mock): + httpx_mock.add_response(url=f"{BASE}/v1/graphs", json=_graphs_body()) + + GraphClient({**CONFIG, "token_provider": lambda: "eyJ.rotated.sig"}).get_graph_info( + GRAPH + ) + + request = httpx_mock.get_requests()[0] + assert request.headers["Authorization"] == "Bearer eyJ.rotated.sig" + + +# ── Reads whose body shape the facade misread ──────────────────────── + + +def _graphs_body() -> dict: + # routers/user — UserGraphsResponse, camelCase on the wire. + return { + "graphs": [ + { + "graphId": GRAPH, + "graphName": "Acme", + "role": "admin", + "isSelected": True, + "createdAt": "2026-01-01T00:00:00Z", + } + ], + "selectedGraphId": GRAPH, + } + + +@pytest.mark.unit +class TestReadShapes: + def test_get_graph_info_finds_the_graph(self, httpx_mock): + httpx_mock.add_response(url=f"{BASE}/v1/graphs", json=_graphs_body()) + + info = GraphClient(CONFIG).get_graph_info(GRAPH) + + assert info.graph_id == GRAPH + assert info.graph_name == "Acme" + assert info.schema_extensions is None # absent on the wire: None, not UNSET + + def test_get_graph_info_reports_a_failed_call(self, httpx_mock): + httpx_mock.add_response( + url=f"{BASE}/v1/graphs", status_code=401, json={"detail": "Invalid API key"} + ) + + with pytest.raises(RuntimeError, match="Failed to get graphs: 401"): + GraphClient(CONFIG).get_graph_info(GRAPH) + + def test_table_list_reads_the_real_fields(self, httpx_mock): + httpx_mock.add_response( + url=f"{BASE}/v1/graphs/{GRAPH}/tables", + json={ + "tables": [ + { + "table_name": "Entity", + "row_count": 3, + "file_count": 1, + "total_size_bytes": 2048, + "s3_location": None, + } + ], + "total_count": 1, + }, + ) + + tables = TableClient(CONFIG).list(GRAPH) + + assert len(tables) == 1 + assert tables[0].table_name == "Entity" + assert tables[0].row_count == 3 + assert tables[0].table_type is None + + def test_file_list_and_get_turn_missing_fields_into_none(self, httpx_mock): + file_row = { + "file_id": "f_1", + "file_name": "a.parquet", + "file_format": "parquet", + "size_bytes": 10, + "upload_status": "uploaded", + "upload_method": "presigned", + "s3_key": "k", + } + httpx_mock.add_response( + url=re.compile(rf"{BASE}/v1/graphs/{GRAPH}/files(\?.*)?$"), + json={ + "graph_id": GRAPH, + "files": [file_row], + "total_files": 1, + "total_size_bytes": 10, + }, + ) + httpx_mock.add_response( + url=f"{BASE}/v1/graphs/{GRAPH}/files/f_1", + json={**file_row, "graph_id": GRAPH, "table_id": "t_1"}, + ) + client = FileClient(CONFIG) + + listed = client.list(GRAPH)[0] + fetched = client.get(GRAPH, "f_1") + + for info in (listed, fetched): + assert info is not None + assert info.row_count is None + assert info.created_at is None + assert info.uploaded_at is None + assert fetched is not None and fetched.layers is None + + +# ── Ledger parity additions ────────────────────────────────────────── + + +@pytest.mark.unit +class TestLedgerParity: + def test_create_report_sends_periods_and_iso_dates(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/extensions/roboledger/{GRAPH}/operations/create-report", + status_code=500, + json={"detail": "stop here"}, + ) + + with pytest.raises(RuntimeError): + LedgerClient(CONFIG).create_report( + GRAPH, + "FY25", + "map_1", + "2025-01-01", + "2025-12-31", + periods=[{"start": "2025-01-01", "end": "2025-03-31", "label": "Q1"}], + ) + + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["period_start"] == "2025-01-01" + assert body["periods"] == [ + {"start": "2025-01-01", "end": "2025-03-31", "label": "Q1"} + ] + + def test_list_information_blocks_sends_scenario_id(self, httpx_mock): + httpx_mock.add_response( + method="POST", + url=f"{BASE}/extensions/{GRAPH}/graphql", + json={"data": {"informationBlocks": []}}, + ) + + blocks = LedgerClient(CONFIG).list_information_blocks(GRAPH, scenario_id="str_fc") + + assert blocks == [] + payload = json.loads(httpx_mock.get_requests()[0].content) + assert payload["variables"] == {"scenarioId": "str_fc"} + assert "$scenarioId: String" in payload["query"] diff --git a/tests/test_graph_client.py b/tests/test_graph_client.py index d0b027f5..e14f50ea 100644 --- a/tests/test_graph_client.py +++ b/tests/test_graph_client.py @@ -507,14 +507,15 @@ def test_sse_uses_correct_url_and_headers(self, mock_httpx_client, mock_config): mock_client = self._create_mock_client(mock_response) mock_httpx_client.return_value = mock_client - client = GraphClient(mock_config) + # API keys are `rfs…`-prefixed; the credential is routed by that shape. + client = GraphClient({**mock_config, "token": "rfs_test_key"}) client._wait_with_sse("op-999", timeout=60, on_progress=None) # Verify stream was called with correct URL and headers mock_client.stream.assert_called_once_with( "GET", "http://localhost:8000/v1/operations/op-999/stream", - headers={"X-API-Key": "test-api-key", "Accept": "text/event-stream"}, + headers={"Accept": "text/event-stream", "X-API-Key": "rfs_test_key"}, ) diff --git a/tests/test_operation_client_ops.py b/tests/test_operation_client_ops.py index 045ba766..3fb4381d 100644 --- a/tests/test_operation_client_ops.py +++ b/tests/test_operation_client_ops.py @@ -300,8 +300,13 @@ class TestCancelOperation: def test_cancel_success(self, mock_cancel, mock_config): """Test successful cancellation.""" mock_resp = Mock() + mock_resp.status_code = 200 mock_resp.parsed = CancelOperationResponseCanceloperation.from_dict( - {"cancelled": True} + { + "operation_id": "op-cancel", + "status": "cancelled", + "message": "Operation has been cancelled", + } ) mock_cancel.return_value = mock_resp diff --git a/tests/test_operation_client_status.py b/tests/test_operation_client_status.py index aba33623..3230117d 100644 --- a/tests/test_operation_client_status.py +++ b/tests/test_operation_client_status.py @@ -11,6 +11,7 @@ ``to_dict()``; these tests pin that the same accessor is used here. """ +from http import HTTPStatus from typing import Any from unittest.mock import MagicMock, patch @@ -41,12 +42,21 @@ def _status_response(body: dict[str, Any]) -> MagicMock: return response -def _cancel_response(body: dict[str, Any]) -> MagicMock: +def _cancel_response(body: dict[str, Any], status_code: int = 200) -> MagicMock: response = MagicMock() + response.status_code = HTTPStatus(status_code) response.parsed = CancelOperationResponseCanceloperation.from_dict(body) return response +# The API's real answer to a cancel that landed (routers/operations.py). +CANCELLED_BODY = { + "operation_id": "op_1", + "status": "cancelled", + "message": "Operation has been cancelled", +} + + @pytest.mark.unit class TestParsedDict: def test_reads_an_additional_properties_model(self): @@ -129,13 +139,17 @@ def test_transport_failure_still_reports_an_error( class TestCancelOperation: @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") def test_returns_true_when_the_cancel_lands(self, mock_cancel, client): - mock_cancel.return_value = _cancel_response({"cancelled": True}) + mock_cancel.return_value = _cancel_response(CANCELLED_BODY) assert client.cancel_operation("op_1") is True @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") def test_returns_false_when_the_server_declines(self, mock_cancel, client): - mock_cancel.return_value = _cancel_response({"cancelled": False}) + # A finished operation is refused with 409. + mock_cancel.return_value = _cancel_response( + {"detail": "Operation cannot be cancelled - current status is completed"}, + status_code=409, + ) assert client.cancel_operation("op_1") is False @@ -143,7 +157,7 @@ def test_returns_false_when_the_server_declines(self, mock_cancel, client): def test_closes_the_stream_on_a_successful_cancel(self, mock_cancel, client): # The cleanup used to sit after an early `return` on the success path, # so the one case that needs it skipped it and leaked the stream. - mock_cancel.return_value = _cancel_response({"cancelled": True}) + mock_cancel.return_value = _cancel_response(CANCELLED_BODY) stream = MagicMock() client.active_operations["op_1"] = stream