From 2d950e827708eb5fd834caf097c0b8ed39e06dad Mon Sep 17 00:00:00 2001 From: Guilherme Martins Crocetti <24530683+gmcrocetti@users.noreply.github.com> Date: Tue, 7 Jun 2022 19:32:35 -0300 Subject: [PATCH] fix: patches #24215. Won't raise KeyError when 'create_job_kwargs' contains the 'Command' key. --- airflow/providers/amazon/aws/hooks/glue.py | 2 +- tests/providers/amazon/aws/hooks/test_glue.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py index be85f2fec4cf3..2fc342675db44 100644 --- 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 index 4e3afd0453834..e8f55c5b81caa 100644 --- 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()