feat: support AWS IAM authentication for RDS (DIA-39107) (#43)

Author: ecotg

.circleci/config.yml

@@ -24,6 +24,7 @@ workflows:
             parameters:
               sqlalchemy_version: ["1.3", "1.4"]
               asyncpg: ["asyncpg", "noasyncpg"]
+              aws_rds_iam: ["aws_rds_iam", "noaws_rds_iam"]
       - release/release:
           name: release
           requires:
@@ -62,13 +63,17 @@ jobs:
         type: enum
         enum: ["asyncpg", "noasyncpg"]
         description: To run tests with and without asyncpg installed.
+      aws_rds_iam:
+        type: enum
+        enum: ["aws_rds_iam", "noaws_rds_iam"]
+        description: To run tests with and without asyncpg installed.
     executor: python-postgres
     working_directory: ~/project/.
     steps:
       - base/setup
       - python/setup
       - utils/with_cache:
-          key: 'sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>-{{ checksum "pyproject.toml" }}-{{ checksum "poetry.lock" }}'
+          key: 'sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>-<<parameters.aws_rds_iam>>-{{ checksum "pyproject.toml" }}-{{ checksum "poetry.lock" }}'
           namespace: tox
           path: ~/project/.tox
           steps:
@@ -77,13 +82,13 @@ jobs:
                 command: |
                   poetry run pip install -U tox
             - run:
-                name: "run tox using sqlalchemy <<parameters.sqlalchemy_version>>.* and -<<parameters.asyncpg>>"
+                name: "run tox using sqlalchemy <<parameters.sqlalchemy_version>>.* and -<<parameters.asyncpg>> and -<<parameters.aws_rds_iam>>"
                 command: |
-                  poetry run tox -e sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>
+                  poetry run tox -e sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>-<<parameters.aws_rds_iam>>
       - store_test_results:
           path: test-reports
       - utils/send_coverage_to_codecov:
-          codecov_flag: sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>
+          codecov_flag: sqlalchemy<<parameters.sqlalchemy_version>>-<<parameters.asyncpg>>-<<parameters.aws_rds_iam>>
 
   publish:
     docker:

README.md

@@ -406,7 +406,7 @@ It returns the path of  `alembic.ini` configuration file. By default, it returns
 ## Setup
 
 ```bash
-$ poetry install --extras tests --extras asyncpg
+$ poetry install --extras tests --extras asyncpg --extras aws_rds_iam
 ```
 
 ## Running tests

fastapi_sqla/__init__.py

@@ -18,6 +18,8 @@
 from sqlalchemy.orm.session import sessionmaker
 from sqlalchemy.sql import Select, func, select
 
+from . import aws_rds_iam_support
+
 try:
     from sqlalchemy.orm import declarative_base
 except ImportError:
@@ -63,6 +65,8 @@ def setup(app: FastAPI):
 def startup():
     lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
     engine = engine_from_config(lowercase_environ, prefix="sqlalchemy_")
+    aws_rds_iam_support.setup(engine.engine)
+
     Base.metadata.bind = engine
     Base.prepare(engine)
     _Session.configure(bind=engine)

fastapi_sqla/_pytest_plugin.py

@@ -17,16 +17,33 @@
 
 
 @fixture(scope="session")
-def db_url():
-    """Default db url used by depending fixtures.
+def db_host():
+    """Default db host used by depending fixtures.
+
+    When CI key is set in environment variables, it uses `postgres` as host name else, host used is `localhost`
+    """
+
+    return "postgres" if "CI" in os.environ else "localhost"
+
+
+@fixture(scope="session")
+def db_user():
+    """Default db user used by depending fixtures.
 
-    When CI key is set in environment variables, it uses `postgres` as host name:
-    postgresql://postgres@posgres/postgres
+    postgres
+    """
+
+    return "postgres"
 
-    Else, host used is `localhost`: postgresql://postgres@localhost/postgres
+
+@fixture(scope="session")
+def db_url(db_host, db_user):
+    """Default db url used by depending fixtures.
+
+    db url example postgresql://{db_user}@{db_host}/postgres
     """
-    host = "postgres" if "CI" in os.environ else "localhost"
-    return f"postgresql://postgres@{host}/postgres"
+
+    return f"postgresql://{db_user}@{db_host}/postgres"
 
 
 @fixture(scope="session")

fastapi_sqla/asyncio_support.py

@@ -8,6 +8,8 @@
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.orm.session import sessionmaker
 
+from . import aws_rds_iam_support
+
 logger = structlog.get_logger(__name__)
 _ASYNC_SESSION_KEY = "fastapi_sqla_async_session"
 _AsyncSession = sessionmaker(class_=SqlaAsyncSession)
@@ -16,6 +18,7 @@
 async def startup():
     async_sqlalchemy_url = os.environ["async_sqlalchemy_url"]
     engine = create_async_engine(async_sqlalchemy_url)
+    aws_rds_iam_support.setup(engine.sync_engine)
     _AsyncSession.configure(bind=engine, expire_on_commit=False)
 
     # Fail early:

fastapi_sqla/aws_rds_iam_support.py

@@ -0,0 +1,38 @@
+try:
+    import boto3
+
+    boto3_installed = True
+except ImportError as err:
+    boto3_installed = False
+    boto3_installed_err = str(err)
+
+from pydantic import BaseSettings
+from sqlalchemy import event
+
+
+def setup(engine):
+    config = Config()
+
+    if config.aws_rds_iam_enabled:
+        assert boto3_installed, boto3_installed_err
+        event.listen(engine, "do_connect", set_connection_token)
+
+
+def get_authentication_token(host, port, user):
+    session = boto3.Session()
+    client = session.client("rds")
+    token = client.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user)
+    return token
+
+
+def set_connection_token(dialect, conn_rec, cargs, cparams):
+    cparams["password"] = get_authentication_token(
+        host=cparams["host"], port=cparams.get("port", 5432), user=cparams["user"]
+    )
+
+
+class Config(BaseSettings):
+    aws_rds_iam_enabled: bool = False
+
+    class Config:
+        env_prefix = "fastapi_sqla_"