DataOps-platform/app/scripts/create_wechat_user_table.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
在Flask应用上下文中创建微信用户表
使用SQLAlchemy创建表结构
"""

import os
import sys
import logging

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from app import create_app, db
from app.core.data_parse.calendar import WechatUser

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def create_wechat_user_table():
    """
    在Flask应用上下文中创建微信用户表
    
    Returns:
        bool: 创建成功返回True，否则返回False
    """
    try:
        # 创建Flask应用
        app = create_app()
        
        with app.app_context():
            logger.info("开始创建微信用户表...")
            
            # 创建表结构
            db.create_all()
            
            # 检查表是否创建成功
            inspector = db.inspect(db.engine)
            tables = inspector.get_table_names(schema='public')
            
            if 'wechat_users' in tables:
                logger.info("微信用户表创建成功")
                
                # 显示表结构信息
                columns = inspector.get_columns('wechat_users', schema='public')
                logger.info("表结构信息:")
                for column in columns:
                    logger.info(f"  - {column['name']}: {column['type']}")
                
                return True
            else:
                logger.error("微信用户表创建失败")
                return False
                
    except Exception as e:
        logger.error(f"创建微信用户表时发生错误: {str(e)}")
        return False


def check_wechat_user_table():
    """
    检查微信用户表是否存在
    
    Returns:
        bool: 表存在返回True，否则返回False
    """
    try:
        # 创建Flask应用
        app = create_app()
        
        with app.app_context():
            inspector = db.inspect(db.engine)
            tables = inspector.get_table_names(schema='public')
            
            if 'wechat_users' in tables:
                logger.info("微信用户表已存在")
                return True
            else:
                logger.info("微信用户表不存在")
                return False
                
    except Exception as e:
        logger.error(f"检查微信用户表时发生错误: {str(e)}")
        return False


def drop_wechat_user_table():
    """
    删除微信用户表
    
    Returns:
        bool: 删除成功返回True，否则返回False
    """
    try:
        # 创建Flask应用
        app = create_app()
        
        with app.app_context():
            logger.info("开始删除微信用户表...")
            
            # 删除表
            with db.engine.connect() as connection:
                connection.execute(db.text("DROP TABLE IF EXISTS public.wechat_users CASCADE;"))
                connection.commit()
            
            logger.info("微信用户表删除成功")
            return True
                
    except Exception as e:
        logger.error(f"删除微信用户表时发生错误: {str(e)}")
        return False


def main():
    """
    主函数，根据命令行参数执行相应操作
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='微信用户表管理脚本')
    parser.add_argument('--action', choices=['create', 'check', 'drop'], default='create',
                        help='执行的操作：create（创建）、check（检查）或 drop（删除）')
    
    args = parser.parse_args()
    
    if args.action == 'create':
        logger.info("开始创建微信用户表...")
        success = create_wechat_user_table()
    elif args.action == 'check':
        logger.info("开始检查微信用户表...")
        success = check_wechat_user_table()
    elif args.action == 'drop':
        logger.info("开始删除微信用户表...")
        success = drop_wechat_user_table()
    else:
        logger.error("未知的操作类型")
        sys.exit(1)
    
    if success:
        logger.info("操作完成")
        sys.exit(0)
    else:
        logger.error("操作失败")
        sys.exit(1)


if __name__ == "__main__":
    main()