pass input data directly to py workflow (#2055)

514-labs · Feb 21, 2025 · 5738bcc · 5738bcc
1 parent 3973558
commit 5738bcc
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/apps/framework-cli/src/framework/python/templates.rs b/apps/framework-cli/src/framework/python/templates.rs
@@ -214,7 +214,7 @@ block = Blocks(teardown=[DROP_MV_QUERY], setup=[MV_QUERY])
 pub static PYTHON_BASE_SCRIPT_TEMPLATE: &str = r#"from moose_lib import task
 
 @task()
-def {{name}}(data: dict):  # The name of your script
+def {{name}}(input: dict):  # The name of your script
     """
     Description of what this script does
     """

diff --git a/apps/framework-cli/src/framework/python/wrappers/scripts/python_worker_wrapper/activity.py b/apps/framework-cli/src/framework/python/wrappers/scripts/python_worker_wrapper/activity.py
@@ -57,12 +57,12 @@ async def dynamic_activity(execution_input: ScriptExecutionInput) -> WorkflowSte
             log.info(f"Task received input: {execution_input.input_data}")
 
             # Pass the input data directly if it exists
-            input_data = execution_input.input_data if execution_input.input_data else {}
+            input_data = execution_input.input_data.get('data', execution_input.input_data) if execution_input.input_data else {}
             log.info(f"Processed input_data for task: {input_data}")
             if asyncio.iscoroutinefunction(task_func):
-                result = await task_func(data=input_data)
+                result = await task_func(input=input_data)
             else:
-                result = task_func(data=input_data)
+                result = task_func(input=input_data)
 
             # Validate and encode result
             if not isinstance(result, dict):

diff --git a/apps/framework-cli/src/infrastructure/olap/clickhouse/inserter.rs b/apps/framework-cli/src/infrastructure/olap/clickhouse/inserter.rs
@@ -138,7 +138,7 @@ impl<C: ClickHouseClientTrait + 'static> Inserter<C> {
 
             let mut queue = self.queue.lock().await;
 
-            if queue.is_empty() || queue.front().map_or(true, |batch| batch.records.is_empty()) {
+            if queue.is_empty() || queue.front().is_none_or(|batch| batch.records.is_empty()) {
                 continue;
             }