feat: support multiple awx tokens (#586)

Author: Alex-Izquierdo

pyproject.toml

@@ -109,6 +109,7 @@ extend-ignore = [
     "S101",   # Asserts allowed in tests
     "ARG",    # Fixtures are not always used explicitly
     "SLF001", # Call private methods in tests
+    "D",      # Docstrings are not required in tests
 ]
 
 [tool.ruff.flake8-tidy-imports]

setup.cfg

@@ -11,3 +11,5 @@ per-file-ignores =
     # Ignore "E501: Lint too long" for migrations, because migrations are generated automatically
     # and the nesting level is high.
     src/aap_eda/core/migrations/*:E501
+    # Ignore docstring errors in tests.
+    tests/*:D

src/aap_eda/api/serializers/activation.py

@@ -48,6 +48,7 @@ class Meta:
             "created_at",
             "modified_at",
             "status_message",
+            "awx_token_id",
         ]
         read_only_fields = [
             "id",
@@ -84,6 +85,7 @@ class Meta:
             "created_at",
             "modified_at",
             "status_message",
+            "awx_token_id",
         ]
         read_only_fields = ["id", "created_at", "modified_at"]
 
@@ -111,12 +113,27 @@ def to_representation(self, activation):
             "created_at": activation.created_at,
             "modified_at": activation.modified_at,
             "status_message": activation.status_message,
+            "awx_token_id": activation.awx_token_id,
         }
 
 
 class ActivationCreateSerializer(serializers.ModelSerializer):
     """Serializer for creating the Activation."""
 
+    class Meta:
+        model = models.Activation
+        fields = [
+            "name",
+            "description",
+            "is_enabled",
+            "decision_environment_id",
+            "rulebook_id",
+            "extra_var_id",
+            "user",
+            "restart_policy",
+            "awx_token_id",
+        ]
+
     rulebook_id = serializers.IntegerField(
         validators=[validators.check_if_rulebook_exists]
     )
@@ -130,25 +147,26 @@ class ActivationCreateSerializer(serializers.ModelSerializer):
     )
     user = serializers.HiddenField(default=serializers.CurrentUserDefault())
 
+    awx_token_id = serializers.IntegerField(
+        allow_null=True,
+        validators=[validators.check_if_awx_token_exists],
+        required=False,
+    )
+
     def validate(self, data):
         user = data["user"]
-        validators.check_awx_tokens(user.id)
+        awx_token = models.AwxToken.objects.filter(
+            id=data.get("awx_token_id"),
+        ).first()
+        if awx_token and awx_token.user != user:
+            raise serializers.ValidationError(
+                "The Awx Token does not belong to the user."
+            )
+        if not awx_token:
+            validate_rulebook_token(data["rulebook_id"])
 
         return data
 
-    class Meta:
-        model = models.Activation
-        fields = [
-            "name",
-            "description",
-            "is_enabled",
-            "decision_environment_id",
-            "rulebook_id",
-            "extra_var_id",
-            "user",
-            "restart_policy",
-        ]
-
     def create(self, validated_data):
         rulebook_id = validated_data["rulebook_id"]
         rulebook = models.Rulebook.objects.get(id=rulebook_id)
@@ -225,6 +243,7 @@ class Meta:
             "modified_at",
             "restarted_at",
             "status_message",
+            "awx_token_id",
         ]
         read_only_fields = ["id", "created_at", "modified_at", "restarted_at"]
 
@@ -290,6 +309,7 @@ def to_representation(self, activation):
             "modified_at": activation.modified_at,
             "restarted_at": restarted_at,
             "status_message": activation.status_message,
+            "awx_token_id": activation.awx_token_id,
         }
 
 
@@ -305,10 +325,17 @@ class PostActivationSerializer(serializers.ModelSerializer):
         allow_null=True,
         validators=[validators.check_if_extra_var_exists],
     )
+    awx_token_id = serializers.IntegerField(
+        allow_null=True,
+        validators=[validators.check_if_awx_token_exists],
+    )
+    rulebook_id = serializers.IntegerField(allow_null=True)
 
     def validate(self, data):
-        user_id = self.initial_data["user_id"]
-        validators.check_awx_tokens(user_id)
+        awx_token = data.get("awx_token_id")
+
+        if not awx_token:
+            validate_rulebook_token(data["rulebook_id"])
 
         return data
 
@@ -323,6 +350,8 @@ class Meta:
             "user_id",
             "created_at",
             "modified_at",
+            "awx_token_id",
+            "rulebook_id",
         ]
         read_only_fields = [
             "id",
@@ -356,3 +385,23 @@ def parse_validation_errors(errors: dict) -> str:
     messages = {key: str(error[0]) for key, error in errors.items() if error}
 
     return str(messages)
+
+
+def validate_rulebook_token(rulebook_id: int) -> None:
+    """Validate if the rulebook requires an Awx Token."""
+    rulebook = models.Rulebook.objects.get(id=rulebook_id)
+
+    # TODO: rulesets are stored as a string in the rulebook model
+    # proper instrospection should require a validation of the
+    # rulesets. https://issues.redhat.com/browse/AAP-19202
+    try:
+        rulesets_data = rulebook.get_rulesets_data()
+    except ValueError:
+        raise serializers.ValidationError("Invalid rulebook data.")
+
+    if validators.check_rulesets_require_token(
+        rulesets_data,
+    ):
+        raise serializers.ValidationError(
+            "The rulebook requires an Awx Token.",
+        )

src/aap_eda/core/migrations/0014_activation_awx_token.py

@@ -0,0 +1,48 @@
+# Generated by Django 4.2.7 on 2023-12-21 18:50
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+def populate_awx_token(apps, schema_editor):
+    """Populate existing activations with AWX tokens.
+
+    Ensure that the awx_token field is populated for all existing
+    activations after the migration is applied. If the user has an
+    AWX token, use it. Always use the first AWX token, previous
+    versions of the code only allowed one AWX token per user.
+    """
+    Activation = apps.get_model("core", "Activation")  # noqa: N806
+    AWXToken = apps.get_model("core", "AWXToken")  # noqa: N806
+    for activation in Activation.objects.all():
+        user = activation.user
+        try:
+            awx_token = AWXToken.objects.filter(user=user).first()
+            if awx_token:
+                activation.awx_token = awx_token
+                activation.save()
+        except AWXToken.DoesNotExist:
+            pass
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("core", "0013_auditaction_status_message"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="activation",
+            name="awx_token",
+            field=models.ForeignKey(
+                default=None,
+                null=True,
+                on_delete=django.db.models.deletion.SET_NULL,
+                to="core.awxtoken",
+            ),
+        ),
+        migrations.RunPython(
+            populate_awx_token,
+            reverse_code=migrations.RunPython.noop,
+        ),
+    ]

src/aap_eda/core/models/activation.py

@@ -27,7 +27,7 @@
     UpdateFieldsRequiredError,
 )
 
-from .user import User
+from .user import AwxToken, User
 
 __all__ = (
     "Activation",
@@ -96,6 +96,12 @@ class Meta:
         on_delete=models.SET_NULL,
         related_name="+",
     )
+    awx_token = models.ForeignKey(
+        AwxToken,
+        on_delete=models.SET_NULL,
+        null=True,
+        default=None,
+    )
 
     def save(self, *args, **kwargs):
         # when creating

src/aap_eda/core/models/rulebook.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import yaml
 from django.db import models
 
 __all__ = (
@@ -36,11 +37,27 @@ class Meta:
     name = models.TextField(null=False)
     description = models.TextField(null=True, default="")
     # TODO: this field should not have a default value.
+    # TODO: should the content of this field be validated?
+    # https://issues.redhat.com/browse/AAP-19202
     rulesets = models.TextField(null=False, default="")
     project = models.ForeignKey("Project", on_delete=models.CASCADE, null=True)
     created_at = models.DateTimeField(auto_now_add=True, null=False)
     modified_at = models.DateTimeField(auto_now=True, null=False)
 
+    # For instrospection purposes we need to return
+    # rulesets data unserialized.
+    def get_rulesets_data(self) -> list[dict]:
+        """Return rulesets data as a list of dicts."""
+        try:
+            return yaml.safe_load(self.rulesets)
+        except yaml.YAMLError as e:
+            raise ValueError(
+                (
+                    "Unable to parse rulesets data for rulebook "
+                    f" {self.id} - {self.name}: Error: {e}"
+                )
+            )
+
 
 class Ruleset(models.Model):
     class Meta:

src/aap_eda/core/validators.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 import logging
+import typing as tp
 
 import yaml
 from rest_framework import serializers
@@ -58,17 +59,43 @@ def check_if_extra_var_exists(extra_var_id: int) -> int:
     return extra_var_id
 
 
-def check_awx_tokens(user_id: int) -> int:
-    tokens = models.AwxToken.objects.filter(user_id=user_id).count()
-    if tokens == 0:
-        raise serializers.ValidationError("No controller token specified")
-    elif tokens > 1:
+def check_if_awx_token_exists(awx_token_id: int) -> int:
+    try:
+        models.AwxToken.objects.get(pk=awx_token_id)
+    except models.AwxToken.DoesNotExist:
         raise serializers.ValidationError(
-            "More than one controller token found, "
-            "currently only 1 token is supported"
+            f"AwxToken with id {awx_token_id} does not exist"
         )
+    return awx_token_id
+
+
+def check_rulesets_require_token(
+    rulesets_data: list[dict[str, tp.Any]],
+) -> bool:
+    """Inspect rulesets data to determine if a token is required.
+
+    Return True if any of the rules has an action that requires a token.
+    """
+    required_actions = {"run_job_template", "run_workflow_template"}
+
+    for ruleset in rulesets_data:
+        for rule in ruleset.get("rules", []):
+            # When it is a single action dict
+            if any(
+                action_key in required_actions
+                for action_key in rule.get("action", {})
+            ):
+                return True
+
+            # When it is a list of actions
+            if any(
+                action_arg in required_actions
+                for action in rule.get("actions", [])
+                for action_arg in action
+            ):
+                return True
 
-    return user_id
+    return False
 
 
 def is_extra_var_dict(extra_var: str):

src/aap_eda/wsapi/consumers.py

@@ -1,6 +1,7 @@
 import base64
 import json
 import logging
+import typing as tp
 from datetime import datetime
 from enum import Enum
 
@@ -111,14 +112,15 @@ async def handle_workers(self, message: WorkerMessage):
             )
             await self.send(text_data=extra_var_message.json())
 
-        controller_info = ControllerInfo(
-            url=settings.EDA_CONTROLLER_URL,
-            token=await self.get_awx_token(message),
-            ssl_verify=settings.EDA_CONTROLLER_SSL_VERIFY,
-        )
-
         await self.send(text_data=rulebook_message.json())
-        await self.send(text_data=controller_info.json())
+        awx_token = await self.get_awx_token(message)
+        if awx_token:
+            controller_info = ControllerInfo(
+                url=settings.EDA_CONTROLLER_URL,
+                token=awx_token,
+                ssl_verify=settings.EDA_CONTROLLER_SSL_VERIFY,
+            )
+            await self.send(text_data=controller_info.json())
         await self.send(text_data=EndOfResponse().json())
 
         # TODO: add broadcasting later by channel groups
@@ -320,15 +322,11 @@ def get_resources(
         return activation.rulebook_rulesets, extra_var
 
     @database_sync_to_async
-    def get_awx_token(self, message):
-        # query for activation
+    def get_awx_token(self, message: WorkerMessage) -> tp.Optional[str]:
+        """Get AWX token from the worker message."""
         activation_instance = models.ActivationInstance.objects.get(
-            id=message.activation_id
+            id=message.activation_id,
         )
+        awx_token = activation_instance.activation.awx_token
 
-        # check/get AWX token
-        awx_token = models.AwxToken.objects.filter(
-            user_id=activation_instance.activation.user_id
-        ).first()
-
-        return awx_token.token.get_secret_value()
+        return awx_token.token.get_secret_value() if awx_token else None