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>

Author: justin-tahara

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

@@ -234,9 +234,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/connectors/jira/connector.py

@@ -84,7 +84,12 @@ def _perform_jql_search(
         f"Fetching Jira issues with JQL: {jql}, "
         f"starting at {start}, max results: {max_results}"
     )
-    issues = jira_client.search_issues(
+
+    # Use custom search method to handle Atlassian's API breaking change
+    # The old /rest/api/{version}/search endpoint has been deprecated
+    # New endpoint is /rest/api/{version}/search/jql
+    issues = _custom_search_issues(
+        jira_client,
         jql_str=jql,
         startAt=start,
         maxResults=max_results,
@@ -98,6 +103,67 @@ def _perform_jql_search(
             raise RuntimeError(f"Found Jira object not of type Issue: {issue}")
 
 
+def _custom_search_issues(
+    jira_client: JIRA,
+    jql_str: str,
+    startAt: int = 0,
+    maxResults: int = 50,
+    fields: str | None = None,
+) -> Iterable[Issue]:
+    """
+    Simple fix for Atlassian's API breaking change.
+
+    The old /rest/api/{version}/search endpoint has been deprecated and removed.
+    New endpoint is /rest/api/{version}/search/jql
+
+    This is a minimal fix to resolve the immediate issue. For performance improvements,
+    see the upgrade instructions in JIRA_API_FIX_SUMMARY.md
+    """
+    if isinstance(fields, str):
+        fields = fields.split(",")
+    elif fields is None:
+        fields = ["*all"]
+
+    # Build search parameters - keep the same interface for backwards compatibility
+    search_params = {
+        "jql": jql_str,
+        "startAt": startAt,
+        "maxResults": maxResults,
+        "fields": fields,
+        "validateQuery": True,
+    }
+
+    # Use the new JQL endpoint
+    url = f"{jira_client.server_url}/rest/api/{jira_client._options.get('rest_api_version', '3')}/search/jql"
+
+    # Make the request directly to the new endpoint
+    response = jira_client._session.post(url, json=search_params)
+
+    if response.status_code == 410:
+        # Fallback to old method if needed (though it should fail now)
+        logger.warning("JQL endpoint returned 410, falling back to old search method")
+        return jira_client.search_issues(
+            jql_str=jql_str,
+            startAt=startAt,
+            maxResults=maxResults,
+            fields=fields,
+        )
+
+    response.raise_for_status()
+    data = response.json()
+
+    # Convert the response to Issue objects
+    issues = []
+    for issue_data in data.get("issues", []):
+        issue = Issue(jira_client, issue_data)
+        # Ensure the issue has the necessary attributes
+        if not hasattr(issue, "key") and "key" in issue_data:
+            issue.key = issue_data["key"]
+        issues.append(issue)
+
+    return issues
+
+
 def process_jira_issue(
     jira_client: JIRA,
     issue: Issue,

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: