xian_algorithm_new/app/services/qgis/qgis_runner.py

#!/usr/bin/env python3
"""
QGIS 专题图生成子进程入口。

由主进程 (Python 3.10) 通过 subprocess 调用，
运行在 QGIS 自带的 Python 3.12 环境中。

支持两种模式：
  - 批量模式（推荐）：单次启动 QgsApplication，顺序处理多个模板
    输入: { "config": {...}, "models": [{...}, {...}, ...] }
  - 单任务模式（兼容）：
    输入: { "config": {...}, "model": {...} }

输出 JSON (stdout):
  批量: { "results": [{"name": "...", "output": "..."}, ...] }
  单任务: { "name": "...", "output": "..." }

错误: stderr + exit code 1
"""
import json
import os
import sys
import time

# ============================================================
# 1. 环境初始化（必须在任何 QGIS/Qt import 之前）
# ============================================================

QGIS_ROOT = os.environ.get("QGIS_ROOT", "D:/QGIS")


def _setup_python_path():
    """将项目根目录和 QGIS Python 路径加入 sys.path"""
    project_root = os.path.dirname(
        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_python = os.path.join(QGIS_ROOT, "apps", "qgis-ltr", "python")
    if os.path.isdir(qgis_python) and qgis_python not in sys.path:
        sys.path.insert(0, qgis_python)


def _setup_environment():
    """设置 QGIS 运行所需的环境变量"""
    # 自动检测 QGIS 应用目录
    qgis_app_dir = os.path.join(QGIS_ROOT, "apps", "qgis-ltr")
    if not os.path.isdir(qgis_app_dir):
        qgis_app_dir = os.path.join(QGIS_ROOT, "apps", "qgis")

    os.environ["QGIS_PREFIX_PATH"] = qgis_app_dir
    os.environ["PYTHONUTF8"] = "1"
    os.environ["GDAL_FILENAME_IS_UTF8"] = "YES"
    os.environ["VSI_CACHE"] = "TRUE"
    os.environ["VSI_CACHE_SIZE"] = "1000000"

    if sys.platform == "win32":
        os.environ["QT_PLUGIN_PATH"] = (
            f"{os.path.join(qgis_app_dir, 'qtplugins')};"
            f"{os.path.join(QGIS_ROOT, 'apps', 'Qt5', 'plugins')}"
        )
        gdal_data = os.path.join(QGIS_ROOT, "apps", "gdal", "share", "gdal")
        if os.path.isdir(gdal_data):
            os.environ["GDAL_DATA"] = gdal_data

        # 强制使用 QGIS 自带的 PROJ，避免 PostgreSQL/PostGIS 的旧 proj.db 干扰
        for proj_dir in [
            os.path.join(qgis_app_dir, "share", "proj"),
            os.path.join(QGIS_ROOT, "share", "proj"),
        ]:
            if os.path.isdir(proj_dir):
                os.environ["PROJ_DATA"] = proj_dir
                break

        import ctypes
        _dll_dirs = [
            os.path.join(qgis_app_dir, "bin"),
            os.path.join(QGIS_ROOT, "apps", "Qt5", "bin"),
            os.path.join(QGIS_ROOT, "apps", "gdal", "lib"),
        ]
        _preload_dlls = [
            "qgis_core.dll", "qgispython.dll",
            "Qt5Core.dll", "Qt5Gui.dll", "Qt5Widgets.dll",
            "Qt5Network.dll", "Qt5Svg.dll", "Qt5Xml.dll",
            "Qt5Concurrent.dll", "Qt5PrintSupport.dll",
        ]
        for dll_dir in _dll_dirs:
            if not os.path.isdir(dll_dir):
                continue
            os.add_dll_directory(dll_dir)
            for dll_name in _preload_dlls:
                dll_path = os.path.join(dll_dir, dll_name)
                if os.path.isfile(dll_path):
                    try:
                        ctypes.WinDLL(dll_path)
                    except OSError:
                        pass


# ============================================================
# 2. 主逻辑
# ============================================================

def _init_qgis():
    """初始化 QgsApplication（只做一次）"""
    from qgis.core import QgsApplication

    qgis_app_dir = os.path.join(QGIS_ROOT, "apps", "qgis-ltr")
    QgsApplication.setPrefixPath(qgis_app_dir, True)
    qgs_app = QgsApplication([], False)
    qgs_app.initQgis()
    return qgs_app


def _process_single(service, model):
    """处理单个模板，返回结果 dict"""
    name = service.generate(model)
    return {"name": name, "output": model["outFile"]}


def main():
    t_start = time.time()

    # 环境初始化
    _setup_environment()
    _setup_python_path()

    # 读取请求 JSON
    if len(sys.argv) > 1 and os.path.isfile(sys.argv[1]):
        with open(sys.argv[1], "r", encoding="utf-8") as f:
            request = json.load(f)
    else:
        request = json.load(sys.stdin)

    config = request["config"]

    # 兼容批量和单任务模式
    models = request.get("models") or [request["model"]]

    # 初始化 QgsApplication（只做一次）
    qgs_app = _init_qgis()

    try:
        from app.services.qgis.map_service import MapService

        service = MapService(config)
        results = []

        for i, model in enumerate(models):
            t_model = time.time()
            try:
                result = _process_single(service, model)
                results.append(result)
                elapsed = time.time() - t_model
                print(
                    f"[qgis_runner] [{i+1}/{len(models)}] 完成: {result['name']}, "
                    f"耗时 {elapsed:.1f}s",
                    file=sys.stderr,
                )
            except Exception as e:
                elapsed = time.time() - t_model
                error_msg = f"{e}"
                print(
                    f"[qgis_runner] [{i+1}/{len(models)}] 失败: {model.get('name', '?')}, "
                    f"耗时 {elapsed:.1f}s — {error_msg}",
                    file=sys.stderr,
                )
                results.append({"name": model.get("name", ""), "output": "", "error": error_msg})

        # 输出结果
        if len(models) == 1 and not request.get("models"):
            # 单任务模式兼容
            json.dump(results[0], sys.stdout, ensure_ascii=False)
        else:
            json.dump({"results": results}, sys.stdout, ensure_ascii=False)
        sys.stdout.flush()

        total = time.time() - t_start
        ok = sum(1 for r in results if not r.get("error"))
        fail = len(results) - ok
        print(
            f"\n[qgis_runner] 批量完成: {ok}成功/{fail}失败, 总耗时 {total:.1f}s",
            file=sys.stderr,
        )
    except Exception as e:
        elapsed = time.time() - t_start
        print(f"[qgis_runner] 致命错误 ({elapsed:.1f}s): {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        qgs_app.exitQgis()


if __name__ == "__main__":
    main()