Support AWS RDS IAM Authentication for PostgreSQK data soruce

Author: winebarrel

poetry.lock

@@ -92,15 +92,15 @@ debugpy = "^1.8.9"
 paramiko = "3.4.1"
 oracledb = "2.5.1"
 ibm-db = { version = "^3.2.7", markers = "platform_machine == 'x86_64' or platform_machine == 'AMD64'" }
+boto3 = "1.28.8"
+botocore = "1.31.8"
 
 [tool.poetry.group.all_ds]
 optional = true
 
 [tool.poetry.group.all_ds.dependencies]
 atsd-client = "3.0.5"
 azure-kusto-data = "5.0.1"
-boto3 = "1.28.8"
-botocore = "1.31.8"
 cassandra-driver = "3.29.3"
 certifi = ">=2019.9.11"
 cmem-cmempy = "21.2.3"

pyproject.toml

@@ -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,
@@ -177,6 +171,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"},
+                "awsIamAuth": {"type": "boolean", "title": "AWS IAM authentication"},
+                "awsRegion": {"type": "string", "title": "AWS Region"},
             },
             "order": ["host", "port", "user", "password"],
             "required": ["dbname"],
@@ -186,6 +182,8 @@ def configuration_schema(cls):
                 "sslrootcertFile",
                 "sslcertFile",
                 "sslkeyFile",
+                "awsIamAuth",
+                "awsRegion",
             ],
         }
 
@@ -255,11 +253,27 @@ def _get_tables(self, schema):
     def _get_connection(self):
         self.ssl_config = _get_ssl_config(self.configuration)
         self.dsn = _parse_dsn(self.configuration)
+
+        user = self.configuration.get("user")
+        password = self.configuration.get("password")
+        host = self.configuration.get("host")
+        port = self.configuration.get("port", 5432)
+
+        if self.configuration.get("awsIamAuth", 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,
@@ -426,7 +440,7 @@ def name(cls):
 
     @classmethod
     def enabled(cls):
-        return IAM_ENABLED
+        return True
 
     def _login_method_selection(self):
         if self.configuration.get("rolename"):