env.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import os
  2. from logging.config import fileConfig
  3. from sqlalchemy import engine_from_config
  4. from sqlalchemy import pool
  5. from alembic import context
  6. from apps.webui.models.auths import Auth
  7. from apps.webui.models.chats import Chat
  8. from apps.webui.models.documents import Document
  9. from apps.webui.models.memories import Memory
  10. from apps.webui.models.models import Model
  11. from apps.webui.models.prompts import Prompt
  12. from apps.webui.models.tags import Tag, ChatIdTag
  13. from apps.webui.models.tools import Tool
  14. from apps.webui.models.users import User
  15. from apps.webui.models.files import File
  16. from apps.webui.models.functions import Function
  17. # this is the Alembic Config object, which provides
  18. # access to the values within the .ini file in use.
  19. config = context.config
  20. # Interpret the config file for Python logging.
  21. # This line sets up loggers basically.
  22. if config.config_file_name is not None:
  23. fileConfig(config.config_file_name)
  24. # add your model's MetaData object here
  25. # for 'autogenerate' support
  26. # from myapp import mymodel
  27. # target_metadata = mymodel.Base.metadata
  28. target_metadata = Auth.metadata
  29. # other values from the config, defined by the needs of env.py,
  30. # can be acquired:
  31. # my_important_option = config.get_main_option("my_important_option")
  32. # ... etc.
  33. database_url = os.getenv("DATABASE_URL", None)
  34. if database_url:
  35. config.set_main_option("sqlalchemy.url", database_url)
  36. def run_migrations_offline() -> None:
  37. """Run migrations in 'offline' mode.
  38. This configures the context with just a URL
  39. and not an Engine, though an Engine is acceptable
  40. here as well. By skipping the Engine creation
  41. we don't even need a DBAPI to be available.
  42. Calls to context.execute() here emit the given string to the
  43. script output.
  44. """
  45. url = config.get_main_option("sqlalchemy.url")
  46. context.configure(
  47. url=url,
  48. target_metadata=target_metadata,
  49. literal_binds=True,
  50. dialect_opts={"paramstyle": "named"},
  51. )
  52. with context.begin_transaction():
  53. context.run_migrations()
  54. def run_migrations_online() -> None:
  55. """Run migrations in 'online' mode.
  56. In this scenario we need to create an Engine
  57. and associate a connection with the context.
  58. """
  59. connectable = engine_from_config(
  60. config.get_section(config.config_ini_section, {}),
  61. prefix="sqlalchemy.",
  62. poolclass=pool.NullPool,
  63. )
  64. with connectable.connect() as connection:
  65. context.configure(connection=connection, target_metadata=target_metadata)
  66. with context.begin_transaction():
  67. context.run_migrations()
  68. if context.is_offline_mode():
  69. run_migrations_offline()
  70. else:
  71. run_migrations_online()