Repository: apache/airflow

Issue:
Unable to specify Python version for AwsGlueJobOperator
### Apache Airflow Provider(s)

amazon

### Versions of Apache Airflow Providers

_No response_

### Apache Airflow version

2.0.2

### Operating System

Amazon Linux

### Deployment

MWAA

### Deployment details

_No response_

### What happened

When a new Glue job is created using the AwsGlueJobOperator, the job is defaulting to Python2. Setting the version in create_job_kwargs fails with key error.

### What you expected to happen

Expected the Glue job to be created with a Python3 runtime. create_job_kwargs are passed to the boto3 glue client create_job method which includes a "Command" parameter that is a dictionary containing the Python version.



### How to reproduce

Create a dag with an AwsGlueJobOperator and pass a "Command" parameter in the create_job_kwargs argument.

```
    create_glue_job_args = {
        "Command": {
            "Name": "abalone-preprocess",
            "ScriptLocation": f"s3://{output_bucket}/code/preprocess.py",
            "PythonVersion": "3"
        }
    }
    glue_etl = AwsGlueJobOperator(  
        task_id="glue_etl",  
        s3_bucket=output_bucket,
        script_args={
                '--S3_INPUT_BUCKET': data_bucket,
                '--S3_INPUT_KEY_PREFIX': 'input/raw',
                '--S3_UPLOADS_KEY_PREFIX': 'input/uploads',
                '--S3_OUTPUT_BUCKET': output_bucket,
                '--S3_OUTPUT_KEY_PREFIX': str(determine_dataset_id.output) +'/input/data' 
            },
        iam_role_name="MLOps",  
        retry_limit=2,
        concurrent_run_limit=3,
        create_job_kwargs=create_glue_job_args,
        dag=dag) 
```

```
[2022-01-04 16:43:42,053] {{logging_mixin.py:104}} INFO - [2022-01-04 16:43:42,053] {{glue.py:190}} ERROR - Failed to create aws glue job, error: 'Command'
[2022-01-04 16:43:42,081] {{logging_mixin.py:104}} INFO - [2022-01-04 16:43:42,081] {{glue.py:112}} ERROR - Failed to run aws glue job, error: 'Command'
[2022-01-04 16:43:42,101] {{taskinstance.py:1482}} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/airflow/providers/amazon/aws/hooks/glue.py", line 166, in get_or_create_glue_job
    get_job_response = glue_client.get_job(JobName=self.job_name)
  File "/usr/local/lib/python3.7/site-packages/botocore/client.py", line 357, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/usr/local/lib/python3.7/site-packages/botocore/client.py", line 676, in _make_api_call
    raise error_class(parsed_response, operation_name)
botocore.errorfactory.EntityNotFoundException: An error occurred (EntityNotFoundException) when calling the GetJob operation: Job with name: abalone-preprocess not found.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1138, in _run_raw_task
    self._prepare_and_execute_task_with_callbacks(context, task)
  File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1311, in _prepare_and_execute_task_with_callbacks
    result = self._execute_task(context, task_copy)
  File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1341, in _execute_task
    result = task_copy.execute(context=context)
  File "/usr/local/lib/python3.7/site-packages/airflow/providers/amazon/aws/operators/glue.py", line 121, in execute
    glue_job_run = glue_job.initialize_job(self.script_args)
  File "/usr/local/lib/python3.7/site-packages/airflow/providers/amazon/aws/hooks/glue.py", line 108, in initialize_job
    job_name = self.get_or_create_glue_job()
  File "/usr/local/lib/python3.7/site-packages/airflow/providers/amazon/aws/hooks/glue.py", line 186, in get_or_create_glue_job
    **self.create_job_kwargs,
KeyError: 'Command'
```

### Anything else

When a new job is being created.

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

I think that this is because the GlueHook is pretty opinionated and hardcodes the value of `Command` when running the `glue_client.create_job` command:

https://github.com/apache/airflow/blob/2ab2ae8849bf6d80a700b1b74cef37eb187161ad/airflow/providers/amazon/aws/hooks/glue.py#L181-L225

So when you provide `Command` as a `create_job_kwargs`, it ends up being supplied twice to that function (Although I suspect that this would be a typeError, not a keyError 🤔)

Anyways, thoughts on just making the `Command.Name` and `Command.PythonVersion` argument configurable in the `GlueJobOperator`?

If y'all think that this is a satisfactory fix, feel free to assign this issue to me and I can put up a quick PR. 
Also for what its worth, I think that the `Command` block in your DAG is invalid, as the `Command.Name` you're using (`abalone-preprocess`) must be one of `glueetl`, `pythonshell` or `gluestreaming`.

https://docs.aws.amazon.com/glue/latest/webapi/API_JobCommand.html

@SamWheating We can make them configurable

I'm just not I'm not 100% why must we create the command. Can't we just leave the `create_job_kwargs` as the user passes? 


https://github.com/apache/airflow/blob/5ca5693b81b129cd34367fe8788d48ed70054f95/airflow/providers/amazon/aws/hooks/glue.py#L206
https://github.com/apache/airflow/blob/5ca5693b81b129cd34367fe8788d48ed70054f95/airflow/providers/amazon/aws/hooks/glue.py#L217

There are other parameters that don't get this special treatment:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue.html#Glue.Client.create_job



