diff --git a/backend/open_webui/internal/db.py b/backend/open_webui/internal/db.py index e0d9c9ede6..aa96d2f2a0 100644 --- a/backend/open_webui/internal/db.py +++ b/backend/open_webui/internal/db.py @@ -202,6 +202,20 @@ def enable_iam_token_auth(connectable) -> None: return engine = getattr(connectable, 'sync_engine', connectable) + url = engine.url + auth = _rds_iam_token_auth + # The token is bound to one host/port/user pair; leave other databases on their own credentials. + if (url.host, url.port or 5432, url.username) != (auth.host, auth.port, auth.username): + log.warning( + 'AWS RDS IAM token auth not applied to %s: the token is issued for %s@%s:%s, ' + 'so this connection uses the password from its own URL', + url.render_as_string(hide_password=True), + auth.username, + auth.host, + auth.port, + ) + return + if not event.contains(engine, 'do_connect', _set_iam_token_password): event.listen(engine, 'do_connect', _set_iam_token_password) diff --git a/backend/open_webui/retrieval/vector/dbs/pgvector.py b/backend/open_webui/retrieval/vector/dbs/pgvector.py index 91158bcc03..75daeccd12 100644 --- a/backend/open_webui/retrieval/vector/dbs/pgvector.py +++ b/backend/open_webui/retrieval/vector/dbs/pgvector.py @@ -17,6 +17,7 @@ from open_webui.config import ( PGVECTOR_POOL_TIMEOUT, PGVECTOR_USE_HALFVEC, ) +from open_webui.internal.db import ScopedSession, enable_iam_token_auth from open_webui.retrieval.vector.main import ( GetResult, SearchResult, @@ -87,8 +88,6 @@ class PgvectorClient(VectorDBBase): def __init__(self) -> None: # if no pgvector uri, use the existing database connection if not PGVECTOR_DB_URL: - from open_webui.internal.db import ScopedSession - self.session = ScopedSession else: if isinstance(PGVECTOR_POOL_SIZE, int): @@ -107,6 +106,7 @@ class PgvectorClient(VectorDBBase): else: engine = create_engine(PGVECTOR_DB_URL, pool_pre_ping=True) + enable_iam_token_auth(engine) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False) self.session = scoped_session(SessionLocal)