xian_algorithm_new/app/services/qgis/qgis_worker.py

#!/usr/bin/env python3
"""
QGIS 长驻 Worker — Docker 容器内运行，常驻进程模式。

与 qgis_runner.py 的区别：
  - qgis_runner.py: 一次性进程，处理完一批任务后退出
  - qgis_worker.py: 长驻进程，初始化 QGIS 一次，循环从 stdin 读取任务处理

通信协议（stdin/stdout JSON Lines）：
  输入: 每行一个 JSON 任务  {"config": {...}, "models": [{...}, ...]}
  输出: 每张图一行 PROGRESS  PROGRESS:{"name": "...", "output": "..."}
        任务结束一行          JOB_DONE:{"total": N, "ok": M, "fail": K, "elapsed": T}
        致命错误              FATAL:{"error": "..."}

用法:
  docker exec -i qgis-server python3 /app/app/services/qgis/qgis_worker.py
"""
import json
import os
import sys
import time


# ============================================================
# 环境初始化（与 qgis_runner.py 保持一致）
# ============================================================

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)

    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"""
    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 _emit(line: str):
    """输出一行到 stdout"""
    print(line, flush=True)


def _emit_progress(result: dict):
    """输出 PROGRESS 行"""
    _emit(f"PROGRESS:{json.dumps(result, ensure_ascii=False)}")


def _process_job(service, job: dict) -> dict:
    """
    处理一个批次任务（多个模板），返回汇总结果。

    每张图输出一行 PROGRESS，最后返回汇总 dict。
    """
    models = job.get("models") or []
    if not models:
        return {"total": 0, "ok": 0, "fail": 0, "elapsed": 0}

    ok = 0
    fail = 0
    t_start = time.time()

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

    total_elapsed = time.time() - t_start
    summary = {"total": len(models), "ok": ok, "fail": fail, "elapsed": round(total_elapsed, 1)}
    print(
        f"[worker] 批次完成: {ok}成功/{fail}失败, 总耗时 {total_elapsed:.1f}s",
        file=sys.stderr,
    )
    return summary


# ============================================================
# 主循环
# ============================================================

def main():
    _setup_environment()
    _setup_python_path()

    print("[worker] 正在初始化 QGIS...", file=sys.stderr)
    t_init = time.time()
    qgs_app = _init_qgis()
    print(f"[worker] QGIS 初始化完成 ({time.time() - t_init:.1f}s)，等待任务...", file=sys.stderr)

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

        # 持有 MapService 实例（config 在第一个任务时设置）
        service = None

        for line_no, line in enumerate(sys.stdin, 1):
            line = line.strip()
            if not line:
                continue

            try:
                job = json.loads(line)
            except json.JSONDecodeError as e:
                print(f"[worker] 第{line_no}行 JSON 解析失败: {e}", file=sys.stderr)
                _emit(f'JOB_DONE:{{"error": "JSON解析失败: {e}"}}')
                continue

            config = job.get("config")
            if config:
                # 更新 config（每个任务可以带不同的 config）
                service = MapService(config)

            if service is None:
                _emit(f'JOB_DONE:{{"error": "未提供 config"}}')
                continue

            summary = _process_job(service, job)
            _emit(f"JOB_DONE:{json.dumps(summary, ensure_ascii=False)}")

            # 清空项目，恢复到初始状态，等待下一个任务
            from qgis.core import QgsProject
            QgsProject.instance().clear()
            print("[worker] 项目已清空，等待下一个任务", file=sys.stderr)

    except KeyboardInterrupt:
        print("[worker] 收到中断信号，退出", file=sys.stderr)
    except Exception as e:
        print(f"[worker] 致命错误: {e}", file=sys.stderr)
        _emit(f'FATAL:{{"error": "{e}"}}')
    finally:
        qgs_app.exitQgis()
        print("[worker] QGIS 已退出", file=sys.stderr)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        # 捕获 main() 未处理的异常，确保输出到 stderr
        import traceback
        print(f"[worker] 未捕获异常: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)