diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 1b752be8a..ea6ec524d 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -12,9 +12,9 @@ is surfaced the same way a bind failure is. * Backpressure uses a *blocking* ``queue.put`` instead of ``put_nowait``+drop: this avoids local drops while the worker pool falls behind and lets the - consumer pause naturally. Because ``aiokafka`` auto-commits fetched offsets, - the current crash semantics are still best-effort / at-most-once rather than - fully durable at-least-once delivery. + consumer pause naturally. Legacy single-record mode keeps auto-commit for + compatibility; opt-in micro-batch mode commits offsets only after a successful + workflow run and therefore provides at-least-once batch delivery. """ from __future__ import annotations @@ -82,6 +82,10 @@ _FETCH_MAX_BYTES = 8 * 1024 * 1024 _MAX_PARTITION_FETCH_BYTES = 4 * 1024 * 1024 _MAX_POLL_RECORDS = 16 +_DEFAULT_BATCH_MAX_RECORDS = 1 +_DEFAULT_BATCH_WAIT_MS = 100 +_MAX_BATCH_RECORDS = 1000 +_MAX_BATCH_WAIT_MS = 60_000 def _worker_count_for_trigger(trigger: TriggerDefinition) -> int: @@ -111,6 +115,52 @@ class _QueuedKafkaMessage: size_bytes: int +@dataclass(frozen=True) +class _QueuedKafkaRecord: + """Kafka record metadata retained until its workflow batch is complete.""" + + topic: str + partition: int + offset: int + timestamp: Optional[int] + key: Optional[bytes] + raw_value: Optional[bytes] + size_bytes: int + + +@dataclass(frozen=True) +class _QueuedKafkaBatch: + """Opt-in micro-batch acknowledged by a worker after workflow execution.""" + + records: tuple[_QueuedKafkaRecord, ...] + completed: asyncio.Future + + +def _batch_settings(data: Dict[str, Any], trigger: TriggerDefinition) -> tuple[int, int]: + """Resolve opt-in micro-batch settings while preserving legacy defaults.""" + + source = trigger.source if isinstance(trigger.source, dict) else {} + raw_max_records = data.get("batchMaxRecords") + if raw_max_records is None: + raw_max_records = source.get("batchMaxRecords") + try: + max_records = int(raw_max_records or _DEFAULT_BATCH_MAX_RECORDS) + except (TypeError, ValueError): + max_records = _DEFAULT_BATCH_MAX_RECORDS + max_records = min(_MAX_BATCH_RECORDS, max(_DEFAULT_BATCH_MAX_RECORDS, max_records)) + if max_records <= 1: + return _DEFAULT_BATCH_MAX_RECORDS, 0 + + raw_wait_ms = data.get("batchWaitMs") + if raw_wait_ms is None: + raw_wait_ms = source.get("batchWaitMs") + try: + wait_ms = int(raw_wait_ms or _DEFAULT_BATCH_WAIT_MS) + except (TypeError, ValueError): + wait_ms = _DEFAULT_BATCH_WAIT_MS + return max_records, min(_MAX_BATCH_WAIT_MS, max(1, wait_ms)) + + def _strip_execution_only_comments(value: Any) -> Any: if isinstance(value, list): return [_strip_execution_only_comments(item) for item in value] @@ -284,6 +334,8 @@ def _default_trigger_from_config(data: Dict[str, Any]) -> TriggerDefinition: "inputTopic": data.get("inputTopic") or "", "inputGroupId": data.get("inputGroupId") or "", "autoOffsetReset": data.get("autoOffsetReset") or "latest", + "batchMaxRecords": data.get("batchMaxRecords") or 1, + "batchWaitMs": data.get("batchWaitMs") or 0, }, "mapping": { str(data.get("inputKey") or "kafka_message"): "$.body", @@ -479,6 +531,7 @@ async def restart_workflow( log.warning("kafka.workflow_plan_failed", {"workflow_id": workflow_id, "error": str(exc)}) return self.get_consumer_status(workflow_id) group_id = str(data.get("inputGroupId") or "").strip() or f"flocks-consumer-{workflow_id}" + batch_max_records, batch_wait_ms = _batch_settings(data, trigger) configured_inputs = _strip_execution_only_comments(trigger.inputs if isinstance(trigger.inputs, dict) else {}) queue_capacity = _queue_size_for_trigger(trigger) @@ -500,6 +553,8 @@ async def restart_workflow( "broker": input_broker, "topic": input_topic, "groupId": group_id, + "batchMaxRecords": batch_max_records, + "batchWaitMs": batch_wait_ms, } # Trigger-configured worker pool, bounded by service safety caps. @@ -532,6 +587,8 @@ async def restart_workflow( queue, abort, ready, + batch_max_records=batch_max_records, + batch_wait_ms=batch_wait_ms, ), name=f"kafka-{workflow_id}", ) @@ -576,6 +633,9 @@ async def _consumer_loop( queue: asyncio.Queue, abort: asyncio.Event, ready: asyncio.Event, + *, + batch_max_records: int = _DEFAULT_BATCH_MAX_RECORDS, + batch_wait_ms: int = 0, ) -> None: try: from aiokafka import AIOKafkaConsumer @@ -592,19 +652,19 @@ async def _consumer_loop( await self._cleanup_runtime_resources(workflow_id) return + batch_enabled = batch_max_records > 1 consumer = AIOKafkaConsumer( topic, bootstrap_servers=broker, group_id=group_id, - # Auto-commit advances based on fetched progress, not worker - # completion. Backpressure narrows the crash window but current - # semantics remain best-effort / at-most-once. - enable_auto_commit=True, + # Preserve legacy auto-commit in single-record mode. Batch mode + # commits explicitly only after the workflow reports success. + enable_auto_commit=not batch_enabled, auto_offset_reset=auto_offset_reset if auto_offset_reset in ("latest", "earliest") else "latest", request_timeout_ms=_REQUEST_TIMEOUT_MS, fetch_max_bytes=_FETCH_MAX_BYTES, max_partition_fetch_bytes=_MAX_PARTITION_FETCH_BYTES, - max_poll_records=_MAX_POLL_RECORDS, + max_poll_records=batch_max_records if batch_enabled else _MAX_POLL_RECORDS, ) try: @@ -641,21 +701,85 @@ async def _consumer_loop( "broker": broker, "topic": topic, "groupId": group_id, + "batchMaxRecords": batch_max_records, + "batchWaitMs": batch_wait_ms if batch_enabled else 0, } ready.set() log.info("kafka.consumer_running", {"workflow_id": workflow_id, "topic": topic}) try: - async for msg in consumer: - if abort.is_set(): - break - raw_value = msg.value - queued = _QueuedKafkaMessage( - raw_value=raw_value, - size_bytes=len(raw_value) if raw_value is not None else 0, - ) - # Blocking put applies backpressure instead of dropping messages. - await queue.put(queued) + if batch_enabled: + while not abort.is_set(): + records_by_partition: Dict[Any, List[Any]] = {} + record_count = 0 + deadline = asyncio.get_running_loop().time() + (batch_wait_ms / 1000) + while record_count < batch_max_records: + remaining_ms = max( + 0, + int((deadline - asyncio.get_running_loop().time()) * 1000), + ) + fetched = await consumer.getmany( + timeout_ms=remaining_ms, + max_records=batch_max_records - record_count, + ) + if not fetched: + break + for topic_partition, partition_records in fetched.items(): + if not partition_records: + continue + records_by_partition.setdefault(topic_partition, []).extend( + partition_records + ) + record_count += len(partition_records) + if asyncio.get_running_loop().time() >= deadline: + break + records = tuple( + _QueuedKafkaRecord( + topic=str(msg.topic), + partition=int(msg.partition), + offset=int(msg.offset), + timestamp=msg.timestamp, + key=msg.key, + raw_value=msg.value, + size_bytes=len(msg.value) if msg.value is not None else 0, + ) + for partition_records in records_by_partition.values() + for msg in partition_records + ) + if not records: + continue + + completed = asyncio.get_running_loop().create_future() + await queue.put(_QueuedKafkaBatch(records=records, completed=completed)) + succeeded = bool(await completed) + if not succeeded: + raise RuntimeError("kafka_batch_workflow_failed") + + commit_offsets = { + topic_partition: partition_records[-1].offset + 1 + for topic_partition, partition_records in records_by_partition.items() + if partition_records + } + if commit_offsets: + await consumer.commit(commit_offsets) + current_status = self._status.get(workflow_id) or {} + self._status[workflow_id] = { + **current_status, + "lastBatchSize": len(records), + "processedRecords": int(current_status.get("processedRecords") or 0) + + len(records), + } + else: + async for msg in consumer: + if abort.is_set(): + break + raw_value = msg.value + queued = _QueuedKafkaMessage( + raw_value=raw_value, + size_bytes=len(raw_value) if raw_value is not None else 0, + ) + # Blocking put applies backpressure instead of dropping messages. + await queue.put(queued) except asyncio.CancelledError: raise except Exception as exc: @@ -700,10 +824,13 @@ async def _worker_loop( continue except asyncio.CancelledError: return + batch = msg if isinstance(msg, _QueuedKafkaBatch) else None try: if isinstance(msg, _QueuedKafkaMessage): msg = _decode_message(msg.raw_value) - await self._trigger_workflow( + elif batch is not None: + msg = [_decode_message(record.raw_value) for record in batch.records] + succeeded = await self._trigger_workflow( workflow_id, workflow_plan, msg, @@ -713,9 +840,15 @@ async def _worker_loop( source=source, generation_cancel_event=run_cancel_event, ) + if batch is not None and not batch.completed.done(): + batch.completed.set_result(bool(succeeded)) except asyncio.CancelledError: + if batch is not None and not batch.completed.done(): + batch.completed.set_result(False) return except Exception as exc: + if batch is not None and not batch.completed.done(): + batch.completed.set_result(False) log.warning( "kafka.worker_dispatch_failed", {"workflow_id": workflow_id, "error": str(exc)}, @@ -732,7 +865,7 @@ async def _trigger_workflow( trigger: Optional[TriggerDefinition] = None, source: Optional[str] = None, generation_cancel_event: Optional[threading.Event] = None, - ) -> None: + ) -> bool: run_cancel_event = generation_cancel_event or threading.Event() trigger = trigger or TriggerDefinition.model_validate( { @@ -864,7 +997,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: "trigger_type": "kafka", }, } - await execute_with_hooks( + dispatch_result = await execute_with_hooks( action_payload, lambda: self._dispatcher.dispatch( trigger=trigger, @@ -874,11 +1007,18 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: before=HookPipeline.run_ingress_before, after=HookPipeline.run_ingress_after, ) + if not isinstance(dispatch_result, dict): + return False + if not dispatch_result.get("executed"): + return True + execution = dispatch_result.get("result") + return isinstance(execution, dict) and execution.get("status") == "success" except TriggerDispatchError as exc: log.warning( "kafka.trigger_dispatch_failed", {"workflow_id": workflow_id, "trigger_id": trigger.id, "error": str(exc)}, ) + return False default_manager = KafkaManager() diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index b838c543a..4876a46a0 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -2703,6 +2703,8 @@ class KafkaConfigRequest(BaseModel): inputGroupId: Optional[str] = None inputKey: str = "kafka_message" autoOffsetReset: str = "latest" + batchMaxRecords: int = Field(1, ge=1, le=1000) + batchWaitMs: int = Field(0, ge=0, le=60_000) inputs: Dict[str, Any] = Field(default_factory=dict) @@ -3401,6 +3403,8 @@ async def save_kafka_config(workflow_id: str, req: KafkaConfigRequest): "inputGroupId": req.inputGroupId, "inputKey": req.inputKey, "autoOffsetReset": req.autoOffsetReset, + "batchMaxRecords": req.batchMaxRecords, + "batchWaitMs": req.batchWaitMs, "inputs": _strip_execution_only_comments(req.inputs), "updatedAt": int(time.time() * 1000), } @@ -3415,6 +3419,8 @@ async def save_kafka_config(workflow_id: str, req: KafkaConfigRequest): "inputTopic": req.inputTopic or "", "inputGroupId": req.inputGroupId or "", "autoOffsetReset": req.autoOffsetReset, + "batchMaxRecords": req.batchMaxRecords, + "batchWaitMs": req.batchWaitMs, }, "mapping": { req.inputKey or "kafka_message": "$.body", diff --git a/flocks/tool/task/workflow_config_manage.py b/flocks/tool/task/workflow_config_manage.py index 90aa078ed..89a601e53 100644 --- a/flocks/tool/task/workflow_config_manage.py +++ b/flocks/tool/task/workflow_config_manage.py @@ -291,6 +291,8 @@ def _normalize_runtime_config( "inputGroupId": req.inputGroupId, "inputKey": req.inputKey, "autoOffsetReset": req.autoOffsetReset, + "batchMaxRecords": req.batchMaxRecords, + "batchWaitMs": req.batchWaitMs, "inputs": routes._strip_execution_only_comments(req.inputs), }, current, diff --git a/flocks/workflow/triggers/compat.py b/flocks/workflow/triggers/compat.py index faf5f5cf4..07a46d620 100644 --- a/flocks/workflow/triggers/compat.py +++ b/flocks/workflow/triggers/compat.py @@ -70,6 +70,8 @@ def legacy_kafka_trigger_from_config(config: Optional[Dict[str, Any]]) -> Option "inputTopic": config.get("inputTopic") or "", "inputGroupId": config.get("inputGroupId") or "", "autoOffsetReset": config.get("autoOffsetReset") or "latest", + "batchMaxRecords": config.get("batchMaxRecords") or 1, + "batchWaitMs": config.get("batchWaitMs") or 0, }, "mapping": { str(config.get("inputKey") or "kafka_message"): "$.body", @@ -124,6 +126,8 @@ def kafka_trigger_to_legacy_config(workflow_id: str, trigger: TriggerDefinition) "inputGroupId": source.get("inputGroupId") or "", "inputKey": input_key, "autoOffsetReset": source.get("autoOffsetReset") or "latest", + "batchMaxRecords": source.get("batchMaxRecords") or 1, + "batchWaitMs": source.get("batchWaitMs") or 0, "inputs": dict(trigger.inputs or {}), "updatedAt": trigger.updatedAt, } diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index a3c2b0432..15e0ea237 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -25,6 +25,10 @@ from flocks.ingest.kafka import manager as kafka_manager from flocks.workflow import execution_store +from flocks.workflow.triggers.compat import ( + kafka_trigger_to_legacy_config, + legacy_kafka_trigger_from_config, +) from flocks.workflow.triggers.models import TriggerDefinition @@ -223,6 +227,275 @@ async def _fake_trigger(workflow_id, workflow_json, msg, input_key, producer=Non assert captured == [{"ok": True}] +@pytest.mark.asyncio +async def test_worker_decodes_batch_as_plain_payload_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Batch delivery must stay invisible to a batch-aware workflow.""" + + manager = kafka_manager.KafkaManager() + workflow_id = "test-wf-batch-queue" + queue: asyncio.Queue = asyncio.Queue(maxsize=8) + abort = asyncio.Event() + completed = asyncio.get_running_loop().create_future() + captured: list[list[dict]] = [] + trigger = TriggerDefinition.model_validate( + {"id": "kafka-default", "type": "kafka", "mapping": {"kafka_message": "$.body"}} + ) + + async def _fake_trigger(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + captured.append(args[2]) + abort.set() + return True + + monkeypatch.setattr(manager, "_trigger_workflow", _fake_trigger) + queue.put_nowait( + kafka_manager._QueuedKafkaBatch( # noqa: SLF001 + records=( + kafka_manager._QueuedKafkaRecord( # noqa: SLF001 + topic="topic-a", + partition=0, + offset=4, + timestamp=1000, + key=None, + raw_value=b'{"id": 1}', + size_bytes=len(b'{"id": 1}'), + ), + kafka_manager._QueuedKafkaRecord( # noqa: SLF001 + topic="topic-a", + partition=0, + offset=5, + timestamp=1001, + key=None, + raw_value=b'{"id": 2}', + size_bytes=len(b'{"id": 2}'), + ), + ), + completed=completed, + ) + ) + + worker = asyncio.create_task( + manager._worker_loop(workflow_id, {}, trigger, {}, queue, abort, "topic-a"), + name="test-worker-batch-queue", + ) + await asyncio.wait_for(worker, timeout=1.0) + + assert captured == [[{"id": 1}, {"id": 2}]] + assert completed.result() is True + + +@pytest.mark.asyncio +async def test_consumer_batch_commits_offsets_only_after_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An enabled micro-batch must commit each partition after workflow success.""" + + manager = kafka_manager.KafkaManager() + queue: asyncio.Queue = asyncio.Queue(maxsize=2) + abort = asyncio.Event() + ready = asyncio.Event() + topic_partition = ("topic-a", 0) + consumer_instances = [] + + class _Consumer: + def __init__(self, *topics, **kwargs): # noqa: ANN002, ANN003 + self.kwargs = kwargs + self.commits = [] + self.getmany_calls = [] + self.stopped = False + consumer_instances.append(self) + + async def start(self) -> None: + return None + + async def stop(self) -> None: + self.stopped = True + + async def getmany(self, *, timeout_ms: int, max_records: int): + self.getmany_calls.append((timeout_ms, max_records)) + if len(self.getmany_calls) == 1: + return { + topic_partition: [ + SimpleNamespace( + topic="topic-a", + partition=0, + offset=10, + timestamp=1000, + key=None, + value=b'{"id": 1}', + ), + ] + } + if len(self.getmany_calls) == 2: + return { + topic_partition: [ + SimpleNamespace( + topic="topic-a", + partition=0, + offset=11, + timestamp=1001, + key=None, + value=b'{"id": 2}', + ), + ] + } + return {} + + async def commit(self, offsets) -> None: # noqa: ANN001 + self.commits.append(offsets) + + monkeypatch.setitem(sys.modules, "aiokafka", SimpleNamespace(AIOKafkaConsumer=_Consumer)) + + consumer_task = asyncio.create_task( + manager._consumer_loop( # noqa: SLF001 + "wf-batch", + "localhost:9092", + "topic-a", + "group-a", + "latest", + queue, + abort, + ready, + batch_max_records=16, + batch_wait_ms=100, + ) + ) + + await asyncio.wait_for(ready.wait(), timeout=1.0) + batch = await asyncio.wait_for(queue.get(), timeout=1.0) + assert isinstance(batch, kafka_manager._QueuedKafkaBatch) # noqa: SLF001 + assert consumer_instances[0].commits == [] + + batch.completed.set_result(True) + abort.set() + await asyncio.wait_for(consumer_task, timeout=1.0) + + consumer = consumer_instances[0] + assert consumer.kwargs["enable_auto_commit"] is False + assert 0 < consumer.getmany_calls[0][0] <= 100 + assert consumer.getmany_calls[0][1] == 16 + assert consumer.getmany_calls[1][1] == 15 + assert consumer.getmany_calls[2][1] == 14 + assert consumer.commits == [{topic_partition: 12}] + assert consumer.stopped is True + + +@pytest.mark.asyncio +async def test_consumer_batch_failure_does_not_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed micro-batch must stop without advancing committed offsets.""" + + manager = kafka_manager.KafkaManager() + queue: asyncio.Queue = asyncio.Queue(maxsize=2) + abort = asyncio.Event() + ready = asyncio.Event() + consumer_instances = [] + + class _Consumer: + def __init__(self, *topics, **kwargs): # noqa: ANN002, ANN003 + self.commits = [] + consumer_instances.append(self) + + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + async def getmany(self, *, timeout_ms: int, max_records: int): + return { + ("topic-a", 0): [ + SimpleNamespace( + topic="topic-a", + partition=0, + offset=20, + timestamp=None, + key=None, + value=b'{"id": 1}', + ) + ] + } + + async def commit(self, offsets) -> None: # noqa: ANN001 + self.commits.append(offsets) + + monkeypatch.setitem(sys.modules, "aiokafka", SimpleNamespace(AIOKafkaConsumer=_Consumer)) + + consumer_task = asyncio.create_task( + manager._consumer_loop( # noqa: SLF001 + "wf-batch-failure", + "localhost:9092", + "topic-a", + "group-a", + "latest", + queue, + abort, + ready, + batch_max_records=8, + batch_wait_ms=50, + ) + ) + + await asyncio.wait_for(ready.wait(), timeout=1.0) + batch = await asyncio.wait_for(queue.get(), timeout=1.0) + batch.completed.set_result(False) + await asyncio.wait_for(consumer_task, timeout=1.0) + + assert consumer_instances[0].commits == [] + assert manager.get_consumer_status("wf-batch-failure")["state"] == "failed" + + +@pytest.mark.asyncio +async def test_consumer_without_batch_config_keeps_legacy_auto_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Missing batch settings must retain the existing single-record path.""" + + manager = kafka_manager.KafkaManager() + queue: asyncio.Queue = asyncio.Queue(maxsize=2) + abort = asyncio.Event() + ready = asyncio.Event() + consumer_instances = [] + + class _Consumer: + def __init__(self, *topics, **kwargs): # noqa: ANN002, ANN003 + self.kwargs = kwargs + self.stopped = False + consumer_instances.append(self) + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def start(self) -> None: + return None + + async def stop(self) -> None: + self.stopped = True + + monkeypatch.setitem(sys.modules, "aiokafka", SimpleNamespace(AIOKafkaConsumer=_Consumer)) + + await manager._consumer_loop( # noqa: SLF001 + "wf-single", + "localhost:9092", + "topic-a", + "group-a", + "latest", + queue, + abort, + ready, + ) + + consumer = consumer_instances[0] + assert consumer.kwargs["enable_auto_commit"] is True + assert consumer.kwargs["max_poll_records"] == kafka_manager._MAX_POLL_RECORDS # noqa: SLF001 + assert consumer.stopped is True + + def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: trigger = TriggerDefinition.model_validate( { @@ -538,6 +811,29 @@ def test_decode_message_variants() -> None: assert kafka_manager._decode_message(b"\xff\xfe") == "fffe" +def test_kafka_batch_config_survives_legacy_trigger_round_trip() -> None: + config = { + "workflowId": "wf-batch", + "enabled": True, + "inputBroker": "broker:9092", + "inputTopic": "alerts", + "inputGroupId": "flocks-alerts", + "inputKey": "kafka_message", + "autoOffsetReset": "earliest", + "batchMaxRecords": 64, + "batchWaitMs": 250, + } + + trigger = legacy_kafka_trigger_from_config(config) + + assert trigger is not None + assert trigger.source["batchMaxRecords"] == 64 + assert trigger.source["batchWaitMs"] == 250 + restored = kafka_trigger_to_legacy_config("wf-batch", trigger) + assert restored["batchMaxRecords"] == 64 + assert restored["batchWaitMs"] == 250 + + @pytest.mark.asyncio async def test_trigger_workflow_compacts_kafka_execution_record( monkeypatch: pytest.MonkeyPatch, @@ -599,13 +895,14 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) - await manager._trigger_workflow( + succeeded = await manager._trigger_workflow( "wf-compact", {"start": "receive_alert", "nodes": [], "edges": []}, {"alarmData": "x" * 50_000}, "kafka_message", ) + assert succeeded is True assert captured_input_params["kafka_message"]["alarmData"]["_type"] == "string" assert captured_input_params["kafka_message"]["alarmData"]["chars"] == 50_000 assert captured_run_kwargs["run_id"] == "exec-compact" @@ -657,7 +954,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) - await manager._trigger_workflow( + succeeded = await manager._trigger_workflow( "wf-merge", {"start": "receive_alert", "nodes": [], "edges": []}, {"alarmData": {"id": 1}}, @@ -670,6 +967,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 }, ) + assert succeeded is True assert captured_run_kwargs["inputs"]["kafka_message"] == {"alarmData": {"id": 1}} assert captured_run_kwargs["inputs"]["kafka_output_enabled"] is True assert captured_run_kwargs["inputs"]["kafka_output_topic"] == "topic_soc_flocks_result_log" @@ -725,7 +1023,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 } ) - await manager._trigger_workflow( + succeeded = await manager._trigger_workflow( "wf-orders", {"start": "receive_alert", "nodes": [], "edges": []}, {"order": {"id": 7, "region": "cn"}}, @@ -734,6 +1032,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 source="orders-topic", ) + assert succeeded is True assert captured_run_kwargs["inputs"]["order_id"] == 7 assert captured_run_kwargs["inputs"]["region"] == "cn" assert captured_run_kwargs["inputs"]["pipeline"] == "orders" @@ -747,7 +1046,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert recorded_exec_data["triggerSource"] == "orders-topic" captured_run_kwargs.clear() - await manager._trigger_workflow( + filtered = await manager._trigger_workflow( "wf-orders", {"start": "receive_alert", "nodes": [], "edges": []}, {"order": {"id": 8, "region": "us"}}, @@ -755,4 +1054,41 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 trigger=trigger, source="orders-topic", ) + assert filtered is True assert captured_run_kwargs == {} + + +@pytest.mark.asyncio +async def test_trigger_workflow_reports_failed_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = kafka_manager.KafkaManager() + + async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + return {"id": "exec-failed", "workflowId": workflow_id, "inputParams": input_params} + + async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + return None + + def _fake_run_workflow(**kwargs): # noqa: ANN003 + return SimpleNamespace( + status="FAILED", + error="workflow failed", + outputs={}, + history=[], + last_node_id="receive_alert", + steps=1, + ) + + monkeypatch.setattr(kafka_manager, "create_execution_record", _fake_create_execution_record) + monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) + monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) + + succeeded = await manager._trigger_workflow( + "wf-failed", + {"start": "receive_alert", "nodes": [], "edges": []}, + [{"id": 1}], + "kafka_message", + ) + + assert succeeded is False diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 588fe227b..53bd9889a 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -335,6 +335,8 @@ async def _fake_persist(workflow_id: str, workflow_data: dict, triggers: list) - inputTopic="workflow-input", inputGroupId="wf-group", inputKey="kafka_message", + batchMaxRecords=16, + batchWaitMs=100, inputs={ "_comment": "remove me", "kafka_output_enabled": True, @@ -354,6 +356,8 @@ async def _fake_persist(workflow_id: str, workflow_data: dict, triggers: list) - assert saved_config["inputTopic"] == "workflow-input" assert saved_config["inputGroupId"] == "wf-group" assert saved_config["inputKey"] == "kafka_message" + assert saved_config["batchMaxRecords"] == 16 + assert saved_config["batchWaitMs"] == 100 assert saved_config["inputs"] == { "kafka_output_enabled": True, "kafka_output_topic": "topic_soc_flocks_result_log", diff --git a/webui/src/api/workflow.ts b/webui/src/api/workflow.ts index dc6ed9230..9a8767555 100644 --- a/webui/src/api/workflow.ts +++ b/webui/src/api/workflow.ts @@ -389,6 +389,8 @@ export interface KafkaConfig { inputGroupId?: string; inputKey?: string; autoOffsetReset?: string; + batchMaxRecords?: number; + batchWaitMs?: number; inputs?: Record; updatedAt?: number; } @@ -403,6 +405,10 @@ export interface KafkaConsumerStatus { queueSize?: number; queueCapacity?: number; workerCount?: number; + batchMaxRecords?: number; + batchWaitMs?: number; + lastBatchSize?: number; + processedRecords?: number; } export interface WorkflowPollerConfig { @@ -578,6 +584,8 @@ export const workflowAPI = { inputGroupId?: string; inputKey?: string; autoOffsetReset?: string; + batchMaxRecords?: number; + batchWaitMs?: number; inputs?: Record; }) => client.post<{ ok: boolean; consumer?: KafkaConsumerStatus }>( diff --git a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx index 31a50cb46..41e555b95 100644 --- a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx +++ b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx @@ -670,6 +670,26 @@ function KafkaTriggerFields({ +
+ + onSourceChange({ batchMaxRecords: Math.max(1, Number(e.target.value) || 1) })} + /> + + + onSourceChange({ batchWaitMs: Math.max(0, Number(e.target.value) || 0) })} + /> + +
); } @@ -883,6 +903,8 @@ function createTriggerDraft( inputTopic: `${workflowId}.events`, inputGroupId: `${workflowId}-group`, autoOffsetReset: 'latest', + batchMaxRecords: 1, + batchWaitMs: 0, }, mapping: { kafka_message: '$.body' }, inputs: workflowSampleInputs,