Add redshift role use option (#4532)

* Add redshift role use option

* Update requirements for SSL socket wrap issue fixes

* Split Redshift class into User and IAM logins

* Update incorrect register

* Change type names

* Correct class name to inherit

* Render IAM redshift image and field order correct

* Update redash/query_runner/pg.py

Co-Authored-By: Arik Fraimovich <arik@arikfr.com>

* Update redash/query_runner/pg.py

Co-Authored-By: Arik Fraimovich <arik@arikfr.com>

* Remove need for specified urllib - specify pyopenssl is enough

* Pyopenssl back down to 19.0.0

Co-authored-by: Arik Fraimovich <arik@arikfr.com>

redash/query_runner/pg.py

@@ -1,6 +1,7 @@
 import os
 import logging
 import select
+from uuid import uuid4
 
 import psycopg2
 from psycopg2.extras import Range
@@ -10,6 +11,13 @@ from redash.utils import JSONEncoder, json_dumps, json_loads
 
 logger = logging.getLogger(__name__)
 
+try:
+    import boto3
+
+    IAM_ENABLED = True
+except ImportError:
+    IAM_ENABLED = False
+
 types_map = {
     20: TYPE_INTEGER,
     21: TYPE_INTEGER,
@@ -227,11 +235,17 @@ class PostgreSQL(BaseSQLQueryRunner):
 
 
 class Redshift(PostgreSQL):
+
     @classmethod
     def type(cls):
         return "redshift"
 
+    @classmethod
+    def name(cls):
+        return "Redshift"
+
     def _get_connection(self):
+
         sslrootcert_path = os.path.join(
             os.path.dirname(__file__), "./files/redshift-ca-bundle.crt"
         )
@@ -331,6 +345,131 @@ class Redshift(PostgreSQL):
 
         return list(schema.values())
 
+class RedshiftIAM(Redshift):
+
+    @classmethod
+    def type(cls):
+        return "redshift_iam"
+
+    @classmethod
+    def name(cls):
+        return "Redshift (with IAM User/Role)"
+
+    @classmethod
+    def enabled(cls):
+        return IAM_ENABLED
+
+    def _login_method_selection(self):
+        if self.configuration.get("rolename"):
+            if not self.configuration.get("aws_access_key_id") or not self.configuration.get("aws_secret_access_key"):
+                return "ASSUME_ROLE_NO_KEYS"
+            else:
+                return "ASSUME_ROLE_KEYS"
+        elif self.configuration.get("aws_access_key_id") and self.configuration.get("aws_secret_access_key"):
+            return "KEYS"
+        elif not self.configuration.get("password"):
+            return "ROLE"
+
+    @classmethod
+    def configuration_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "rolename": {"type": "string", "title": "IAM Role Name"},
+                "aws_region": {"type": "string", "title": "AWS Region"},
+                "aws_access_key_id": {"type": "string", "title": "AWS Access Key ID"},
+                "aws_secret_access_key": {"type": "string", "title": "AWS Secret Access Key"},
+                "clusterid": {"type": "string", "title": "Redshift Cluster ID"},
+                "user": {"type": "string"},
+                "host": {"type": "string"},
+                "port": {"type": "number"},
+                "dbname": {"type": "string", "title": "Database Name"},
+                "sslmode": {"type": "string", "title": "SSL Mode", "default": "prefer"},
+                "adhoc_query_group": {
+                    "type": "string",
+                    "title": "Query Group for Adhoc Queries",
+                    "default": "default",
+                },
+                "scheduled_query_group": {
+                    "type": "string",
+                    "title": "Query Group for Scheduled Queries",
+                    "default": "default",
+                },
+            },
+            "order": [
+                "rolename",
+                "aws_region",
+                "aws_access_key_id",
+                "aws_secret_access_key",
+                "clusterid",
+                "host",
+                "port",
+                "user",
+                "dbname",
+                "sslmode",
+                "adhoc_query_group",
+                "scheduled_query_group",
+            ],
+            "required": ["dbname", "user", "host", "port", "aws_region"],
+            "secret": ["aws_secret_access_key"],
+        }
+
+    def _get_connection(self):
+
+        sslrootcert_path = os.path.join(
+            os.path.dirname(__file__), "./files/redshift-ca-bundle.crt"
+        )
+
+        login_method = self._login_method_selection()
+
+
+        if login_method == "KEYS":
+            client = boto3.client("redshift",
+                                  region_name=self.configuration.get("aws_region"),
+                                  aws_access_key_id=self.configuration.get("aws_access_key_id"),
+                                  aws_secret_access_key=self.configuration.get("aws_secret_access_key"))
+        elif login_method == "ROLE":
+            client = boto3.client("redshift",
+                                  region_name=self.configuration.get("aws_region"))
+        else:
+            if login_method == "ASSUME_ROLE_KEYS":
+                assume_client = client = boto3.client('sts',
+                                                      region_name=self.configuration.get("aws_region"),
+                                                      aws_access_key_id=self.configuration.get("aws_access_key_id"),
+                                                      aws_secret_access_key=self.configuration.get(
+                                                          "aws_secret_access_key"))
+            else:
+                assume_client = client = boto3.client('sts',
+                                                      region_name=self.configuration.get("aws_region"))
+            role_session = f"redash_{uuid4().hex}"
+            session_keys = assume_client.assume_role(
+                RoleArn=self.configuration.get("rolename"),
+                RoleSessionName=role_session)["Credentials"]
+            client = boto3.client("redshift",
+                                  region_name=self.configuration.get("aws_region"),
+                                  aws_access_key_id=session_keys["AccessKeyId"],
+                                  aws_secret_access_key=session_keys["SecretAccessKey"],
+                                  aws_session_token=session_keys["SessionToken"]
+                                  )
+        credentials = client.get_cluster_credentials(
+            DbUser=self.configuration.get("user"),
+            DbName=self.configuration.get("dbname"),
+            ClusterIdentifier=self.configuration.get("clusterid"))
+        db_user = credentials["DbUser"]
+        db_password = credentials["DbPassword"]
+        connection = psycopg2.connect(
+            user=db_user,
+            password=db_password,
+            host=self.configuration.get("host"),
+            port=self.configuration.get("port"),
+            dbname=self.configuration.get("dbname"),
+            sslmode=self.configuration.get("sslmode", "prefer"),
+            sslrootcert=sslrootcert_path,
+            async_=True,
+        )
+
+        return connection
+
 
 class CockroachDB(PostgreSQL):
     @classmethod
@@ -340,4 +479,5 @@ class CockroachDB(PostgreSQL):
 
 register(PostgreSQL)
 register(Redshift)
+register(RedshiftIAM)
 register(CockroachDB)

requirements.txt

@@ -47,7 +47,7 @@ xlsxwriter==1.2.2
 pystache==0.5.4
 parsedatetime==2.4
 PyJWT==1.7.1
-cryptography==2.7
+cryptography==2.8
 simplejson==3.16.0
 ua-parser==0.8.0
 user-agents==2.0

requirements_all_ds.txt

@@ -11,8 +11,8 @@ td-client==1.0.0
 pymssql==2.1.4
 dql==0.5.26
 dynamo3==0.4.10
-boto3>=1.9.241,<1.10.0
-botocore>=1.12.241,<1.13.0
+boto3>=1.10.0,<1.11.0
+botocore>=1.13,<1.14.0
 sasl>=0.1.3
 thrift>=0.8.0
 thrift_sasl>=0.1.0

requirements_dev.txt

@@ -6,7 +6,7 @@ mock==3.0.5
 # PyMongo and Athena dependencies are needed for some of the unit tests:
 # (this is not perfect and we should resolve this in a different way)
 pymongo[srv,tls]==3.9.0
-botocore>=1.12.241,<1.13.0
+botocore>=1.13,<1.14.0
 PyAthena>=1.5.0
 ptvsd==4.3.2
 freezegun==0.3.12