SuperAGI
116 строк · 3.4 Кб
1from logging.config import fileConfig
2
3from sqlalchemy import engine_from_config
4from sqlalchemy import pool
5from alembic import context
6from urllib.parse import urlparse
7
8# this is the Alembic Config object, which provides
9# access to the values within the .ini file in use.
10config = context.config
11
12# Interpret the config file for Python logging.
13# This line sets up loggers basically.
14if config.config_file_name is not None:
15fileConfig(config.config_file_name)
16
17# add your model's MetaData object here
18# for 'autogenerate' support
19# from myapp import mymodel
20# target_metadata = mymodel.Base.metadata
21from superagi.models.base_model import DBBaseModel
22target_metadata = DBBaseModel.metadata
23from superagi.models import *
24from superagi.config.config import get_config
25
26# other values from the config, defined by the needs of env.py,
27# can be acquired:
28# my_important_option = config.get_main_option("my_important_option")
29# ... etc.
30
31db_host = get_config('DB_HOST', 'super__postgres')
32db_username = get_config('DB_USERNAME')
33db_password = get_config('DB_PASSWORD')
34db_name = get_config('DB_NAME')
35database_url = get_config('DB_URL', None)
36
37def run_migrations_offline() -> None:
38"""Run migrations in 'offline' mode.
39
40This configures the context with just a URL
41and not an Engine, though an Engine is acceptable
42here as well. By skipping the Engine creation
43we don't even need a DBAPI to be available.
44
45Calls to context.execute() here emit the given string to the
46script output.
47
48"""
49
50db_url = database_url
51if db_url is None:
52if db_username is None:
53db_url = f'postgresql://{db_host}/{db_name}'
54else:
55db_url = f'postgresql://{db_username}:{db_password}@{db_host}/{db_name}'
56else:
57db_url = urlparse(db_url)
58db_url = db_url.scheme + "://" + db_url.netloc + db_url.path
59
60config.set_main_option("sqlalchemy.url", db_url)
61
62url = config.get_main_option("sqlalchemy.url")
63context.configure(
64url=url,
65target_metadata=target_metadata,
66literal_binds=True,
67dialect_opts={"paramstyle": "named"},
68)
69
70with context.begin_transaction():
71context.run_migrations()
72
73
74def run_migrations_online() -> None:
75"""Run migrations in 'online' mode.
76
77In this scenario we need to create an Engine
78and associate a connection with the context.
79
80"""
81
82db_host = get_config('DB_HOST', 'super__postgres')
83db_username = get_config('DB_USERNAME')
84db_password = get_config('DB_PASSWORD')
85db_name = get_config('DB_NAME')
86db_url = get_config('DB_URL', None)
87
88if db_url is None:
89if db_username is None:
90db_url = f'postgresql://{db_host}/{db_name}'
91else:
92db_url = f'postgresql://{db_username}:{db_password}@{db_host}/{db_name}'
93else:
94db_url = urlparse(db_url)
95db_url = db_url.scheme + "://" + db_url.netloc + db_url.path
96
97config.set_main_option('sqlalchemy.url', db_url)
98connectable = engine_from_config(
99config.get_section(config.config_ini_section, {}),
100prefix="sqlalchemy.",
101poolclass=pool.NullPool,
102)
103
104with connectable.connect() as connection:
105context.configure(
106connection=connection, target_metadata=target_metadata
107)
108
109with context.begin_transaction():
110context.run_migrations()
111
112
113if context.is_offline_mode():
114run_migrations_offline()
115else:
116run_migrations_online()
117