Merge pull request #3595 from open-webui/dev-migration

feat: db migration
open-webui · Jul 6, 2024 · 8f6f766 · 8f6f766
2 parents 97a8491 4e75150
commit 8f6f766
Show file tree

Hide file tree

Showing 44 changed files with 2,720 additions and 957 deletions.
diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml
@@ -67,6 67,28 @@ jobs:
  path: compose-logs.txt
  if-no-files-found: ignore
 
  pytest:
  name: Run backend tests
  runs-on: ubuntu-latest
  steps:
  - uses: actions/checkout@v4
 
  - name: Set up Python
  uses: actions/setup-python@v4
  with:
  python-version: ${{ matrix.python-version }}
 
  - name: Install dependencies
  run: |
  python -m pip install --upgrade pip
  pip install -r backend/requirements.txt
 
  - name: pytest run
  run: |
  ls -al
  cd backend
  PYTHONPATH=. pytest . -o log_cli=true -o log_cli_level=INFO
 
  migration_test:
  name: Run Migration Tests
  runs-on: ubuntu-latest
@@ -171,7 193,7 @@ jobs:
  fi
 
  # Check that service will reconnect to postgres when connection will be closed
- status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health)
  status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db)
  if [[ "$status_code" -ne 200 ]] ; then
  echo "Server has failed before postgres reconnect check"
  exit 1
@@ -183,7 205,7 @@ jobs:
  cur = conn.cursor(); \
  cur.execute('SELECT pg_terminate_backend(psa.pid) FROM pg_stat_activity psa WHERE datname = current_database() AND pid <> pg_backend_pid();')"
 
- status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health)
  status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db)
  if [[ "$status_code" -ne 200 ]] ; then
  echo "Server has not reconnected to postgres after connection was closed: returned status $status_code"
  exit 1

diff --git a/backend/alembic.ini b/backend/alembic.ini
@@ -0,0 1,114 @@
 # A generic, single database configuration.
 
 [alembic]
 # path to migration scripts
 script_location = migrations
 
 # template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
 # Uncomment the line below if you want the files to be prepended with date and time
 # file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
 
 # sys.path path, will be prepended to sys.path if present.
 # defaults to the current working directory.
 prepend_sys_path = .
 
 # timezone to use when rendering the date within the migration file
 # as well as the filename.
 # If specified, requires the python>=3.9 or backports.zoneinfo library.
 # Any required deps can installed by adding `alembic[tz]` to the pip requirements
 # string value is passed to ZoneInfo()
 # leave blank for localtime
 # timezone =
 
 # max length of characters to apply to the
 # "slug" field
 # truncate_slug_length = 40
 
 # set to 'true' to run the environment during
 # the 'revision' command, regardless of autogenerate
 # revision_environment = false
 
 # set to 'true' to allow .pyc and .pyo files without
 # a source .py file to be detected as revisions in the
 # versions/ directory
 # sourceless = false
 
 # version location specification; This defaults
 # to migrations/versions. When using multiple version
 # directories, initial revisions must be specified with --version-path.
 # The path separator used here should be the separator specified by "version_path_separator" below.
 # version_locations = %(here)s/bar:%(here)s/bat:migrations/versions
 
 # version path separator; As mentioned above, this is the character used to split
 # version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
 # If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
 # Valid values for version_path_separator are:
 #
 # version_path_separator = :
 # version_path_separator = ;
 # version_path_separator = space
 version_path_separator = os # Use os.pathsep. Default configuration used for new projects.
 
 # set to 'true' to search source files recursively
 # in each "version_locations" directory
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = utf-8
 
 # sqlalchemy.url = REPLACE_WITH_DATABASE_URL
 
 
 [post_write_hooks]
 # post_write_hooks defines scripts or Python functions that are run
 # on newly generated revision scripts. See the documentation for further
 # detail and examples
 
 # format using "black" - use the console_scripts runner, against the "black" entrypoint
 # hooks = black
 # black.type = console_scripts
 # black.entrypoint = black
 # black.options = -l 79 REVISION_SCRIPT_FILENAME
 
 # lint with attempts to fix using "ruff" - use the exec runner, execute a binary
 # hooks = ruff
 # ruff.type = exec
 # ruff.executable = %(here)s/.venv/bin/ruff
 # ruff.options = --fix REVISION_SCRIPT_FILENAME
 
 # Logging configuration
 [loggers]
 keys = root,sqlalchemy,alembic
 
 [handlers]
 keys = console
 
 [formatters]
 keys = generic
 
 [logger_root]
 level = WARN
 handlers = console
 qualname =
 
 [logger_sqlalchemy]
 level = WARN
 handlers =
 qualname = sqlalchemy.engine
 
 [logger_alembic]
 level = INFO
 handlers =
 qualname = alembic
 
 [handler_console]
 class = StreamHandler
 args = (sys.stderr,)
 level = NOTSET
 formatter = generic
 
 [formatter_generic]
 format = %(levelname)-5.5s [%(name)s] %(message)s
 datefmt = %H:%M:%S
diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py
@@ -1,18 1,39 @@
 import os
 import logging
 import json
 from contextlib import contextmanager
 
-from peewee import *
 from peewee_migrate import Router
-
 from apps.webui.internal.wrappers import register_connection
 
 from typing import Optional, Any
 from typing_extensions import Self
 
 from sqlalchemy import create_engine, types, Dialect
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, scoped_session
 from sqlalchemy.sql.type_api import _T
 
 from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["DB"])
 
 
-class JSONField(TextField):
 class JSONField(types.TypeDecorator):
  impl = types.Text
  cache_ok = True
 
  def process_bind_param(self, value: Optional[_T], dialect: Dialect) -> Any:
  return json.dumps(value)
 
  def process_result_value(self, value: Optional[_T], dialect: Dialect) -> Any:
  if value is not None:
  return json.loads(value)
 
  def copy(self, **kw: Any) -> Self:
  return JSONField(self.impl.length)
 
  def db_value(self, value):
  return json.dumps(value)
 
@@ -30,25 51,57 @@ def python_value(self, value):
  pass
 
 
-# The `register_connection` function encapsulates the logic for setting up
-# the database connection based on the connection string, while `connect`
-# is a Peewee-specific method to manage the connection state and avoid errors
-# when a connection is already open.
-try:
- DB = register_connection(DATABASE_URL)
- log.info(f"Connected to a {DB.__class__.__name__} database.")
-except Exception as e:
- log.error(f"Failed to initialize the database connection: {e}")
- raise
-
-router = Router(
- DB,
- migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations",
- logger=log,
 # Workaround to handle the peewee migration
 # This is required to ensure the peewee migration is handled before the alembic migration
 def handle_peewee_migration():
  try:
  db = register_connection(DATABASE_URL)
  migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations"
  router = Router(db, logger=log, migrate_dir=migrate_dir)
  router.run()
  db.close()
 
  # check if db connection has been closed
 
  except Exception as e:
  log.error(f"Failed to initialize the database connection: {e}")
  raise
 
  finally:
  # Properly closing the database connection
  if db and not db.is_closed():
  db.close()
 
  # Assert if db connection has been closed
  assert db.is_closed(), "Database connection is still open."
 
 
 handle_peewee_migration()
 
 
 SQLALCHEMY_DATABASE_URL = DATABASE_URL
 if "sqlite" in SQLALCHEMY_DATABASE_URL:
  engine = create_engine(
  SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
  )
 else:
  engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
 
 
 SessionLocal = sessionmaker(
  autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
 )
-router.run()
-try:
- DB.connect(reuse_if_open=True)
-except OperationalError as e:
- log.info(f"Failed to connect to database again due to: {e}")
- pass
 Base = declarative_base()
 Session = scoped_session(SessionLocal)
 
 
 # Dependency
 def get_session():
  db = SessionLocal()
  try:
  yield db
  finally:
  db.close()
 
 
 get_db = contextmanager(get_session)
diff --git a/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py b/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py
@@ -1,10 1,7 @@
 """Peewee migrations -- 017_add_user_oauth_sub.py.
-
 Some examples (model - class or model name)::
-
  > Model = migrator.orm['table_name'] # Return model in current state by name
  > Model = migrator.ModelClass # Return model in current state by name
-
  > migrator.sql(sql) # Run custom SQL
  > migrator.run(func, *args, **kwargs) # Run python function with the given args
  > migrator.create_model(Model) # Create a model (could be used as decorator)
@@ -21,7 18,6 @@
  > migrator.drop_index(model, *col_names)
  > migrator.drop_not_null(model, *field_names)
  > migrator.drop_constraints(model, *constraints)
-
 """
 
 from contextlib import suppress

diff --git a/backend/apps/webui/internal/migrations/README.md b/backend/apps/webui/internal/migrations/README.md
diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py
@@ -3,7 3,7 @@
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.middleware.sessions import SessionMiddleware
-
 from sqlalchemy.orm import Session
 from apps.webui.routers import (
  auths,
  users,