horsies 0.1.0a4__py3-none-any.whl → 0.1.0a5__py3-none-any.whl

horsies/core/models/workflow_pg.py

@@ -236,7 +236,9 @@ class WorkflowTaskModel(Base):
 
     # Unique constraint: one task per index per workflow
     __table_args__ = (
-        UniqueConstraint('workflow_id', 'task_index', name='uq_horsies_workflow_task_index'),
+        UniqueConstraint(
+            'workflow_id', 'task_index', name='uq_horsies_workflow_task_index'
+        ),
     )
 
 

horsies/core/scheduler/service.py

@@ -21,6 +21,10 @@ from horsies.core.worker.worker import import_by_path
 
 logger = get_logger('scheduler')
 
+SCHEDULE_ADVISORY_LOCK_SQL = text(
+    """SELECT pg_advisory_xact_lock(CAST(:key AS BIGINT))"""
+)
+
 
 class Scheduler:
     """
@@ -257,7 +261,7 @@ class Scheduler:
         async with self.broker.session_factory() as session:
             # Acquire transaction-scoped advisory lock for this specific schedule
             await session.execute(
-                text('SELECT pg_advisory_xact_lock(CAST(:key AS BIGINT))'),
+                SCHEDULE_ADVISORY_LOCK_SQL,
                 {'key': lock_key},
             )
 

horsies/core/scheduler/state.py

@@ -9,6 +9,39 @@ from horsies.core.logging import get_logger
 
 logger = get_logger('scheduler.state')
 
+UPDATE_SCHEDULE_AFTER_RUN_SQL = text("""
+    UPDATE horsies_schedule_state
+    SET last_run_at = :executed_at,
+        next_run_at = :next_run_at,
+        last_task_id = :task_id,
+        run_count = run_count + 1,
+        updated_at = :now
+    WHERE schedule_name = :schedule_name
+""")
+
+UPDATE_SCHEDULE_NEXT_RUN_WITH_HASH_SQL = text("""
+    UPDATE horsies_schedule_state
+    SET next_run_at = :next_run_at,
+        config_hash = :config_hash,
+        updated_at = :now
+    WHERE schedule_name = :schedule_name
+""")
+
+UPDATE_SCHEDULE_NEXT_RUN_SQL = text("""
+    UPDATE horsies_schedule_state
+    SET next_run_at = :next_run_at,
+        updated_at = :now
+    WHERE schedule_name = :schedule_name
+""")
+
+DELETE_SCHEDULE_STATE_SQL = text("""
+    DELETE FROM horsies_schedule_state WHERE schedule_name = :schedule_name
+""")
+
+GET_ALL_SCHEDULE_STATES_SQL = text("""
+    SELECT * FROM horsies_schedule_state ORDER BY schedule_name
+""")
+
 
 class ScheduleStateManager:
     """
