add tag support to primary API

Author: snopoke

taskbadger/cli/basics.py

@@ -11,7 +11,7 @@
     configure_api,
     err_console,
     get_actions,
-    get_metadata,
+    merge_kv_json,
 )
 
 
@@ -69,12 +69,23 @@ def create(
         show_default=False,
         help="Metadata to associate with the task. Must be valid JSON.",
     ),
+    tag: list[str] = typer.Option(
+        None,
+        show_default=False,
+        help="Metadata 'key=value' pair to associate with the task. Can be specified multiple times.",
+    ),
+    tags_json: str = typer.Option(
+        None,
+        show_default=False,
+        help="Tags to associate with the task. Must be valid JSON mapping name -> value.",
+    ),
     quiet: bool = typer.Option(False, "--quiet", "-q", help="Minimal output. Only the Task ID."),
 ):
     """Create a task."""
     configure_api(ctx)
     actions = get_actions(action_def)
-    metadata = get_metadata(metadata, metadata_json)
+    metadata = merge_kv_json(metadata, metadata_json)
+    tags = merge_kv_json(tag, tags_json)
 
     try:
         task = create_task(
@@ -84,6 +95,7 @@ def create(
             data=metadata,
             actions=actions,
             monitor_id=monitor_id,
+            tags=tags,
         )
     except Exception as e:
         err_console.print(f"Error creating task: {e}")
@@ -119,12 +131,23 @@ def update(
         show_default=False,
         help="Metadata to associate with the task. Must be valid JSON.",
     ),
+    tag: list[str] = typer.Option(
+        None,
+        show_default=False,
+        help="Metadata 'key=value' pair to associate with the task. Can be specified multiple times.",
+    ),
+    tags_json: str = typer.Option(
+        None,
+        show_default=False,
+        help="Tags to associate with the task. Must be valid JSON mapping name -> value.",
+    ),
     quiet: bool = typer.Option(False, "--quiet", "-q", help="No output."),
 ):
     """Update a task."""
     configure_api(ctx)
     actions = get_actions(action_def)
-    metadata = get_metadata(metadata, metadata_json)
+    metadata = merge_kv_json(metadata, metadata_json)
+    tags = merge_kv_json(tag, tags_json)
 
     try:
         task = update_task(
@@ -135,6 +158,7 @@ def update(
             value_max=value_max,
             data=metadata,
             actions=actions,
+            tags=tags,
         )
     except Exception as e:
         err_console.print(f"Error creating task: {e}")

taskbadger/cli/utils.py

@@ -28,7 +28,7 @@ def get_actions(action_def: tuple[str, str, str]) -> list[Action]:
     return []
 
 
-def get_metadata(metadata_kv: list[str], metadata_json: str) -> dict:
+def merge_kv_json(metadata_kv: list[str], metadata_json: str) -> dict:
     metadata = {}
     for kv in metadata_kv:
         k, v = kv.strip().split("=", 1)

taskbadger/cli/wrapper.py

@@ -2,7 +2,7 @@
 from rich import print
 
 from taskbadger import DefaultMergeStrategy, Session, StatusEnum, Task
-from taskbadger.cli.utils import configure_api, err_console, get_actions
+from taskbadger.cli.utils import configure_api, err_console, get_actions, merge_kv_json
 from taskbadger.process import ProcessRunner
 
 
@@ -19,6 +19,11 @@ def run(
         show_default=False,
         help="Action definition e.g. 'success,error email to:me@email.com'",
     ),
+    tag: list[str] = typer.Option(
+        None,
+        show_default=False,
+        help="Tags: 'name=value' pair to associate with the task. Can be specified multiple times.",
+    ),
     capture_output: bool = typer.Option(False, help="Capture stdout and stderr."),
 ):
     """Execute a command using the CLI and create a Task to track its outcome.
@@ -34,6 +39,7 @@ def run(
     """
     configure_api(ctx)
     actions = get_actions(action_def)
