context.py

from flask import current_app, has_app_context
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine

from app.config import config

db = SQLAlchemy()


def exchange_db(db_name: str):
    """
    切换数据库
    :param db:数据
    :return:
    """
    if has_app_context():
        app = current_app
        rel = app.config['SQLALCHEMY_BINDS'][db_name]
        app.config['SQLALCHEMY_DATABASE_URI'] = rel
        return db
    else:
        return create_engine_db(db_name)


def create_engine_db(site_name='us') -> Engine:
    if has_app_context():
        return db.engines.get(site_name)
    else:
        url = config['default'].SQLALCHEMY_BINDS.get(site_name)
        return create_engine(url)