db.py 738 B

12345678910111213141516171819202122232425
  1. from peewee import *
  2. from peewee_migrate import Router
  3. from playhouse.db_url import connect
  4. from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
  5. import os
  6. import logging
  7. log = logging.getLogger(__name__)
  8. log.setLevel(SRC_LOG_LEVELS["DB"])
  9. # Check if the file exists
  10. if os.path.exists(f"{DATA_DIR}/ollama.db"):
  11. # Rename the file
  12. os.rename(f"{DATA_DIR}/ollama.db", f"{DATA_DIR}/webui.db")
  13. log.info("Database migrated from Ollama-WebUI successfully.")
  14. else:
  15. pass
  16. DB = connect(DATABASE_URL)
  17. log.info(f"Connected to a {DB.__class__.__name__} database.")
  18. router = Router(
  19. DB, migrate_dir=BACKEND_DIR / "apps" / "web" / "internal" / "migrations", logger=log
  20. )
  21. router.run()
  22. DB.connect(reuse_if_open=True)