AcidMap/scripts/demos/reset_db.py

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

"""
数据库重置脚本

这个脚本用于清理数据库并准备新的迁移
"""

import os
import sys
from sqlalchemy import text, inspect, create_engine
import re
import logging
from sqlalchemy.ext.declarative import declarative_base

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 加载环境变量配置
from dotenv import load_dotenv
load_dotenv("config.env")

# 从环境变量获取数据库连接信息，如果没有则使用默认值
DB_USER = os.getenv("DB_USER", "postgres")
DB_PASSWORD = os.getenv("DB_PASSWORD", "123456789Qq")
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_NAME = os.getenv("DB_NAME", "soilgd")
print(f"连接到数据库: {DB_USER}@{DB_HOST}:{DB_PORT}/{DB_NAME}")

# 构建数据库连接URL
SQLALCHEMY_DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# 创建数据库引擎
engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    pool_size=5,
    max_overflow=10,
    pool_timeout=30,
    pool_recycle=1800
)

# 要删除的表列表（按顺序）
TABLES_TO_DROP = [
    'fifty_thousand_survey_data',
    'surveydata',
    'unit_ceil',
    'raster_table',
    'Conventional_land_information',
    'Point_information',
    'alembic_version'
]

# 不应该删除的系统表
SYSTEM_TABLES = [
    'spatial_ref_sys',  # PostGIS系统表
    'geography_columns',
    'geometry_columns',
    'raster_columns',
    'raster_overviews'
]

def get_all_tables():
    """获取数据库中所有表的名称（排除系统表）"""
    inspector = inspect(engine)
    tables = inspector.get_table_names()
    # 排除系统表
    return [t for t in tables if t not in SYSTEM_TABLES]

def get_table_case_insensitive(table_name):
    """不区分大小写查找表名"""
    all_tables = get_all_tables()
    return [t for t in all_tables if t.lower() == table_name.lower()]

def reset_database():
    """重置数据库，删除所有表"""
    print("\n==== 开始重置数据库 ====")
    
    # 获取删除前的所有表
    before_tables = get_all_tables()
    print(f"\n删除前的表: {before_tables}")
    
    try:
        for table_name in TABLES_TO_DROP:
            # 查找该表在数据库中的实际名称（考虑大小写）
            actual_tables = get_table_case_insensitive(table_name)
            
            if not actual_tables:
                print(f"表 '{table_name}' 不存在，跳过")
                continue
            
            for actual_table in actual_tables:
                try:
                    with engine.begin() as conn:
                        print(f"\n尝试删除表 '{actual_table}'...")
                        # 使用双引号包裹表名以保持大小写
                        sql = f'DROP TABLE IF EXISTS "{actual_table}" CASCADE'
                        conn.execute(text(sql))
                        print(f"执行SQL: {sql}")
                        
                        # 验证表是否已删除
                        check_tables = get_table_case_insensitive(actual_table)
                        if not check_tables:
                            print(f"✓ 表 '{actual_table}' 已成功删除")
                        else:
                            print(f"✗ 表 '{actual_table}' 删除失败，仍然存在")
                except Exception as e:
                    print(f"删除表 '{actual_table}' 时出错: {str(e)}")
        
        # 获取删除后的所有表
        after_tables = get_all_tables()
        print(f"\n==== 删除操作完成 ====")
        print(f"删除前的表: {len(before_tables)}个 - {before_tables}")
        print(f"删除后的表: {len(after_tables)}个 - {after_tables}")
        
        # 找出仍然存在的表
        remaining_tables = []
        for table in before_tables:
            # 检查表是否应该被删除
            should_be_deleted = any(re.match(f"^{t}$", table, re.IGNORECASE) for t in TABLES_TO_DROP)
            if should_be_deleted and table in after_tables:
                remaining_tables.append(table)
        
        if remaining_tables:
            print(f"\n警告：以下表应该被删除但仍然存在: {remaining_tables}")
            print("这可能是由于权限问题或者表依赖关系导致")
        else:
            print("\n所有目标表已成功删除")
        
        print("\n==== 数据库重置完成 ====")
        return 0
    except Exception as e:
        print(f"\n重置数据库时出错: {str(e)}")
        return 1

if __name__ == '__main__':
    sys.exit(reset_database())