xian_algorithm_new/app/services/qgis/map_service.py

"""
地图生成主流程控制器。
模板加载 → 图层过滤 → 缩放 → 文本更新 → 导出。
"""
import os
import time

from qgis.core import QgsProject, QgsDataSourceUri

from app.config.paths import get_logger
from app.config.qgis_mappings import TABLE_RENAMES, SCHEMA_REPLACEMENTS
from .template_modifier import TemplateModifier
from .layer_filter import LayerFilter
from .map_exporter import MapExporter
from app.utils.map_zoom import MapZoom

logger = get_logger("qgis.service")


class MapService:
    def __init__(self, config: dict):
        self.config = config

    def generate(self, model: dict) -> str:
        _timing = {}
        t_total = time.time()

        template_path = model["path"]
        template_name = os.path.basename(template_path)
        project = QgsProject.instance()

        # ── 步骤 1：加载模板 ──
        t0 = time.time()
        project.clear()
        modifier = TemplateModifier(self.config)
        actual_path = modifier.modify(template_path)
        project.read(actual_path)
        if actual_path != template_path:
            try:
                os.remove(actual_path)
            except OSError:
                pass
        _timing["1.load"] = time.time() - t0

        # ── 步骤 2：图层连接修正 ──
        t0 = time.time()
        self._fix_invalid_layers(project)
        _timing["2.connections"] = time.time() - t0

        # ── 步骤 3：图层过滤 ──
        t0 = time.time()
        LayerFilter().apply(project, model)
        _timing["3.filter"] = time.time() - t0

        # ── 步骤 4：地图缩放 ──
        t0 = time.time()
        layout = project.layoutManager().layoutByName(model["mapLayout"])
        if layout is None:
            available = [l.name() for l in project.layoutManager().layouts()]
            raise RuntimeError(f"模板中未找到布局 '{model['mapLayout']}'，可用布局：{available}")
        map_item = layout.itemById("Map")
        zoom = MapZoom(project, layout, map_item)
        zoom.execute(model["zoomRule"], {
            "X": model["centerX"], "Y": model["centerY"], "value": model["zoomValue"],
        })
        _timing["4.zoom"] = time.time() - t0

        # ── 步骤 5：文本 + 比例尺 + 导出 ──
        t0 = time.time()
        exporter = MapExporter(self.config, layout)
        exporter.update_texts(model)
        exporter.update_scale_bar()
        exporter.export(model["outFile"])
        _timing["5.export"] = time.time() - t0

        total = time.time() - t_total
        steps = ", ".join(f"{k}={v:.1f}s" for k, v in _timing.items())
        logger.info(f"{template_name} → {total:.1f}s ({steps})")
        return model["name"]

    def _fix_invalid_layers(self, project: QgsProject) -> None:
        """修复 TemplateModifier 处理后仍无效的 PostgreSQL 图层。

        批量查询 DB 确认表存在性，只对存在的表执行 setDataSource 修复。
        """
        t0 = time.time()
        db_config = self.config["db"]
        actual_schema = "qgis"

        invalid_pg = []
        valid_count = 0
        total = 0
        for layer in project.mapLayers().values():
            if layer.providerType() != "postgres":
                continue
            total += 1
            if layer.isValid():
                valid_count += 1
                continue
            invalid_pg.append(layer)

        if not invalid_pg:
            elapsed = time.time() - t0
            logger.info(f"  图层修正: 总计{total}个PG层, 全部有效, 耗时{elapsed:.3f}s")
            return

        # 批量 DB 查询确认表存在性
        existing = self._batch_check_tables(db_config, actual_schema, invalid_pg)

        fixed = 0
        skipped_missing = 0
        failed = 0
        for layer in invalid_pg:
            uri = layer.dataProvider().uri()
            key = (uri.schema() or actual_schema, uri.table())
            if key not in existing:
                logger.debug(f"  跳过不存在的表: {key[0]}.{key[1]} (图层: {layer.name()})")
                skipped_missing += 1
                continue
            try:
                self._fix_postgres_layer(layer, db_config, actual_schema)
                if layer.isValid():
                    fixed += 1
                else:
                    failed += 1
            except Exception as e:
                logger.error(f"  修复图层 {layer.name()} 失败: {e}")
                failed += 1

        elapsed = time.time() - t0
        if total:
            logger.info(
                f"  图层修正: 总计{total}个PG层, 有效{valid_count}, "
                f"跳过{skipped_missing}(表不存在), "
                f"修复{fixed}, 仍失败{failed}, 耗时{elapsed:.3f}s"
            )

    @staticmethod
    def _batch_check_tables(db_config, schema, layers):
        """批量检查表是否存在（单次 DB 查询），返回存在的表集合"""
        tables = set()
        for layer in layers:
            uri = layer.dataProvider().uri()
            tables.add((uri.schema() or schema, uri.table()))

        if not tables:
            return set()

        try:
            import psycopg2
            conn = psycopg2.connect(
                host=db_config["host"],
                port=int(db_config["port"]),
                dbname=db_config["database"],
                user=db_config["username"],
                password=db_config["password"],
                connect_timeout=3,
            )
            cur = conn.cursor()
            conditions = " OR ".join(
                "(table_schema = %s AND table_name = %s)"
                for _ in tables
            )
            params = []
            for s, t in tables:
                params.extend([s, t])
            cur.execute(
                f"SELECT table_schema, table_name FROM information_schema.tables "
                f"WHERE {conditions}",
                params,
            )
            existing = {(row[0], row[1]) for row in cur.fetchall()}
            cur.close()
            conn.close()
            return existing
        except Exception as e:
            logger.warning(f"  表存在性检查失败，将尝试修复所有图层: {e}")
            return set(tables)

    @staticmethod
    def _fix_postgres_layer(layer, db_config, actual_schema):
        """修正单个 PostgreSQL 图层的连接参数"""
        try:
            uri = layer.dataProvider().uri()
            uri.setConnection(
                db_config["host"], str(db_config["port"]),
                db_config["database"], db_config["username"], db_config["password"],
            )
            uri_str = uri.uri()
            for old_schema in SCHEMA_REPLACEMENTS:
                if f'table="{old_schema}".' in uri_str:
                    uri_str = uri_str.replace(f'table="{old_schema}".', f'table="{actual_schema}".')
                    uri = QgsDataSourceUri(uri_str)
                    break
            uri_str = uri.uri()
            for old_name, new_name in TABLE_RENAMES.items():
                full_old = f'table="{actual_schema}"."{old_name}"'
                full_new = f'table="{actual_schema}"."{new_name}"'
                if full_old in uri_str:
                    uri_str = uri_str.replace(full_old, full_new)
                    uri = QgsDataSourceUri(uri_str)
            uri_str = uri.uri()
            if " srid=0 " in uri_str:
                uri_str = uri_str.replace(" srid=0 ", " srid=4326 ")
                uri = QgsDataSourceUri(uri_str)
            layer.setDataSource(uri.uri(), layer.name(), "postgres")
            if layer.isValid():
                logger.debug(f"  图层 {layer.name()} 连接更新成功")
            else:
                logger.error(f"  图层 {layer.name()} 更新后仍无效")
        except Exception as e:
            logger.error(f"  更新图层 {layer.name()} 连接失败: {e}")