feat: support AWS IAM database authentication

pyproject.toml

@@ -28,6 +28,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
 ]
 dependencies = [
+    "boto3>=1.26.0",
     "paramiko>=4",
     "psycopg2-binary==2.9.12",
     "sqlalchemy==2.0.51",
@@ -82,6 +83,12 @@ python_version = "3.13"
 warn_unused_configs = true
 warn_unused_ignores = true
 
+[[tool.mypy.overrides]]
+module = [
+  "boto3.*",
+]
+ignore_missing_imports = true
+
 [build-system]
 requires = [
   "hatchling==1.31.0",

tap_postgres/client.py

@@ -12,8 +12,10 @@
 import typing as t
 from types import MappingProxyType
 
+import boto3
 import psycopg2
 import sqlalchemy as sa
+import sqlalchemy.event
 import sqlalchemy.types
 from psycopg2 import extras
 from singer_sdk.helpers.conform import TypeConformanceLevel
@@ -192,6 +194,41 @@ def get_schema_names(self, engine: Engine, inspected: Inspector) -> list[str]:
             return self.config["filter_schemas"]
         return super().get_schema_names(engine, inspected)
 
+    def create_engine(self) -> sa.Engine:
+        """Create a SQLAlchemy engine with dynamic RDS IAM authentication support."""
+        engine = super().create_engine()
+
+        # If IAM authentication is enabled, hook into the do_connect event
+        if self.config.get("aws_iam_auth"):
+
+            @sa.event.listens_for(engine, "do_connect")
+            def provide_token(dialect, conn_rec, cargs, cparams):
+                self.provide_token(dialect, conn_rec, cargs, cparams)
+
+        return engine
+
+    def provide_token(self, dialect, conn_rec, cargs, cparams):
+        """Inject a fresh RDS IAM authentication token into connection parameters."""
+        host = cparams.get("host")
+        port = cparams.get("port", 5432)
+        user = cparams.get("user")
+
+        # Get AWS credentials region and profile from tap configuration
+        region = self.config.get("aws_region")
+        profile = self.config.get("aws_profile")
+
+        # Initialize AWS session and RDS client
+        session = boto3.Session(profile_name=profile) if profile else boto3.Session()
+        rds_client = session.client("rds", region_name=region)
+
+        # Dynamically generate a fresh token
+        token = rds_client.generate_db_auth_token(
+            DBHostname=host, Port=int(port), DBUsername=user, Region=region
+        )
+
+        # Overwrite the password with the generated token
+        cparams["password"] = token
+
 
 class PostgresStream(SQLStream):
     """Stream class for Postgres streams."""

tap_postgres/connection_parameters.py

@@ -71,7 +71,7 @@ def from_tap_config(cls, config: Mapping[str, Any]) -> ConnectionParameters:
             port=int(config["port"]),
             database=config["database"],
             user=config["user"],
-            password=config["password"],
+            password=config.get("password") or "",
             options=_build_options_from_tap_config(config),
         )
 

tap_postgres/tap.py

@@ -230,10 +230,10 @@ def __init__(
             self.config.get("host") is not None
             and self.config.get("port") is not None
             and self.config.get("user") is not None
-            and self.config.get("password") is not None
+            and (self.config.get("password") is not None or self.config.get("aws_iam_auth") is True)
         ), (
             "Need either the sqlalchemy_url to be set or host, port, user,"
-            + " and password to be set"
+            + " and password (or aws_iam_auth) to be set"
         )
 
         # If sqlalchemy_url is not being used and ssl_enable is on, ssl_mode must have
@@ -535,6 +535,22 @@ def __init__(
                 "this choice. One of `FULL_TABLE`, `INCREMENTAL`, or `LOG_BASED`."
             ),
         ),
+        th.Property(
+            "aws_iam_auth",
+            th.BooleanType,
+            default=False,
+            description=("Whether to use AWS IAM database authentication."),
+        ),
+        th.Property(
+            "aws_region",
+            th.StringType,
+            description=("The AWS region where the RDS instance is located."),
+        ),
+        th.Property(
+            "aws_profile",
+            th.StringType,
+            description=("Optional AWS CLI profile name to use for credential resolution."),
+        ),
         th.Property(
             "log_based_single_connection",
             th.BooleanType,

tests/test_aws_iam_auth.py

@@ -0,0 +1,54 @@
+from unittest.mock import MagicMock, patch
+
+from tap_postgres.tap import TapPostgres
+
+
+@patch("tap_postgres.client.boto3.Session")
+def test_aws_iam_auth_token_injection(mock_session_class):
+    # Mock the boto3 Session and RDS client
+    mock_session = MagicMock()
+    mock_client = MagicMock()
+    mock_session.client.return_value = mock_client
+    mock_session_class.return_value = mock_session
+
+    mock_client.generate_db_auth_token.return_value = "mocked-aws-token"
+
+    # Configuration with AWS IAM enabled and password omitted
+    config = {
+        "host": "localhost",
+        "port": 5432,
+        "user": "test_user",
+        "database": "test_db",
+        "aws_iam_auth": True,
+        "aws_region": "us-east-1",
+    }
+
+    # 1. Verify Tap initialization succeeds without a password config
+    tap = TapPostgres(config=config, setup_mapper=False)
+    assert tap.config.get("aws_iam_auth") is True
+
+    # 2. Retrieve the connector
+    connector = tap.connector
+
+    # 3. Simulate a database connection parameters payload
+    cparams = {
+        "host": "localhost",
+        "port": 5432,
+        "user": "test_user",
+        "password": "",  # Empty initial password
+    }
+
+    # 4. Directly test the provide_token method (completely bypassing database connections)
+    connector.provide_token(None, None, None, cparams)
+
+    # 5. Assertions:
+    # Verify that the connection password was replaced with our mocked AWS token
+    assert cparams["password"] == "mocked-aws-token"
+
+    # Verify that the boto3 rds client generated the token with correct configs
+    mock_client.generate_db_auth_token.assert_called_once_with(
+        DBHostname="localhost",
+        Port=5432,
+        DBUsername="test_user",
+        Region="us-east-1",
+    )