73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
from contextlib import contextmanager
|
|
from sqlalchemy.schema import CreateTable
|
|
|
|
from db.models import Base
|
|
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.dialects.mysql import insert
|
|
from sqlalchemy import create_engine, text, inspect
|
|
from env_provider import DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD, DB_SCHEMA_UPDATED, ENVIRONMENT
|
|
engine = create_engine(f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}")
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
|
|
@contextmanager
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
def clear_db():
|
|
with engine.connect() as conn:
|
|
conn.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
|
|
inspector = inspect(engine)
|
|
for table_name in inspector.get_table_names():
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {table_name}"))
|
|
conn.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))
|
|
|
|
|
|
def create_all():
|
|
with engine.begin() as conn:
|
|
Base.metadata.create_all(bind=conn)
|
|
|
|
def run_scripts():
|
|
from db.models import table_models
|
|
with get_db() as session:
|
|
for model in table_models:
|
|
if hasattr(model, "__scripts__"):
|
|
scripts = model.__scripts__
|
|
for script in scripts:
|
|
session.execute(script)
|
|
session.commit()
|
|
|
|
def init_payload():
|
|
from db.models import table_models
|
|
with get_db() as session:
|
|
session.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
|
|
for model in table_models:
|
|
print(str(CreateTable(model.__table__)))
|
|
|
|
print(f"MODEL -- {model}, {hasattr(model, '__payload__')}")
|
|
if hasattr(model, "__payload__"):
|
|
payload = model.__payload__[ENVIRONMENT]
|
|
print(f"- - [ - ] hasattr, {ENVIRONMENT} - {payload}")
|
|
stmt = insert(model.__table__).values(payload).prefix_with("IGNORE")
|
|
print(f"- - [ - ] {stmt}\n")
|
|
session.execute(stmt)
|
|
session.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))
|
|
session.commit()
|
|
|
|
|
|
def setup_db():
|
|
if DB_SCHEMA_UPDATED:
|
|
clear_db()
|
|
print("[ x ] db cleared")
|
|
create_all()
|
|
print("[ x ] db created")
|
|
run_scripts()
|
|
print("[ x ] db scripts executed")
|
|
init_payload()
|
|
print("[ x ] payload loaded")
|
|
|