[executors] chore: wrap server response into ApiExecutorResponse

Author: AndreyMarkinPPC

libs/garf_executors/garf_executors/__init__.py

@@ -15,12 +15,13 @@
 
 from __future__ import annotations
 
-from garf_executors.api_executor import ApiQueryExecutor
+from garf_executors.api_executor import ApiExecutionContext, ApiQueryExecutor
 from garf_executors.fetchers import FETCHERS
 
 __all__ = [
   'FETCHERS',
   'ApiQueryExecutor',
+  'ApiExecutionContext',
 ]
 
 __version__ = '0.0.7'

libs/garf_executors/garf_executors/entrypoints/server.py

@@ -40,7 +40,7 @@ class ApiExecutorRequest(pydantic.BaseModel):
   title: str | None = None
   query: str | None = None
   query_path: str | list[str] | None = None
-  context: garf_executors.api_executor.ApiExecutionContext
+  context: garf_executors.ApiExecutionContext
 
   @pydantic.model_validator(mode='after')
   def check_query_specified(self):
@@ -57,35 +57,42 @@ def model_post_init(self, __context__) -> None:
       self.title = str(self.query_path)
 
 
+class ApiExecutorResponse(pydantic.BaseModel):
+  """Response after executing a query.
+
+  Attributes:
+    results: Results of query execution.
+  """
+
+  results: list[str]
+
+
 router = fastapi.APIRouter(prefix='/api')
 
 
 @router.post('/execute')
-async def execute(request: ApiExecutorRequest) -> dict[str, str]:
+async def execute(request: ApiExecutorRequest) -> ApiExecutorResponse:
   if not (concrete_api_fetcher := garf_executors.FETCHERS.get(request.source)):
     raise exceptions.GarfExecutorError(
       f'Source {request.source} is not available.'
     )
 
-  query_executor = garf_executors.api_executor.ApiQueryExecutor(
+  query_executor = garf_executors.ApiQueryExecutor(
     concrete_api_fetcher(**request.context.fetcher_parameters)
   )
 
   result = query_executor.execute(request.query, request.title, request.context)
-
-  return fastapi.responses.JSONResponse(
-    content=fastapi.encoders.jsonable_encoder({'result': result})
-  )
+  return ApiExecutorResponse(results=[result])
 
 
 @router.post('/execute:batch')
-async def execute_batch(request: ApiExecutorRequest) -> dict[str, str]:
+async def execute_batch(request: ApiExecutorRequest) -> ApiExecutorResponse:
   if not (concrete_api_fetcher := garf_executors.FETCHERS.get(request.source)):
     raise exceptions.GarfExecutorError(
       f'Source {request.source} is not available.'
     )
 
-  query_executor = garf_executors.api_executor.ApiQueryExecutor(
+  query_executor = garf_executors.ApiQueryExecutor(
     concrete_api_fetcher(**request.context.fetcher_parameters)
   )
   file_reader = reader.FileReader()
@@ -102,10 +109,7 @@ async def execute_batch(request: ApiExecutorRequest) -> dict[str, str]:
     }
     for future in futures.as_completed(future_to_query):
       results.append(future.result())
-
-  return fastapi.responses.JSONResponse(
-    content=fastapi.encoders.jsonable_encoder({'result': results})
-  )
+  return ApiExecutorResponse(results=results)
 
 
 if __name__ == '__main__':

libs/garf_executors/tests/end-to-end/test_server.py

@@ -50,7 +50,7 @@ def test_fake_source_from_query_text(self, tmp_path):
     }
     response = client.post('/api/execute', json=request)
     assert response.status_code == fastapi.status.HTTP_200_OK
-    expected_output = {'result': f'[CSV] - at {tmp_path}/test.csv'}
+    expected_output = {'results': [f'[CSV] - at {tmp_path}/test.csv']}
     assert response.json() == expected_output
 
   def test_fake_source_from_query_path(self, tmp_path):
@@ -71,7 +71,7 @@ def test_fake_source_from_query_path(self, tmp_path):
     }
     response = client.post('/api/execute', json=request)
     assert response.status_code == fastapi.status.HTTP_200_OK
-    expected_output = {'result': f'[CSV] - at {tmp_path}/query.csv'}
+    expected_output = {'results': [f'[CSV] - at {tmp_path}/query.csv']}
     assert response.json() == expected_output
 
   def test_batch_fake_source_from_query_path(self, tmp_path):
@@ -96,7 +96,7 @@ def test_batch_fake_source_from_query_path(self, tmp_path):
     response = client.post('/api/execute:batch', json=request)
     assert response.status_code == fastapi.status.HTTP_200_OK
     expected_output = {
-      'result': [
+      'results': [
         f'[CSV] - at {tmp_path}/query.csv',
         f'[CSV] - at {tmp_path}/query.csv',
       ]