feat(ingest): allow max_workers=1 with ASYNC_BATCH rest sink (#12088)

datahub-project · Dec 10, 2024 · d953718 · d953718
1 parent 00f0ee8
commit d953718
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
@@ -65,11 +65,11 @@ class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
 
     # These only apply in async modes.
-    max_threads: int = _DEFAULT_REST_SINK_MAX_THREADS
-    max_pending_requests: int = 2000
+    max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
+    max_pending_requests: pydantic.PositiveInt = 2000
 
     # Only applies in async batch mode.
-    max_per_batch: int = 100
+    max_per_batch: pydantic.PositiveInt = 100
 
 
 @dataclasses.dataclass

diff --git a/metadata-ingestion/src/datahub/utilities/partition_executor.py b/metadata-ingestion/src/datahub/utilities/partition_executor.py
@@ -268,7 +268,7 @@ def __init__(
         self.process_batch = process_batch
         self.min_process_interval = min_process_interval
         self.read_from_pending_interval = read_from_pending_interval
-        assert self.max_workers > 1
+        assert self.max_workers >= 1
 
         self._state_lock = threading.Lock()
         self._executor = ThreadPoolExecutor(

diff --git a/metadata-ingestion/tests/unit/utilities/test_partition_executor.py b/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
@@ -80,7 +80,8 @@ def task(id: str) -> str:
         assert len(done_tasks) == 16
 
 
-def test_batch_partition_executor_sequential_key_execution():
+@pytest.mark.parametrize("max_workers", [1, 2, 10])
+def test_batch_partition_executor_sequential_key_execution(max_workers: int) -> None:
     executing_tasks = set()
     done_tasks = set()
     done_task_batches = set()
@@ -99,7 +100,7 @@ def process_batch(batch):
         done_task_batches.add(tuple(id for _, id in batch))
 
     with BatchPartitionExecutor(
-        max_workers=2,
+        max_workers=max_workers,
         max_pending=10,
         max_per_batch=2,
         process_batch=process_batch,