From ca61f5ae8c322c2428c8839f5443283bc2c2aed0 Mon Sep 17 00:00:00 2001 From: Brendan Shanahan Date: Wed, 8 Apr 2026 17:03:33 -0400 Subject: [PATCH] add support for rds iam authentication add support for rds iam authentication add support for rds iam authentication add support for rds iam authentication fix arg name rm async engine import, return type annotation fix Path coercion of input ca path from env var handle iam token auth prior to peewee migration pass ssl args to peewee migration connection use sslmode without cert verification if ca path can't be resolved debug log log level info -> warning use rds iam config engine for alembic migration when using iam auth rm debug log statements fix formatting add support for rds iam authentication add support for rds iam authentication fix arg name rm async engine import, return type annotation fix Path coercion of input ca path from env var handle iam token auth prior to peewee migration pass ssl args to peewee migration connection use sslmode without cert verification if ca path can't be resolved debug log log level info -> warning use rds iam config engine for alembic migration when using iam auth rm debug log statements fix formatting --- backend/open_webui/env.py | 4 + backend/open_webui/internal/db.py | 144 +++++++++++++++++++++++- backend/open_webui/internal/wrappers.py | 9 +- backend/open_webui/migrations/env.py | 5 +- 4 files changed, 155 insertions(+), 7 deletions(-) diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index 53737a1f2f..76084f048d 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -335,6 +335,10 @@ elif DATABASE_TYPE == 'sqlite+sqlcipher' and not os.environ.get('DATABASE_URL'): if 'postgres://' in DATABASE_URL: DATABASE_URL = DATABASE_URL.replace('postgres://', 'postgresql://') +DATABASE_ENABLE_IAM_TOKEN_AUTH = os.environ.get('DATABASE_ENABLE_IAM_TOKEN_AUTH', 'False').lower() == 'true' + +DATABASE_CA_PATH = os.getenv('DATABASE_CA_PATH', None) + DATABASE_SCHEMA = os.environ.get('DATABASE_SCHEMA', None) DATABASE_POOL_SIZE = os.environ.get('DATABASE_POOL_SIZE', None) diff --git a/backend/open_webui/internal/db.py b/backend/open_webui/internal/db.py index b0545255a6..6f6de40671 100644 --- a/backend/open_webui/internal/db.py +++ b/backend/open_webui/internal/db.py @@ -3,11 +3,19 @@ import json import logging from contextlib import contextmanager from typing import Any, Optional +from pathlib import Path +from datetime import datetime, timedelta, timezone +from urllib.parse import parse_qs, quote_plus, urlsplit +import boto3 from open_webui.internal.wrappers import register_connection from open_webui.env import ( OPEN_WEBUI_DIR, + DB_VARS, + DATABASE_CA_PATH, + DATABASE_ENABLE_IAM_TOKEN_AUTH, DATABASE_URL, + DATABASE_USER, DATABASE_SCHEMA, DATABASE_POOL_MAX_OVERFLOW, DATABASE_POOL_RECYCLE, @@ -18,7 +26,10 @@ from open_webui.env import ( ENABLE_DB_MIGRATIONS, ) from peewee_migrate import Router +from pydantic import BaseModel, SecretStr from sqlalchemy import Dialect, create_engine, MetaData, event, types +from sqlalchemy.engine import Engine +from sqlalchemy.engine.url import URL from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import scoped_session, sessionmaker, Session from sqlalchemy.pool import QueuePool, NullPool @@ -28,6 +39,107 @@ from typing_extensions import Self log = logging.getLogger(__name__) +class IAMToken(BaseModel): + token: SecretStr + expiration: datetime + + +class RDSIAMConfig: + def __init__( + self, db_type: str, db_name: str, db_host: str, db_port: int, db_user: str, ca_path: str | None = None + ) -> None: + self.url = URL.create( + drivername=db_type, + username=db_user, + password=None, + host=db_host, + port=db_port, + database=db_name, + ) + + self.endpoint = db_host + self.port = db_port + self.db_user = db_user + + self.connect_args = {"sslmode": "require"} + if ca_path is None: + log.info("No CA provided; using sslmode=require without certificate verification.") + elif isinstance(ca_path, str): + _ca_path = Path(ca_path).resolve() + if _ca_path.exists() and _ca_path.is_file(): + self.connect_args["sslmode"] = "verify-full" + self.connect_args["sslrootcert"] = str(_ca_path) + else: + log.warning(f"CA file not found at {_ca_path}; using sslmode=require without certificate verification.") + else: + log.warning("Unable to verify CA file path; using sslmode=require without certificate verification.") + + self.client = boto3.client("rds", region_name=os.getenv("AWS_DEFAULT_REGION", "us-east-1")) + self.token = self.get_token() + self.engine = self.create_engine() + + def get_token(self) -> IAMToken: + """Generate and return an RDS IAM authorization token. + + RDS tokens are presigned URLs; expiration time is calculated from the token's + query string parameters. + """ + try: + token = self.client.generate_db_auth_token(self.endpoint, self.port, self.db_user) + expiration = self.set_token_expiration(token) + log.info(f"AWS RDS token for {self.endpoint} expires {expiration} UTC.") + return IAMToken(token=SecretStr(token), expiration=expiration) + except Exception as e: + log.error(f"AWS RDS token error: {e.__class__.__qualname__} {e}") + raise + + def set_token_expiration(self, token: str) -> datetime: + """ + Calculate authorization token expiration time from a presigned RDS URL. + + RDS tokens are presigned URLs. The query string contains the time at which the + token was generated, and the number of seconds until it expires. This function + converts these values and returns the token expiration as a UTC-aware datetime. + """ + split_token = urlsplit(token) + parsed_qs: dict = parse_qs(split_token.query) + generation_time = datetime.strptime(parsed_qs["X-Amz-Date"][0], "%Y%m%dT%H%M%SZ") + duration = timedelta(seconds=int(parsed_qs["X-Amz-Expires"][0])) + expiration = generation_time + duration + return expiration.replace(tzinfo=timezone.utc) + + def create_engine(self) -> Engine: + if isinstance(DATABASE_POOL_SIZE, int): + if DATABASE_POOL_SIZE > 0: + engine = create_engine( + self.url, + pool_size=DATABASE_POOL_SIZE, + max_overflow=DATABASE_POOL_MAX_OVERFLOW, + pool_timeout=DATABASE_POOL_TIMEOUT, + pool_recycle=DATABASE_POOL_RECYCLE, + pool_pre_ping=True, + poolclass=QueuePool, + connect_args=self.connect_args, + ) + else: + engine = create_engine( + self.url, + pool_pre_ping=True, + poolclass=NullPool, + connect_args=self.connect_args, + ) + else: + engine = create_engine(self.url, pool_pre_ping=True, connect_args=self.connect_args) + + return engine + + def check_token(self, dialect, connection_record, connection_args, connection_kwargs) -> None: + now = datetime.now(tz=timezone.utc) + if not self.token or now >= self.token.expiration: + self.token = self.get_token() + connection_kwargs["password"] = self.token.token.get_secret_value() + + class JSONField(types.TypeDecorator): impl = types.Text cache_ok = True @@ -52,11 +164,11 @@ class JSONField(types.TypeDecorator): # Workaround to handle the peewee migration # This is required to ensure the peewee migration is handled before the alembic migration -def handle_peewee_migration(DATABASE_URL): - # db = None +def handle_peewee_migration(db_url: str = DATABASE_URL, connect_args: dict | None = None): + db = None try: # Replace the postgresql:// with postgres:// to handle the peewee migration - db = register_connection(DATABASE_URL.replace('postgresql://', 'postgres://')) + db = register_connection(db_url.replace('postgresql://', 'postgres://'), connect_args=connect_args) migrate_dir = OPEN_WEBUI_DIR / 'internal' / 'migrations' router = Router(db, logger=log, migrate_dir=migrate_dir) router.run() @@ -72,11 +184,30 @@ def handle_peewee_migration(DATABASE_URL): db.close() # Assert if db connection has been closed - assert db.is_closed(), 'Database connection is still open.' + if db is not None: + assert db.is_closed(), 'Database connection is still open.' +if DATABASE_ENABLE_IAM_TOKEN_AUTH: + rds_iam_config = RDSIAMConfig( + db_type=DB_VARS["db_type"], + db_name=DB_VARS["db_name"], + db_host=DB_VARS["db_host"], + db_port=DB_VARS["db_port"], + db_user=DATABASE_USER, + ca_path=DATABASE_CA_PATH, + ) + if ENABLE_DB_MIGRATIONS: - handle_peewee_migration(DATABASE_URL) + if DATABASE_ENABLE_IAM_TOKEN_AUTH: + _token_val = rds_iam_config.token.token.get_secret_value() + _migration_url = ( + f'{DB_VARS["db_type"]}://{DATABASE_USER}:{quote_plus(_token_val)}' + f'@{DB_VARS["db_host"]}:{DB_VARS["db_port"]}/{DB_VARS["db_name"]}' + ) + handle_peewee_migration(_migration_url, connect_args=rds_iam_config.connect_args) + else: + handle_peewee_migration(DATABASE_URL) SQLALCHEMY_DATABASE_URL = DATABASE_URL @@ -137,6 +268,9 @@ elif 'sqlite' in SQLALCHEMY_DATABASE_URL: cursor.close() event.listen(engine, 'connect', on_connect) +elif DATABASE_ENABLE_IAM_TOKEN_AUTH: + engine = rds_iam_config.engine + event.listen(engine, 'do_connect', rds_iam_config.check_token) else: if isinstance(DATABASE_POOL_SIZE, int): if DATABASE_POOL_SIZE > 0: diff --git a/backend/open_webui/internal/wrappers.py b/backend/open_webui/internal/wrappers.py index 3d54d02e3a..4fcfe5be56 100644 --- a/backend/open_webui/internal/wrappers.py +++ b/backend/open_webui/internal/wrappers.py @@ -41,7 +41,7 @@ class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase): pass -def register_connection(db_url): +def register_connection(db_url: str, connect_args: dict | None = None): # Check if using SQLCipher protocol if db_url.startswith('sqlite+sqlcipher://'): database_password = os.environ.get('DATABASE_PASSWORD') @@ -71,7 +71,14 @@ def register_connection(db_url): # Get the connection details connection = parse(db_url, unquote_user=True, unquote_password=True) + if connect_args: + if "sslmode" in connect_args: + connection["sslmode"] = connect_args["sslmode"] + if "sslrootcert" in connect_args: + connection["sslrootcert"] = connect_args["sslrootcert"] + # Use our custom database class that supports reconnection + log.debug(f"Peewee connection params (password redacted): { {k: ('***' if k == 'password' else v) for k, v in connection.items()} }") db = ReconnectingPostgresqlDatabase(**connection) db.connect(reuse_if_open=True) elif isinstance(db, SqliteDatabase): diff --git a/backend/open_webui/migrations/env.py b/backend/open_webui/migrations/env.py index 9ee6c2dceb..4e6a2b9b25 100644 --- a/backend/open_webui/migrations/env.py +++ b/backend/open_webui/migrations/env.py @@ -3,7 +3,7 @@ from logging.config import fileConfig from alembic import context from open_webui.models.auths import Auth -from open_webui.env import DATABASE_URL, DATABASE_PASSWORD, LOG_FORMAT +from open_webui.env import DATABASE_URL, DATABASE_ENABLE_IAM_TOKEN_AUTH, DATABASE_PASSWORD, LOG_FORMAT from sqlalchemy import engine_from_config, pool, create_engine # this is the Alembic Config object, which provides @@ -93,6 +93,9 @@ def run_migrations_online() -> None: creator=create_sqlcipher_connection, echo=False, ) + elif DATABASE_ENABLE_IAM_TOKEN_AUTH: + from open_webui.internal.db import rds_iam_config + connectable = rds_iam_config.engine else: # Standard database connection (existing logic) connectable = engine_from_config(