+    tags = merge_kv_json(tag, "")
     stale_timeout = update_frequency * 2
     with Session():
         try:
@@ -43,6 +49,7 @@ def run(
                 stale_timeout=stale_timeout,
                 actions=actions,
                 monitor_id=monitor_id,
+                tags=tags,
             )
         except Exception as e:
             err_console.print(f"Error creating task: {e}")

taskbadger/sdk.py

@@ -18,8 +18,10 @@
 )
 from taskbadger.internal.models import (
     PatchedTaskRequest,
+    PatchedTaskRequestTags,
     StatusEnum,
     TaskRequest,
+    TaskRequestTags,
 )
 from taskbadger.internal.types import UNSET
 from taskbadger.mug import Badger, Session, Settings
@@ -95,6 +97,7 @@ def create_task(
     stale_timeout: int = None,
     actions: list[Action] = None,
     monitor_id: str = None,
+    tags: dict[str, str] = None,
 ) -> "Task":
     """Create a Task.
 
@@ -108,6 +111,7 @@ def create_task(
         stale_timeout: Maximum allowed time between updates (seconds).
         actions: Task actions.
         monitor_id: ID of the monitor to associate this task with.
+        tags: Dictionary of namespace -> value tags.
 
     Returns:
         Task: The created Task object.
@@ -132,6 +136,8 @@ def create_task(
         task.data = {**scope_data, **data}
     if actions:
         task.additional_properties = {"actions": [a.to_dict() for a in actions]}
+    if tags:
+        task.tags = TaskRequestTags.from_dict(tags)
     kwargs = _make_args(body=task)
     if monitor_id:
         kwargs["x_taskbadger_monitor"] = monitor_id
@@ -151,6 +157,7 @@ def update_task(
     max_runtime: int = None,
     stale_timeout: int = None,
     actions: list[Action] = None,
+    tags: dict[str, str] = None,
 ) -> "Task":
     """Update a task.
     Requires only the task ID and fields to update.
