zrb 1.0.0a21__py3-none-any.whl → 1.0.0b2__py3-none-any.whl

zrb/runner/web_util/user.py

@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING
+
+from zrb.runner.web_config.config import WebConfig
+from zrb.runner.web_schema.user import User
+
+if TYPE_CHECKING:
+    # Import Request only for type checking to reduce runtime dependencies
+    from fastapi import Request
+
+
+def get_user_by_credentials(
+    web_config: WebConfig, username: str, password: str
+) -> User | None:
+    user = web_config.find_user_by_username(username)
+    if user is None or not user.is_password_match(password):
+        return None
+    return user
+
+
+async def get_user_from_request(
+    web_config: WebConfig, request: "Request"
+) -> User | None:
+    from fastapi.security import OAuth2PasswordBearer
+
+    if not web_config.enable_auth:
+        return web_config.default_user
+    # Normally we use "Depends"
+    get_bearer_token = OAuth2PasswordBearer(tokenUrl="/api/v1/login", auto_error=False)
+    bearer_token = await get_bearer_token(request)
+    token_user = _get_user_from_token(web_config, bearer_token)
+    if token_user is not None:
+        return token_user
+    cookie_user = _get_user_from_cookie(web_config, request)
+    if cookie_user is not None:
+        return cookie_user
+    return web_config.default_user
+
+
+def _get_user_from_cookie(web_config: WebConfig, request: "Request") -> User | None:
+    token = request.cookies.get(web_config.access_token_cookie_name)
+    if token:
+        return _get_user_from_token(web_config, token)
+    return None
+
+
+def _get_user_from_token(web_config: WebConfig, token: str) -> User | None:
+    try:
+        from jose import jwt
+
+        payload = jwt.decode(
+            token,
+            web_config.secret_key,
+            options={"require_sub": True, "require_exp": True},
+        )
+        username: str = payload.get("sub")
+        if username is None:
+            return None
+        user = web_config.find_user_by_username(username)
+        if user is None:
+            return None
+        return user
+    except Exception:
+        return None

zrb/session/session.py

@@ -11,9 +11,7 @@ from zrb.session_state_log.session_state_log import (
     TaskStatusStateLog,
 )
 from zrb.session_state_logger.any_session_state_logger import AnySessionStateLogger
