xian_algorithm_new/app/services/qgis/qgis_runner.py

#!/usr/bin/env python3
"""
QGIS 专题图生成子进程入口 — Docker 容器内运行。

由主进程通过 docker exec 调用，运行在 QGIS Docker 容器内。

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

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

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

# ============================================================
# 环境初始化（Docker 容器内，QGIS 通过 apt 安装在 /usr）
# ============================================================

def _setup_python_path():
    """将项目根目录和 QGIS dist-packages 加入 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 包路径（从配置读取，避免硬编码）
    try:
        from config import settings
        pythonpath = getattr(settings, "QGIS_DOCKER_PYTHONPATH", None) or []
    except Exception:
        pythonpath = []

    for p in pythonpath:
        if os.path.isdir(p) and p not in sys.path:
            sys.path.insert(0, p)


def _setup_environment():
    """设置 QGIS 运行所需的环境变量"""
    os.environ["PYTHONUTF8"] = "1"
    os.environ["GDAL_FILENAME_IS_UTF8"] = "YES"
    os.environ["VSI_CACHE"] = "TRUE"
    os.environ["VSI_CACHE_SIZE"] = "1000000"


# ============================================================
# 主逻辑
# ============================================================

def _init_qgis():
    """初始化 QgsApplication（从配置读取 prefixPath，避免硬编码）"""
    from qgis.core import QgsApplication

    try:
        from config import settings
        prefix_path = getattr(settings, "QGIS_DOCKER_PREFIX_PATH", "/usr")
    except Exception:
        prefix_path = "/usr"

    QgsApplication.setPrefixPath(prefix_path, True)
    qgs_app = QgsApplication([], False)
    qgs_app.initQgis()
    return qgs_app


def _process_single(service, model):
    """处理单个模板，返回结果 dict。成功/失败均输出 PROGRESS 行。"""
    name = service.generate(model)
    result = {"name": name, "output": model["outFile"]}
    print(f"PROGRESS:{json.dumps(result, ensure_ascii=False)}", flush=True)
    return result


def _emit_progress(result: dict):
    """输出 PROGRESS 行（成功或失败均调用）"""
    print(f"PROGRESS:{json.dumps(result, ensure_ascii=False)}", flush=True)


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,
                )
                fail_result = {"name": model.get("name", ""), "output": "", "error": error_msg}
                _emit_progress(fail_result)
                results.append(fail_result)

        # 输出结果
        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()