|
25 | 25 | from onyx.configs.app_configs import REDIS_SSL
|
26 | 26 | from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
|
27 | 27 | from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
|
| 28 | +from onyx.configs.app_configs import USE_REDIS_IAM_AUTH |
28 | 29 | from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
|
29 | 30 | from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
|
| 31 | +from onyx.redis.iam_auth import configure_redis_iam_auth |
| 32 | +from onyx.redis.iam_auth import create_redis_ssl_context_if_iam |
30 | 33 | from onyx.utils.logger import setup_logger
|
31 | 34 | from shared_configs.configs import DEFAULT_REDIS_PREFIX
|
32 | 35 | from shared_configs.contextvars import get_current_tenant_id
|
@@ -186,12 +189,41 @@ def create_pool(
|
186 | 189 | ssl_cert_reqs: str = REDIS_SSL_CERT_REQS,
|
187 | 190 | ssl: bool = False,
|
188 | 191 | ) -> redis.BlockingConnectionPool:
|
189 |
| - """We use BlockingConnectionPool because it will block and wait for a connection |
| 192 | + """ |
| 193 | + Create a Redis connection pool with appropriate SSL configuration. |
| 194 | + SSL Configuration Priority: |
| 195 | + 1. IAM Authentication (USE_REDIS_IAM_AUTH=true): Uses system CA certificates |
| 196 | + 2. Regular SSL (REDIS_SSL=true): Uses custom SSL configuration |
| 197 | + 3. No SSL: Standard connection without encryption |
| 198 | + Note: IAM authentication automatically enables SSL and takes precedence |
| 199 | + over regular SSL configuration to ensure proper security. |
| 200 | +
|
| 201 | + We use BlockingConnectionPool because it will block and wait for a connection |
190 | 202 | rather than error if max_connections is reached. This is far more deterministic
|
191 | 203 | behavior and aligned with how we want to use Redis."""
|
192 | 204 |
|
193 | 205 | # Using ConnectionPool is not well documented.
|
194 | 206 | # Useful examples: https://github.com/redis/redis-py/issues/780
|
| 207 | + |
| 208 | + # Handle IAM authentication |
| 209 | + if USE_REDIS_IAM_AUTH: |
| 210 | + # For IAM authentication, we don't use password |
| 211 | + # and ensure SSL is enabled with proper context |
| 212 | + ssl_context = create_redis_ssl_context_if_iam() |
| 213 | + return redis.BlockingConnectionPool( |
| 214 | + host=host, |
| 215 | + port=port, |
| 216 | + db=db, |
| 217 | + password=None, # No password with IAM auth |
| 218 | + max_connections=max_connections, |
| 219 | + timeout=None, |
| 220 | + health_check_interval=REDIS_HEALTH_CHECK_INTERVAL, |
| 221 | + socket_keepalive=True, |
| 222 | + socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS, |
| 223 | + connection_class=redis.SSLConnection, |
| 224 | + ssl_context=ssl_context, # Use IAM auth SSL context |
| 225 | + ) |
| 226 | + |
195 | 227 | if ssl:
|
196 | 228 | return redis.BlockingConnectionPool(
|
197 | 229 | host=host,
|
@@ -363,7 +395,9 @@ async def get_async_redis_connection() -> aioredis.Redis:
|
363 | 395 | "socket_keepalive_options": REDIS_SOCKET_KEEPALIVE_OPTIONS,
|
364 | 396 | }
|
365 | 397 |
|
366 |
| - if REDIS_SSL: |
| 398 | + if USE_REDIS_IAM_AUTH: |
| 399 | + configure_redis_iam_auth(connection_kwargs) |
| 400 | + elif REDIS_SSL: |
367 | 401 | ssl_context = ssl.create_default_context()
|
368 | 402 |
|
369 | 403 | if REDIS_SSL_CA_CERTS:
|
|
0 commit comments