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

zrb/runner/web_schema/user.py

@@ -0,0 +1,32 @@
+from pydantic import BaseModel, ConfigDict
+
+from zrb.group.any_group import AnyGroup
+from zrb.task.any_task import AnyTask
+from zrb.util.group import get_all_subtasks
+
+
+class User(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    username: str
+    password: str = ""
+    is_super_admin: bool = False
+    is_guest: bool = False
+    accessible_tasks: list[AnyTask | str] = []
+
+    def is_password_match(self, password: str) -> bool:
+        return self.password == password
+
+    def can_access_group(self, group: AnyGroup) -> bool:
+        if self.is_super_admin:
+            return True
+        all_tasks = get_all_subtasks(group, web_only=True)
+        if any(self.can_access_task(task) for task in all_tasks):
+            return True
+        return False
+
+    def can_access_task(self, task: AnyTask) -> bool:
+        if self.is_super_admin:
+            return True
+        if task.name in self.accessible_tasks or task in self.accessible_tasks:
+            return True
+        return False

zrb/runner/web_util/cookie.py

@@ -0,0 +1,29 @@
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from zrb.runner.web_config.config import WebConfig
+from zrb.runner.web_schema.token import Token
+
+if TYPE_CHECKING:
+    # We want fastapi to only be loaded when necessary to decrease footprint
+    from fastapi import Response
+
+
+def set_auth_cookie(web_config: WebConfig, response: "Response", token: Token):
+    access_token_max_age = web_config.access_token_expire_minutes * 60
+    refresh_token_max_age = web_config.refresh_token_expire_minutes * 60
+    now = datetime.now(timezone.utc)
+    response.set_cookie(
+        key=web_config.access_token_cookie_name,
+        value=token.access_token,
+        httponly=True,
+        max_age=access_token_max_age,
+        expires=now + timedelta(seconds=access_token_max_age),
+    )
+    response.set_cookie(
+        key=web_config.refresh_token_cookie_name,
+        value=token.refresh_token,
+        httponly=True,
+        max_age=refresh_token_max_age,
+        expires=now + timedelta(seconds=refresh_token_max_age),
+    )

zrb/runner/{web_util.py → web_util/html.py}

@@ -1,23 +1,9 @@
-import os
-
 from zrb.group.any_group import AnyGroup
-from zrb.runner.web_config import User
+from zrb.runner.web_schema.user import User
 from zrb.task.any_task import AnyTask
-from zrb.util.file import read_file
 from zrb.util.group import get_non_empty_subgroups, get_subtasks
 
 
-def url_to_args(url: str) -> list[str]:
-    stripped_url = url.strip("/")
-    return [part for part in stripped_url.split("/") if part.strip() != ""]
-
-
-def node_path_to_url(args: list[str]) -> str:
-    pruned_args = [part for part in args if part.strip() != ""]
-    stripped_url = "/".join(pruned_args)
-    return f"/{stripped_url}/"
-
-
 def get_html_auth_link(user: User) -> str:
     if user.is_guest and user.is_super_admin:
         return f"Hi, {user.username}"
@@ -26,14 +12,6 @@ def get_html_auth_link(user: User) -> str:
     return f'Hi, {user.username} <a href="/logout">Logout 🚪</a>'
 
 
-def get_refresh_token_js(refresh_interval_seconds: int):
-    _DIR = os.path.dirname(__file__)
-    return read_file(
-        os.path.join(_DIR, "refresh-token.template.js"),
-        {"refreshIntervalSeconds": f"{refresh_interval_seconds}"},
-    )
-
-
 def get_html_subtask_info(user: User, parent_url: str, parent_group: AnyGroup) -> str:
     subtasks = get_subtasks(parent_group, web_only=True)
     task_li = "\n".join(

zrb/runner/web_util/token.py

@@ -0,0 +1,72 @@
+from datetime import datetime, timedelta, timezone
+
+from zrb.runner.web_config.config import WebConfig
+from zrb.runner.web_schema.token import Token
+from zrb.runner.web_util.user import get_user_by_credentials
+
+
+def generate_tokens_by_credentials(
+    web_config: WebConfig, username: str, password: str
+) -> Token | None:
+    if not web_config.enable_auth:
+        user = web_config.default_user
+    else:
+        user = get_user_by_credentials(web_config, username, password)
+    if user is None:
+        return None
+    access_token = _generate_access_token(web_config, user.username)
+    refresh_token = _generate_refresh_token(web_config, user.username)
+    return Token(
+        access_token=access_token, refresh_token=refresh_token, token_type="bearer"
+    )
+
+
+def regenerate_tokens(web_config: WebConfig, refresh_token: str) -> Token:
+    from fastapi import HTTPException
+    from jose import jwt
+
+    # Decode and validate token
+    try:
+        payload = jwt.decode(
+            refresh_token,
+            web_config.secret_key,
+            options={"require_exp": True, "require_sub": True},
+        )
+    except Exception:
+        raise HTTPException(status_code=401, detail="Invalid JWT token")
+    if payload.get("type") != "refresh":
+        raise HTTPException(status_code=401, detail="Invalid token type")
+    username: str = payload.get("sub")
+    if username is None:
+        raise HTTPException(status_code=401, detail="Invalid refresh token")
+    user = web_config.find_user_by_username(username)
+    if user is None:
+        raise HTTPException(status_code=401, detail="User not found")
+    # Create new token
+    new_access_token = _generate_access_token(web_config, username)
+    new_refresh_token = _generate_refresh_token(web_config, username)
+    return Token(
+        access_token=new_access_token,
+        refresh_token=new_refresh_token,
+        token_type="bearer",
+    )
+
+
+def _generate_access_token(web_config: WebConfig, username: str) -> str:
+    from jose import jwt
+
+    expire = datetime.now(timezone.utc) + timedelta(
+        minutes=web_config.access_token_expire_minutes
+    )
+    to_encode = {"sub": username, "exp": expire, "type": "access"}
+    return jwt.encode(to_encode, web_config.secret_key)
+
+
+def _generate_refresh_token(web_config: WebConfig, username: str) -> str:
+    from jose import jwt
+
+    expire = datetime.now(timezone.utc) + timedelta(
+        minutes=web_config.refresh_token_expire_minutes
+    )
+    to_encode = {"sub": username, "exp": expire, "type": "refresh"}
+    return jwt.encode(to_encode, web_config.secret_key)

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))
+        return asyncio.run(self._run_and_cleanup(session, str_kwargs))
+
+    async def _run_and_cleanup(
+        self,
+        session: AnySession | None = None,
+        str_kwargs: dict[str, str] = {},
+    ) -> Any:
+        current_task = asyncio.create_task(self.async_run(session, str_kwargs))
+        try:
+            result = await current_task
+        finally:
+            if session and not session.is_terminated:
+                session.terminate()
+            # Cancel all running tasks except the current one
+            pending = [task for task in asyncio.all_tasks() if task is not current_task]
+            for task in pending:
+                task.cancel()
+            if pending:
+                try:
+                    await asyncio.wait(pending, timeout=5)
+                except asyncio.CancelledError:
+                    pass
+        return result
 
     async def async_run(
         self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
@@ -252,7 +274,8 @@ class BaseTask(AnyTask):
         # Update session
         self.__fill_shared_context_inputs(session.shared_ctx, str_kwargs)
         self.__fill_shared_context_envs(session.shared_ctx)
-        return await run_async(self.exec_root_tasks(session))
+        result = await run_async(self.exec_root_tasks(session))
+        return result
 
     def __fill_shared_context_inputs(
         self, shared_context: AnySharedContext, str_kwargs: dict[str, str] = {}
@@ -265,11 +288,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)
@@ -332,7 +353,7 @@ class BaseTask(AnyTask):
             session.get_task_status(self).mark_as_skipped()
             return
         # Wait for task to be ready
-        await run_async(self.__exec_action_until_ready(session))
+        return await run_async(self.__exec_action_until_ready(session))
 
     def __get_execute_condition(self, session: Session) -> bool:
         ctx = self.get_ctx(session)
@@ -416,6 +437,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 +471,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 +486,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 +497,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 +508,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