Update query endpoints to use FastApiReadCypherQueryResponse model; enhance response handling in service methods

Author: rafaljanicki

src/sandbox_api_mcp_server/sandbox/routes.py

@@ -4,7 +4,7 @@
 from fastapi import Query, Path, status
 from fastapi.responses import PlainTextResponse
 
-from .models import StartSandboxBody, StopSandboxBody, ExtendSandboxBody, AuraUploadBody, BackupDownloadUrlBody, FastApiReadCypherQueryBody, FastApiWriteCypherQueryBody
+from .models import StartSandboxBody, StopSandboxBody, ExtendSandboxBody, AuraUploadBody, BackupDownloadUrlBody, FastApiReadCypherQueryBody, FastApiWriteCypherQueryBody, FastApiReadCypherQueryResponse
 from ..helpers import get_logger
 from .service import call_sandbox_api, SandboxApiClient, get_sandbox_client
 
@@ -160,23 +160,23 @@ async def get_aura_upload_result_ep(
             logger.error(f"Error getting Aura upload result: {e}")
             raise e
 
-    @router.get("/query/schema", operation_id="get_schema", tags=["Query"], response_model=Dict)
+    @router.get("/query/schema", operation_id="get_schema", tags=["Query"], response_model=FastApiReadCypherQueryResponse)
     async def get_schema(hash_key: str, client: Annotated[SandboxApiClient, Depends(get_sandbox_client)]):
         try:
             return await call_sandbox_api("get_schema", client, hash_key=hash_key)
         except Exception as e:
             logger.error(f"Error getting schema: {e}")
             raise e
 
-    @router.post("/query/read", operation_id="read_query", tags=["Query"], response_model=Dict)
+    @router.post("/query/read", operation_id="read_query", tags=["Query"], response_model=FastApiReadCypherQueryResponse)
     async def read(cypher_query: FastApiReadCypherQueryBody, client: Annotated[SandboxApiClient, Depends(get_sandbox_client)]):
         try:
             return await call_sandbox_api("read_query", client, hash_key=cypher_query.hash_key, query=cypher_query.query, params=cypher_query.params)
         except Exception as e:
             logger.error(f"Error reading query: {e}")
             raise e
 
-    @router.post("/query/write", operation_id="write_query", tags=["Query"], response_model=Dict)
+    @router.post("/query/write", operation_id="write_query", tags=["Query"], response_model=FastApiReadCypherQueryResponse)
     async def write(cypher_query: FastApiWriteCypherQueryBody, client: Annotated[SandboxApiClient, Depends(get_sandbox_client)]):
         try:
             return await call_sandbox_api("write_query", client, hash_key=cypher_query.hash_key, query=cypher_query.query, params=cypher_query.params)

src/sandbox_api_mcp_server/sandbox/service.py

@@ -265,14 +265,19 @@ async def read_query(self, hash_key: str, query: str, params: Optional[Dict[str,
         Executes a read query on the Neo4j database.
         Corresponds to POST /SandboxQuery in swagger.
         """
-        return await self._request("POST", "/SandboxRunQuery", json_data={"hash_key": hash_key, "statement": query, "params": params, "accessMode": "Read"})
+        result = await self._request("POST", "/SandboxRunQuery",
+                                    json_data={"hash_key": hash_key, "statement": query, "params": params,
+                                               "accessMode": "Read"})
+        return FastApiReadCypherQueryResponse(data=result, count=len(result))
 
     async def write_query(self, hash_key: str, query: str, params: Optional[Dict[str, Any]] = None) -> FastApiReadCypherQueryResponse:
         """
         Executes a write query on the Neo4j database.
         Corresponds to POST /SandboxQuery in swagger.
         """
-        return await self._request("POST", "/SandboxRunQuery", json_data={"hash_key": hash_key, "statement": query, "params": params})
+        result = await self._request("POST", "/SandboxRunQuery",
+                                      json_data={"hash_key": hash_key, "statement": query, "params": params})
+        return FastApiReadCypherQueryResponse(data=result, count=len(result))
 
 
 def get_sandbox_client(user: Annotated[Dict[str, Any], Depends(verify_auth)]) -> SandboxApiClient: