diff --git a/poetry.lock b/poetry.lock index 3b9659455d..e99424b62a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "adal" @@ -5493,4 +5493,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"] [metadata] lock-version = "2.0" python-versions = ">=3.8,<3.11" -content-hash = "971596e47325293cbc984bb5a8aabd88a211f4ff4bbd72323f5eb6a168643feb" +content-hash = "a60a856644c1282f9e2c1cc98949e47164cce6683c84b7e92de1224991014ab6" diff --git a/pyproject.toml b/pyproject.toml index b5c706a18c..514ac87d15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,6 +88,8 @@ tzlocal = "4.3.1" pyodbc = "5.1.0" debugpy = "^1.8.9" paramiko = "3.4.1" +boto3 = "1.28.8" +botocore = "1.31.8" [tool.poetry.group.all_ds] optional = true @@ -95,8 +97,6 @@ optional = true [tool.poetry.group.all_ds.dependencies] atsd-client = "3.0.5" azure-kusto-data = "0.0.35" -boto3 = "1.28.8" -botocore = "1.31.8" cassandra-driver = "3.21.0" certifi = ">=2019.9.11" cmem-cmempy = "21.2.3" diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index c7ddef1eb7..ab54ae2f88 100644 --- a/redash/query_runner/pg.py +++ b/redash/query_runner/pg.py @@ -5,6 +5,7 @@ from tempfile import NamedTemporaryFile from uuid import uuid4 +import boto3 import psycopg2 from psycopg2.extras import Range @@ -23,13 +24,6 @@ logger = logging.getLogger(__name__) -try: - import boto3 - - IAM_ENABLED = True -except ImportError: - IAM_ENABLED = False - types_map = { 20: TYPE_INTEGER, 21: TYPE_INTEGER, @@ -167,6 +161,8 @@ def configuration_schema(cls): "sslrootcertFile": {"type": "string", "title": "SSL Root Certificate"}, "sslcertFile": {"type": "string", "title": "SSL Client Certificate"}, "sslkeyFile": {"type": "string", "title": "SSL Client Key"}, + "iamAuth": {"type": "boolean", "title": "IAM authentication"}, + "awsRegion": {"type": "string", "title": "AWS Region"}, }, "order": ["host", "port", "user", "password"], "required": ["dbname"], @@ -176,6 +172,8 @@ def configuration_schema(cls): "sslrootcertFile", "sslcertFile", "sslkeyFile", + "iamAuth", + "awsRegion", ], } @@ -251,11 +249,27 @@ def _get_tables(self, schema): def _get_connection(self): self.ssl_config = _get_ssl_config(self.configuration) + + user = self.configuration.get("user") + password = self.configuration.get("password") + host = self.configuration.get("host") + port = self.configuration.get("port") + + if self.configuration.get("iamAuth", False): + region_name = self.configuration.get("awsRegion") + rds_client = boto3.client("rds", region_name=region_name) + auth_token = rds_client.generate_db_auth_token( + DBHostname=host, + Port=port, + DBUsername=user, + ) + password = auth_token + connection = psycopg2.connect( - user=self.configuration.get("user"), - password=self.configuration.get("password"), - host=self.configuration.get("host"), - port=self.configuration.get("port"), + user=user, + password=password, + host=host, + port=port, dbname=self.configuration.get("dbname"), async_=True, **self.ssl_config, @@ -421,7 +435,7 @@ def name(cls): @classmethod def enabled(cls): - return IAM_ENABLED + return True def _login_method_selection(self): if self.configuration.get("rolename"):