refac

backend/open_webui/internal/wrappers.py

@@ -10,13 +10,13 @@ from playhouse.shortcuts import ReconnectMixin
 
 log = logging.getLogger(__name__)
 
-db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
-db_state = ContextVar("db_state", default=db_state_default.copy())
+db_state_default = {'closed': None, 'conn': None, 'ctx': None, 'transactions': None}
+db_state = ContextVar('db_state', default=db_state_default.copy())
 
 
 class PeeweeConnectionState(object):
     def __init__(self, **kwargs):
-        super().__setattr__("_state", db_state)
+        super().__setattr__('_state', db_state)
         super().__init__(**kwargs)
 
     def __setattr__(self, name, value):
@@ -30,10 +30,10 @@ class PeeweeConnectionState(object):
 class CustomReconnectMixin(ReconnectMixin):
     reconnect_errors = (
         # psycopg2
-        (OperationalError, "termin"),
-        (InterfaceError, "closed"),
+        (OperationalError, 'termin'),
+        (InterfaceError, 'closed'),
         # peewee
-        (PeeWeeInterfaceError, "closed"),
+        (PeeWeeInterfaceError, 'closed'),
     )
 
 
@@ -43,23 +43,21 @@ class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase):
 
 def register_connection(db_url):
     # 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"
-            )
+    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')
         from playhouse.sqlcipher_ext import SqlCipherDatabase
 
         # 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://", "")
+        db_path = db_url.replace('sqlite+sqlcipher://', '')
 
         # 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 encrypted SQLite database using SQLCipher")
+        log.info('Connected to encrypted SQLite database using SQLCipher')
 
     else:
         # Standard database connection (existing logic)
@@ -68,7 +66,7 @@ def register_connection(db_url):
             # Enable autoconnect for SQLite databases, managed by Peewee
             db.autoconnect = True
             db.reuse_if_open = True
-            log.info("Connected to PostgreSQL database")
+            log.info('Connected to PostgreSQL database')
 
             # Get the connection details
             connection = parse(db_url, unquote_user=True, unquote_password=True)
@@ -80,7 +78,7 @@ def register_connection(db_url):
             # Enable autoconnect for SQLite databases, managed by Peewee
             db.autoconnect = True
             db.reuse_if_open = True
-            log.info("Connected to SQLite database")
+            log.info('Connected to SQLite database')
         else:
-            raise ValueError("Unsupported database connection")
+            raise ValueError('Unsupported database connection')
     return db