[GitHub] kaxil closed pull request #4418: [AIRFLOW-3610] Add region param for EMR jobflow creation

[GitBox <gi...@apache.org>]

kaxil closed pull request #4418: [AIRFLOW-3610] Add region param for EMR jobflow creation
URL: https://github.com/apache/airflow/pull/4418
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/airflow/contrib/hooks/emr_hook.py b/airflow/contrib/hooks/emr_hook.py
index fcdf4ac848..7571d0988f 100644
--- a/airflow/contrib/hooks/emr_hook.py
+++ b/airflow/contrib/hooks/emr_hook.py
@@ -27,12 +27,13 @@ class EmrHook(AwsHook):
     create_job_flow method.
     """
 
-    def __init__(self, emr_conn_id=None, *args, **kwargs):
+    def __init__(self, emr_conn_id=None, region_name=None, *args, **kwargs):
         self.emr_conn_id = emr_conn_id
+        self.region_name = region_name
         super(EmrHook, self).__init__(*args, **kwargs)
 
     def get_conn(self):
-        self.conn = self.get_client_type('emr')
+        self.conn = self.get_client_type('emr', self.region_name)
         return self.conn
 
     def create_job_flow(self, job_flow_overrides):
diff --git a/airflow/contrib/operators/emr_create_job_flow_operator.py b/airflow/contrib/operators/emr_create_job_flow_operator.py
index 60256235c6..bcb625b00a 100644
--- a/airflow/contrib/operators/emr_create_job_flow_operator.py
+++ b/airflow/contrib/operators/emr_create_job_flow_operator.py
@@ -46,6 +46,7 @@ def __init__(
             aws_conn_id='s3_default',
             emr_conn_id='emr_default',
             job_flow_overrides=None,
+            region_name=None,
             *args, **kwargs):
         super(EmrCreateJobFlowOperator, self).__init__(*args, **kwargs)
         self.aws_conn_id = aws_conn_id
@@ -53,9 +54,10 @@ def __init__(
         if job_flow_overrides is None:
             job_flow_overrides = {}
         self.job_flow_overrides = job_flow_overrides
+        self.region_name = region_name
 
     def execute(self, context):
-        emr = EmrHook(aws_conn_id=self.aws_conn_id, emr_conn_id=self.emr_conn_id)
+        emr = EmrHook(aws_conn_id=self.aws_conn_id, emr_conn_id=self.emr_conn_id, region_name=self.region_name)
 
         self.log.info(
             'Creating JobFlow using aws-conn-id: %s, emr-conn-id: %s',
diff --git a/tests/contrib/hooks/test_emr_hook.py b/tests/contrib/hooks/test_emr_hook.py
index 27f2c2250d..03f23f6906 100644
--- a/tests/contrib/hooks/test_emr_hook.py
+++ b/tests/contrib/hooks/test_emr_hook.py
@@ -38,7 +38,7 @@ def setUp(self):
 
     @mock_emr
     def test_get_conn_returns_a_boto3_connection(self):
-        hook = EmrHook(aws_conn_id='aws_default')
+        hook = EmrHook(aws_conn_id='aws_default', region_name='ap-southeast-2')
         self.assertIsNotNone(hook.get_conn().list_clusters())
 
     @mock_emr
diff --git a/tests/contrib/operators/test_emr_create_job_flow_operator.py b/tests/contrib/operators/test_emr_create_job_flow_operator.py
index 704f5982d8..fd95f78fd2 100644
--- a/tests/contrib/operators/test_emr_create_job_flow_operator.py
+++ b/tests/contrib/operators/test_emr_create_job_flow_operator.py
@@ -71,12 +71,14 @@ def setUp(self):
             aws_conn_id='aws_default',
             emr_conn_id='emr_default',
             job_flow_overrides=self._config,
+            region_name='ap-southeast-2',
             dag=DAG('test_dag_id', default_args=args)
         )
 
     def test_init(self):
         self.assertEqual(self.operator.aws_conn_id, 'aws_default')
         self.assertEqual(self.operator.emr_conn_id, 'emr_default')
+        self.assertEqual(self.operator.region_name, 'ap-southeast-2')
 
     def test_render_template(self):
         ti = TaskInstance(self.operator, DEFAULT_DATE)


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services