feat: preflight-check batch:TagResource when tags are configured

When tags are set (via --aws-batch-tags or SNAKEMAKE_AWS_BATCH_JOB_TAGS), every
job is tagged at submit time, so a missing batch:TagResource surfaces as an
opaque AccessDenied an hour into the run. Extend _preflight_validate with a
non-destructive tag/untag round-trip on the job queue ARN so the permission gap
fails fast at startup instead. Only a permissions error raises; any other error
(or no tags / no queue ARN) degrades to a no-op. Adds tag_resource/untag_resource
wrappers to BatchClient and documents the batch:TagResource/UntagResource needs.

Author: nh13

docs/further.md

@@ -261,6 +261,19 @@ misconfiguration (a disabled/invalid queue or compute environment, `maxvCpus=0`,
 or a non-existent job role) fails fast with a clear error, so you don't wait for
 jobs that could never start.
 
+When tags are configured (via `--aws-batch-tags` or the
+`SNAKEMAKE_AWS_BATCH_JOB_TAGS` environment variable), the executor also runs a
+tag/untag round-trip on the job queue as a best-effort *proxy* for the
+`batch:TagResource` permission — otherwise a missing tag permission would only
+surface as an opaque `AccessDenied` at submission time. This probes the queue,
+whereas jobs are tagged on the job and job-definition resources, so an IAM
+policy that scopes `batch:TagResource` per resource cannot be fully verified
+this way — treat a pass or failure as a strong hint, not a guarantee. The check
+needs `batch:TagResource` and `batch:UntagResource` (and possibly
+`ecs:TagResource`, depending on your account's tag-authorization settings); if
+the cleanup untag is denied, the throwaway `snakemake-preflight` tag is left on
+the queue.
+
 The check is deliberately conservative about *uncertainty*: a transient API
 error, a queue mid-update (`status` `CREATING`/`UPDATING`), or a missing
 `batch:Describe*` / `iam:GetRole` permission is treated as a warning and the

snakemake_executor_plugin_aws_batch/__init__.py

@@ -3,6 +3,7 @@
 __email__ = "jake.vancampen7@gmail.com"
 __license__ = "MIT"
 
+import os
 from dataclasses import dataclass, field
 from pprint import pformat
 from typing import List, AsyncGenerator, Optional
@@ -11,7 +12,10 @@
 from botocore.exceptions import ClientError
 
 from snakemake_executor_plugin_aws_batch.batch_client import BatchClient
-from snakemake_executor_plugin_aws_batch.batch_job_builder import BatchJobBuilder
+from snakemake_executor_plugin_aws_batch.batch_job_builder import (
+    BatchJobBuilder,
+    SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR,
+)
 from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo
 from snakemake_interface_executor_plugins.executors.remote import RemoteExecutor
 from snakemake_interface_executor_plugins.settings import (
@@ -30,6 +34,18 @@
 FATAL_BATCH_STATUSES = frozenset({"INVALID", "DELETING", "DELETED"})
 
 
+def _is_access_denied(error: ClientError) -> bool:
+    """True if a botocore ClientError is a permissions failure (vs. transient)."""
+    response = getattr(error, "response", {}) or {}
+    code = response.get("Error", {}).get("Code", "")
+    status = response.get("ResponseMetadata", {}).get("HTTPStatusCode")
+    return code in {
+        "AccessDenied",
+        "AccessDeniedException",
+        "UnauthorizedOperation",
+    } or status in (401, 403)
+
+
 # Optional:
 # Define additional settings for your executor.
 # They will occur in the Snakemake CLI as --<executor-name>-<param-name>
@@ -362,6 +378,72 @@ def _preflight_validate(self) -> None:
                 "environment(s)."
             )
         self._validate_job_role()
+        self._preflight_check_tags()
+
+    def _preflight_check_tags(self) -> None:
+        """Best-effort probe of ``batch:TagResource`` when tags are configured.
+
+        When tags are set (via ``--aws-batch-tags`` or the
+        ``SNAKEMAKE_AWS_BATCH_JOB_TAGS`` env var) every job is tagged at submit
+        time, so a missing ``batch:TagResource`` would otherwise surface as an
+        opaque ``AccessDenied`` an hour into the run. This runs a tag/untag
+        round-trip on the job queue ARN as a *proxy* for that permission and
+        fails fast if it is denied.
+
+        The probe is a heuristic, not a proof: jobs are tagged on the *job* and
+        *job-definition* resources (via ``submit_job`` / ``register_job_definition``),
+        not the queue, so an IAM policy that scopes ``batch:TagResource`` per
+        resource-type/ARN could grant it on the queue while denying it on the
+        job/job-definition resources (or the reverse). Treat a pass as "the
+        permission is very likely present" and a denial as "very likely
+        missing", not a guarantee either way.
+
+        Best-effort: only a permissions denial raises; any other error (or a
+        missing queue ARN / no tags configured) is a no-op. If the cleanup
+        untag fails, the throwaway ``snakemake-preflight`` tag is left on the
+        job queue (logged at debug) — so the round-trip is not guaranteed to be
+        fully non-destructive.
+        """
+        if not self._tags_configured():
+            return
+        queue_arn = getattr(self.settings, "job_queue", None)
+        if not queue_arn:
+            return
+
+        try:
+            self.batch_client.tag_resource(
+                resourceArn=queue_arn, tags={"snakemake-preflight": "1"}
+            )
+            # Clean up the throwaway tag; failure to untag is non-fatal.
+            try:
+                self.batch_client.untag_resource(
+                    resourceArn=queue_arn, tagKeys=["snakemake-preflight"]
+                )
+            except Exception as e:
+                self.logger.debug(
+                    "could not remove throwaway 'snakemake-preflight' tag from "
+                    f"the job queue (left in place): {e}"
+                )
+        except ClientError as e:
+            if _is_access_denied(e):
+                raise WorkflowError(
+                    "AWS Batch preflight check failed — the executor role was "
+                    "denied batch:TagResource on the job queue, but tags are "
+                    "configured. Job tagging at submit time would very likely "
+                    "fail with AccessDenied. Grant batch:TagResource (and "
+                    "ecs:TagResource if your account requires it) on the job / "
+                    "job-definition resources, or unset the tags."
+                ) from e
+            self.logger.debug(f"could not verify batch:TagResource: {e}")
+        except Exception as e:
+            self.logger.debug(f"could not verify batch:TagResource: {e}")
+
+    def _tags_configured(self) -> bool:
+        """True if any tags are set via the setting or the env var."""
+        settings_tags = getattr(self.settings, "tags", None)
+        if isinstance(settings_tags, dict) and settings_tags:
+            return True
+        return bool(os.environ.get(SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR, "").strip())
 
     def _queue_problems(self) -> Optional[List[str]]:
         """Return definitive job-queue / compute-environment misconfigurations.

snakemake_executor_plugin_aws_batch/batch_client.py

@@ -91,3 +91,21 @@ def describe_compute_environments(self, **kwargs):
         :return: The response from the describe_compute_environments method.
         """
         return self.client.describe_compute_environments(**kwargs)
+
+    def tag_resource(self, **kwargs):
+        """
+        Add tags to an AWS Batch resource.
+
+        :param kwargs: The keyword arguments to pass to the tag_resource method.
+        :return: The response from the tag_resource method.
+        """
+        return self.client.tag_resource(**kwargs)
+
+    def untag_resource(self, **kwargs):
+        """
+        Remove tags from an AWS Batch resource.
+
+        :param kwargs: The keyword arguments to pass to the untag_resource method.
+        :return: The response from the untag_resource method.
+        """
+        return self.client.untag_resource(**kwargs)

tests/test_preflight.py

@@ -7,14 +7,24 @@
 permission).
 """
 
+import os
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 from botocore.exceptions import ClientError
 from snakemake_interface_common.exceptions import WorkflowError
 
-from snakemake_executor_plugin_aws_batch import Executor
+from snakemake_executor_plugin_aws_batch import Executor, _is_access_denied
+from snakemake_executor_plugin_aws_batch.batch_job_builder import (
+    SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR,
+)
+
+
+def _env_without_tags() -> dict:
+    return {
+        k: v for k, v in os.environ.items() if k != SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR
+    }
 
 
 def _executor(**settings) -> Executor:
@@ -295,6 +305,86 @@ def test_no_queue_configured_does_not_raise(self):
         ex._validate_job_role.assert_called_once()
 
 
+class TestPreflightCheckTags:
+    """The tag/untag round-trip that verifies batch:TagResource when tags are set."""
+
+    def _access_denied(self) -> ClientError:
+        return ClientError({"Error": {"Code": "AccessDenied"}}, "TagResource")
+
+    def test_no_tags_configured_skips(self):
+        ex = _executor()
+        with patch.dict(os.environ, _env_without_tags(), clear=True):
+            ex._preflight_check_tags()
+        ex.batch_client.tag_resource.assert_not_called()
+
+    def test_no_queue_skips(self):
+        ex = _executor(job_queue=None, tags={"Env": "prod"})
+        ex._preflight_check_tags()
+        ex.batch_client.tag_resource.assert_not_called()
+
+    def test_success_tags_then_untags(self):
+        ex = _executor(tags={"Env": "prod"})
+        ex._preflight_check_tags()  # must not raise
+        ex.batch_client.tag_resource.assert_called_once()
+        ex.batch_client.untag_resource.assert_called_once()
+
+    def test_env_var_tags_trigger_check(self):
+        ex = _executor()
+        with patch.dict(
+            os.environ, {SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR: "Team=data"}
+        ):
+            ex._preflight_check_tags()
+        ex.batch_client.tag_resource.assert_called_once()
+
+    def test_missing_permission_raises(self):
+        ex = _executor(tags={"Env": "prod"})
+        ex.batch_client.tag_resource.side_effect = self._access_denied()
+        with pytest.raises(WorkflowError, match="batch:TagResource"):
+            ex._preflight_check_tags()
+
+    def test_non_permission_error_degrades(self):
+        ex = _executor(tags={"Env": "prod"})
+        ex.batch_client.tag_resource.side_effect = ClientError(
+            {"Error": {"Code": "ThrottlingException"}}, "TagResource"
+        )
+        ex._preflight_check_tags()  # must not raise
+
+    @pytest.mark.parametrize("code", ["AccessDeniedException", "UnauthorizedOperation"])
+    def test_alternate_access_denied_codes_raise(self, code):
+        ex = _executor(tags={"Env": "prod"})
+        ex.batch_client.tag_resource.side_effect = ClientError(
+            {"Error": {"Code": code}}, "TagResource"
+        )
+        with pytest.raises(WorkflowError, match="batch:TagResource"):
+            ex._preflight_check_tags()
+
+    def test_http_forbidden_status_is_access_denied(self):
+        # A permissions failure signalled only by HTTP 403 (unrecognized code).
+        err = ClientError(
+            {
+                "Error": {"Code": "Forbidden"},
+                "ResponseMetadata": {"HTTPStatusCode": 403},
+            },
+            "TagResource",
+        )
+        assert _is_access_denied(err) is True
+
+    def test_untag_failure_is_non_fatal(self):
+        # Tag succeeds but the cleanup untag fails: must not raise (the stray
+        # snakemake-preflight tag is left behind and only logged at debug).
+        ex = _executor(tags={"Env": "prod"})
+        ex.batch_client.untag_resource.side_effect = Exception("no UntagResource")
+        ex._preflight_check_tags()  # must not raise
+        ex.batch_client.tag_resource.assert_called_once()
+
+    def test_untag_not_called_when_tagging_denied(self):
+        ex = _executor(tags={"Env": "prod"})
+        ex.batch_client.tag_resource.side_effect = self._access_denied()
+        with pytest.raises(WorkflowError, match="batch:TagResource"):
+            ex._preflight_check_tags()
+        ex.batch_client.untag_resource.assert_not_called()
+
+
 class TestValidateJobRole:
     def _iam(self, **get_role_kwargs) -> MagicMock:
         return MagicMock(get_role=MagicMock(**get_role_kwargs))