db.py 693 B

1234567891011121314151617181920212223
  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
  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(DB, migrate_dir="apps/web/internal/migrations", logger=log)
  19. router.run()
  20. DB.connect(reuse_if_open=True)