AcidMap/app/scripts/import_counties.py

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

"""
导入县级GeoJSON数据脚本
"""

import json
import logging
from pathlib import Path
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.county import County

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

def import_counties_from_geojson(file_path: str, db: Session):
    """从GeoJSON文件导入县级数据
    
    Args:
        file_path: GeoJSON文件路径
        db: 数据库会话
    """
    try:
        # 读取GeoJSON文件
        with open(file_path, 'r', encoding='utf-8') as f:
            geojson_data = json.load(f)
        
        # 检查是否为FeatureCollection
        if geojson_data['type'] != 'FeatureCollection':
            raise ValueError("输入文件必须是GeoJSON FeatureCollection格式")
        
        # 导入每个Feature
        for feature in geojson_data['features']:
            try:
                # 使用模型类方法创建County实例
                county = County.from_geojson_feature(feature)
                db.add(county)
                logger.info(f"成功导入: {county.name}")
            except Exception as e:
                logger.error(f"导入失败: {feature.get('properties', {}).get('县名', '未知')} - {str(e)}")
                continue
        
        # 提交事务
        db.commit()
        logger.info("数据导入完成")
        
    except Exception as e:
        db.rollback()
        logger.error(f"导入过程出错: {str(e)}")
        raise

def main():
    """主函数"""
    # 获取数据文件路径
    data_dir = Path(__file__).parent.parent.parent / 'data'
    geojson_file = data_dir / 'counties.geojson'
    
    if not geojson_file.exists():
        logger.error(f"找不到数据文件: {geojson_file}")
        return
    
    # 创建数据库会话
    db = SessionLocal()
    try:
        import_counties_from_geojson(str(geojson_file), db)
    finally:
        db.close()

if __name__ == '__main__':
    main()