-from zrb.session_state_logger.default_session_state_logger import (
-    default_session_state_logger,
-)
+from zrb.session_state_logger.session_state_logger_factory import session_state_logger
 from zrb.task.any_task import AnyTask
 from zrb.task_status.task_status import TaskStatus
 from zrb.util.cli.style import (
@@ -126,7 +124,7 @@ class Session(AnySession):
     @property
     def state_logger(self) -> AnySessionStateLogger:
         if self._state_logger is None:
-            return default_session_state_logger
+            return session_state_logger
         return self._state_logger
 
     def set_main_task(self, main_task: AnyTask):
@@ -215,6 +213,10 @@ class Session(AnySession):
         self._register_single_task(task)
         for readiness_check in task.readiness_checks:
             self.register_task(readiness_check)
+        for successor in task.successors:
+            self.register_task(successor)
+        for fallback in task.fallbacks:
+            self.register_task(fallback)
         for upstream in task.upstreams:
             self.register_task(upstream)
             if task not in self._downstreams[upstream]:

zrb/session_state_logger/{default_session_state_logger.py → session_state_logger_factory.py}

@@ -1,4 +1,4 @@
 from zrb.config import SESSION_LOG_DIR
 from zrb.session_state_logger.file_session_state_logger import FileSessionStateLogger
 
-default_session_state_logger = FileSessionStateLogger(SESSION_LOG_DIR)
+session_state_logger = FileSessionStateLogger(SESSION_LOG_DIR)

zrb/task/base_task.py

@@ -242,7 +242,29 @@ class BaseTask(AnyTask):
     def run(
         self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
     ) -> Any:
-        return asyncio.run(self.async_run(session, str_kwargs))
+        loop = asyncio.new_event_loop()
+        try:
+            return loop.run_until_complete(self._run_and_cleanup(session, str_kwargs))
+        finally:
+            loop.close()
+
+    async def _run_and_cleanup(
+        self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
+    ) -> Any:
+        try:
+            result = await self.async_run(session, str_kwargs)
+        finally:
+            if not session.is_terminated:
+                session.terminate()
+            # Cancel all running tasks except the current one
+            current_task = asyncio.current_task()
+            pending = [task for task in asyncio.all_tasks() if task is not current_task]
+            for task in pending:
+                task.cancel()
+            # Wait for all tasks to complete with a timeout
+            if pending:
+                await asyncio.wait(pending, timeout=5)
+        return result
 
     async def async_run(
         self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
@@ -265,11 +287,9 @@ class BaseTask(AnyTask):
     def __fill_shared_context_envs(self, shared_context: AnySharedContext):
         # Inject os environ
         os_env_map = {
-            key: val
-            for key, val in os.environ.items()
-            if key not in shared_context._env
+            key: val for key, val in os.environ.items() if key not in shared_context.env
         }
-        shared_context._env.update(os_env_map)
+        shared_context.env.update(os_env_map)
 
     async def exec_root_tasks(self, session: AnySession):
         session.set_main_task(self)
@@ -416,6 +436,23 @@ class BaseTask(AnyTask):
             ctx.log_info("Continue monitoring")
 
     async def __exec_action_and_retry(self, session: AnySession) -> Any:
+        """
+        Executes an action with retry logic.
+
+        This method attempts to execute the action defined in `_exec_action` with a specified number of retries.
+        If the action fails, it will retry after a specified period until the maximum number of attempts is reached.
+        If the action succeeds, it marks the task as completed and executes any successors.
+        If the action fails permanently, it marks the task as permanently failed and executes any fallbacks.
+
+        Args:
+            session (AnySession): The session object containing the task status and context.
+
+        Returns:
+            Any: The result of the executed action if successful.
+
+        Raises:
+            Exception: If the action fails permanently after all retry attempts.
+        """
         ctx = self.get_ctx(session)
         max_attempt = self._retries + 1
         ctx.set_max_attempt(max_attempt)
@@ -433,13 +470,14 @@ class BaseTask(AnyTask):
                 # Put result on xcom
                 task_xcom: Xcom = ctx.xcom.get(self.name)
                 task_xcom.push(result)
+                self.__skip_fallbacks(session)
                 await run_async(self.__exec_successors(session))
                 return result
             except (asyncio.CancelledError, KeyboardInterrupt):
                 ctx.log_info("Marked as failed")
                 session.get_task_status(self).mark_as_failed()
                 return
-            except Exception as e:
+            except BaseException as e:
                 ctx.log_error(e)
                 if attempt < max_attempt - 1:
                     ctx.log_info("Marked as failed")
@@ -447,6 +485,7 @@ class BaseTask(AnyTask):
                     continue
                 ctx.log_info("Marked as permanently failed")
                 session.get_task_status(self).mark_as_permanently_failed()
+                self.__skip_successors(session)
                 await run_async(self.__exec_fallbacks(session))
                 raise e
 
@@ -457,6 +496,10 @@ class BaseTask(AnyTask):
         ]
         await asyncio.gather(*successor_coros)
 
+    def __skip_successors(self, session: AnySession) -> Any:
+        for successor in self.successors:
+            session.get_task_status(successor).mark_as_skipped()
+
     async def __exec_fallbacks(self, session: AnySession) -> Any:
         fallbacks: list[AnyTask] = self.fallbacks
         fallback_coros = [
@@ -464,6 +507,10 @@ class BaseTask(AnyTask):
         ]
         await asyncio.gather(*fallback_coros)
 
+    def __skip_fallbacks(self, session: AnySession) -> Any:
+        for fallback in self.fallbacks:
+            session.get_task_status(fallback).mark_as_skipped()
+
     async def _exec_action(self, ctx: AnyContext) -> Any:
         """Execute the main action of the task.
         By default will render and run the _action attribute.

zrb/task/base_trigger.py

@@ -42,6 +42,7 @@ class BaseTrigger(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -63,6 +64,7 @@ class BaseTrigger(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._callbacks = callback
         self._queue_name = queue_name

zrb/task/cmd_task.py

@@ -56,6 +56,7 @@ class CmdTask(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -76,6 +77,7 @@ class CmdTask(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._shell = shell
         self._render_shell = render_shell
@@ -233,17 +235,18 @@ class CmdTask(BaseTask):
 
     def _render_cmd_val(self, ctx: AnyContext, cmd_val: CmdVal) -> str:
         if isinstance(cmd_val, list):
+            cmd_val_list = [
+                self.__render_single_cmd_val(ctx, single_cmd_val)
+                for single_cmd_val in cmd_val
+            ]
             return "\n".join(
-                [
-                    self.__render_single_cmd_val(ctx, single_cmd_val)
-                    for single_cmd_val in cmd_val
-                ]
+                [cmd_val for cmd_val in cmd_val_list if cmd_val is not None]
             )
         return self.__render_single_cmd_val(ctx, cmd_val)
 
     def __render_single_cmd_val(
         self, ctx: AnyContext, single_cmd_val: SingleCmdVal
-    ) -> str:
+    ) -> str | None:
         if callable(single_cmd_val):
             return single_cmd_val(ctx)
         if isinstance(single_cmd_val, str):
@@ -252,6 +255,7 @@ class CmdTask(BaseTask):
             return single_cmd_val
         if isinstance(single_cmd_val, AnyCmdVal):
             return single_cmd_val.to_str(ctx)
+        return None
 
     def __get_multiline_repr(self, text: str) -> str:
         lines_repr: list[str] = []

zrb/task/http_check.py

@@ -28,6 +28,7 @@ class HttpCheck(BaseTask):
         execute_condition: bool | str | Callable[[Context], bool] = True,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -41,6 +42,7 @@ class HttpCheck(BaseTask):
             retries=0,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._url = url
         self._render_url = render_url

zrb/task/llm_task.py

@@ -67,6 +67,7 @@ class LLMTask(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -87,6 +88,7 @@ class LLMTask(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._model = model
         self._render_model = render_model

zrb/task/make_task.py

@@ -29,6 +29,7 @@ def make_task(
     monitor_readiness: bool = False,
     upstream: list[AnyTask] | AnyTask | None = None,
     fallback: list[AnyTask] | AnyTask | None = None,
+    successor: list[AnyTask] | AnyTask | None = None,
     group: AnyGroup | None = None,
     alias: str | None = None,
 ) -> Callable[[Callable[[AnyContext], Any]], AnyTask]:
@@ -53,6 +54,7 @@ def make_task(
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         if group is not None:
             return group.add_task(task, alias=alias)

zrb/task/rsync_task.py

@@ -49,6 +49,7 @@ class RsyncTask(CmdTask):
         readiness_check: list[AnyTask] | AnyTask | None = None,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -80,6 +81,7 @@ class RsyncTask(CmdTask):
             readiness_check=readiness_check,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._remote_source_path = remote_source_path
         self._render_remote_source_path = render_remote_source_path

zrb/task/scaffolder.py

@@ -11,6 +11,7 @@ from zrb.input.any_input import AnyInput
 from zrb.task.any_task import AnyTask
 from zrb.task.base_task import BaseTask
 from zrb.util.attr import get_str_attr
+from zrb.util.cli.style import stylize_faint
 
 TransformConfig = dict[str, str] | Callable[[AnyContext, str], str]
 
@@ -46,6 +47,7 @@ class Scaffolder(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -66,6 +68,7 @@ class Scaffolder(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._source_path = source_path
         self._render_source_path = render_source_path
@@ -83,13 +86,12 @@ class Scaffolder(BaseTask):
         return get_str_attr(ctx, self._destination_path, "", auto_render=True)
 
     def _get_content_transformers(self) -> list[AnyContentTransformer]:
-        if callable(self._content_transformers):
-            return [
-                ContentTransformer(match=".*", transform=self._content_transformers)
-            ]
-        if isinstance(self._content_transformers, dict):
+        if callable(self._content_transformers) or isinstance(
+            self._content_transformers, dict
+        ):
             return [
                 ContentTransformer(
+                    name="default-transform",
                     match=".*",
                     transform=self._content_transformers,
                     auto_render=self._render_content_transformers,
@@ -109,6 +111,7 @@ class Scaffolder(BaseTask):
             for transformer in transformers:
                 if transformer.match(ctx, file_path):
                     try:
+                        ctx.print(stylize_faint(f"{transformer.name}: {file_path}"))
                         transformer.transform_file(ctx, file_path)
                     except UnicodeDecodeError:
                         pass

zrb/task/scheduler.py

@@ -39,6 +39,7 @@ class Scheduler(BaseTrigger):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -61,6 +62,7 @@ class Scheduler(BaseTrigger):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._cron_pattern = schedule
 

zrb/task/tcp_check.py

@@ -28,6 +28,7 @@ class TcpCheck(BaseTask):
         execute_condition: bool | str | Callable[[Context], bool] = True,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -41,6 +42,7 @@ class TcpCheck(BaseTask):
             retries=0,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._host = host
         self._render_host = render_host

zrb/task_status/task_status.py

@@ -62,9 +62,10 @@ class TaskStatus:
         self._history.append((TASK_PERMANENTLY_FAILED, datetime.datetime.now()))
 
     def mark_as_terminated(self):
-        self._is_terminated = True
-        if not self.is_completed and not self.is_permanently_failed:
-            self._history.append((TASK_TERMINATED, datetime.datetime.now()))
+        if not self._is_terminated:
+            self._is_terminated = True
+            if not (self.is_skipped or self.is_completed or self.is_permanently_failed):
+                self._history.append((TASK_TERMINATED, datetime.datetime.now()))
 
     @property
     def is_started(self) -> bool:

zrb/util/cmd/command.py

@@ -62,6 +62,7 @@ async def run_command(
             log_method(line)
         return "\n".join(lines)
 
+    cmd_process = None
     try:
         if cwd is None:
             cwd = os.getcwd()

zrb/util/file.py

@@ -1,4 +1,5 @@
 import os
+import re
 
 
 def read_file(file_path: str, replace_map: dict[str, str] = {}) -> str:
@@ -14,5 +15,10 @@ def write_file(file_path: str, content: str | list[str]):
         content = "\n".join([line for line in content if line is not None])
     dir_path = os.path.dirname(file_path)
     os.makedirs(dir_path, exist_ok=True)
+    content = re.sub(r"\n{3,}$", "\n\n", content)
+    # Remove trailing newlines, but keep one if it exists
+    content = content.rstrip("\n")
+    if content.endswith("\n"):
+        content += "\n"
     with open(file_path, "w") as f:
-        f.write(content.strip())
+        f.write(content)

{zrb-1.0.0a21.dist-info → zrb-1.0.0b2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0a21
+Version: 1.0.0b2
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later