feat(infra): Add IAM support for Redis

backend/onyx/background/celery/configs/base.py

@@ -12,6 +12,7 @@
 from onyx.configs.app_configs import REDIS_SSL
 from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
 from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
+from onyx.configs.app_configs import USE_REDIS_IAM_AUTH
 from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
 
@@ -25,7 +26,7 @@
 
 # SSL-specific query parameters for Redis URL
 SSL_QUERY_PARAMS = ""
-if REDIS_SSL:
+if REDIS_SSL and not USE_REDIS_IAM_AUTH:
     REDIS_SCHEME = "rediss"
     SSL_QUERY_PARAMS = f"?ssl_cert_reqs={REDIS_SSL_CERT_REQS}"
     if REDIS_SSL_CA_CERTS:

backend/onyx/configs/app_configs.py

@@ -231,9 +231,12 @@
 except ValueError:
     POSTGRES_POOL_RECYCLE = POSTGRES_POOL_RECYCLE_DEFAULT
 
+# RDS IAM authentication - enables IAM-based authentication for PostgreSQL
 USE_IAM_AUTH = os.getenv("USE_IAM_AUTH", "False").lower() == "true"
 
-
+# Redis IAM authentication - enables IAM-based authentication for Redis ElastiCache
+# Note: This is separate from RDS IAM auth as they use different authentication mechanisms
+USE_REDIS_IAM_AUTH = os.getenv("USE_REDIS_IAM_AUTH", "False").lower() == "true"
 REDIS_SSL = os.getenv("REDIS_SSL", "").lower() == "true"
 REDIS_HOST = os.environ.get("REDIS_HOST") or "localhost"
 REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))

backend/onyx/redis/iam_auth.py

@@ -0,0 +1,38 @@
+"""
+Redis IAM Authentication Module
+This module provides Redis IAM authentication functionality for AWS ElastiCache.
+Unlike RDS IAM auth, Redis IAM auth relies on IAM roles and policies rather than
+generating authentication tokens.
+Key functions:
+- configure_redis_iam_auth: Configure Redis connection parameters for IAM auth
+- create_redis_ssl_context_if_iam: Create SSL context for secure connections
+"""
+
+import ssl
+from typing import Any
+
+
+def configure_redis_iam_auth(connection_kwargs: dict[str, Any]) -> None:
+    """
+    Configure Redis connection parameters for IAM authentication.
+    Modifies the connection_kwargs dict in-place to:
+    1. Remove password (not needed with IAM)
+    2. Enable SSL with system CA certificates
+    3. Set proper SSL context for secure connections
+    """
+    # Remove password as it's not needed with IAM authentication
+    if "password" in connection_kwargs:
+        del connection_kwargs["password"]
+
+    # Ensure SSL is enabled for IAM authentication
+    connection_kwargs["ssl"] = True
+    connection_kwargs["ssl_context"] = create_redis_ssl_context_if_iam()
+
+
+def create_redis_ssl_context_if_iam() -> ssl.SSLContext:
+    """Create an SSL context for Redis IAM authentication using system CA certificates."""
+    # Use system CA certificates by default - no need for additional CA files
+    ssl_context = ssl.create_default_context()
+    ssl_context.check_hostname = True
+    ssl_context.verify_mode = ssl.CERT_REQUIRED
+    return ssl_context

backend/onyx/redis/redis_pool.py

@@ -25,8 +25,11 @@
 from onyx.configs.app_configs import REDIS_SSL
 from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
 from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
+from onyx.configs.app_configs import USE_REDIS_IAM_AUTH
 from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
+from onyx.redis.iam_auth import configure_redis_iam_auth
+from onyx.redis.iam_auth import create_redis_ssl_context_if_iam
 from onyx.utils.logger import setup_logger
 from shared_configs.configs import DEFAULT_REDIS_PREFIX
 from shared_configs.contextvars import get_current_tenant_id
@@ -186,12 +189,41 @@ def create_pool(
         ssl_cert_reqs: str = REDIS_SSL_CERT_REQS,
         ssl: bool = False,
     ) -> redis.BlockingConnectionPool:
-        """We use BlockingConnectionPool because it will block and wait for a connection
+        """
+        Create a Redis connection pool with appropriate SSL configuration.
+        SSL Configuration Priority:
+        1. IAM Authentication (USE_REDIS_IAM_AUTH=true): Uses system CA certificates
+        2. Regular SSL (REDIS_SSL=true): Uses custom SSL configuration
+        3. No SSL: Standard connection without encryption
+        Note: IAM authentication automatically enables SSL and takes precedence
+        over regular SSL configuration to ensure proper security.
+
+        We use BlockingConnectionPool because it will block and wait for a connection
         rather than error if max_connections is reached. This is far more deterministic
         behavior and aligned with how we want to use Redis."""
 
         # Using ConnectionPool is not well documented.
         # Useful examples: https://github.com/redis/redis-py/issues/780
+
+        # Handle IAM authentication
+        if USE_REDIS_IAM_AUTH:
+            # For IAM authentication, we don't use password
+            # and ensure SSL is enabled with proper context
+            ssl_context = create_redis_ssl_context_if_iam()
+            return redis.BlockingConnectionPool(
+                host=host,
+                port=port,
+                db=db,
+                password=None,  # No password with IAM auth
+                max_connections=max_connections,
+                timeout=None,
+                health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,
+                socket_keepalive=True,
+                socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,
+                connection_class=redis.SSLConnection,
+                ssl_context=ssl_context,  # Use IAM auth SSL context
+            )
+
         if ssl:
             return redis.BlockingConnectionPool(
                 host=host,
@@ -363,7 +395,9 @@ async def get_async_redis_connection() -> aioredis.Redis:
                     "socket_keepalive_options": REDIS_SOCKET_KEEPALIVE_OPTIONS,
                 }
 
-                if REDIS_SSL:
+                if USE_REDIS_IAM_AUTH:
+                    configure_redis_iam_auth(connection_kwargs)
+                elif REDIS_SSL:
                     ssl_context = ssl.create_default_context()
 
                     if REDIS_SSL_CA_CERTS: