xian_algorithm_new/app/services/qgis/template_modifier.py

"""
模板 XML 修改模块。在 project.read() 之前：
1. 替换 PostgreSQL 连接参数（host/port/dbname/schema）
2. 将静态底图的 datasource 替换为本地 GeoPackage 路径
"""
import os
import re
import zipfile
import tempfile

from app.config.paths import get_logger
from app.config.qgis_mappings import OGR_TO_POSTGRES, TABLE_RENAMES, SCHEMA_REPLACEMENTS

logger = get_logger("qgis.modifier")


class TemplateModifier:
    def __init__(self, config: dict):
        self.config = config
        self._static_map = self._build_static_map()

    def _build_static_map(self) -> dict[str, str]:
        """构建 table_key → gpkg_abs_path 的映射"""
        sl = self.config.get("static_layers")
        if not sl or not sl.get("enabled", False):
            return {}

        gpkg_dir = sl["gpkg_dir"]
        mapping = {}
        for info in sl["layers"].values():
            schema, table = info["table"].split(".", 1)
            xml_key = f'table="{schema}"."{table}"'
            gpkg_path = os.path.join(gpkg_dir, info["file"]).replace("\\", "/")
            mapping[xml_key] = gpkg_path
        return mapping

    @staticmethod
    def _fix_provider_tags(content: str) -> str:
        """
        将 GPKG 文件路径对应的 <provider ...>postgres</provider> 改为 ogr。

        策略：按 <maplayer> 块处理，若块内 datasource 是文件路径（盘符开头），
        则该块内的 provider 改为 ogr。避免跨层误改。
        """
        maplayer_re = re.compile(r'(<maplayer[^>]*>.*?</maplayer>)', re.DOTALL)
        provider_re = re.compile(r'(<provider[^>]*>)postgres(</provider>)')
        file_ds_re = re.compile(r'<datasource>([A-Za-z]:/[^<]+)</datasource>')

        def _fix_layer(m):
            layer_xml = m.group(1)
            if file_ds_re.search(layer_xml):
                layer_xml = provider_re.sub(r'\1ogr\2', layer_xml)
            return layer_xml

        return maplayer_re.sub(_fix_layer, content)

    def modify(self, template_path: str) -> str:
        """修改模板文件，返回修改后的临时 .qgz 路径"""
        override = self.config.get("template_override")
        has_override = override and override.get("enabled", False)
        has_static = bool(self._static_map)

        if not has_override and not has_static:
            return template_path

        orig = override["original"] if has_override else None
        actual = override["actual"] if has_override else None

        try:
            tmp = tempfile.NamedTemporaryFile(
                suffix=".qgz", delete=False,
                dir=tempfile.gettempdir(),
            )
            tmp_path = tmp.name
            tmp.close()

            datasource_re = re.compile(r"(<datasource>)(.*?)(</datasource>)", re.DOTALL)
            table_re = re.compile(r'table="(\w+)"\."(\w+)"')

            with zipfile.ZipFile(template_path, "r") as zin, \
                 zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as zout:

                for item in zin.infolist():
                    data = zin.read(item.filename)

                    if item.filename.endswith(".qgs"):
                        content = data.decode("utf-8")

                        # ★ 先替换静态底图 datasource（在 schema 替换之前）
                        # 否则 table="base"."xxx" 会被改为 table="qgis"."xxx"，导致匹配失败
                        if self._static_map:
                            def _replace(m):
                                tbl = table_re.search(m.group(2))
                                if tbl:
                                    key = f'table="{tbl.group(1)}"."{tbl.group(2)}"'
                                    if key in self._static_map:
                                        return f"{m.group(1)}{self._static_map[key]}{m.group(3)}"
                                return m.group(0)

                            content = datasource_re.sub(_replace, content)

                        # ★ 转换避难场所 OGR 图层为 PostgreSQL（模板引用本地 JSON 文件）
                        _target_schema = (actual or {}).get("schema") or "qgis"
                        _db = self.config.get("db", {})
                        content = self._convert_ogr_to_postgres(content, _db, _target_schema)

                        # 替换所有 PostgreSQL 连接参数（统一指向目标库）
                        db = self.config.get("db", {})

                        # host: 模板中 host=localhost 无引号，匹配带/不带引号
                        if db.get("host"):
                            content = re.sub(
                                r"(host\s*=\s*)(?:['\"])?(?:localhost|127\.0\.0\.1)(?:['\"])?(?=\s|$)",
                                rf"\g<1>{db['host']}",
                                content,
                            )
                        # port: 模板中 port=5432 无引号，匹配带/不带引号
                        if db.get("port"):
                            content = re.sub(
                                r"(port\s*=\s*)(?:['\"])?5432(?:['\"])?(?=\s|$)",
                                rf"\g<1>{str(db['port'])}",
                                content,
                            )
                        if db.get("database"):
                            content = re.sub(
                                r"(dbname\s*=\s*['\"])([^'\"]+)(['\"])",
                                rf"\g<1>{db['database']}\3",
                                content,
                            )
                        # 移除 authcfg（QGIS 本地认证配置，子进程环境无效）
                        content = re.sub(r'\s*authcfg\s*=\s*\S+', '', content)
                        # 移除可能残留的 user=/password=（防止重复注入）
                        if db.get("username"):
                            content = re.sub(r"\s*user\s*=\s*'[^']*'", '', content)
                        if db.get("password"):
                            content = re.sub(r"\s*password\s*=\s*'[^']*'", '', content)
                        # 在 port 后注入显式 user + password
                        if db.get("username") and db.get("password"):
                            content = re.sub(
                                r'(port\s*=\s*\d+)',
                                rf"\1 user='{db['username']}' password='{db['password']}'",
                                content,
                            )
                        # schema 替换：统一指向目标 schema
                        target_schema = (actual or {}).get("schema") or "qgis"
                        for old_schema in SCHEMA_REPLACEMENTS:
                            content = content.replace(
                                f'table="{old_schema}".',
                                f'table="{target_schema}".',
                            )

                        # 表名映射（模板表名 ≠ 目标库表名）
                        for old_name, new_name in TABLE_RENAMES.items():
                            content = content.replace(
                                f'table="{target_schema}"."{old_name}"',
                                f'table="{target_schema}"."{new_name}"',
                            )

                        # tid 列已通过 add_tid_column.py 添加到所有 qgis 动态表中
                        # 模板中的 key='tid' 可正常工作，无需移除

                        # 修正 srid=0 → srid=4326（QGIS 无法从 srid=0 自动检测 SRID）
                        content = content.replace(" srid=0 ", " srid=4326 ")

                        # ★ 修复 provider 标签：GPKG 文件路径的图层必须用 ogr provider
                        # 否则 QGIS 会将 .gpkg 路径当作 PostgreSQL 连接字符串解析，导致图层无效
                        if self._static_map:
                            content = self._fix_provider_tags(content)

                        data = content.encode("utf-8")

                    zout.writestr(item, data)

            logger.info(f"模板已修改并写入: {tmp_path}")
            return tmp_path

        except Exception as e:
            logger.error(f"模板修改失败，将使用原始模板: {e}")
            return template_path

    @staticmethod
    def _convert_ogr_to_postgres(content: str, db: dict, schema: str) -> str:
        """
        将模板中引用本地文件的避难场所 OGR 图层转换为 PostgreSQL 图层。

        设计师在本地用 JSON 文件制作了避难场所图层（公园/学校/文化馆等），
        这些数据已迁移到目标库 qgis schema 中。
        """
        _ogr_shelter_map = OGR_TO_POSTGRES

        maplayer_re = re.compile(r'(<maplayer[^>]*>.*?</maplayer>)', re.DOTALL)
        provider_re = re.compile(r'(<provider[^>]*>)ogr(</provider>)')
        datasource_re = re.compile(r'(<datasource>).*?(</datasource>)', re.DOTALL)
        file_ds_re = re.compile(r'<datasource>.*?\|layername=([^<]+)</datasource>', re.DOTALL)

        def _convert_layer(m):
            layer_xml = m.group(1)
            ds_match = file_ds_re.search(layer_xml)
            if not ds_match:
                return layer_xml
            layer_name = ds_match.group(1)
            table = _ogr_shelter_map.get(layer_name)
            if not table:
                return layer_xml

            # 构建 PostgreSQL URI
            pg_uri = (
                f"dbname='{db.get('database', 'xian_new')}' "
                f"host={db.get('host', '47.92.216.173')} "
                f"port={db.get('port', 7654)} "
                f"user='{db.get('username', 'postgres')}' "
                f"password='{db.get('password', '')}' "
                f"key='tid' srid=4326 type=Point "
                f"table=\"{schema}\".\"{table}\" (geometry)"
            )
            new_ds = f"<datasource>{pg_uri}</datasource>"

            layer_xml = datasource_re.sub(new_ds, layer_xml)
            layer_xml = provider_re.sub(r'\1postgres\2', layer_xml)
            logger.info(f"避难场所图层已转换: {layer_name} → qgis.{table}")
            return layer_xml

        return maplayer_re.sub(_convert_layer, content)