[AIRFLOW-3197] EMRHook is missing new parameters of the AWS API (#4044)

Allow passing any params to the CreateJobFlow API, so that we don't have
to stay up to date with AWS api changes.

Author: ashb

UPDATING.md

@@ -73,6 +73,16 @@ then you need to change it like this
     def is_active(self):
       return self.active
 
+### EMRHook now passes all of connection's extra to CreateJobFlow API
+
+EMRHook.create_job_flow has been changed to pass all keys to the create_job_flow API, rather than
+just specific known keys for greater flexibility.
+
+However prior to this release the "emr_default" sample connection that was created had invalid
+configuration, so creating EMR clusters might fail until your connection is updated. (Ec2KeyName,
+Ec2SubnetId, TerminationProtection and KeepJobFlowAliveWhenNoSteps were all top-level keys when they
+should be inside the "Instances" dict)
+
 ## Airflow 1.10
 
 Installation and upgrading requires setting `SLUGIFY_USES_TEXT_UNIDECODE=yes` in your environment or

airflow/contrib/hooks/emr_hook.py

@@ -51,19 +51,6 @@ def create_job_flow(self, job_flow_overrides):
         config = emr_conn.extra_dejson.copy()
         config.update(job_flow_overrides)
 
-        response = self.get_conn().run_job_flow(
-            Name=config.get('Name'),
-            LogUri=config.get('LogUri'),
-            ReleaseLabel=config.get('ReleaseLabel'),
-            Instances=config.get('Instances'),
-            Steps=config.get('Steps', []),
-            BootstrapActions=config.get('BootstrapActions', []),
-            Applications=config.get('Applications'),
-            Configurations=config.get('Configurations', []),
-            VisibleToAllUsers=config.get('VisibleToAllUsers'),
-            JobFlowRole=config.get('JobFlowRole'),
-            ServiceRole=config.get('ServiceRole'),
-            Tags=config.get('Tags')
-        )
+        response = self.get_conn().run_job_flow(**config)
 
         return response

airflow/utils/db.py

@@ -225,6 +225,8 @@ def initdb(rbac=False):
                     "LogUri": "s3://my-emr-log-bucket/default_job_flow_location",
                     "ReleaseLabel": "emr-4.6.0",
                     "Instances": {
+                        "Ec2KeyName": "mykey",
+                        "Ec2SubnetId": "somesubnet",
                         "InstanceGroups": [
                             {
                                 "Name": "Master nodes",
@@ -240,12 +242,10 @@ def initdb(rbac=False):
                                 "InstanceType": "r3.2xlarge",
                                 "InstanceCount": 1
                             }
-                        ]
+                        ],
+                        "TerminationProtected": false,
+                        "KeepJobFlowAliveWhenNoSteps": false
                     },
-                    "Ec2KeyName": "mykey",
-                    "KeepJobFlowAliveWhenNoSteps": false,
-                    "TerminationProtected": false,
-                    "Ec2SubnetId": "somesubnet",
                     "Applications":[
                         { "Name": "Spark" }
                     ],

tests/contrib/hooks/test_emr_hook.py

@@ -31,28 +31,49 @@
     mock_emr = None
 
 
+@unittest.skipIf(mock_emr is None, 'moto package not present')
 class TestEmrHook(unittest.TestCase):
     @mock_emr
     def setUp(self):
         configuration.load_test_config()
 
-    @unittest.skipIf(mock_emr is None, 'mock_emr package not present')
     @mock_emr
     def test_get_conn_returns_a_boto3_connection(self):
         hook = EmrHook(aws_conn_id='aws_default')
         self.assertIsNotNone(hook.get_conn().list_clusters())
 
-    @unittest.skipIf(mock_emr is None, 'mock_emr package not present')
     @mock_emr
     def test_create_job_flow_uses_the_emr_config_to_create_a_cluster(self):
         client = boto3.client('emr', region_name='us-east-1')
-        if len(client.list_clusters()['Clusters']):
-            raise ValueError('AWS not properly mocked')
 
         hook = EmrHook(aws_conn_id='aws_default', emr_conn_id='emr_default')
         cluster = hook.create_job_flow({'Name': 'test_cluster'})
 
-        self.assertEqual(client.list_clusters()['Clusters'][0]['Id'], cluster['JobFlowId'])
+        self.assertEqual(client.list_clusters()['Clusters'][0]['Id'],
+                         cluster['JobFlowId'])
+
+    @mock_emr
+    def test_create_job_flow_extra_args(self):
+        """
+        Test that we can add extra arguments to the launch call.
+
+        This is useful for when AWS add new options, such as
+        "SecurityConfiguration" so that we don't have to change our code
+        """
+        client = boto3.client('emr', region_name='us-east-1')
+
+        hook = EmrHook(aws_conn_id='aws_default', emr_conn_id='emr_default')
+        # AmiVersion is really old and almost no one will use it anymore, but
+        # it's one of the "optional" request params that moto supports - it's
+        # coverage of EMR isn't 100% it turns out.
+        cluster = hook.create_job_flow({'Name': 'test_cluster',
+                                        'ReleaseLabel': '',
+                                        'AmiVersion': '3.2'})
+
+        cluster = client.describe_cluster(ClusterId=cluster['JobFlowId'])['Cluster']
+
+        # The AmiVersion comes back as {Requested,Running}AmiVersion fields.
+        self.assertEqual(cluster['RequestedAmiVersion'], '3.2')
 
 if __name__ == '__main__':
     unittest.main()