Restrict ansible-playbook's --limit option on specific operations

tdp/core/collections/collection_reader.py

@@ -7,11 +7,12 @@
 import logging
 from collections.abc import Generator
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import yaml
 from pydantic import BaseModel, ConfigDict, ValidationError
 
+from tdp.core.collections.playbook_validate import validate_playbook
 from tdp.core.constants import (
     DAG_DIRECTORY_NAME,
     DEFAULT_VARS_DIRECTORY_NAME,
@@ -20,7 +21,7 @@
     SCHEMA_VARS_DIRECTORY_NAME,
     YML_EXTENSION,
 )
-from tdp.core.entities.operation import Playbook
+from tdp.core.entities.operation import Playbook, PlaybookMeta
 from tdp.core.inventory_reader import InventoryReader
 from tdp.core.repository.utils.get_repository_version import get_repository_version
 from tdp.core.types import PathLike
@@ -32,6 +33,9 @@
 except ImportError:
     from yaml import Loader
 
+if TYPE_CHECKING:
+    from tdp.core.collections.playbook_validate import PlaybookIn
+
 MANDATORY_DIRECTORIES = [
     DAG_DIRECTORY_NAME,
     DEFAULT_VARS_DIRECTORY_NAME,
@@ -168,10 +172,12 @@ def read_dag_nodes(self) -> Generator[TDPLibDagNodeModel, None, None]:
     def read_playbooks(self) -> Generator[Playbook, None, None]:
         """Read the playbooks stored in the playbooks_directory."""
         for playbook_path in (self.playbooks_directory).glob("*" + YML_EXTENSION):
+            playbook: PlaybookIn = validate_playbook(playbook_path)
             yield Playbook(
                 path=playbook_path,
                 collection_name=self.name,
-                hosts=read_hosts_from_playbook(playbook_path, self._inventory_reader),
+                hosts=self._inventory_reader.get_hosts_from_playbook(playbook),
+                meta=_get_playbook_meta(playbook, playbook_path),
             )
 
     def read_schemas(self) -> list[ServiceCollectionSchema]:
@@ -216,25 +222,28 @@ def _check_collection_structure(self, path: Path) -> None:
                 )
 
 
-def read_hosts_from_playbook(
-    playbook_path: Path, inventory_reader: Optional[InventoryReader]
-) -> frozenset[str]:
-    """Read the hosts from a playbook.
-
-    Args:
-        playbook_path: Path to the playbook.
-        inventory_reader: Inventory reader.
+def _get_playbook_meta(playbook: PlaybookIn, playbook_path: Path) -> PlaybookMeta:
+    can_limit = True
+    can_limit_true_plays = list[str]()
+
+    for play_nb, play in enumerate(playbook):
+        play_name = f"{play.name}[{play_nb}]" if play.name else f"play[{play_nb}]"
+        if vars := play.vars:
+            if tdp_lib := vars.tdp_lib:
+                if tdp_lib.can_limit == True:
+                    can_limit_true_plays.append(play_name)
+                elif can_limit == True and tdp_lib.can_limit == False:
+                    can_limit = False
+
+    if can_limit == False and len(can_limit_true_plays) > 0:
+        logger.warning(
+            f"Playbook '{playbook_path}': tdp_lib.can_limit is both true and false "
+            "accross plays. Because a play sets 'can_limit: false', the playbook "
+            "can_limit is false; the 'can_limit: true' flags on these plays are "
+            "ignored: " + ", ".join(can_limit_true_plays)
+        )
 
-    Returns:
-        Set of hosts.
-    """
-    if not inventory_reader:
-        inventory_reader = InventoryReader()
-    try:
-        with playbook_path.open() as fd:
-            return inventory_reader.get_hosts_from_playbook(fd)
-    except Exception as e:
-        raise ValueError(f"Can't parse playbook {playbook_path}.") from e
+    return PlaybookMeta(can_limit=can_limit)
 
 
 def _get_galaxy_version(

tdp/core/collections/playbook_validate.py

@@ -0,0 +1,73 @@
+# Copyright 2025 TOSIT.IO
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Extract and validate custom `tdp_lib` metadata from an Ansible playbook.
+"""
+
+from pathlib import Path
+from typing import Optional, Union
+
+import yaml
+from pydantic import BaseModel, ConfigDict, Field, RootModel, ValidationError, conlist
+
+
+class _PlaybookPlayVarsMetaIn(BaseModel):
+    """Pydantic model describing the expected structure of a playbook's play[].vars.tdp_lib."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    can_limit: Optional[bool] = Field(
+        None,
+        description="Can the task be limited to specific host. false will be applied to the whole playbook.",
+    )
+
+
+class _PlaybookPlayVarsIn(BaseModel):
+    """Pydantic model describing the expected structure of a playbook's play[].vars."""
+
+    model_config = ConfigDict(frozen=True)
+
+    tdp_lib: Optional[_PlaybookPlayVarsMetaIn] = Field(None)
+
+
+class _PlaybookPlayIn(BaseModel):
+    """Pydantic model describing the expected structure of a playbook's play."""
+
+    model_config = ConfigDict(frozen=True)
+
+    hosts: Union[str, list[str]]
+    name: Optional[str] = Field(None)
+    vars: Optional[_PlaybookPlayVarsIn] = Field(None)
+
+
+class PlaybookIn(RootModel[conlist(_PlaybookPlayIn, min_length=1)]):
+    """Pydantic model describing the expected structure of a playbook."""
+
+    model_config = ConfigDict(frozen=True)
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+
+def validate_playbook(playbook_path: Path) -> PlaybookIn:
+    """Validate the content of a playbook file."""
+    # Open playbook file and get content
+    try:
+        with playbook_path.open() as f:
+            data = yaml.safe_load(f)
+    except Exception as exc:
+        raise ValueError(
+            f"Parsing error for playbook file: '{playbook_path}':\n{exc}"
+        ) from exc
+
+    # Validate the file
+    try:
+        return PlaybookIn.model_validate(data)
+    except ValidationError as exc:
+        raise ValueError(
+            f"Validation error for playbook file: '{playbook_path}':\n{exc}"
+        ) from exc

