diff --git a/pyproject.toml b/pyproject.toml index 288c6896..62c1fa6b 100644 --- a/pyproject.toml +++ b/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", diff --git a/tap_postgres/client.py b/tap_postgres/client.py index d71c5a85..d2a01396 100644 --- a/tap_postgres/client.py +++ b/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.""" diff --git a/tap_postgres/connection_parameters.py b/tap_postgres/connection_parameters.py index 8fd130e4..ec8ed638 100644 --- a/tap_postgres/connection_parameters.py +++ b/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), ) diff --git a/tap_postgres/tap.py b/tap_postgres/tap.py index 89c45856..e6b08515 100644 --- a/tap_postgres/tap.py +++ b/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, diff --git a/tests/test_aws_iam_auth.py b/tests/test_aws_iam_auth.py new file mode 100644 index 00000000..22147bfd --- /dev/null +++ b/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", + ) diff --git a/uv.lock b/uv.lock index 9724f9d8..3ad82b1e 100644 --- a/uv.lock +++ b/uv.lock @@ -161,6 +161,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e4/f8/972c96f5a2b6c4b3deca57009d93e946bbdbe2241dca9806d502f29dd3ee/bcrypt-5.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:6b8f520b61e8781efee73cba14e3e8c9556ccfb375623f4f97429544734545b4", size = 273375, upload-time = "2025-09-25T19:50:45.43Z" }, ] +[[package]] +name = "boto3" +version = "1.43.29" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/85/a6/9c02ff00d08ea87908934351d244e35bb6fb5cbc169e1a14fc5bd80d124b/boto3-1.43.29.tar.gz", hash = "sha256:354006c512cdb87ef8214a095f2ade961c8145734475cd7a7e6b39260ff5494a", size = 113198, upload-time = "2026-06-12T19:32:23.442Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/07/5c/f12a9978526c7068c873ccf9788161fc6af338c6a025f1354a46134a6e46/boto3-1.43.29-py3-none-any.whl", hash = "sha256:77c27ada27cdbf619a3bbc41fa9e991caef818d3a2988cf92ea722e107d90108", size = 140537, upload-time = "2026-06-12T19:32:21.682Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.29" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6f/54/df99c5ca5c9ef275e34b87e177782e3ca054fc35f1f462c40fe180936c81/botocore-1.43.29.tar.gz", hash = "sha256:dce39d33b707aa162aa3820975f99d7f8f746d46576169fb42ce4f2b3b56b261", size = 15512384, upload-time = "2026-06-12T19:32:13.754Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/b1/aa410c22355f8f6c4ac2433db6a1c557dd959acf2953ccae4bfc37488119/botocore-1.43.29-py3-none-any.whl", hash = "sha256:5d62f2a03ed279a50207ca2824e009313df15f082b6bb591a095a4f04c7faef3", size = 15194135, upload-time = "2026-06-12T19:32:09.463Z" }, +] + [[package]] name = "certifi" version = "2026.7.22" @@ -669,6 +697,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/de/bbc12563bbf979618d17625a4e753ff7a078523e28d870d3626daa97261a/invoke-3.0.3-py3-none-any.whl", hash = "sha256:f11327165e5cbb89b2ad1d88d3292b5113332c43b8553b494da435d6ec6f5053", size = 160958, upload-time = "2026-04-07T15:17:46.875Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joblib" version = "1.5.3" @@ -806,6 +843,7 @@ wheels = [ name = "meltanolabs-tap-postgres" source = { editable = "." } dependencies = [ + { name = "boto3" }, { name = "paramiko" }, { name = "psycopg2-binary" }, { name = "singer-sdk", extra = ["faker", "sql"] }, @@ -853,6 +891,7 @@ typing = [ [package.metadata] requires-dist = [ + { name = "boto3", specifier = ">=1.26.0" }, { name = "msgspec", marker = "extra == 'msgspec'", specifier = ">=0.19.0" }, { name = "paramiko", specifier = ">=4" }, { name = "psycopg2-binary", specifier = "==2.9.12" }, @@ -1212,6 +1251,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3c/ea/a3adc8d2cda57480cd7020ae6857f57b17ce1720bb6a652bbb0516fa4dff/python_backoff-2.3.1-py3-none-any.whl", hash = "sha256:d3cfcc0938ec78683edc530bbeb5f5c08252477a40270516497d2bb500da392b", size = 15106, upload-time = "2025-12-18T21:30:08.601Z" }, ] +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + [[package]] name = "python-dotenv" version = "1.2.2" @@ -1592,6 +1643,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/9a/8415f2657cbe200f41a4531ccededf135505a92d4a012229121f885b26f9/ruff-0.16.0-py3-none-win_arm64.whl", hash = "sha256:14296fedcd2705c77ab8235439278bbb38f285cf7da5528b00b3e330c3d4872d", size = 11273407, upload-time = "2026-07-23T19:11:28.705Z" }, ] +[[package]] +name = "s3transfer" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/1f/12417f7f493fc45e1f9fd5d4a9b6c125cf8d2cf3f8ddbdfab3e76406e9d6/s3transfer-0.18.0.tar.gz", hash = "sha256:3760b8b7ec1315da54048b2d626276732bee4300d054d492d4e1d43e20d4ecbd", size = 160560, upload-time = "2026-05-28T19:39:09.124Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/58/a58fc997655386daa2e25784e30c288aa3e3819e401f77029ee4899fb55a/s3transfer-0.18.0-py3-none-any.whl", hash = "sha256:239c13b09e65ad0346e1be7348b8a202dcad44ac7ea7c6eb858fc881dce739b6", size = 88572, upload-time = "2026-05-28T19:39:07.999Z" }, +] + [[package]] name = "simpleeval" version = "1.0.7" @@ -1718,6 +1781,15 @@ testing = [ { name = "pytest" }, ] +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + [[package]] name = "sortedcontainers" version = "2.4.0"