Resolving patch:
diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
--- a/airflow/providers/amazon/aws/hooks/glue.py
+++ b/airflow/providers/amazon/aws/hooks/glue.py
@@ -181,6 +181,12 @@ def get_or_create_glue_job(self) -> str:
             s3_log_path = f's3://{self.s3_bucket}/{self.s3_glue_logs}{self.job_name}'
             execution_role = self.get_iam_execution_role()
             try:
+                default_command = {
+                    "Name": "glueetl",
+                    "ScriptLocation": self.script_location,
+                }
+                command = self.create_job_kwargs.get("Command", default_command)
+
                 if "WorkerType" in self.create_job_kwargs and "NumberOfWorkers" in self.create_job_kwargs:
                     create_job_response = glue_client.create_job(
                         Name=self.job_name,
@@ -188,7 +194,7 @@ def get_or_create_glue_job(self) -> str:
                         LogUri=s3_log_path,
                         Role=execution_role['Role']['Arn'],
                         ExecutionProperty={"MaxConcurrentRuns": self.concurrent_run_limit},
-                        Command={"Name": "glueetl", "ScriptLocation": self.script_location},
+                        Command=command,
                         MaxRetries=self.retry_limit,
                         **self.create_job_kwargs,
                     )
@@ -199,7 +205,7 @@ def get_or_create_glue_job(self) -> str:
                         LogUri=s3_log_path,
                         Role=execution_role['Role']['Arn'],
                         ExecutionProperty={"MaxConcurrentRuns": self.concurrent_run_limit},
-                        Command={"Name": "glueetl", "ScriptLocation": self.script_location},
+                        Command=command,
                         MaxRetries=self.retry_limit,
                         MaxCapacity=self.num_of_dpus,
                         **self.create_job_kwargs,


Test patch:
diff --git a/tests/providers/amazon/aws/hooks/test_glue.py b/tests/providers/amazon/aws/hooks/test_glue.py
--- a/tests/providers/amazon/aws/hooks/test_glue.py
+++ b/tests/providers/amazon/aws/hooks/test_glue.py
@@ -22,9 +22,9 @@
 from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
 
 try:
-    from moto import mock_iam
+    from moto import mock_glue, mock_iam
 except ImportError:
-    mock_iam = None
+    mock_iam = mock_glue = None
 
 
 class TestGlueJobHook(unittest.TestCase):
@@ -57,23 +57,56 @@ def test_get_iam_execution_role(self):
         assert "Arn" in iam_role['Role']
         assert iam_role['Role']['Arn'] == "arn:aws:iam::123456789012:role/my_test_role"
 
-    @mock.patch.object(GlueJobHook, "get_iam_execution_role")
     @mock.patch.object(GlueJobHook, "get_conn")
-    def test_get_or_create_glue_job(self, mock_get_conn, mock_get_iam_execution_role):
-        mock_get_iam_execution_role.return_value = mock.MagicMock(Role={'RoleName': 'my_test_role'})
+    def test_get_or_create_glue_job_get_existing_job(self, mock_get_conn):
+        """
+        Calls 'get_or_create_glue_job' with a existing job.
+        Should retrieve existing one.
+        """
+        expected_job_name = "simple-job"
+        mock_get_conn.return_value.get_job.return_value = {"Job": {"Name": expected_job_name}}
+
         some_script = "s3:/glue-examples/glue-scripts/sample_aws_glue_job.py"
         some_s3_bucket = "my-includes"
 
-        mock_glue_job = mock_get_conn.return_value.get_job()['Job']['Name']
-        glue_job = GlueJobHook(
-            job_name='aws_test_glue_job',
-            desc='This is test case job from Airflow',
+        hook = GlueJobHook(
+            job_name="aws_test_glue_job",
+            desc="This is test case job from Airflow",
             script_location=some_script,
-            iam_role_name='my_test_role',
+            iam_role_name="my_test_role",
             s3_bucket=some_s3_bucket,
             region_name=self.some_aws_region,
-        ).get_or_create_glue_job()
-        assert glue_job == mock_glue_job
+        )
+
+        result = hook.get_or_create_glue_job()
+
+        mock_get_conn.assert_called_once()
+        mock_get_conn.return_value.get_job.assert_called_once_with(JobName=hook.job_name)
+        assert result == expected_job_name
+
+    @unittest.skipIf(mock_glue is None, "mock_glue package not present")
+    @mock_glue
+    @mock.patch.object(GlueJobHook, "get_iam_execution_role")
+    def test_get_or_create_glue_job_create_new_job(self, mock_get_iam_execution_role):
+        """
+        Calls 'get_or_create_glue_job' with no existing job.
+        Should create a new job.
+        """
+        mock_get_iam_execution_role.return_value = {"Role": {"RoleName": "my_test_role", "Arn": "test_role"}}
+        expected_job_name = "aws_test_glue_job"
+
+        hook = GlueJobHook(
+            job_name=expected_job_name,
+            desc="This is test case job from Airflow",
+            iam_role_name="my_test_role",
+            script_location="s3://bucket",
+            s3_bucket="bucket",
+            region_name=self.some_aws_region,
+        )
+
+        result = hook.get_or_create_glue_job()
+
+        assert result == expected_job_name
 
     @mock.patch.object(GlueJobHook, "get_iam_execution_role")
     @mock.patch.object(GlueJobHook, "get_conn")