feat: compatibility with BatchGetJobEntity API changes for jobRunAsUser

- provides forwards/backwards compatibility with the current API usage of the BatchGetJobEntity "jobDetails" entity responses - Updates some developer test code and configuration to accomodate API changes to CreateQueue and to be able to end-to-end test these changes Signed-off-by: Josh Usiskin <56369778+jusiskin@users.noreply.github.com>
aws-deadline · Jan 12, 2024 · 40fcd7c · 40fcd7c
1 parent f607954
commit 40fcd7c
Show file tree

Hide file tree

Showing 12 changed files with 340 additions and 26 deletions.
diff --git a/scripts/create_service_resources.sh b/scripts/create_service_resources.sh
@@ -10,10 +10,10 @@ if test $# -lt 1
 then
     echo "Usage: $0 <Fleet Role Arn> <s3 bucket name> [<Queue#1 Role Arn> [<Queue#2 Role Arn]]"
     echo "where:"
-    echo "  <s3 bucketname>: The name of an S3 bucket to configure on the Queues for use"
-    echo "                   with the Job Attachments feature."
     echo "  <Fleet Role Arn>: The ARN of the Worker Role to attach to the Fleet;"
     echo "                    this is used by the worker agent during operations."
+    echo "  <s3 bucketname>: The name of an S3 bucket to configure on the Queues for use"
+    echo "                   with the Job Attachments feature."
     echo "  <Queue* Role Arn>: The ARN of the Role whose credentials will be provided"
     echo "                      to the running jobs."
     exit 1
@@ -50,9 +50,10 @@ then
     "status": "IDLE",
     "jobRunAsUser": {
         "posix": {
-            "user": "",
-            "group": ""
-        }
+            "user": "jobuser",
+            "group": "jobuser"
+        },
+        "runAs": "QUEUE_CONFIGURED_USER"
     },
     "jobAttachmentSettings": {
         "s3BucketName": "${assets_s3_bucket}",
@@ -69,9 +70,10 @@ EOF
     "roleArn": "$queue_1_iam_role",
     "jobRunAsUser": {
         "posix": {
-            "user": "",
-            "group": ""
-        }
+            "user": "jobuser",
+            "group": "jobuser"
+        },
+        "runAs": "QUEUE_CONFIGURED_USER"
     },
     "jobAttachmentSettings": {
         "s3BucketName": "${assets_s3_bucket}",
@@ -113,9 +115,10 @@ then
     "status": "IDLE",
     "jobRunAsUser": {
         "posix": {
-            "user": "",
-            "group": ""
-        }
+            "user": "jobuser",
+            "group": "jobuser"
+        },
+        "runAs": "QUEUE_CONFIGURED_USER"
     },
     "jobAttachmentSettings": {
         "s3BucketName": "${assets_s3_bucket}",
@@ -132,9 +135,10 @@ EOF
     "roleArn": "$queue_2_iam_role",
     "jobRunAsUser": {
         "posix": {
-            "user": "",
-            "group": ""
-        }
+            "user": "jobuser",
+            "group": "jobuser"
+        },
+        "runAs": "QUEUE_CONFIGURED_USER"
     },
     "jobAttachmentSettings": {
         "s3BucketName": "${assets_s3_bucket}",

diff --git a/scripts/run_posix_docker.sh b/scripts/run_posix_docker.sh
@@ -14,6 +14,7 @@ fi
 USE_LDAP="False"
 DO_BUILD="False"
 BUILD_ONLY="False"
+OVERRIDE_JOB_USER="False"
 while [[ "${1:-}" != "" ]]; do
     case $1 in
         -h|--help)
@@ -27,6 +28,10 @@ while [[ "${1:-}" != "" ]]; do
             echo "Using the LDAP client container image for testing."
             USE_LDAP="True"
             ;;
+        --override-job-user)
+            echo "Using the LDAP client container image which overrides the jobRunAsUSer"
+            OVERRIDE_JOB_USER="True"
+            ;;
         --build-only)
             BUILD_ONLY="True"
             ;;
@@ -38,6 +43,11 @@ while [[ "${1:-}" != "" ]]; do
     shift
 done
 
+if test "${OVERRIDE_JOB_USER}" == "True" && test "${USE_LDAP}" == "True"; then
+    echo "ERROR: Cannot use --ldap and --override-job-user together"
+    exit 1
+fi
+
 if ! test -d ${HOME}/.aws/models/deadline
 then
     echo "ERROR: Amazon Deadline Cloud service model must be installed to ~/.aws/models/deadline"
@@ -52,8 +62,13 @@ if test "${USE_LDAP}" == "True"; then
     CONTAINER_IMAGE_DIR="posix_ldap_multiuser"
 else
     ARGS="${ARGS} -h localuser.environment.internal"
-    CONTAINER_IMAGE_TAG="agent_posix_local_multiuser"
-    CONTAINER_IMAGE_DIR="posix_local_multiuser"
+    if test "${OVERRIDE_JOB_USER}" == "True"; then
+        CONTAINER_IMAGE_TAG="agent_posix_local_multiuser"
+        CONTAINER_IMAGE_DIR="posix_local_multiuser"
+    else
+        CONTAINER_IMAGE_TAG="agent_posix_local_multiuser_jobrunasuser"
+        CONTAINER_IMAGE_DIR="posix_local_multiuser_jobRunAsUser"
+    fi
 fi
 
 if test "${DO_BUILD}" == "True"; then

diff --git a/src/deadline_worker_agent/api_models.py b/src/deadline_worker_agent/api_models.py
@@ -241,8 +241,9 @@ class PosixUser(TypedDict):
 
 
 class JobRunAsUser(TypedDict):
-    posix: PosixUser
+    posix: NotRequired[PosixUser]
     # TODO: windows support
+    runAs: NotRequired[Literal["QUEUE_CONFIGURED_USER", "WORKER_AGENT_USER"]]
 
 
 class JobDetailsData(JobDetailsIdentifierFields):

diff --git a/src/deadline_worker_agent/sessions/job_entities/job_details.py b/src/deadline_worker_agent/sessions/job_entities/job_details.py
@@ -90,15 +90,22 @@ def job_run_as_user_api_model_to_worker_agent(
     """Converts the 'JobRunAsUser' api model to the 'JobRunAsUser' dataclass
     expected by the Worker Agent.
     """
+    if "runAs" in job_run_as_user_data and job_run_as_user_data["runAs"] == "WORKER_AGENT_USER":
+        return None
+
     if os.name == "posix":
         job_run_as_user_posix = job_run_as_user_data.get("posix", {})
-        user = job_run_as_user_posix.get("user", "")
-        group = job_run_as_user_posix.get("group", "")
-        if not (user and group):
+        user = job_run_as_user_posix["user"]
+        group = job_run_as_user_posix["group"]
+
+        if "runAs" not in job_run_as_user_data and not group and not user:
             return None
 
         job_run_as_user = JobRunAsUser(
-            posix=PosixSessionUser(user=user, group=group),
+            posix=PosixSessionUser(
+                user=user,
+                group=group,
+            ),
         )
     else:
         # TODO: windows support
@@ -267,6 +274,11 @@ def validate_entity_data(cls, entity_data: dict[str, Any]) -> JobDetailsData:
                                 Field(key="group", expected_type=str, required=True),
                             ),
                         ),
+                        Field(
+                            key="runAs",
+                            expected_type=str,
+                            required=False,
+                        ),
                     ),
                 ),
                 Field(
@@ -306,6 +318,30 @@ def validate_entity_data(cls, entity_data: dict[str, Any]) -> JobDetailsData:
                     ),
                 )
 
+        # Validate jobRunAsUser -> runAs is one of ("QUEUE_CONFIGURED_USER" / "WORKER_AGENT_USER")
+        if run_as_value := entity_data["jobRunAsUser"].get("runAs", None):
+            if run_as_value not in ("QUEUE_CONFIGURED_USER", "WORKER_AGENT_USER"):
+                raise ValueError(
+                    f'Expected "jobRunAs" -> "runAs" to be one of "QUEUE_CONFIGURED_USER", "WORKER_AGENT_USER" but got "{run_as_value}"'
+                )
+            elif run_as_value == "QUEUE_CONFIGURED_USER":
+                if not (run_as_posix := entity_data["jobRunAsUser"].get("posix", None)):
+                    raise ValueError(
+                        'Expected "jobRunAs" -> "posix" to exist when "jobRunAs" -> "runAs" is "QUEUE_CONFIGURED_USER" but it was not present'
+                    )
+                if run_as_posix["user"] == "":
+                    raise ValueError(
+                        'Got empty "jobRunAs" -> "posix" -> "user" but "jobRunAs" -> "runAs" is "QUEUE_CONFIGURED_USER"'
+                    )
+                if run_as_posix["group"] == "":
+                    raise ValueError(
+                        'Got empty "jobRunAs" -> "posix" -> "group" but "jobRunAs" -> "runAs" is "QUEUE_CONFIGURED_USER"'
+                    )
+            elif run_as_value == "WORKER_AGENT_USER" and "posix" in entity_data["jobRunAsUser"]:
+                raise ValueError(
+                    f'Expected "jobRunAs" -> "posix" is not valid when "jobRunAs" -> "runAs" is "WORKER_AGENT_USER" but got {entity_data["jobRunAsUser"]["posix"]}'
+                )
+
         return cast(JobDetailsData, entity_data)
 
     @classmethod

diff --git a/test/integ/conftest.py b/test/integ/conftest.py
@@ -68,7 +68,7 @@ def queue_with_job_run_as_user(
         client=deadline_client,
         display_name=f"Queue with jobsRunAsUser {job_run_as_user.user}",
         farm=farm,
-        job_run_as_user=JobRunAsUser(posix=job_run_as_user),
+        job_run_as_user=JobRunAsUser(posix=job_run_as_user, runAs="QUEUE_CONFIGURED_USER"),
     )
 
     qfa = QueueFleetAssociation.create(