fix: patches #24215. Won't raise KeyError when 'create_job_kwargs' co…

…ntains the 'Command' key. (#24308)
apache · Jun 8, 2022 · dc9c083 · dc9c083
1 parent a1a9a8f
commit dc9c083
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
@@ -185,7 +185,7 @@ def get_or_create_glue_job(self) -> str:
                     "Name": "glueetl",
                     "ScriptLocation": self.script_location,
                 }
-                command = self.create_job_kwargs.get("Command", default_command)
+                command = self.create_job_kwargs.pop("Command", default_command)
 
                 if "WorkerType" in self.create_job_kwargs and "NumberOfWorkers" in self.create_job_kwargs:
                     create_job_response = glue_client.create_job(

diff --git a/tests/providers/amazon/aws/hooks/test_glue.py b/tests/providers/amazon/aws/hooks/test_glue.py
@@ -102,6 +102,7 @@ def test_get_or_create_glue_job_create_new_job(self, mock_get_iam_execution_role
             script_location="s3://bucket",
             s3_bucket="bucket",
             region_name=self.some_aws_region,
+            create_job_kwargs={"Command": {}},
         )
 
         result = hook.get_or_create_glue_job()