feat: AWS IAM auth for DocumentDB and pooled Postgres (v0.2.10)

README.md

@@ -276,6 +276,16 @@ client = get_mongodb(
     max_pool_size=200,
 )
 
+# AWS DocumentDB via IAM auth (MONGODB-AWS) — credentials from the AWS chain
+# (env vars, ECS task role, EC2 instance profile); requires cloudrift[aws].
+client = get_mongodb(
+    "documentdb",
+    auth="iam",
+    host="cluster.docdb.amazonaws.com",
+    port=27017,
+    tls_ca_file="/etc/ssl/rds-ca-bundle.pem",
+)
+
 # Azure Cosmos DB (MongoDB API)
 client = get_mongodb("cosmos", connection_string="mongodb://...")
 client = get_mongodb("cosmos", account="myacct", account_key="...")
@@ -322,6 +332,9 @@ from cloudrift.document import get_mongodb_sync
 
 client = get_mongodb_sync("documentdb", uri="mongodb://...")
 client = get_mongodb_sync("cosmos", account="myacct", account_key="...")
+client = get_mongodb_sync(
+    "documentdb", auth="iam",
+    host="cluster.docdb.amazonaws.com", tls_ca_file="/etc/ssl/rds-ca-bundle.pem")
 
 users = client["lyzr"]["users"]
 users.insert_one({"name": "Alice"})

cloudrift/__init__.py

@@ -8,7 +8,7 @@
 from cloudrift.pubsub import get_pubsub
 from cloudrift.email import get_email
 
