Skip to content

Commit 2dc564c

Browse files
feat(infra): Add IAM support for Redis (#5267)
* feat: JIRA support for custom JQL filter (#5164) * jira jql support * jira jql fixes * Address comment --------- Co-authored-by: sktbcpraha <131408565+sktbcpraha@users.noreply.github.com>
1 parent b259f53 commit 2dc564c

File tree

4 files changed

+80
-4
lines changed

4 files changed

+80
-4
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from onyx.configs.app_configs import REDIS_SSL
1313
from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
1414
from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
15+
from onyx.configs.app_configs import USE_REDIS_IAM_AUTH
1516
from onyx.configs.constants import OnyxCeleryPriority
1617
from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
1718

@@ -25,7 +26,7 @@
2526

2627
# SSL-specific query parameters for Redis URL
2728
SSL_QUERY_PARAMS = ""
28-
if REDIS_SSL:
29+
if REDIS_SSL and not USE_REDIS_IAM_AUTH:
2930
REDIS_SCHEME = "rediss"
3031
SSL_QUERY_PARAMS = f"?ssl_cert_reqs={REDIS_SSL_CERT_REQS}"
3132
if REDIS_SSL_CA_CERTS:

backend/onyx/configs/app_configs.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,9 +234,12 @@
234234
except ValueError:
235235
POSTGRES_POOL_RECYCLE = POSTGRES_POOL_RECYCLE_DEFAULT
236236

237+
# RDS IAM authentication - enables IAM-based authentication for PostgreSQL
237238
USE_IAM_AUTH = os.getenv("USE_IAM_AUTH", "False").lower() == "true"
238239

239-
240+
# Redis IAM authentication - enables IAM-based authentication for Redis ElastiCache
241+
# Note: This is separate from RDS IAM auth as they use different authentication mechanisms
242+
USE_REDIS_IAM_AUTH = os.getenv("USE_REDIS_IAM_AUTH", "False").lower() == "true"
240243
REDIS_SSL = os.getenv("REDIS_SSL", "").lower() == "true"
241244
REDIS_HOST = os.environ.get("REDIS_HOST") or "localhost"
242245
REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))

backend/onyx/redis/iam_auth.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
"""
2+
Redis IAM Authentication Module
3+
This module provides Redis IAM authentication functionality for AWS ElastiCache.
4+
Unlike RDS IAM auth, Redis IAM auth relies on IAM roles and policies rather than
5+
generating authentication tokens.
6+
Key functions:
7+
- configure_redis_iam_auth: Configure Redis connection parameters for IAM auth
8+
- create_redis_ssl_context_if_iam: Create SSL context for secure connections
9+
"""
10+
11+
import ssl
12+
from typing import Any
13+
14+
15+
def configure_redis_iam_auth(connection_kwargs: dict[str, Any]) -> None:
16+
"""
17+
Configure Redis connection parameters for IAM authentication.
18+
Modifies the connection_kwargs dict in-place to:
19+
1. Remove password (not needed with IAM)
20+
2. Enable SSL with system CA certificates
21+
3. Set proper SSL context for secure connections
22+
"""
23+
# Remove password as it's not needed with IAM authentication
24+
if "password" in connection_kwargs:
25+
del connection_kwargs["password"]
26+
27+
# Ensure SSL is enabled for IAM authentication
28+
connection_kwargs["ssl"] = True
29+
connection_kwargs["ssl_context"] = create_redis_ssl_context_if_iam()
30+
31+
32+
def create_redis_ssl_context_if_iam() -> ssl.SSLContext:
33+
"""Create an SSL context for Redis IAM authentication using system CA certificates."""
34+
# Use system CA certificates by default - no need for additional CA files
35+
ssl_context = ssl.create_default_context()
36+
ssl_context.check_hostname = True
37+
ssl_context.verify_mode = ssl.CERT_REQUIRED
38+
return ssl_context

backend/onyx/redis/redis_pool.py

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,11 @@
2525
from onyx.configs.app_configs import REDIS_SSL
2626
from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
2727
from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
28+
from onyx.configs.app_configs import USE_REDIS_IAM_AUTH
2829
from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
2930
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
3033
from onyx.utils.logger import setup_logger
3134
from shared_configs.configs import DEFAULT_REDIS_PREFIX
3235
from shared_configs.contextvars import get_current_tenant_id
@@ -186,12 +189,41 @@ def create_pool(
186189
ssl_cert_reqs: str = REDIS_SSL_CERT_REQS,
187190
ssl: bool = False,
188191
) -> 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
190202
rather than error if max_connections is reached. This is far more deterministic
191203
behavior and aligned with how we want to use Redis."""
192204

193205
# Using ConnectionPool is not well documented.
194206
# 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+
195227
if ssl:
196228
return redis.BlockingConnectionPool(
197229
host=host,
@@ -363,7 +395,9 @@ async def get_async_redis_connection() -> aioredis.Redis:
363395
"socket_keepalive_options": REDIS_SOCKET_KEEPALIVE_OPTIONS,
364396
}
365397

366-
if REDIS_SSL:
398+
if USE_REDIS_IAM_AUTH:
399+
configure_redis_iam_auth(connection_kwargs)
400+
elif REDIS_SSL:
367401
ssl_context = ssl.create_default_context()
368402

369403
if REDIS_SSL_CA_CERTS:

0 commit comments

Comments
 (0)