UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-02-22 13:18:25 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

add logging for user upon db connection

Browse Source
This commit is contained in:
perf3ct 2024-06-17 10:34:19 -07:00
parent 5c655f298b
commit 48e1356ed9
No known key found for this signature in database
GPG Key ID: 569C4EEC436F5232
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
backend/apps/webui/internal/wrappers.py
View File

@ -1,8 +1,12 @@
from contextvars import ContextVar from contextvars import ContextVar
from peewee import * from peewee import *
import logging
from playhouse.db_url import connect from playhouse.db_url import connect
from playhouse.pool import PooledPostgresqlExtDatabase
from playhouse.pool import PooledSqliteDatabase from config import SRC_LOG_LEVELS
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["DB"])
db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
db_state = ContextVar("db_state", default=db_state_default.copy()) db_state = ContextVar("db_state", default=db_state_default.copy())
@ -22,11 +26,13 @@ class PeeweeConnectionState(object):
def register_connection(db_url): def register_connection(db_url):
db = connect(db_url) db = connect(db_url)
if isinstance(db, PostgresqlDatabase): if isinstance(db, PostgresqlDatabase):
# Directly use PostgresqlDatabase without pooling # Enable autoconnect for SQLite databases, managed by Peewee
db.autoconnect = True db.autoconnect = True
log.info("Connected to PostgreSQL database")
elif isinstance(db, SqliteDatabase): elif isinstance(db, SqliteDatabase):
# Directly use SqliteDatabase without pooling # Enable autoconnect for SQLite databases, managed by Peewee
db.autoconnect = True db.autoconnect = True
log.info("Connected to SQLite database")
else: else:
raise ValueError('Unsupported database connection') raise ValueError('Unsupported database connection')
return db return db