|
@@ -1,4 +1,5 @@
|
|
|
import logging
|
|
|
+import os
|
|
|
from contextvars import ContextVar
|
|
|
|
|
|
from open_webui.env import SRC_LOG_LEVELS
|
|
@@ -7,6 +8,7 @@ from peewee import InterfaceError as PeeWeeInterfaceError
|
|
|
from peewee import PostgresqlDatabase
|
|
|
from playhouse.db_url import connect, parse
|
|
|
from playhouse.shortcuts import ReconnectMixin
|
|
|
+from playhouse.sqlcipher_ext import SqlCipherDatabase
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
log.setLevel(SRC_LOG_LEVELS["DB"])
|
|
@@ -43,24 +45,44 @@ class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase):
|
|
|
|
|
|
|
|
|
def register_connection(db_url):
|
|
|
- db = connect(db_url, unquote_user=True, unquote_password=True)
|
|
|
- if isinstance(db, PostgresqlDatabase):
|
|
|
- # Enable autoconnect for SQLite databases, managed by Peewee
|
|
|
+ # Check if using SQLCipher protocol
|
|
|
+ if db_url.startswith('sqlite+sqlcipher://'):
|
|
|
+ database_password = os.environ.get("DATABASE_PASSWORD")
|
|
|
+ if not database_password or database_password.strip() == "":
|
|
|
+ raise ValueError("DATABASE_PASSWORD is required when using sqlite+sqlcipher:// URLs")
|
|
|
+
|
|
|
+ # Parse the database path from SQLCipher URL
|
|
|
+ # Convert sqlite+sqlcipher:///path/to/db.sqlite to /path/to/db.sqlite
|
|
|
+ db_path = db_url.replace('sqlite+sqlcipher://', '')
|
|
|
+ if db_path.startswith('/'):
|
|
|
+ db_path = db_path[1:] # Remove leading slash for relative paths
|
|
|
+
|
|
|
+ # Use Peewee's native SqlCipherDatabase with encryption
|
|
|
+ db = SqlCipherDatabase(db_path, passphrase=database_password)
|
|
|
db.autoconnect = True
|
|
|
db.reuse_if_open = True
|
|
|
- log.info("Connected to PostgreSQL database")
|
|
|
+ log.info("Connected to encrypted SQLite database using SQLCipher")
|
|
|
+
|
|
|
+ else:
|
|
|
+ # Standard database connection (existing logic)
|
|
|
+ db = connect(db_url, unquote_user=True, unquote_password=True)
|
|
|
+ if isinstance(db, PostgresqlDatabase):
|
|
|
+ # Enable autoconnect for SQLite databases, managed by Peewee
|
|
|
+ db.autoconnect = True
|
|
|
+ db.reuse_if_open = True
|
|
|
+ log.info("Connected to PostgreSQL database")
|
|
|
|
|
|
- # Get the connection details
|
|
|
- connection = parse(db_url, unquote_user=True, unquote_password=True)
|
|
|
+ # Get the connection details
|
|
|
+ connection = parse(db_url, unquote_user=True, unquote_password=True)
|
|
|
|
|
|
- # Use our custom database class that supports reconnection
|
|
|
- db = ReconnectingPostgresqlDatabase(**connection)
|
|
|
- db.connect(reuse_if_open=True)
|
|
|
- elif isinstance(db, SqliteDatabase):
|
|
|
- # Enable autoconnect for SQLite databases, managed by Peewee
|
|
|
- db.autoconnect = True
|
|
|
- db.reuse_if_open = True
|
|
|
- log.info("Connected to SQLite database")
|
|
|
- else:
|
|
|
- raise ValueError("Unsupported database connection")
|
|
|
+ # Use our custom database class that supports reconnection
|
|
|
+ db = ReconnectingPostgresqlDatabase(**connection)
|
|
|
+ db.connect(reuse_if_open=True)
|
|
|
+ elif isinstance(db, SqliteDatabase):
|
|
|
+ # Enable autoconnect for SQLite databases, managed by Peewee
|
|
|
+ db.autoconnect = True
|
|
|
+ db.reuse_if_open = True
|
|
|
+ log.info("Connected to SQLite database")
|
|
|
+ else:
|
|
|
+ raise ValueError("Unsupported database connection")
|
|
|
return db
|