@@ -127,15 +160,7 @@ class ScheduleStateManager:
         async with self.session_factory() as session:
             # Use raw SQL for atomic update with increment
             result = await session.execute(
-                text("""
-                    UPDATE horsies_schedule_state
-                    SET last_run_at = :executed_at,
-                        next_run_at = :next_run_at,
-                        last_task_id = :task_id,
-                        run_count = run_count + 1,
-                        updated_at = :now
-                    WHERE schedule_name = :schedule_name
-                """),
+                UPDATE_SCHEDULE_AFTER_RUN_SQL,
                 {
                     'schedule_name': schedule_name,
                     'executed_at': executed_at,
@@ -174,13 +199,7 @@ class ScheduleStateManager:
         async with self.session_factory() as session:
             # Build UPDATE query dynamically based on whether config_hash is provided
             if config_hash is not None:
-                query = """
-                    UPDATE horsies_schedule_state
-                    SET next_run_at = :next_run_at,
-                        config_hash = :config_hash,
-                        updated_at = :now
-                    WHERE schedule_name = :schedule_name
-                """
+                query = UPDATE_SCHEDULE_NEXT_RUN_WITH_HASH_SQL
                 params = {
                     'schedule_name': schedule_name,
                     'next_run_at': next_run_at,
@@ -188,19 +207,14 @@ class ScheduleStateManager:
                     'now': datetime.now(timezone.utc),
                 }
             else:
-                query = """
-                    UPDATE horsies_schedule_state
-                    SET next_run_at = :next_run_at,
-                        updated_at = :now
-                    WHERE schedule_name = :schedule_name
-                """
+                query = UPDATE_SCHEDULE_NEXT_RUN_SQL
                 params = {
                     'schedule_name': schedule_name,
                     'next_run_at': next_run_at,
                     'now': datetime.now(timezone.utc),
                 }
 
-            result = await session.execute(text(query), params)
+            result = await session.execute(query, params)
             await session.commit()
 
             rows_updated = getattr(result, 'rowcount', 0)
@@ -223,7 +237,7 @@ class ScheduleStateManager:
         """
         async with self.session_factory() as session:
             result = await session.execute(
-                text('DELETE FROM horsies_schedule_state WHERE schedule_name = :schedule_name'),
+                DELETE_SCHEDULE_STATE_SQL,
                 {'schedule_name': schedule_name},
             )
             await session.commit()
@@ -244,9 +258,7 @@ class ScheduleStateManager:
             List of all ScheduleStateModel records
         """
         async with self.session_factory() as session:
-            result = await session.execute(
-                text('SELECT * FROM horsies_schedule_state ORDER BY schedule_name')
-            )
+            result = await session.execute(GET_ALL_SCHEDULE_STATES_SQL)
             rows = result.fetchall()
             columns = result.keys()
 

horsies/core/task_decorator.py

@@ -6,7 +6,9 @@ from typing import (
     get_origin,
     get_type_hints,
     get_args,
+    Literal,
     ParamSpec,
+    Sequence,
     TypeVar,
     Generic,
     Protocol,
@@ -24,6 +26,11 @@ if TYPE_CHECKING:
     from horsies.core.models.tasks import TaskOptions
     from horsies.core.models.tasks import TaskError, TaskResult
     from horsies.core.models.tasks import TaskInfo
+    from horsies.core.models.workflow import (
+        TaskNode,
+        SubWorkflowNode,
+        WorkflowContext,
+    )
 
 from horsies.core.models.tasks import TaskResult, TaskError, LibraryErrorCode
 from horsies.core.models.workflow import WorkflowContextMissingIdError
@@ -241,6 +248,87 @@ class TaskHandle(Generic[T]):
         self._result_fetched = True
 
 
+class NodeFactory(Generic[P, T]):
+    """
+    Factory for creating TaskNode instances with typed arguments.
+
+    Returned by TaskFunction.node(). Call with the task's arguments
+    to create a TaskNode with full static type checking.
+
+    Example:
+        node = my_task.node(waits_for=[dep])(value='test')
+        # Type checker validates 'value' against my_task's signature
+    """
+
+    _fn: 'TaskFunction[P, T]'
+    _waits_for: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None
+    _workflow_ctx_from: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None
+    _args_from: dict[str, 'TaskNode[Any] | SubWorkflowNode[Any]'] | None
+    _queue: str | None
+    _priority: int | None
+    _allow_failed_deps: bool
+    _run_when: Callable[['WorkflowContext'], bool] | None
+    _skip_when: Callable[['WorkflowContext'], bool] | None
+    _join: Literal['all', 'any', 'quorum']
+    _min_success: int | None
+    _good_until: datetime | None
+    _node_id: str | None
+
+    def __init__(
+        self,
+        fn: 'TaskFunction[P, T]',
+        *,
+        waits_for: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None,
+        workflow_ctx_from: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None,
+        args_from: dict[str, 'TaskNode[Any] | SubWorkflowNode[Any]'] | None,
+        queue: str | None,
+        priority: int | None,
+        allow_failed_deps: bool,
+        run_when: Callable[['WorkflowContext'], bool] | None,
+        skip_when: Callable[['WorkflowContext'], bool] | None,
+        join: Literal['all', 'any', 'quorum'],
+        min_success: int | None,
+        good_until: datetime | None,
+        node_id: str | None,
+    ) -> None:
+        self._fn = fn
+        self._waits_for = waits_for
+        self._workflow_ctx_from = workflow_ctx_from
+        self._args_from = args_from
+        self._queue = queue
+        self._priority = priority
+        self._allow_failed_deps = allow_failed_deps
+        self._run_when = run_when
+        self._skip_when = skip_when
+        self._join = join
+        self._min_success = min_success
+        self._good_until = good_until
+        self._node_id = node_id
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> 'TaskNode[T]':
+        from horsies.core.models.workflow import TaskNode
+
+        return TaskNode(
+            fn=self._fn,
+            args=args,
+            kwargs=dict(kwargs),
+            waits_for=list(self._waits_for) if self._waits_for else [],
+            workflow_ctx_from=list(self._workflow_ctx_from)
+            if self._workflow_ctx_from
+            else None,
+            args_from=dict(self._args_from) if self._args_from else {},
+            queue=self._queue,
+            priority=self._priority,
+            allow_failed_deps=self._allow_failed_deps,
+            run_when=self._run_when,
+            skip_when=self._skip_when,
+            join=self._join,
+            min_success=self._min_success,
+            good_until=self._good_until,
+            node_id=self._node_id,
+        )
+
+
 class TaskFunction(Protocol[P, T]):
     """
     A TaskFunction is a function that gets a @task decorator applied to it.
@@ -281,6 +369,24 @@ class TaskFunction(Protocol[P, T]):
         **kwargs: P.kwargs,
     ) -> 'TaskHandle[T]': ...
 
+    @abstractmethod
+    def node(
+        self,
+        *,
+        waits_for: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+        workflow_ctx_from: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+        args_from: dict[str, 'TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+        queue: str | None = None,
+        priority: int | None = None,
+        allow_failed_deps: bool = False,
+        run_when: Callable[['WorkflowContext'], bool] | None = None,
+        skip_when: Callable[['WorkflowContext'], bool] | None = None,
+        join: Literal['all', 'any', 'quorum'] = 'all',
+        min_success: int | None = None,
+        good_until: datetime | None = None,
+        node_id: str | None = None,
+    ) -> 'NodeFactory[P, T]': ...
+
 
 def create_task_wrapper(
     fn: Callable[P, TaskResult[T, TaskError]],
@@ -647,6 +753,38 @@ def create_task_wrapper(
         ) -> TaskHandle[T]:
             return schedule(delay, *args, **kwargs)
 
+        def node(
+            self,
+            *,
+            waits_for: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+            workflow_ctx_from: Sequence['TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+            args_from: dict[str, 'TaskNode[Any] | SubWorkflowNode[Any]'] | None = None,
+            queue: str | None = None,
+            priority: int | None = None,
+            allow_failed_deps: bool = False,
+            run_when: Callable[['WorkflowContext'], bool] | None = None,
+            skip_when: Callable[['WorkflowContext'], bool] | None = None,
+            join: Literal['all', 'any', 'quorum'] = 'all',
+            min_success: int | None = None,
+            good_until: datetime | None = None,
+            node_id: str | None = None,
+        ) -> NodeFactory[P, T]:
+            return NodeFactory(
+                fn=self,  # type: ignore[arg-type]
+                waits_for=waits_for,
+                workflow_ctx_from=workflow_ctx_from,
+                args_from=args_from,
+                queue=queue,
+                priority=priority,
+                allow_failed_deps=allow_failed_deps,
+                run_when=run_when,
+                skip_when=skip_when,
+                join=join,
+                min_success=min_success,
+                good_until=good_until,
+                node_id=node_id,
+            )
+
         # Copy metadata
         def __getattr__(self, name: str) -> Any:
             return getattr(wrapped_function, name)

horsies/core/types/status.py

@@ -31,8 +31,10 @@ class TaskStatus(Enum):
         return self in TASK_TERMINAL_STATES
 
 
-TASK_TERMINAL_STATES: frozenset[TaskStatus] = frozenset({
-    TaskStatus.COMPLETED,
-    TaskStatus.FAILED,
-    TaskStatus.CANCELLED,
-})
+TASK_TERMINAL_STATES: frozenset[TaskStatus] = frozenset(
+    {
+        TaskStatus.COMPLETED,
+        TaskStatus.FAILED,
+        TaskStatus.CANCELLED,
+    }
+)

horsies/core/utils/imports.py

@@ -20,7 +20,7 @@ from typing import Any
 
 from horsies.core.logging import get_logger
 
-logger = get_logger("imports")
+logger = get_logger('imports')
 
 
 def find_project_root(start_dir: str) -> str | None:
@@ -34,7 +34,7 @@ def find_project_root(start_dir: str) -> str | None:
     Returns start_dir if it contains a marker file, None otherwise.
     """
     start_dir = os.path.abspath(start_dir)
-    for marker in ("pyproject.toml", "setup.cfg", "setup.py"):
+    for marker in ('pyproject.toml', 'setup.cfg', 'setup.py'):
         if os.path.exists(os.path.join(start_dir, marker)):
             return start_dir
     return None
@@ -54,7 +54,7 @@ def setup_sys_path_from_cwd() -> str | None:
     cwd = os.getcwd()
     if find_project_root(cwd) and cwd not in sys.path:
         sys.path.insert(0, cwd)
-        logger.debug(f"Added cwd to sys.path: {cwd}")
+        logger.debug(f'Added cwd to sys.path: {cwd}')
         return cwd
     return None
 
@@ -90,7 +90,7 @@ def _compute_synthetic_module_name(path: str) -> str:
     """
     realpath = os.path.realpath(path)
     hash_prefix = hashlib.sha256(realpath.encode()).hexdigest()[:12]
-    return f"horsies._dynamic.{hash_prefix}"
+    return f'horsies._dynamic.{hash_prefix}'
 
 
 def import_file_path(
@@ -121,11 +121,11 @@ def import_file_path(
     file_path = os.path.realpath(file_path)
 
     if not os.path.exists(file_path):
-        raise FileNotFoundError(f"Module file not found: {file_path}")
+        raise FileNotFoundError(f'Module file not found: {file_path}')
 
     # Check if already loaded
     for name, mod in list(sys.modules.items()):
-        mod_file = getattr(mod, "__file__", None)
+        mod_file = getattr(mod, '__file__', None)
         if mod_file and os.path.realpath(mod_file) == file_path:
             return mod
 
@@ -142,7 +142,7 @@ def import_file_path(
     # Load the module
     spec = importlib.util.spec_from_file_location(module_name, file_path)
     if spec is None or spec.loader is None:
-        raise ImportError(f"Could not load module from path: {file_path}")
+        raise ImportError(f'Could not load module from path: {file_path}')
 
     mod = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = mod
@@ -164,7 +164,7 @@ def import_by_path(path: str, module_name: str | None = None) -> Any:
     For file paths: delegates to import_file_path()
     For module paths: delegates to import_module_path()
     """
-    if path.endswith(".py") or os.path.sep in path:
+    if path.endswith('.py') or os.path.sep in path:
         return import_file_path(path, module_name)
     else:
         return import_module_path(path)
@@ -187,7 +187,7 @@ def compute_package_path_from_fs(file_path: str) -> tuple[str | None, str | None
 
     components = [module_name]
     while True:
-        init_path = os.path.join(current_dir, "__init__.py")
+        init_path = os.path.join(current_dir, '__init__.py')
         if not os.path.exists(init_path):
             break
         package_name = os.path.basename(current_dir)
@@ -198,6 +198,6 @@ def compute_package_path_from_fs(file_path: str) -> tuple[str | None, str | None
         return (None, None)
 
     components.reverse()
-    dotted_name = ".".join(components)
+    dotted_name = '.'.join(components)
     package_root = current_dir
     return (dotted_name, package_root)