tdp/core/deployment/deployment_runner.py

@@ -54,20 +54,21 @@ def _run_operation(self, operation_rec: OperationModel) -> None:
 
         operation = self._collections.operations[operation_rec.operation]
 
-        # Check if the operation is available for the given host
-        if operation_rec.host and (
-            not isinstance(operation, PlaybookOperation)
-            or operation_rec.host not in operation.playbook.hosts
-        ):
-            logs = (
-                f"Operation '{operation_rec.operation}' not available for host "
-                + f"'{operation_rec.host}'"
-            )
-            logger.error(logs)
-            operation_rec.state = OperationStateEnum.FAILURE
-            operation_rec.logs = logs.encode("utf-8")
-            operation_rec.end_time = datetime.utcnow()
-            return
+        if operation_rec.host:
+            logs = None
+            if not isinstance(operation, PlaybookOperation):
+                logs = f"Operation '{operation_rec.operation}' is not related to any playbook (noop). It can't be limited to any host."
+            elif not operation.playbook.meta.can_limit:
+                logs = f"Operation '{operation_rec.operation}' can't be limited to specific host."
+            elif operation_rec.host not in operation.playbook.hosts:
+                logs = f"Operation '{operation_rec.operation}' is not available on host '{operation_rec.host}'."
+
+            if logs:
+                logger.error(logs)
+                operation_rec.state = OperationStateEnum.FAILURE
+                operation_rec.logs = logs.encode("utf-8")
+                operation_rec.end_time = datetime.utcnow()
+                return
 
         # Execute the operation
         playbook_file = self._collections.playbooks[operation.name.name].path

tdp/core/entities/operation.py

@@ -68,11 +68,22 @@ def __str__(self):
         return self.name
 
 
+@dataclass(frozen=True)
+class PlaybookMeta:
+    """Playbook metadata.
+
+    Aggregate data from each playbook's play.
+    """
+
+    can_limit: bool = True
+
+
 @dataclass(frozen=True)
 class Playbook:
     path: Path
     collection_name: str
     hosts: frozenset[str]
+    meta: PlaybookMeta
 
     def __post_init__(self):
         for host_name in self.hosts:

tdp/core/inventory_reader.py

@@ -3,16 +3,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, TextIO
-
-import yaml
+from typing import TYPE_CHECKING, Optional
 
 from tdp.core.ansible_loader import AnsibleLoader
-
-try:
-    from yaml import CLoader as Loader
-except ImportError:
-    from yaml import Loader
+from tdp.core.collections.playbook_validate import PlaybookIn
 
 if TYPE_CHECKING:
     from ansible.inventory.manager import InventoryManager
@@ -39,29 +33,13 @@ def get_hosts(self, *args, **kwargs) -> list[str]:
         # so we convert them to "str"
         return [str(host) for host in self.inventory.get_hosts(*args, **kwargs)]
 
-    def get_hosts_from_playbook(self, fd: TextIO) -> frozenset[str]:
+    def get_hosts_from_playbook(self, playbook: PlaybookIn) -> frozenset[str]:
         """Takes a playbook content, read all plays inside and return a set
         of matching host like "ansible-playbook --list-hosts playbook.yml".
-
-        Args:
-            fd: file-like object from which the playbook content must be read.
-
-        Returns:
-            Set of hosts.
         """
-        plays = yaml.load(fd, Loader=Loader)
-        if not isinstance(plays, list):
-            raise TypeError(f"Playbook content is not a list, given {type(plays)}")
-
         hosts: set[str] = set()
 
-        for play in plays:
-            if not isinstance(play, dict):
-                raise TypeError(f"A play must be a dict, given {type(play)}")
-            if "hosts" not in play:
-                raise ValueError(
-                    f"'hosts' key is mandatory for a play, keys are {play.keys()}"
-                )
-            hosts.update(self.get_hosts(play["hosts"]))
+        for play in playbook:
+            hosts.update(self.get_hosts(play.hosts))
 
         return frozenset(hosts)

tests/unit/core/collections/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2025 TOSIT.IO
+# SPDX-License-Identifier: Apache-2.0