[executors] feat: add execute_batch method for executing queries in parallel

* Add common Executor class with execute_batch method
* Inherit Api, Bq and SqlAlchemy executors from it
* Replace parallel execution in CLI and server with execute_batch call

Author: AndreyMarkinPPC

libs/garf_executors/garf_executors/api_executor.py

@@ -23,7 +23,7 @@
 import logging
 
 from garf_core import report_fetcher
-from garf_executors import exceptions, execution_context
+from garf_executors import exceptions, execution_context, executor
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +34,7 @@ class ApiExecutionContext(execution_context.ExecutionContext):
   writer: str = 'console'
 
 
-class ApiQueryExecutor:
+class ApiQueryExecutor(executor.Executor):
   """Gets data from API and writes them to local/remote storage.
 
   Attributes:

libs/garf_executors/garf_executors/bq_executor.py

@@ -28,7 +28,7 @@
 from google.cloud import exceptions as google_cloud_exceptions
 
 from garf_core import query_editor, report
-from garf_executors import exceptions, execution_context
+from garf_executors import exceptions, execution_context, executor
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +37,7 @@ class BigQueryExecutorError(exceptions.GarfExecutorError):
   """Error when BigQueryExecutor fails to run query."""
 
 
-class BigQueryExecutor(query_editor.TemplateProcessorMixin):
+class BigQueryExecutor(executor.Executor, query_editor.TemplateProcessorMixin):
   """Handles query execution in BigQuery.
 
   Attributes:

libs/garf_executors/garf_executors/entrypoints/cli.py

@@ -21,10 +21,9 @@
 
 import argparse
 import sys
-from concurrent import futures
 
 import garf_executors
-from garf_executors import bq_executor, exceptions, sql_executor
+from garf_executors import exceptions
 from garf_executors.entrypoints import utils
 from garf_io import reader
 
@@ -86,18 +85,8 @@ def main():
   )
   if args.parallel_queries:
     logger.info('Running queries in parallel')
-    with futures.ThreadPoolExecutor(args.parallel_threshold) as executor:
-      future_to_query = {
-        executor.submit(
-          query_executor.execute,
-          reader_client.read(query),
-          query,
-          context,
-        ): query
-        for query in args.query
-      }
-      for future in futures.as_completed(future_to_query):
-        future.result()
+    batch = {query: reader_client.read(query) for query in args.query}
+    query_executor.execute_batch(batch, context, args.parallel_queries)
   else:
     logger.info('Running queries sequentially')
     for query in args.query:

libs/garf_executors/garf_executors/entrypoints/server.py

@@ -14,15 +14,12 @@
 
 """FastAPI endpoint for executing queries."""
 
-from concurrent import futures
-
 import fastapi
 import pydantic
 import uvicorn
 
 import garf_executors
 from garf_executors import exceptions
-from garf_executors.entrypoints import utils
 from garf_io import reader
 
 
@@ -85,20 +82,9 @@ async def execute_batch(request: ApiExecutorRequest) -> ApiExecutorResponse:
   query_executor = garf_executors.setup_executor(
     request.source, 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())
+  reader_client = reader.FileReader()
+  batch = {query: reader_client.read(query) for query in request.query_path}
+  results = query_executor.execute_batch(batch, request.context)
   return ApiExecutorResponse(results=results)
 
 

libs/garf_executors/garf_executors/executor.py

@@ -0,0 +1,54 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines common functionality between executors."""
+
+from concurrent import futures
+
+from garf_executors import execution_context
+
+
+class Executor:
+  """Defines common functionality between executors."""
+
+  def execute_batch(
+    self,
+    batch: dict[str, str],
+    context: execution_context.ExecutionContext,
+    parallel_threshold: int = 10,
+  ) -> list[str]:
+    """Executes batch of queries for a common context.
+
+    Args:
+      batch: Mapping between query_title and its text.
+      context: Execution context.
+      parallel_threshold: Number of queries to execute in parallel.
+
+    Returns:
+      Results of execution.
+    """
+    results = []
+    with futures.ThreadPoolExecutor(max_workers=parallel_threshold) as executor:
+      future_to_query = {
+        executor.submit(
+          self.execute,
+          query,
+          title,
+          context,
+        ): query
+        for title, query in batch.items()
+      }
+      for future in futures.as_completed(future_to_query):
+        results.append(future.result())
+    return results

libs/garf_executors/garf_executors/sql_executor.py

@@ -29,7 +29,7 @@
 import pandas as pd
 
 from garf_core import query_editor, report
-from garf_executors import exceptions, execution_context
+from garf_executors import exceptions, execution_context, executor
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +38,9 @@ class SqlAlchemyQueryExecutorError(exceptions.GarfExecutorError):
   """Error when SqlAlchemyQueryExecutor fails to run query."""
 
 
-class SqlAlchemyQueryExecutor(query_editor.TemplateProcessorMixin):
+class SqlAlchemyQueryExecutor(
+  executor.Executor, query_editor.TemplateProcessorMixin
+):
   """Handles query execution via SqlAlchemy.
 
   Attributes:
@@ -57,6 +59,10 @@ def __init__(self, engine: sqlalchemy.engine.base.Engine) -> None:
   def from_connection_string(
     cls, connection_string: str
   ) -> SqlAlchemyQueryExecutor:
+    """Creates executor from SqlAlchemy connection string.
+
+    https://docs.sqlalchemy.org/en/20/core/engines.html
+    """
     engine = sqlalchemy.create_engine(connection_string)
     return cls(engine)
 

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

@@ -75,15 +75,18 @@ def test_fake_source_from_query_path(self, tmp_path):
     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:
+    query_path1 = tmp_path / 'query1.sql'
+    with pathlib.Path.open(query_path1, 'w', encoding='utf-8') as f:
+      f.write(self.query)
+    query_path2 = tmp_path / 'query2.sql'
+    with pathlib.Path.open(query_path2, '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),
+        str(query_path1),
+        str(query_path2),
       ],
       'context': {
         'fetcher_parameters': {
@@ -97,8 +100,8 @@ def test_batch_fake_source_from_query_path(self, tmp_path):
     assert response.status_code == fastapi.status.HTTP_200_OK
     expected_output = {
       'results': [
-        f'[CSV] - at {tmp_path}/query.csv',
-        f'[CSV] - at {tmp_path}/query.csv',
+        f'[CSV] - at {tmp_path}/query1.csv',
+        f'[CSV] - at {tmp_path}/query2.csv',
       ]
     }
     assert response.json() == expected_output