[executors] feat: add support for batch execution of queries

Author: AndreyMarkinPPC

libs/garf_executors/garf_executors/entrypoints/server.py

@@ -14,6 +14,8 @@
 
 """FastAPI endpoint for executing queries."""
 
+from concurrent import futures
+
 import fastapi
 import pydantic
 import uvicorn
@@ -37,7 +39,7 @@ class ApiExecutorRequest(pydantic.BaseModel):
   source: str
   title: str | None = None
   query: str | None = None
-  query_path: str | None = None
+  query_path: str | list[str] | None = None
   context: garf_executors.api_executor.ApiExecutionContext
 
   @pydantic.model_validator(mode='after')
@@ -49,7 +51,7 @@ def check_query_specified(self):
     return self
 
   def model_post_init(self, __context__) -> None:
-    if self.query_path:
+    if self.query_path and isinstance(self.query_path, str):
       self.query = reader.FileReader().read(self.query_path)
     if not self.title:
       self.title = str(self.query_path)
@@ -76,6 +78,36 @@ async def execute(request: ApiExecutorRequest) -> dict[str, str]:
   )
 
 
+@router.post('/execute:batch')
+async def execute_batch(request: ApiExecutorRequest) -> dict[str, str]:
+  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(
+    concrete_api_fetcher(**request.context.fetcher_parameters)
+  )
+  file_reader = reader.FileReader()
+  results = []
+  with futures.ThreadPoolExecutor() as executor:
+    future_to_query = {
+      executor.submit(
+        query_executor.execute,
+        file_reader.read(query),
+        query,
+        request.context,
+      ): query
+      for query in request.query_path
+    }
+    for future in futures.as_completed(future_to_query):
+      results.append(future.result())
+
+  return fastapi.responses.JSONResponse(
+    content=fastapi.encoders.jsonable_encoder({'result': results})
+  )
+
+
 if __name__ == '__main__':
   app = fastapi.FastAPI()
   app.include_router(router)

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

@@ -73,3 +73,32 @@ def test_fake_source_from_query_path(self, tmp_path):
     assert response.status_code == fastapi.status.HTTP_200_OK
     expected_output = {'result': f'[CSV] - at {tmp_path}/query.csv'}
     assert response.json() == expected_output
+
+  def test_batch_fake_source_from_query_path(self, tmp_path):
+    query_path = tmp_path / 'query.sql'
+    with pathlib.Path.open(query_path, 'w', encoding='utf-8') as f:
+      f.write(self.query)
+    fake_data = _SCRIPT_PATH / 'test.json'
+    request = {
+      'source': 'fake',
+      'query_path': [
+        str(query_path),
+        str(query_path),
+      ],
+      'context': {
+        'fetcher_parameters': {
+          'data_location': str(fake_data),
+        },
+        'writer': 'csv',
+        'writer_parameters': {'destination_folder': str(tmp_path)},
+      },
+    }
+    response = client.post('/api/execute:batch', json=request)
+    assert response.status_code == fastapi.status.HTTP_200_OK
+    expected_output = {
+      'result': [
+        f'[CSV] - at {tmp_path}/query.csv',
+        f'[CSV] - at {tmp_path}/query.csv',
+      ]
+    }
+    assert response.json() == expected_output