xian_algorithm_new/app/script/export_static_layers.py

"""
一次性脚本：从 PostgreSQL 导出静态底图为 GeoPackage 文件。
优先使用 ogr2ogr（QGIS 自带），回退到 geopandas。

运行一次即可，之后服务直接读本地 GPKG。

用法: python -m app.script.export_static_layers
"""
import os
import subprocess
import sys
import time

# 确保项目根目录在 sys.path 中
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# ============================================================
# 配置
# ============================================================

QGIS_ROOT = os.environ.get("QGIS_ROOT", "D:/QGIS")
ogr2ogr_path = os.path.join(QGIS_ROOT, "bin", "ogr2ogr.exe")

GPKG_DIR = os.path.join(project_root, "app", "data", "gpkg")

# 静态图层定义: {显示名: (schema.table, gpkg文件名)}
STATIC_LAYERS = [
    ("水库",               "qgis.rivers",                  "rivers.gpkg"),
    ("市州驻地",           "qgis.sx_capital",              "sx_capital.gpkg"),
    ("河流",               "qgis.river",                   "river.gpkg"),
    ("active_fault",       "qgis.active_fault",            "active_fault.gpkg"),
    ("陕西省",             "qgis.sx",                      "sx.gpkg"),
    ("乡镇驻地",           "qgis.sx_street",               "sx_street.gpkg"),
    ("区县驻地",           "qgis.sx_xa_county",            "sx_xa_county.gpkg"),
    ("县界",               "qgis.sx_xa_county_boundary",   "sx_xa_county_boundary.gpkg"),
    ("周边区县",           "qgis.sx_zb_county_boundary",   "sx_zb_county_boundary.gpkg"),
    ("周边市州",           "qgis.sx_zb_city",              "sx_zb_city.gpkg"),
    ("周边县区",           "qgis.sx_zb_county",            "sx_zb_county.gpkg"),
    ("traffic_expressway", "qgis.traffic_expressway",      "traffic_expressway.gpkg"),
    ("traffic_provincial", "qgis.traffic_provincial",      "traffic_provincial.gpkg"),
    ("traffic_railway",    "qgis.traffic_railway",         "traffic_railway.gpkg"),
    ("traffic_township",   "qgis.traffic_township",        "traffic_township.gpkg"),
    ("traffic_trunk_line", "qgis.traffic_trunk_line",      "traffic_trunk_line.gpkg"),
    ("积水点",             "qgis.hazard_hydrops",          "hazard_hydrops.gpkg"),
    ("排水口",             "qgis.lifeline_outfall",        "lifeline_outfall.gpkg"),
    ("供水管网",           "qgis.lifeline_pipe",           "lifeline_pipe.gpkg"),
    ("风险人口",           "qgis.risk_census_population",  "risk_census_population.gpkg"),
    ("西安乡镇",           "qgis.sx_xa_towns",             "sx_xa_towns.gpkg"),
]


# ============================================================
# 方法一: ogr2ogr（推荐，QGIS 自带）
# ============================================================

def _setup_gdal_env():
    """设置 GDAL/OGR 运行环境"""
    gdal_data = os.path.join(QGIS_ROOT, "apps", "gdal", "share", "gdal")
    gdal_lib = os.path.join(QGIS_ROOT, "apps", "gdal", "lib")
    gdal_bin = os.path.join(QGIS_ROOT, "apps", "gdal", "bin")

    if os.path.isdir(gdal_data):
        os.environ["GDAL_DATA"] = gdal_data
    os.environ["GDAL_FILENAME_IS_UTF8"] = "YES"

    paths_to_add = [p for p in [gdal_bin, gdal_lib] if os.path.isdir(p)]
    os.environ["PATH"] = ";".join(paths_to_add) + ";" + os.environ.get("PATH", "")


def _export_with_ogr2ogr(host, port, dbname, user, password, schema, table, gpkg_path):
    """用 ogr2ogr 导出单个图层"""
    conn = f"PG:host={host} port={port} dbname={dbname} user={user} password={password}"
    cmd = [
        ogr2ogr_path,
        "-f", "GPKG",
        gpkg_path,
        conn,
        "-sql", f'SELECT * FROM "{schema}"."{table}"',
        "-nln", table,
        "-overwrite",
        "-t_srs", "EPSG:4326",
    ]
    result = subprocess.run(cmd, capture_output=True, timeout=120)
    if result.returncode != 0:
        stderr = result.stderr.decode("utf-8", errors="replace").strip()
        raise RuntimeError(stderr[:300])
    return True


# ============================================================
# 方法二: geopandas（回退）
# ============================================================

def _export_with_geopandas(host, port, dbname, user, password, schema, table, gpkg_path):
    """用 geopandas 导出单个图层"""
    import geopandas as gpd
    from sqlalchemy import create_engine

    conn_str = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
    engine = create_engine(conn_str)

    gdf = gpd.read_postgis(
        f'SELECT * FROM "{schema}"."{table}"',
        engine,
        geom_col="Geometry",
    )
    if gdf.crs is None:
        gdf = gdf.set_crs(epsg=4326)

    gdf.to_file(gpkg_path, driver="GPKG")
    engine.dispose()
    return len(gdf)


# ============================================================
# 主入口
# ============================================================

def main():
    # 从 config.settings 读取数据库配置
    try:
        from config import settings
        host = getattr(settings, "DB_HOST", "47.92.216.173")
        port = str(getattr(settings, "DB_PORT", 7654))
        dbname = getattr(settings, "DB_NAME", "xian_new")
        user = getattr(settings, "DB_USER", "postgres")
        password = getattr(settings, "DB_PASSWORD", "zhangsan")
    except Exception:
        host, port, dbname, user, password = "47.92.216.173", "7654", "xian_new", "postgres", "zhangsan"

    os.makedirs(GPKG_DIR, exist_ok=True)

    # 选择导出方法
    use_ogr2ogr = os.path.isfile(ogr2ogr_path)
    method = "ogr2ogr" if use_ogr2ogr else "geopandas"

    if use_ogr2ogr:
        _setup_gdal_env()

    print(f"数据库: {host}:{port}/{dbname}")
    print(f"输出目录: {GPKG_DIR}")
    print(f"导出方法: {method}")
    print(f"共 {len(STATIC_LAYERS)} 个图层\n")

    success = 0
    failed = 0

    for name, table_ref, gpkg_file in STATIC_LAYERS:
        schema, table = table_ref.split(".", 1)
        gpkg_path = os.path.join(GPKG_DIR, gpkg_file)

        print(f"[{success + failed + 1}/{len(STATIC_LAYERS)}] {name} ({table_ref})", end=" ... ", flush=True)
        t0 = time.time()

        try:
            if use_ogr2ogr:
                _export_with_ogr2ogr(host, port, dbname, user, password, schema, table, gpkg_path)
                size_kb = os.path.getsize(gpkg_path) / 1024
                print(f"✓ {size_kb:.0f} KB, {time.time() - t0:.1f}s")
            else:
                count = _export_with_geopandas(host, port, dbname, user, password, schema, table, gpkg_path)
                print(f"✓ {count} 行, {time.time() - t0:.1f}s")
            success += 1
        except Exception as e:
            print(f"✗ 失败: {e}")
            failed += 1

    print(f"\n{'='*50}")
    print(f"完成: 成功={success}, 失败={failed}, 共={len(STATIC_LAYERS)}")
    print(f"输出目录: {GPKG_DIR}")


if __name__ == "__main__":
    main()