Bläddra i källkod

Merge pull request #8510 from MadsLang/main

Add DATABASE_SCHEMA as env var
Timothy Jaeryang Baek 3 månader sedan
förälder
incheckning
d83acb2f03
2 ändrade filer med 6 tillägg och 2 borttagningar
  1. 2 0
      backend/open_webui/env.py
  2. 4 2
      backend/open_webui/internal/db.py

+ 2 - 0
backend/open_webui/env.py

@@ -274,6 +274,8 @@ DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite:///{DATA_DIR}/webui.db")
 if "postgres://" in DATABASE_URL:
     DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://")
 
+DATABASE_SCHEMA = os.environ.get("DATABASE_SCHEMA", None)
+
 DATABASE_POOL_SIZE = os.environ.get("DATABASE_POOL_SIZE", 0)
 
 if DATABASE_POOL_SIZE == "":

+ 4 - 2
backend/open_webui/internal/db.py

@@ -7,6 +7,7 @@ from open_webui.internal.wrappers import register_connection
 from open_webui.env import (
     OPEN_WEBUI_DIR,
     DATABASE_URL,
+    DATABASE_SCHEMA,
     SRC_LOG_LEVELS,
     DATABASE_POOL_MAX_OVERFLOW,
     DATABASE_POOL_RECYCLE,
@@ -14,7 +15,7 @@ from open_webui.env import (
     DATABASE_POOL_TIMEOUT,
 )
 from peewee_migrate import Router
-from sqlalchemy import Dialect, create_engine, types
+from sqlalchemy import Dialect, create_engine, MetaData, types
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import scoped_session, sessionmaker
 from sqlalchemy.pool import QueuePool, NullPool
@@ -99,7 +100,8 @@ else:
 SessionLocal = sessionmaker(
     autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
 )
-Base = declarative_base()
+metadata_obj = MetaData(schema=DATABASE_SCHEMA)
+Base = declarative_base(metadata=metadata_obj)
 Session = scoped_session(SessionLocal)