From 9ce64514986062a66de96f9d55b73e2e1b58114c Mon Sep 17 00:00:00 2001
From: Toshimitsu Takahashi <toshi@tilfin.com>
Date: Fri, 26 Oct 2018 00:57:30 +0900
Subject: [PATCH] Support AWS IAM profile for Amazon Elasticsearch

---
 redash/query_runner/amazon_elasticsearch.py | 21 +++++++++++++++++----
 requirements_all_ds.txt                     |  2 +-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/redash/query_runner/amazon_elasticsearch.py b/redash/query_runner/amazon_elasticsearch.py
index 8ea42c05a9..7d465de863 100644
--- a/redash/query_runner/amazon_elasticsearch.py
+++ b/redash/query_runner/amazon_elasticsearch.py
@@ -2,7 +2,8 @@
 from . import register
 
 try:
-    from requests_aws4auth import AWS4Auth
+    from requests_aws_sign import AWSV4Sign
+    from botocore import session, credentials
     enabled = True
 except ImportError:
     enabled = False
@@ -40,17 +41,29 @@ def configuration_schema(cls):
                 'secret_key': {
                     'type': 'string',
                     'title': 'Secret Key'
+                },
+                'use_aws_iam_profile': {
+                    'type': 'boolean',
+                    'title': 'Use AWS IAM Profile'
                 }
             },
             "secret": ["secret_key"],
-            "order": ["server", "region", "access_key", "secret_key"],
-            "required": ['server', 'region', 'access_key', 'secret_key']
+            "order": ["server", "region", "access_key", "secret_key", "use_aws_iam_profile"],
+            "required": ['server', 'region']
         }
 
     def __init__(self, configuration):
         super(AmazonElasticsearchService, self).__init__(configuration)
 
-        self.auth = AWS4Auth(configuration['access_key'], configuration['secret_key'], configuration['region'], 'es')
+        region = configuration['region']
+        cred = None
+        if configuration.get('use_aws_iam_profile', False):
+            cred = credentials.get_credentials(session.Session())
+        else:
+            cred = credentials.Credentials(access_key=configuration.get('access_key', ''),
+                                           secret_key=configuration.get('secret_key', ''))
+
+        self.auth = AWSV4Sign(cred, region, 'es')
 
 
 register(AmazonElasticsearchService)
diff --git a/requirements_all_ds.txt b/requirements_all_ds.txt
index 0dc951cb6a..3630c31ff4 100644
--- a/requirements_all_ds.txt
+++ b/requirements_all_ds.txt
@@ -25,7 +25,7 @@ pymapd>=0.2.1
 qds-sdk>=1.9.6
 ibm-db>=2.0.9
 pydruid==0.4
-requests-aws4auth==0.9
+requests_aws_sign==0.1.4
 # certifi is needed to support MongoDB and SSL:
 certifi
 # We don't install snowflake connector by default, as it's causing conflicts with