@@ -165,6 +172,7 @@ def update_task(
         max_runtime: Maximum expected runtime (seconds).
         stale_timeout: Maximum allowed time between updates (seconds).
         actions: Task actions.
+        tags: Dictionary of namespace -> value tags.
 
     Returns:
         Task: The updated Task object.
@@ -189,6 +197,8 @@ def update_task(
     )
     if actions:
         body.additional_properties = {"actions": [a.to_dict() for a in actions]}
+    if tags:
+        body.tags = PatchedTaskRequestTags.from_dict(tags)
     kwargs = _make_args(id=task_id, body=body)
     with Session() as client:
         response = task_partial_update.sync_detailed(client=client, **kwargs)
@@ -245,6 +255,7 @@ def create(
         stale_timeout: int = None,
         actions: list[Action] = None,
         monitor_id: str = None,
+        tags: dict[str, str] = None,
     ) -> "Task":
         """Create a new task"""
         return create_task(
@@ -257,6 +268,7 @@ def create(
             stale_timeout=stale_timeout,
             actions=actions,
             monitor_id=monitor_id,
+            tags=tags,
         )
 
     def __init__(self, task):
@@ -321,6 +333,7 @@ def update(
         max_runtime: int = None,
         stale_timeout: int = None,
         actions: list[Action] = None,
+        tags: dict[str, str] = None,
         data_merge_strategy: Any = None,
     ):
         """Generic update method used to update any of the task fields.
@@ -345,13 +358,18 @@ def update(
             max_runtime=max_runtime,
             stale_timeout=stale_timeout,
             actions=actions,
+            tags=tags,
         )
         self._task = task._task
 
     def add_actions(self, actions: list[Action]):
         """Add actions to the task."""
         self.update(actions=actions)
 
+    def tag(self, tags: dict[str, str]):
+        """Add tags to the task."""
+        self.update(tags=tags)
+
     def ping(self):
         """Update the task without changing any values. This can be used in conjunction
         with 'stale_timeout' to indicate that the task is still running."""

tests/test_cli_create_update.py

@@ -8,8 +8,10 @@
 from taskbadger.cli_main import app
 from taskbadger.internal.models import (
     PatchedTaskRequest,
+    PatchedTaskRequestTags,
     StatusEnum,
     TaskRequest,
+    TaskRequestTags,
 )
 from taskbadger.internal.types import Response
 from tests.utils import task_for_test
@@ -47,6 +49,10 @@ def test_cli_create():
             "b=2",
             "--metadata",
             "a=3",
+            "--tag",
+            "name=foo",
+            "--tag",
+            "bar=baz",
         ]
         result = runner.invoke(app, args, catch_exceptions=False)
         assert result.exit_code == 0, result.output
@@ -56,6 +62,7 @@ def test_cli_create():
             status=StatusEnum.PROCESSING,
             value_max=100,
             data={"b": "2", "a": 1, "c": 1},
+            tags=TaskRequestTags.from_dict({"name": "foo", "bar": "baz"}),
         )
         create.assert_called_with(
             client=mock.ANY,
@@ -72,12 +79,14 @@ def test_cli_update():
 
         result = runner.invoke(
             app,
-            ["update", "task123", "--status=success", "--value", "100"],
+            ["update", "task123", "--status=success", "--value", "100", "--tag", "name=foo"],
             catch_exceptions=False,
         )
         assert result.exit_code == 0, result.output
 
-        body = PatchedTaskRequest(status=StatusEnum.SUCCESS, value=100)
+        body = PatchedTaskRequest(
+            status=StatusEnum.SUCCESS, value=100, tags=PatchedTaskRequestTags.from_dict({"name": "foo"})
+        )
 
         update.assert_called_with(
             client=mock.ANY,

tests/test_cli_run.py

@@ -10,6 +10,7 @@
     PatchedTaskRequest,
     StatusEnum,
     TaskRequest,
+    TaskRequestTags,
 )
 from taskbadger.internal.types import UNSET, Response
 from taskbadger.mug import Badger
@@ -37,6 +38,15 @@ def test_cli_run_success():
     _test_cli_run(["echo", "test"], 0, args=["task_name"])
 
 
+def test_cli_run_tags():
+    _test_cli_run(
+        ["echo", "test"],
+        0,
+        ["task_name", "--tag", "name=value", "--tag", "name1=value1"],
+        tags={"name": "value", "name1": "value1"},
+    )
+
+
 def test_cli_long_run():
     def _should_update_task(last_update, update_frequency_seconds):
         return True
@@ -106,7 +116,7 @@ def test_cli_run_webhook():
     )
 
 
-def _test_cli_run(command, return_code, args=None, action=None, update_call_count=1):
+def _test_cli_run(command, return_code, args=None, action=None, tags=None, update_call_count=1):
     update_mock = mock.MagicMock()
 
     def _update(*args, **kwargs):
@@ -133,6 +143,8 @@ def _update(*args, **kwargs):
         request = TaskRequest(name="task_name", status=StatusEnum.PROCESSING, stale_timeout=10)
         if action:
             request.additional_properties = {"actions": [action]}
+        if tags:
+            request.tags = TaskRequestTags.from_dict(tags)
         create.assert_called_with(
             client=mock.ANY,
             organization_slug="org",

tests/test_sdk_primatives.py

@@ -8,7 +8,7 @@
 
 
 @pytest.fixture(autouse=True)
-def init_skd():
+def _init_skd():
     init("org", "project", "token")
 
 
@@ -75,6 +75,7 @@ def test_update_task(httpx_mock):
         "value": 150,
         "value_max": 150,
         "data": {"custom": "value"},
+        "tags": {"tag": "value"},
     }
     httpx_mock.add_response(
         url="https://taskbadger.net/api/org/project/tasks/test_id/",