From 7a40f12b47b6e60350767a720b76cc48dfe5e4cf Mon Sep 17 00:00:00 2001 From: Ievgen Aleinikov Date: Fri, 15 Mar 2019 13:18:46 +1100 Subject: [PATCH 1/5] allowing to specify a custom work group for AWS Athena queries --- redash/query_runner/athena.py | 6 ++++++ requirements_all_ds.txt | 6 +++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index 8d4f3dfcff..819c7be246 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -78,6 +78,11 @@ def configuration_schema(cls): 'type': 'boolean', 'title': 'Use Glue Data Catalog', }, + 'work_group': { + 'type': 'string', + 'title': 'Athena work group', + 'default': 'primary' + }, }, 'required': ['region', 's3_staging_dir'], 'order': ['region', 'aws_access_key', 'aws_secret_key', 's3_staging_dir', 'schema'], @@ -170,6 +175,7 @@ def run_query(self, query, user): schema_name=self.configuration.get('schema', 'default'), encryption_option=self.configuration.get('encryption_option', None), kms_key=self.configuration.get('kms_key', None), + work_group=self.configuration.get('work_group', 'primary'), formatter=SimpleFormatter()).cursor() try: diff --git a/requirements_all_ds.txt b/requirements_all_ds.txt index adfbd171bd..4f7a49cae0 100644 --- a/requirements_all_ds.txt +++ b/requirements_all_ds.txt @@ -11,8 +11,8 @@ td-client==0.8.0 pymssql==2.1.3 dql==0.5.24 dynamo3==0.4.7 -boto3==1.9.85 -botocore==1.12.85 +boto3==1.9.115 +botocore==1.12.115 sasl>=0.1.3 thrift>=0.8.0 thrift_sasl>=0.1.0 @@ -20,7 +20,7 @@ cassandra-driver==3.11.0 memsql==2.16.0 atsd_client==2.0.12 simple_salesforce==0.72.2 -PyAthena>=1.0.0 +PyAthena>=1.5.0 pymapd==0.7.1 qds-sdk>=1.9.6 ibm-db>=2.0.9 From 08eb3fae8850ff8657bc2108a87668ebb5721025 Mon Sep 17 00:00:00 2001 From: Ievgen Aleinikov Date: Mon, 18 Mar 2019 10:08:03 +1100 Subject: [PATCH 2/5] Fixing title + adding correct position in the UI --- redash/query_runner/athena.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index 819c7be246..d9d2736531 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -80,12 +80,12 @@ def configuration_schema(cls): }, 'work_group': { 'type': 'string', - 'title': 'Athena work group', + 'title': 'Athena Work Group', 'default': 'primary' }, }, 'required': ['region', 's3_staging_dir'], - 'order': ['region', 'aws_access_key', 'aws_secret_key', 's3_staging_dir', 'schema'], + 'order': ['region', 'aws_access_key', 'aws_secret_key', 's3_staging_dir', 'schema', 'work_group'], 'secret': ['aws_secret_key'] } From c2b481d48d355978f86049a963da4c4d6bbe2984 Mon Sep 17 00:00:00 2001 From: Ievgen Aleinikov Date: Thu, 1 Aug 2019 11:34:52 +1000 Subject: [PATCH 3/5] Adding assume role configuration to Athena query runner. --- redash/query_runner/athena.py | 67 ++++++++++++++++++++++++++++------- 1 file changed, 54 insertions(+), 13 deletions(-) diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index d9d2736531..4cdb9ea8a9 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -8,6 +8,7 @@ logger = logging.getLogger(__name__) ANNOTATE_QUERY = parse_boolean(os.environ.get('ATHENA_ANNOTATE_QUERY', 'true')) SHOW_EXTRA_SETTINGS = parse_boolean(os.environ.get('ATHENA_SHOW_EXTRA_SETTINGS', 'true')) +ASSUME_ROLE = parse_boolean(os.environ.get('ATHENA_ASSUME_ROLE', 'false')) OPTIONAL_CREDENTIALS = parse_boolean(os.environ.get('ATHENA_OPTIONAL_CREDENTIALS', 'true')) try: @@ -85,7 +86,7 @@ def configuration_schema(cls): }, }, 'required': ['region', 's3_staging_dir'], - 'order': ['region', 'aws_access_key', 'aws_secret_key', 's3_staging_dir', 'schema', 'work_group'], + 'order': ['region', 's3_staging_dir', 'schema', 'work_group'], 'secret': ['aws_secret_key'] } @@ -101,8 +102,30 @@ def configuration_schema(cls): }, }) - if not OPTIONAL_CREDENTIALS: - schema['required'] += ['aws_access_key', 'aws_secret_key'] + if ASSUME_ROLE: + del schema['properties']['aws_access_key'] + del schema['properties']['aws_secret_key'] + schema['secret'] = [] + + schema['order'].insert(1, 'iam_role') + schema['order'].insert(2, 'external_id') + schema['properties'].update({ + 'iam_role': { + 'type': 'string', + 'title': 'IAM role to assume', + }, + 'external_id': { + 'type': 'string', + 'title': 'External ID to be used while STS assume role', + }, + }) + else: + schema['order'].insert(1, 'aws_access_key') + schema['order'].insert(2, 'aws_secret_key') + + + if not OPTIONAL_CREDENTIALS and not ASSUME_ROLE: + schema['required'] += ['aws_access_key', 'aws_secret_key'] return schema @@ -118,13 +141,33 @@ def annotate_query(cls): def type(cls): return "athena" - def __get_schema_from_glue(self): - client = boto3.client( - 'glue', - aws_access_key_id=self.configuration.get('aws_access_key', None), - aws_secret_access_key=self.configuration.get('aws_secret_key', None), - region_name=self.configuration['region'] + def _get_iam_credentials(self, user=None): + if ASSUME_ROLE: + role_session_name = 'redash' if user is None else str(user.email) + sts = boto3.client('sts') + creds = sts.assume_role( + RoleArn=self.configuration.get('iam_role'), + RoleSessionName=role_session_name, + ExternalId=self.configuration.get('external_id') ) + return { + 'aws_access_key_id': creds['Credentials']['AccessKeyId'], + 'aws_secret_access_key': creds['Credentials']['SecretAccessKey'], + 'aws_session_token': creds['Credentials']['SessionToken'], + 'region_name': self.configuration['region'] + } + else: + return { + 'aws_access_key_id': self.configuration.get('aws_access_key', None), + 'aws_secret_access_key': self.configuration.get('aws_secret_key', None), + 'aws_session_token': None, + 'region_name': self.configuration['region'] + } + + + + def __get_schema_from_glue(self): + client = boto3.client('glue', **self._get_iam_credentials()) schema = {} database_paginator = client.get_paginator('get_databases') @@ -169,14 +212,12 @@ def get_schema(self, get_stats=False): def run_query(self, query, user): cursor = pyathena.connect( s3_staging_dir=self.configuration['s3_staging_dir'], - region_name=self.configuration['region'], - aws_access_key_id=self.configuration.get('aws_access_key', None), - aws_secret_access_key=self.configuration.get('aws_secret_key', None), schema_name=self.configuration.get('schema', 'default'), encryption_option=self.configuration.get('encryption_option', None), kms_key=self.configuration.get('kms_key', None), work_group=self.configuration.get('work_group', 'primary'), - formatter=SimpleFormatter()).cursor() + formatter=SimpleFormatter(), + **self._get_iam_credentials(user=user)).cursor() try: cursor.execute(query) From c8b2ec1d0a6cdfe04a205556810850cb83b983e3 Mon Sep 17 00:00:00 2001 From: Ievgen Aleinikov Date: Fri, 2 Aug 2019 13:02:20 +1000 Subject: [PATCH 4/5] removing extra blank lines --- redash/query_runner/athena.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index 4cdb9ea8a9..ec7319ce3c 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -123,7 +123,6 @@ def configuration_schema(cls): schema['order'].insert(1, 'aws_access_key') schema['order'].insert(2, 'aws_secret_key') - if not OPTIONAL_CREDENTIALS and not ASSUME_ROLE: schema['required'] += ['aws_access_key', 'aws_secret_key'] @@ -164,8 +163,6 @@ def _get_iam_credentials(self, user=None): 'region_name': self.configuration['region'] } - - def __get_schema_from_glue(self): client = boto3.client('glue', **self._get_iam_credentials()) schema = {} From f72842c3e838cb978bb2e4bb5e2b7a64314491a8 Mon Sep 17 00:00:00 2001 From: Ievgen Aleinikov Date: Mon, 12 Aug 2019 13:39:14 +1000 Subject: [PATCH 5/5] fixes based on comments to the PR --- redash/query_runner/athena.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index ec7319ce3c..db13297caa 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -142,7 +142,7 @@ def type(cls): def _get_iam_credentials(self, user=None): if ASSUME_ROLE: - role_session_name = 'redash' if user is None else str(user.email) + role_session_name = 'redash' if user is None else user.email sts = boto3.client('sts') creds = sts.assume_role( RoleArn=self.configuration.get('iam_role'), @@ -159,7 +159,6 @@ def _get_iam_credentials(self, user=None): return { 'aws_access_key_id': self.configuration.get('aws_access_key', None), 'aws_secret_access_key': self.configuration.get('aws_secret_key', None), - 'aws_session_token': None, 'region_name': self.configuration['region'] }