-__version__ = "0.2.7"
+__version__ = "0.2.10"
 __all__ = [
     "get_storage",
     "get_queue",

cloudrift/document/__init__.py

@@ -39,6 +39,9 @@ def get_mongodb(provider: str, **kwargs) -> AsyncIOMotorClient:
     if provider == "documentdb":
         from cloudrift.document import documentdb
 
+        if kwargs.pop("auth", None) == "iam":
+            return documentdb.connect_iam_auth(**kwargs)
+
         if "uri" in kwargs:
             return documentdb.connect_uri(**kwargs)
         if "tls_cert_key_file" in kwargs:
@@ -79,6 +82,9 @@ def get_mongodb_sync(provider: str, **kwargs) -> MongoClient:
     if provider == "documentdb":
         from cloudrift.document import documentdb_sync
 
+        if kwargs.pop("auth", None) == "iam":
+            return documentdb_sync.connect_iam_auth(**kwargs)
+
         if "uri" in kwargs:
             return documentdb_sync.connect_uri(**kwargs)
         if "tls_cert_key_file" in kwargs:

cloudrift/document/documentdb.py

@@ -107,3 +107,40 @@ def connect_tls_cert(
         return AsyncIOMotorClient(uri, **kwargs)
     except Exception as e:
         raise DocumentConnectionError(f"Failed to connect to DocumentDB: {e}") from e
+
+
+def connect_iam_auth(
+    host: str,
+    port: int = 27017,
+    *,
+    tls_ca_file: str | None = None,
+    max_pool_size: int = 100,
+    min_pool_size: int = 0,
+    **client_kwargs,
+) -> AsyncIOMotorClient:
+    """Connect using AWS IAM authentication (``MONGODB-AWS`` mechanism).
+
+    Credentials are resolved from the standard AWS provider chain (environment,
+    ECS/EC2 instance role, etc.) by ``pymongo-auth-aws`` — install via
+    ``cloudrift[aws]``. IAM auth requires TLS, and the DocumentDB cluster must
+    have IAM authentication enabled.
+
+    Args:
+        host: DocumentDB cluster endpoint hostname.
+        port: Port number (DocumentDB default: 27017).
+        tls_ca_file: Optional path to the CA certificate bundle (PEM).
+        max_pool_size: Max connection pool size.
+        min_pool_size: Min connection pool size.
+    """
+    uri = (
+        f"mongodb://{host}:{port}/?tls=true&retryWrites=false"
+        "&authMechanism=MONGODB-AWS&authSource=%24external"
+    )
+    kwargs: dict = {"maxPoolSize": max_pool_size, "minPoolSize": min_pool_size}
+    if tls_ca_file:
+        kwargs["tlsCAFile"] = tls_ca_file
+    kwargs.update(client_kwargs)
+    try:
+        return AsyncIOMotorClient(uri, **kwargs)
+    except Exception as e:
+        raise DocumentConnectionError(f"Failed to connect to DocumentDB: {e}") from e

cloudrift/document/documentdb_sync.py

@@ -108,3 +108,40 @@ def connect_tls_cert(
         return MongoClient(uri, **kwargs)
     except Exception as e:
         raise DocumentConnectionError(f"Failed to connect to DocumentDB: {e}") from e
+
+
+def connect_iam_auth(
+    host: str,
+    port: int = 27017,
+    *,
+    tls_ca_file: str | None = None,
+    max_pool_size: int = 100,
+    min_pool_size: int = 0,
+    **client_kwargs,
+) -> MongoClient:
+    """Connect using AWS IAM authentication (``MONGODB-AWS`` mechanism).
+
+    Credentials are resolved from the standard AWS provider chain (environment,
+    ECS/EC2 instance role, etc.) by ``pymongo-auth-aws`` — install via
+    ``cloudrift[aws]``. IAM auth requires TLS, and the DocumentDB cluster must
+    have IAM authentication enabled.
+
+    Args:
+        host: DocumentDB cluster endpoint hostname.
+        port: Port number (DocumentDB default: 27017).
+        tls_ca_file: Optional path to the CA certificate bundle (PEM).
+        max_pool_size: Max connection pool size.
+        min_pool_size: Min connection pool size.
+    """
+    uri = (
+        f"mongodb://{host}:{port}/?tls=true&retryWrites=false"
+        "&authMechanism=MONGODB-AWS&authSource=%24external"
+    )
+    kwargs: dict = {"maxPoolSize": max_pool_size, "minPoolSize": min_pool_size}
+    if tls_ca_file:
+        kwargs["tlsCAFile"] = tls_ca_file
+    kwargs.update(client_kwargs)
+    try:
+        return MongoClient(uri, **kwargs)
+    except Exception as e:
+        raise DocumentConnectionError(f"Failed to connect to DocumentDB: {e}") from e

cloudrift/sql/postgresql.py

@@ -129,13 +129,19 @@
         user: str,
         database: str,
         region: str,
+        pool: bool = False,
+        pool_min_size: int = 0,
+        pool_max_size: int = 10,
         **connect_kwargs,
     ) -> "PostgresSQLBackend":
         """Authenticate to AWS RDS/Aurora PostgreSQL using an IAM auth token.
 
         A short-lived (15 min) token is generated on every :meth:`connect` call
         and used in place of a password. IAM auth requires TLS, so ``sslmode``
         defaults to ``require`` unless overridden in ``connect_kwargs``.
+
+        Set ``pool=True`` to enable a ``psycopg_pool`` pool used by
+        :meth:`acquire`; each pooled connection mints its own fresh token.
         """
         connect_kwargs.setdefault("sslmode", "require")
         return cls(
@@ -146,6 +152,9 @@
             iam=True,
             region=region,
             connect_kwargs=connect_kwargs,
+            pool=pool,
+            pool_min_size=pool_min_size,
+            pool_max_size=pool_max_size,
         )
 
     @classmethod
@@ -156,6 +165,9 @@
         user: str,
         database: str,
         client_id: str | None = None,
+        pool: bool = False,
+        pool_min_size: int = 0,
+        pool_max_size: int = 10,
         **connect_kwargs,
     ) -> "PostgresSQLBackend":
         """Authenticate to Azure Database for PostgreSQL via a Microsoft Entra
@@ -176,6 +188,9 @@
             entra=True,
             client_id=client_id,
             connect_kwargs=connect_kwargs,
+            pool=pool,
+            pool_min_size=pool_min_size,
+            pool_max_size=pool_max_size,
         )
 
     # ------------------------------------------------------------------
@@ -277,21 +292,45 @@
                 "host": self._host,
                 "port": self._port,
                 "user": self._user,
-                "password": self._password,
                 "dbname": self._database,
                 **self._connect_kwargs,
             }
-            pool = AsyncConnectionPool(
+            pool_kwargs: dict = dict(
                 conninfo="",
                 kwargs=kwargs,
                 min_size=self._pool_min_size,
                 max_size=self._pool_max_size,
                 open=False,
             )
+            if self._iam or self._entra:
+                # Token auth: no static password. Each physical connection the
+                # pool opens mints its own fresh short-lived token (RDS IAM
+                # tokens last ~15 min) via a custom connection class.
+                pool_kwargs["connection_class"] = self._token_connection_class()
+            else:
+                kwargs["password"] = self._password
+            pool = AsyncConnectionPool(**pool_kwargs)
             await pool.open()
             self._pool = pool
         return self._pool
 
+    def _token_connection_class(self):
+        """Build a ``psycopg.AsyncConnection`` subclass that authenticates each
+        new physical pool connection with a freshly minted IAM/Entra token."""
+        import psycopg
+
+        backend = self
+
+        class _TokenConnection(psycopg.AsyncConnection):
+            @classmethod
+            async def connect(cls, conninfo="", **kwargs):
+                kwargs["password"] = await (
+                    backend._rds_token() if backend._iam else backend._entra_token()
+                )
+                return await super().connect(conninfo, **kwargs)
+
+        return _TokenConnection
+
     @asynccontextmanager
     async def acquire(self, timeout: float | None = None):
         if not self._pool_enabled:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "lyzr-cloudrift"
-version = "0.2.9"
+version = "0.2.10"
 description = "Cloud-agnostic abstraction for storage, messaging, document databases, cache, secrets, SQL, crypto (KMS), pub/sub, and email"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -36,7 +36,7 @@ Documentation = "https://github.com/LYZR-OSS/cloudrift#readme"
 aws = [
     "aioboto3>=13.0.0",       # native async S3 / SQS / SNS / Secrets Manager / SES
     "motor>=3.3.0",           # async MongoDB driver for DocumentDB
-    "pymongo>=4.6.3",         # sync MongoDB driver (optional sync client)
+    "pymongo[aws]>=4.6.3",    # sync MongoDB driver + MONGODB-AWS IAM auth (pymongo-auth-aws)
     "redis[hiredis]>=5.0.0",  # ElastiCache (also used by standalone)
 ]
 azure = [
@@ -89,7 +89,7 @@ sql = [
 all = [
     "aioboto3>=13.0.0",
     "motor>=3.3.0",
-    "pymongo>=4.6.3",
+    "pymongo[aws]>=4.6.3",
     "azure-storage-blob>=12.19.0",
     "azure-servicebus>=7.11.0",
     "azure-identity>=1.15.0",
@@ -112,7 +112,7 @@ dev = [
     "moto[s3,sqs,sns,ses,secretsmanager,kms,server]>=5.0",
     "aioboto3>=13.0.0",
     "motor>=3.3.0",
-    "pymongo>=4.6.3",
+    "pymongo[aws]>=4.6.3",
     "fakeredis>=2.20.0",
     "httpx>=0.25.0",
     "ruff>=0.4.0",

tests/test_document.py

@@ -91,6 +91,37 @@ def test_documentdb_tls_cert_passes_cert_path(recorder):
     assert inst.kwargs["tlsCAFile"] == "/secrets/ca.pem"
 
 
+def test_documentdb_iam_auth_builds_mongodb_aws_uri(recorder):
+    get_mongodb(
+        "documentdb",
+        auth="iam",
+        host="cluster.docdb.amazonaws.com",
+        port=27017,
+        tls_ca_file="/etc/ssl/rds-ca-bundle.pem",
+    )
+    inst = recorder.instances[-1]
+    uri = inst.args[0]
+    assert uri.startswith("mongodb://cluster.docdb.amazonaws.com:27017/")
+    assert "authMechanism=MONGODB-AWS" in uri
+    assert "authSource=%24external" in uri
+    assert "tls=true" in uri
+    assert "retryWrites=false" in uri
+    # IAM URI carries no embedded credentials
+    assert "@" not in uri
+    assert inst.kwargs["tlsCAFile"] == "/etc/ssl/rds-ca-bundle.pem"
+    assert inst.kwargs["maxPoolSize"] == 100
+    assert inst.kwargs["minPoolSize"] == 0
+
+
+def test_documentdb_iam_auth_default_port_and_pool(recorder):
+    get_mongodb("documentdb", auth="iam", host="h", max_pool_size=250, min_pool_size=25)
+    inst = recorder.instances[-1]
+    assert inst.args[0].startswith("mongodb://h:27017/")
+    assert "tlsCAFile" not in inst.kwargs
+    assert inst.kwargs["maxPoolSize"] == 250
+    assert inst.kwargs["minPoolSize"] == 25
+
+
 def test_cosmos_account_key_builds_mongo_uri(recorder):
     get_mongodb("cosmos", account="myacct", account_key="raw+key/with=special")
     inst = recorder.instances[-1]

tests/test_document_sync.py

@@ -57,6 +57,27 @@ def test_documentdb_uri_passes_pool_kwargs(recorder):
     assert inst.kwargs["tlsCAFile"] == "/etc/ssl/ca.pem"
 
 
+def test_documentdb_iam_auth_builds_mongodb_aws_uri(recorder):
+    get_mongodb_sync(
+        "documentdb",
+        auth="iam",
+        host="cluster.docdb.amazonaws.com",
+        port=27017,
+        tls_ca_file="/etc/ssl/rds-ca-bundle.pem",
+    )
+    inst = recorder.instances[-1]
+    uri = inst.args[0]
+    assert uri.startswith("mongodb://cluster.docdb.amazonaws.com:27017/")
+    assert "authMechanism=MONGODB-AWS" in uri
+    assert "authSource=%24external" in uri
+    assert "tls=true" in uri
+    assert "retryWrites=false" in uri
+    assert "@" not in uri
+    assert inst.kwargs["tlsCAFile"] == "/etc/ssl/rds-ca-bundle.pem"
+    assert inst.kwargs["maxPoolSize"] == 100
+    assert inst.kwargs["minPoolSize"] == 0
+
+
 def test_documentdb_credentials_url_encodes_password(recorder):
     get_mongodb_sync(
         "documentdb",