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

Author: hsheth2

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

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(

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,