cube-studio/myapp/check_tables.py

#!/usr/bin/env python
import os


# @pysnooper.snoop()
def check_tables():
    SQLALCHEMY_DATABASE_URI = os.getenv('MYSQL_SERVICE', '')
    if SQLALCHEMY_DATABASE_URI:
        import sqlalchemy.engine.url as url
        uri = url.make_url(SQLALCHEMY_DATABASE_URI)
        """Inits the Myapp application"""
        import pymysql
        # 创建连接
        conn = pymysql.connect(host=uri.host, port=uri.port, user=uri.username, password=uri.password, charset='utf8')
        # 创建游标
        cursor = conn.cursor()

        # 创建数据库的sql(如果数据库存在就不创建，防止异常)
        sql = f"SELECT table_name FROM information_schema.tables  WHERE table_schema='{uri.database}'"
        cursor.execute(sql)
        results = list(cursor.fetchall())
        results = [item[0] for item in results]
        print(results)
        for table_name in ['ab_permission', 'ab_permission_view', 'ab_permission_view_role', 'ab_register_user',
                           'ab_role', 'ab_user', 'ab_user_role', 'ab_view_menu', 'alembic_version', 'dimension',
                           'docker', 'images', 'inferenceservice', 'job_template', 'logs', 'metadata_metric', 'model',
                           'nni', 'notebook', 'pipeline', 'project', 'project_user', 'repository', 'run',
                           'service', 'task', 'metadata_table', 'workflow']:
            if table_name not in results:
                print('%s db下，table %s不完整，请\n1、mysql: drop dabatase %s\n2、重启当前pod' % (uri.database,table_name,uri.database))
                exit(1)


check_tables()