moose-lib 0.6.18__py3-none-any.whl → 0.6.20__py3-none-any.whl

moose_lib/__init__.py

@@ -4,8 +4,6 @@ from .blocks import *
 
 from .commons import *
 
-from .tasks import *
-
 from .data_models import *
 
 from .dmv2 import *

moose_lib/dmv2/__init__.py

@@ -65,6 +65,7 @@ from .materialized_view import (
 )
 
 from .workflow import (
+    TaskContext,
     TaskConfig,
     Task,
     WorkflowConfig,
@@ -140,6 +141,7 @@ __all__ = [
     'MaterializedView',
 
     # Workflow
+    'TaskContext',
     'TaskConfig',
     'Task',
     'WorkflowConfig',

moose_lib/dmv2/workflow.py

@@ -11,16 +11,17 @@ from pydantic import BaseModel
 from .types import TypedMooseResource, T_none, U_none
 from ._registry import _workflows
 
-type TaskRunFunc[T_none, U_none] = Union[
-    # Case 1: No input, no output
-    Callable[[], None],
-    # Case 2: No input, with output
-    Callable[[], Union[U_none, Awaitable[U_none]]],
-    # Case 3: With input, no output
-    Callable[[T_none], None],
-    # Case 4: With input, with output
-    Callable[[T_none], Union[U_none, Awaitable[U_none]]]
-]
+@dataclasses.dataclass
+class TaskContext(Generic[T_none]):
+    """Context object passed to task handlers.
+
+    - When a task declares an input model `T`, `input` is of type `T` (not Optional).
+    - For no-input tasks (`T` is `None`), `input` is exactly `None`.
+    """
+    state: Dict[str, Any]
+    input: T_none
+
+type TaskRunFunc[T_none, U_none] = Callable[[TaskContext[T_none]], Union[U_none, Awaitable[U_none]]]
 
 @dataclasses.dataclass
 class TaskConfig(Generic[T_none, U_none]):
@@ -28,12 +29,16 @@ class TaskConfig(Generic[T_none, U_none]):
 
     Attributes:
         run: The handler function that executes the task logic.
+             Signature: run(context: TaskContext[T]) -> U
         on_complete: Optional list of tasks to run after this task completes.
-        timeout: Optional timeout string (e.g. "5m", "1h").
+        on_cancel: Optional function to call when the task is cancelled.
+                  Signature: on_cancel(context: TaskContext[T]) -> None
+        timeout: Optional timeout string (e.g. "5m", "1h", "never").
         retries: Optional number of retry attempts.
     """
     run: TaskRunFunc[T_none, U_none]
     on_complete: Optional[list["Task[U_none, Any]"]] = None
+    on_cancel: Optional[Callable[[TaskContext[T_none]], Union[None, Awaitable[None]]]] = None
     timeout: Optional[str] = None
     retries: Optional[int] = None
 

moose_lib/main.py

@@ -220,7 +220,7 @@ class QueryClient:
 
 
 class WorkflowClient:
-    """Client for interacting with Temporal workflows.
+    """Client for interacting with Temporal DMv2 workflows.
 
     Args:
         temporal_client: An instance of the Temporal client.
@@ -228,8 +228,6 @@ class WorkflowClient:
 
     def __init__(self, temporal_client: TemporalClient):
         self.temporal_client = temporal_client
-        self.configs = self.load_consolidated_configs()
-        print(f"WorkflowClient - configs: {self.configs}")
 
     # Test workflow executor in rust if this changes significantly
     def execute(self, name: str, input_data: Any) -> Dict[str, Any]:
@@ -278,42 +276,43 @@ class WorkflowClient:
         run_timeout = self.parse_timeout_to_timedelta(config['timeout_str'])
 
         print(
-            f"WorkflowClient - starting {'DMv2 ' if config['is_dmv2'] else ''}workflow: {name} with retry policy: {retry_policy} and timeout: {run_timeout}")
+            f"WorkflowClient - starting DMv2 workflow: {name} with retry policy: {retry_policy} and timeout: {run_timeout}")
 
         # Start workflow with appropriate args
-        workflow_args = self._build_workflow_args(name, processed_input, config['is_dmv2'])
+        workflow_args = self._build_workflow_args(name, processed_input)
+
+        # Handle "never" timeout by omitting run_timeout parameter
+        workflow_kwargs = {
+            "args": workflow_args,
+            "id": workflow_id,
+            "task_queue": "python-script-queue",
+            "id_conflict_policy": WorkflowIDConflictPolicy.FAIL,
+            "id_reuse_policy": WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+            "retry_policy": retry_policy,
+        }
+
+        if run_timeout is not None:
+            workflow_kwargs["run_timeout"] = run_timeout
 
         workflow_handle = await self.temporal_client.start_workflow(
             "ScriptWorkflow",
-            args=workflow_args,
-            id=workflow_id,
-            task_queue="python-script-queue",
-            id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
-            id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE,
-            retry_policy=retry_policy,
-            run_timeout=run_timeout
+            **workflow_kwargs
         )
 
         return workflow_id, workflow_handle.result_run_id
 
     def _get_workflow_config(self, name: str) -> Dict[str, Any]:
-        """Extract workflow configuration from DMv2 or legacy config."""
+        """Extract workflow configuration from DMv2 workflow."""
         from moose_lib.dmv2 import get_workflow
 
         dmv2_workflow = get_workflow(name)
-        if dmv2_workflow is not None:
-            return {
-                'retry_count': dmv2_workflow.config.retries or 3,
-                'timeout_str': dmv2_workflow.config.timeout or "1h",
-                'is_dmv2': True
-            }
-        else:
-            config = self.configs.get(name, {})
-            return {
-                'retry_count': config.get('retries', 3),
-                'timeout_str': config.get('timeout', "1h"),
-                'is_dmv2': False
-            }
+        if dmv2_workflow is None:
+            raise ValueError(f"DMv2 workflow '{name}' not found")
+            
+        return {
+            'retry_count': dmv2_workflow.config.retries or 3,
+            'timeout_str': dmv2_workflow.config.timeout or "1h",
+        }
 
     def _process_input_data(self, name: str, input_data: Any) -> tuple[Any, str]:
         """Process input data and generate workflow ID."""
@@ -339,26 +338,16 @@ class WorkflowClient:
 
         return input_data, workflow_id
 
-    def _build_workflow_args(self, name: str, input_data: Any, is_dmv2: bool) -> list:
-        """Build workflow arguments based on workflow type."""
-        if is_dmv2:
-            return [f"{name}", input_data]
-        else:
-            return [f"{os.getcwd()}/app/scripts/{name}", input_data]
+    def _build_workflow_args(self, name: str, input_data: Any) -> list:
+        """Build workflow arguments for DMv2 workflow."""
+        return [{"workflow_name": name, "execution_mode": "start"}, input_data]
+
 
-    # TODO: Remove when workflows dmv1 is removed
-    def load_consolidated_configs(self):
-        try:
-            file_path = os.path.join(os.getcwd(), ".moose", "workflow_configs.json")
-            with open(file_path, 'r') as file:
-                data = json.load(file)
-                config_map = {config['name']: config for config in data}
-                return config_map
-        except Exception as e:
-            print(f"Could not load configs for workflows v1: {e}")
 
-    def parse_timeout_to_timedelta(self, timeout_str: str) -> timedelta:
-        if timeout_str.endswith('h'):
+    def parse_timeout_to_timedelta(self, timeout_str: str) -> Optional[timedelta]:
+        if timeout_str == "never":
+            return None  # Unlimited execution timeout
+        elif timeout_str.endswith('h'):
             return timedelta(hours=int(timeout_str[:-1]))
         elif timeout_str.endswith('m'):
             return timedelta(minutes=int(timeout_str[:-1]))

{moose_lib-0.6.18.dist-info → moose_lib-0.6.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: moose_lib
-Version: 0.6.18
+Version: 0.6.20
 Home-page: https://www.fiveonefour.com/moose
 Author: Fiveonefour Labs Inc.
 Author-email: support@fiveonefour.com

{moose_lib-0.6.18.dist-info → moose_lib-0.6.20.dist-info}/RECORD

@@ -1,18 +1,17 @@
-moose_lib/__init__.py,sha256=_XXfHrzRl-M5M1e3CwBQGGblVQYKnU5QJkipghe4bXU,184
+moose_lib/__init__.py,sha256=LiUdVqmtASPqbMF53dXVZzCdU1jtFVx62_tiSbW65a0,162
 moose_lib/blocks.py,sha256=_wdvC2NC_Y3MMEnB71WTgWbeQ--zPNHk19xjToJW0C0,3185
 moose_lib/commons.py,sha256=FUpRv8D3-LeGcjhcqtDyiimz5izwpCq53h50ydxC_uA,3711
 moose_lib/data_models.py,sha256=PgoFXTzf4C66FlKUowMT4VOrwSRR9_bk36P43ub4LzU,10274
 moose_lib/dmv2-serializer.py,sha256=CL_Pvvg8tJOT8Qk6hywDNzY8MYGhMVdTOw8arZi3jng,49
 moose_lib/internal.py,sha256=UYF42oFujLH1g9GDo_HdP-AYipKTf_lReyVi6qs945A,14397
-moose_lib/main.py,sha256=WAxrM39rp5leAtHxZQZlC8-a8pA8Ksy1TjVQnq12YeA,17354
+moose_lib/main.py,sha256=XcVX_sTnt5QbrPXKNLCKZGCvpFpE8oiqSG2S1So9ztI,16713
 moose_lib/query_param.py,sha256=kxcR09BMIsEg4o2qetjKrVu1YFRaLfMEzwzyGsKUpvA,6474
-moose_lib/tasks.py,sha256=6MXA0j7nhvQILAJVTQHCAsquwrSOi2zAevghAc_7kXs,1554
 moose_lib/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 moose_lib/clients/redis_client.py,sha256=UBCdxwgZpIOIOy2EnPyxJIAYjw_qmNwGsJQCQ66SxUI,8117
 moose_lib/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 moose_lib/config/config_file.py,sha256=NyjY6YFraBel7vBk18lLkpGaPR1viKMAEv4ZldyfLIA,2585
 moose_lib/config/runtime.py,sha256=gF1wrUCbp0MxRU92k0vyItvWX2x6P1FuBRzM61nJE_U,2776
-moose_lib/dmv2/__init__.py,sha256=MXVSaMw8V0h06JlA_uomSjl6MXPerjC9tVo2KvIvLgE,2732
+moose_lib/dmv2/__init__.py,sha256=3DVAtNMZUoP94CMJBFhuXfYEQXDbQUNKSgg9XnKqae0,2768
 moose_lib/dmv2/_registry.py,sha256=RweMQWHSvhxgSceZ5mfd0NGnJ-49VZORbaFmw-O0bFY,611
 moose_lib/dmv2/consumption.py,sha256=KJV9MKgrnVLIRhlFVZMDhTsXvOaZoMiHk4hnRePdPWA,8745
 moose_lib/dmv2/ingest_api.py,sha256=Snek9NGwaJl_BuImSWGtQq91m9D3AJ4qBoGiKZ-9yTQ,2323
@@ -25,14 +24,14 @@ moose_lib/dmv2/sql_resource.py,sha256=kUZoGqxhZMHMthtBZGYJBxTFjXkspXiWLXhJRYXgGU
 moose_lib/dmv2/stream.py,sha256=jiUWBsjFalLLP63mikOxyHRdieiDAlzf9lXfLye-Wjc,10761
 moose_lib/dmv2/types.py,sha256=5FsB0HLHFkYB-8cjJ0rtRUjqahVA-ToLr2JXT1lFiss,3276
 moose_lib/dmv2/view.py,sha256=fVbfbJgc2lvhjpGvpfKcFUqZqxKuLD4X59jdupxIe94,1350
-moose_lib/dmv2/workflow.py,sha256=ZNEMaYWGCkOw_1qropfl2m553aq5YG16Y0rOJjo8eak,5916
+moose_lib/dmv2/workflow.py,sha256=_FY4-VRo7uWxRtoipxGSo04qzBb4pbP30iQei1W0Ios,6287
 moose_lib/streaming/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 moose_lib/streaming/streaming_function_runner.py,sha256=toide6mpUIuHWv75Iv8PmKL6HlsPnBB2XRq2aIjOWBY,23421
 tests/__init__.py,sha256=0Gh4yzPkkC3TzBGKhenpMIxJcRhyrrCfxLSfpTZnPMQ,53
 tests/conftest.py,sha256=ZVJNbnr4DwbcqkTmePW6U01zAzE6QD0kNAEZjPG1f4s,169
 tests/test_moose.py,sha256=mBsx_OYWmL8ppDzL_7Bd7xR6qf_i3-pCIO3wm2iQNaA,2136
 tests/test_redis_client.py,sha256=d9_MLYsJ4ecVil_jPB2gW3Q5aWnavxmmjZg2uYI3LVo,3256
-moose_lib-0.6.18.dist-info/METADATA,sha256=ewael_1KEnFWIZtC7fMG8sY1dTNkn_B1PjgfiSaDcdo,730
-moose_lib-0.6.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-moose_lib-0.6.18.dist-info/top_level.txt,sha256=XEns2-4aCmGp2XjJAeEH9TAUcGONLnSLy6ycT9FSJh8,16
-moose_lib-0.6.18.dist-info/RECORD,,
+moose_lib-0.6.20.dist-info/METADATA,sha256=5oRl34L_r8CXqY2lKc__Tx2CX3M4pcQcbQCuojDbAS4,730
+moose_lib-0.6.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+moose_lib-0.6.20.dist-info/top_level.txt,sha256=XEns2-4aCmGp2XjJAeEH9TAUcGONLnSLy6ycT9FSJh8,16
+moose_lib-0.6.20.dist-info/RECORD,,

moose_lib/tasks.py

@@ -1,45 +0,0 @@
-from functools import wraps
-from typing import TypeVar, Callable, Any
-from .commons import Logger
-import asyncio
-
-T = TypeVar('T')
-
-def task(func: Callable[..., T] = None, *, retries: int = 3) -> Callable[..., T]:
-    """Decorator to mark a function as a Moose task
-    
-    Args:
-        func: The function to decorate
-        retries: Number of times to retry the task if it fails
-    """
-    def validate_result(result: Any) -> None:
-        """Ensure proper return format"""
-        if not isinstance(result, dict):
-            raise ValueError("Task must return a dictionary with 'task' and 'data' keys")
-        if "task" not in result or "data" not in result:
-            raise ValueError("Task result must contain 'task' and 'data' keys")
-
-    def decorator(f: Callable[..., T]) -> Callable[..., T]:
-        if asyncio.iscoroutinefunction(f):
-            @wraps(f)
-            async def wrapper(*args, **kwargs) -> T:
-                result = await f(*args, **kwargs)
-                validate_result(result)
-                return result
-        else:
-            @wraps(f)
-            def wrapper(*args, **kwargs) -> T:
-                result = f(*args, **kwargs)
-                validate_result(result)
-                return result
-
-        # Add the markers to the wrapper
-        wrapper._is_moose_task = True
-        wrapper._retries = retries
-        wrapper.logger = Logger(is_moose_task=True)
-        return wrapper
-    
-    # Handle both @task and @task() syntax
-    if func is None:
-        return decorator
-    return decorator(func)