fix: Cleanup and fix lint/format issues (#731)

alembic/versions/0a85ea95e68c_add_schedule_timeout.py

@@ -5,6 +5,7 @@
 Create Date: 2024-11-27 11:46:05.761782
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/307c3bb9fad5_add_user_last_login_at.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-04 10:43:38.280178
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/3bc0a0970817_add_tags_and_workflow_tags_tables.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-23 00:34:49.841129
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/496c162975d7_add_organization_settings_table.py

@@ -5,6 +5,7 @@
 Create Date: 2025-01-09 00:02:47.484038
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/8be6393b2ee0_add_workflow_alias.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-24 18:10:17.717675
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/9194fb66b4ea_add_doc_url_and_author_to_.py

@@ -5,6 +5,7 @@
 Create Date: 2025-01-06 20:10:46.871277
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa
@@ -13,21 +14,27 @@
 from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = '9194fb66b4ea'
-down_revision: str | None = 'f2e840c9fb2c'
+revision: str = "9194fb66b4ea"
+down_revision: str | None = "f2e840c9fb2c"
 branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column('registryaction', sa.Column('doc_url', sqlmodel.sql.sqltypes.AutoString(), nullable=True))
-    op.add_column('registryaction', sa.Column('author', sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    op.add_column(
+        "registryaction",
+        sa.Column("doc_url", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    )
+    op.add_column(
+        "registryaction",
+        sa.Column("author", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    )
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_column('registryaction', 'author')
-    op.drop_column('registryaction', 'doc_url')
+    op.drop_column("registryaction", "author")
+    op.drop_column("registryaction", "doc_url")
     # ### end Alembic commands ###

alembic/versions/b7a3a2146bac_add_accesstoken_id.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-04 13:00:57.251229
 
 """
+
 # Import uuid for generating IDs
 from collections.abc import Sequence
 

alembic/versions/da03a0dfa2db_add_repo_last_synced_at_and_commit_sha.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-15 19:01:06.608544
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

alembic/versions/f2e840c9fb2c_add_workflow_error_handler.py

@@ -5,6 +5,7 @@
 Create Date: 2024-12-31 15:39:57.045107
 
 """
+
 from collections.abc import Sequence
 
 import sqlalchemy as sa

justfile

@@ -22,13 +22,20 @@ build:
 	docker compose build --no-cache
 
 lint-ui:
-	cd frontend && pnpm lint:fix && pnpm format:write && cd ..
+	cd frontend && pnpm lint:fix && cd ..
 lint-app:
-	ruff check .
+	ruff check
+
+lint-fix-ui:
+	cd frontend && pnpm lint:fix && pnpm format:write && cd ..
+lint-fix-app:
+	ruff check . && ruff format .
 
 lint: lint-ui lint-app
+lint-fix: lint-fix-ui lint-fix-app
+
 mypy path:
-	mypy --ignore-missing-imports --enable-incomplete-feature=NewGenericSyntax {{path}}
+	mypy --ignore-missing-imports {{path}}
 gen-client:
 	cd frontend && pnpm generate-client && cd ..
 # Update version number. If no version is provided, increments patch version.

pyproject.toml

@@ -72,6 +72,7 @@ dev = [
     "python-dotenv",
     "psycopg>=3.1.19",
     "respx",
+    "ruff==0.9.1",
 ]
 
 [tool.hatch.version]

tests/unit/test_executor.py

@@ -6,9 +6,8 @@
 import pytest
 from pydantic import SecretStr
 
-from tracecat.contexts import RunContext
 from tracecat.dsl.common import create_default_execution_context
-from tracecat.dsl.models import ActionStatement, RunActionInput
+from tracecat.dsl.models import ActionStatement, RunActionInput, RunContext
 from tracecat.executor.models import ExecutorActionErrorInfo
 from tracecat.executor.service import run_action_from_input, sync_executor_entrypoint
 from tracecat.expressions.expectations import ExpectedField

tests/unit/test_functions.py

@@ -746,9 +746,9 @@ def test_set_timezone(
     assert offset is not None
     offset_hours = offset.total_seconds() / 3600
     min_offset, max_offset = expected_range
-    assert (
-        min_offset <= offset_hours <= max_offset
-    ), f"Offset {offset_hours} not in expected range [{min_offset}, {max_offset}]"
+    assert min_offset <= offset_hours <= max_offset, (
+        f"Offset {offset_hours} not in expected range [{min_offset}, {max_offset}]"
+    )
 
 
 @pytest.mark.parametrize(

tests/unit/test_sandbox.py

@@ -6,7 +6,7 @@
 
 from tracecat.auth.credentials import TemporaryRole
 from tracecat.auth.sandbox import AuthSandbox
-from tracecat.contexts import ctx_env, ctx_role
+from tracecat.contexts import ctx_role, get_env
 from tracecat.db.schemas import BaseSecret
 from tracecat.secrets import secrets_manager
 from tracecat.secrets.encryption import encrypt_keyvalues
@@ -175,13 +175,13 @@ def test_env_sandbox_initial_env():
 
 def test_env_sandbox_isolation():
     """Test that changes inside the sandbox don't affect the outer environment."""
-    outer_env = ctx_env.get()
+    outer_env = get_env()
     with secrets_manager.env_sandbox({"TEMP_KEY": "temp_value"}):
         secrets_manager.set("NEW_KEY", "new_value")
         assert secrets_manager.get("NEW_KEY") == "new_value"
 
-    assert "NEW_KEY" not in ctx_env.get()
-    assert ctx_env.get() == outer_env
+    assert "NEW_KEY" not in get_env()
+    assert get_env() == outer_env
 
 
 def test_env_sandbox_nested():
@@ -208,14 +208,14 @@ async def async_function():
 
 def test_env_sandbox_exception():
     """Test that env_sandbox resets the environment even if an exception occurs."""
-    outer_env = ctx_env.get()
+    outer_env = get_env()
     try:
         with secrets_manager.env_sandbox({"EXCEPTION_KEY": "exception_value"}):
             raise ValueError("Test exception")
     except ValueError:
         pass
 
-    assert ctx_env.get() == outer_env
+    assert get_env() == outer_env
     assert secrets_manager.get("EXCEPTION_KEY") is None
 
 

tests/unit/test_workflows.py

@@ -2655,7 +2655,9 @@ async def test_workflow_error_handler_success(
         < eh_init_evt.event_id
         < eh_start_evt.event_id
         < eh_complete_evt.event_id
-    ), f"Event order is not correct: {fail_evt.event_id} < {eh_init_evt.event_id} < {eh_start_evt.event_id} < {eh_complete_evt.event_id}"
+    ), (
+        f"Event order is not correct: {fail_evt.event_id} < {eh_init_evt.event_id} < {eh_start_evt.event_id} < {eh_complete_evt.event_id}"
+    )
 
 
 @pytest.mark.parametrize(

tracecat/contexts.py

@@ -3,24 +3,16 @@
 from contextvars import ContextVar
 
 import loguru
-from pydantic import BaseModel
 
-from tracecat.identifiers import WorkflowExecutionID, WorkflowID, WorkflowRunID
+from tracecat.dsl.models import RunContext
 from tracecat.types.auth import Role
 
-
-class RunContext(BaseModel):
-    """This is the runtime context model for a workflow run. Passed into activities."""
-
-    wf_id: WorkflowID
-    wf_exec_id: WorkflowExecutionID
-    wf_run_id: WorkflowRunID
-    environment: str
-
-
 ctx_run: ContextVar[RunContext] = ContextVar("run", default=None)
 ctx_role: ContextVar[Role] = ContextVar("role", default=None)
 ctx_logger: ContextVar[loguru.Logger] = ContextVar("logger", default=None)
 
-SecretContextEnv = dict[str, dict[str, str]]
-ctx_env: ContextVar[SecretContextEnv] = ContextVar("env", default={})
+ctx_env: ContextVar[dict[str, str] | None] = ContextVar("env", default=None)
+
+
+def get_env() -> dict[str, str]:
+    return ctx_env.get() or {}

tracecat/db/engine.py

@@ -153,7 +153,7 @@ def get_session_context_manager() -> contextlib.AbstractContextManager[Session]:
     return contextlib.contextmanager(get_session)()
 
 
-def get_async_session_context_manager() -> (
-    contextlib.AbstractAsyncContextManager[AsyncSession]
-):
+def get_async_session_context_manager() -> contextlib.AbstractAsyncContextManager[
+    AsyncSession
+]:
     return contextlib.asynccontextmanager(get_async_session)()

tracecat/dsl/common.py

@@ -18,7 +18,6 @@
 )
 from pydantic_core import PydanticCustomError
 
-from tracecat.contexts import RunContext
 from tracecat.db.schemas import Action
 from tracecat.dsl.enums import EdgeType, FailStrategy, LoopStrategy
 from tracecat.dsl.models import (
@@ -27,6 +26,7 @@
     DSLEnvironment,
     DSLExecutionError,
     ExecutionContext,
+    RunContext,
     Trigger,
     TriggerInputs,
 )

tracecat/dsl/models.py

@@ -6,11 +6,11 @@
 
 from pydantic import BaseModel, Field
 
-from tracecat.contexts import RunContext
 from tracecat.dsl.constants import DEFAULT_ACTION_TIMEOUT
 from tracecat.dsl.enums import JoinStrategy
 from tracecat.expressions.common import ExprContext
 from tracecat.expressions.validation import ExpressionStr, TemplateValidator
+from tracecat.identifiers import WorkflowExecutionID, WorkflowID, WorkflowRunID
 from tracecat.secrets.constants import DEFAULT_SECRETS_ENVIRONMENT
 
 SLUG_PATTERN = r"^[a-z0-9_]+$"
@@ -174,6 +174,15 @@ class DSLEnvironment(TypedDict, total=False):
     """The registry version to use for the workflow."""
 
 
+class RunContext(BaseModel):
+    """This is the runtime context model for a workflow run. Passed into activities."""
+
+    wf_id: WorkflowID
+    wf_exec_id: WorkflowExecutionID
+    wf_run_id: WorkflowRunID
+    environment: str
+
+
 class RunActionInput(BaseModel):
     """This object contains all the information needed to execute an action."""
 

tracecat/dsl/workflow.py

@@ -26,7 +26,7 @@
 
     from tracecat import identifiers
     from tracecat.concurrency import GatheringTaskGroup
-    from tracecat.contexts import RunContext, ctx_logger, ctx_role, ctx_run
+    from tracecat.contexts import ctx_logger, ctx_role, ctx_run
     from tracecat.dsl.action import (
         DSLActivities,
         ValidateActionActivityInput,
@@ -48,6 +48,7 @@
         DSLNodeResult,
         ExecutionContext,
         RunActionInput,
+        RunContext,
         TriggerInputs,
     )
     from tracecat.dsl.scheduler import DSLScheduler
@@ -341,7 +342,7 @@ async def _run_workflow(self, args: DSLRunArgs) -> Any:
         if task_exceptions:
             n_exc = len(task_exceptions)
             formatted_exc = "\n".join(
-                f"{'='*20} ({i+1}/{n_exc}) {details.expr_context}.{ref} {'='*20}\n\n{info.exception!s}"
+                f"{'=' * 20} ({i + 1}/{n_exc}) {details.expr_context}.{ref} {'=' * 20}\n\n{info.exception!s}"
                 for i, (ref, info) in enumerate(task_exceptions.items())
                 if (details := info.details)
             )

tracecat/executor/models.py

@@ -33,7 +33,7 @@ class ExecutorActionErrorInfo(BaseModel):
     def __str__(self) -> str:
         return (
             f"{self.type}: {self.message}"
-            f"\n\n{'-'*30}"
+            f"\n\n{'-' * 30}"
             f"\nFile: {self.filename}"
             f"\nFunction: {self.function}"
             f"\nLine: {self.lineno}"

tracecat/expressions/core.py

@@ -94,7 +94,7 @@ def validate(
         except TracecatExpressionError as e:
             return visitor.add(
                 status="error",
-                msg=f"[{loc or "parser"}]\n\nError parsing expression `{self._expr}`\n\n{e}",
+                msg=f"[{loc or 'parser'}]\n\nError parsing expression `{self._expr}`\n\n{e}",
             )
 
         # 2) Validate the AST
@@ -103,7 +103,7 @@ def validate(
         except TracecatExpressionError as e:
             return visitor.add(
                 status="error",
-                msg=f"[{loc or "validator"}]\n\nError validating expression `{self._expr}`\n\n{e}",
+                msg=f"[{loc or 'validator'}]\n\nError validating expression `{self._expr}`\n\n{e}",
             )
 
     def extract(self, visitor: ExprExtractor) -> Mapping[ExprContext, set[str]]:

tracecat/expressions/parser/evaluator.py

@@ -166,7 +166,7 @@ def function(self, fn_name: str, fn_args: Sequence[Any]):
         fn = functions.FUNCTION_MAPPING.get(fn_name)
         if fn is None:
             raise TracecatExpressionError(
-                f"Unknown function {fn_name!r}." f" ({is_mapped=})"
+                f"Unknown function {fn_name!r}. ({is_mapped=})"
             )
         final_fn = fn.map if is_mapped else fn
         result = final_fn(*fn_args)

tracecat/expressions/parser/validator.py

@@ -238,7 +238,7 @@ def iterator(self, node: Tree):
         if collection.data in blacklist:
             self.add(
                 status="error",
-                msg=f"You cannot use {", ".join(repr(e) for e in blacklist)} expressions in the `for_each` collection.",
+                msg=f"You cannot use {', '.join(repr(e) for e in blacklist)} expressions in the `for_each` collection.",
                 type=ExprType.ITERATOR,
             )
 

tracecat/registry/repository.py

@@ -617,15 +617,13 @@ def _enforce_restrictions(fn: F) -> F:
     # Check for direct access to os.environ
     if "os" in names and "environ" in names:
         raise ValueError(
-            "`os.environ` usage is not allowed in user-defined code."
-            f" Found in: {path}"
+            f"`os.environ` usage is not allowed in user-defined code. Found in: {path}"
         )
 
     # Check for invocations of os.getenv
     if "os" in names and "getenv" in names:
         raise ValueError(
-            "`os.getenv()` usage is not allowed in user-defined code."
-            f" Found in: {path}"
+            f"`os.getenv()` usage is not allowed in user-defined code. Found in: {path}"
         )
 
     return fn