sanya-dbms/processing/algs/gdal/gdbtopostgislist.py

"""
***************************************************************************
    gdbtopostgislist.py
    ---------------------
    Date                 : November 2012
    Copyright            : (C) 2012 by Victor Olaya
    Email                : volayaf at gmail dot com
***************************************************************************
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

__author__ = 'wanger'
__date__ = 'November 2024'
__copyright__ = '(C) 2024, wanger'

import os
import subprocess

from PyQt5.QtCore import QCoreApplication
from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QMessageBox
from osgeo import ogr

from qgis._core import QgsProcessingParameterDefinition, QgsProcessingAlgorithm, QgsProcessingParameterFolderDestination
from qgis.core import (QgsProcessing,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterString,
                       QgsProcessingParameterDateTime,
                       QgsProcessingParameterEnum,
                       QgsProcessingParameterCrs,
                       QgsProcessingParameterField,
                       QgsProcessingParameterExtent,
                       QgsProcessingParameterBoolean,
                       QgsProcessingParameterProviderConnection,
                       QgsProcessingParameterDatabaseSchema,
                       QgsProcessingParameterDatabaseTable,
                       QgsProviderRegistry,
                       QgsProcessingException,
                       QgsCoordinateReferenceSystem,
                       QgsProviderConnectionException,
                       QgsDataSourceUri)

from processing.algs.gdal.GdalAlgorithm import GdalAlgorithm
from processing.algs.gdal.GdalUtils import GdalUtils
from processing.tools.PrintUtils import printStr
from processing.tools.PostgreSQL.PostgreSQL import PostgreSQL
from processing.tools.system import isWindows
from processing.tools.PrintUtils import getLastPrint
from processing.tools.FileListPrintUtils import getFileListPrint

pluginPath = os.path.normpath(os.path.join(
    os.path.split(os.path.dirname(__file__))[0], os.pardir))


class GdbToPostGisList(GdalAlgorithm):
    UID = 'UID'
    DATABASE = 'DATABASE'
    INPUTGDB = 'INPUTGDB'
    SHAPE_ENCODING = 'SHAPE_ENCODING'
    GTYPE = 'GTYPE'
    GEOMTYPE = ['', 'NONE', 'GEOMETRY', 'POINT', 'LINESTRING', 'POLYGON', 'GEOMETRYCOLLECTION', 'MULTIPOINT',
                'MULTIPOLYGON', 'MULTILINESTRING', 'CIRCULARSTRING', 'COMPOUNDCURVE', 'CURVEPOLYGON', 'MULTICURVE',
                'MULTISURFACE']
    S_SRS = 'S_SRS'
    T_SRS = 'T_SRS'
    A_SRS = 'A_SRS'
    HOST = 'HOST'
    PORT = 'PORT'
    USER = 'USER'
    DBNAME = 'DBNAME'
    PASSWORD = 'PASSWORD'
    SCHEMA = 'SCHEMA'
    TABLE = 'TABLE'
    TABLE_ALIAS = 'TABLE_ALIAS'
    PK = 'PK'
    PRIMARY_KEY = 'PRIMARY_KEY'
    GEOCOLUMN = 'GEOCOLUMN'
    DIM = 'DIM'
    DIMLIST = ['2', '3', '4']
    SIMPLIFY = 'SIMPLIFY'
    SEGMENTIZE = 'SEGMENTIZE'
    SPAT = 'SPAT'
    CLIP = 'CLIP'
    WHERE = 'WHERE'
    GT = 'GT'
    OVERWRITE = 'OVERWRITE'
    APPEND = 'APPEND'
    ADDFIELDS = 'ADDFIELDS'
    LAUNDER = 'LAUNDER'
    INDEX = 'INDEX'
    SKIPFAILURES = 'SKIPFAILURES'
    PRECISION = 'PRECISION'
    PROMOTETOMULTI = 'PROMOTETOMULTI'
    OPTIONS = 'OPTIONS'
    # TODO 三亚项目新追加入库属性
    VECTOR_YEAR = 'VECTOR_YEAR'
    VECTOR_XZQH = 'VECTOR_XZQH'
    VECTOR_XMLX = 'VECTOR_XMLX'
    VECTOR_SJLY = 'VECTOR_SJLY'
    VECTOR_YWLX = 'VECTOR_YWLX'
    SOURCE_TYPE = 'SOURCE_TYPE'
    XZQH_FIELD = 'XZQH_FIELD'
    VECTOR_GLBM = 'VECTOR_GLBM'
    VECTOR_ZYML = 'VECTOR_ZYML'
    Metadata_storage = 'Metadata_storage'
    XLK = 'XLK'
    XLKXZX = ['选择一', '选择二', '选择三', '选择四', '选择五']
    selectedValue = "selectedValue"

    def __init__(self):
        super().__init__()

    def initAlgorithm(self, config=None):
        self.addParameter(QgsProcessingParameterString(self.UID,
                                                       self.tr('唯一数据标识前缀，用于更新和删除关联使用'), "",
                                                       optional=True))
        folderParameter = QgsProcessingParameterFolderDestination(self.INPUTGDB, self.tr('GDB数据文件夹'))
        print(folderParameter)
        self.addParameter(folderParameter)
        db_param = QgsProcessingParameterProviderConnection(
            self.DATABASE,
            self.tr('数据库'), provider='postgres',
            defaultValue=self.pgcoon["host"])
        self.addParameter(db_param)
        schema_param = QgsProcessingParameterDatabaseSchema(
            self.SCHEMA,
            self.tr('模式'), defaultValue=self.pgcoon["schema"], connectionParameterName=self.DATABASE,
            optional=False)
        self.addParameter(schema_param)
        # table_param = QgsProcessingParameterString(
        #     self.TABLE,
        #     self.tr('表名'), defaultValue="", optional=True)
        # self.addParameter(table_param)
        # tablealias_param = QgsProcessingParameterString(
        #     self.TABLE_ALIAS,
        #     self.tr('别名'), defaultValue="", optional=True)
        # self.addParameter(tablealias_param)
        # 行政区划字段选择
        # self.addParameter(QgsProcessingParameterField(self.XZQH_FIELD,
        #                                               self.tr('行政区划字段')))
        # 数据来源
        self.addParameter(QgsProcessingParameterString(self.VECTOR_SJLY,
                                                       self.tr('数据来源'), "",
                                                       optional=False))

        self.addParameter(
            QgsProcessingParameterDateTime(self.VECTOR_YEAR, '数据时效', type=QgsProcessingParameterDateTime.Type.Date,
                                           defaultValue=None))
        pgconn = PostgreSQL(schema='base')
        rows = pgconn.getVectorYwlx()
        self.ywlxs = []
        for row in rows:
            self.ywlxs.append(row[0])
        self.addParameter(QgsProcessingParameterEnum(name=self.VECTOR_YWLX,
                                                     description=self.tr('业务类型'), options=self.ywlxs))

        rows = pgconn.getDeptList()
        self.depts = []
        for row in rows:
            self.depts.append(row[0])
        self.addParameter(QgsProcessingParameterEnum(name=self.VECTOR_GLBM,
                                                     description=self.tr('管理部门'), options=self.depts))
        rows = pgconn.getVectorZyml()
        self.zymls = []
        for row in rows:
            self.zymls.append(row[1])
        self.addParameter(QgsProcessingParameterEnum(name=self.VECTOR_ZYML,
                                                     description=self.tr('资源目录'), options=self.zymls))

        crs = QgsCoordinateReferenceSystem("EPSG:4525")
        crs_parameter = QgsProcessingParameterCrs(self.T_SRS,
                                                  self.tr('指定入库坐标系'),
                                                  defaultValue=crs,
                                                  optional=False)
        self.addParameter(crs_parameter)
        string = QgsProcessingParameterString(self.SOURCE_TYPE,
                                              self.tr('数据源类型'), "vector",
                                              optional=False)
        self.addParameter(string)
        self.addParameter(QgsProcessingParameterString(self.GEOCOLUMN,
                                                       self.tr('Geometry column name'), defaultValue='geom',
                                                       optional=True))
        self.addParameter(QgsProcessingParameterString(self.PK,
                                                       self.tr('Primary key (new field)'), defaultValue='id',
                                                       optional=True))
        metadata = QgsProcessingParameterBoolean(self.Metadata_storage,
                                                 self.tr('元数据入库'), defaultValue=True)
        self.addParameter(metadata)
        self.addParameter(QgsProcessingParameterBoolean(self.OVERWRITE,
                                                        self.tr('Overwrite existing table'), defaultValue=True))
        self.addParameter(QgsProcessingParameterBoolean(self.APPEND,
                                                        self.tr('Append to existing table'), defaultValue=False))
        self.addParameter(QgsProcessingParameterBoolean(self.ADDFIELDS,
                                                        self.tr('Append and add new fields to existing table'),
                                                        defaultValue=False))
        self.addParameter(QgsProcessingParameterBoolean(self.LAUNDER,
                                                        self.tr('Do not launder columns/table names'),
                                                        defaultValue=False))
        self.addParameter(QgsProcessingParameterBoolean(self.INDEX,
                                                        self.tr('Do not create spatial index'), defaultValue=False))
        self.addParameter(QgsProcessingParameterBoolean(self.SKIPFAILURES,
                                                        self.tr(
                                                            'Continue after a failure, skipping the failed feature'),
                                                        defaultValue=False))
        self.addParameter(QgsProcessingParameterBoolean(self.PROMOTETOMULTI,
                                                        self.tr('Promote to Multipart'),
                                                        defaultValue=True))
        self.addParameter(QgsProcessingParameterBoolean(self.PRECISION,
                                                        self.tr('Keep width and precision of input attributes'),
                                                        defaultValue=False))
        self.addParameter(QgsProcessingParameterString(self.OPTIONS,
                                                       self.tr('Additional creation options'), defaultValue='',
                                                       optional=True))
        pgconn.close()

    def name(self):
        return 'gdbtopostgis'

    def icon(self):
        return QIcon(os.path.join(pluginPath, 'images', 'dbms', 'importvector.png'))

    def displayName(self):
        return self.tr('GDB数据入库')

    def shortDescription(self):
        return self.tr('Exports a vector layer to an existing PostgreSQL database connection')

    def tags(self):
        t = self.tr('import,into,postgis,database,vector').split(',')
        t.extend(super().tags())
        return t

    def group(self):
        return self.tr('数据入库工具')

    def groupId(self):
        return 'vector2postgis'

    def setSelectedValue(self, v):
        printStr(v)
        self.selectedValue = v

    def getSelectedValue(self):
        return self.selectedValue

    def gdbimport(self, parameters, context, feedback, executing=True):
        print("############进入自身组件的processAlgorithm################")
        gdb_folder = self.parameterAsString(parameters, self.INPUTGDB, context)
        uid = self.parameterAsString(parameters, self.UID, context)
        driver = ogr.GetDriverByName('OpenFileGDB')
        if not driver:
            return {
                "状态": "OpenFileGDB 驱动不可用"
            }
        gdb = driver.Open(gdb_folder, 0)  # 0 表示只读模式
        if not gdb:
            return {
                "状态": f"无法打开 {gdb_folder} 文件夹"
            }
        layer_count = gdb.GetLayerCount()
        if layer_count > 0:
            # 数据库连接参数
            db_params = self.pgcoon
            # 将每个图层导入到 PostGIS
            connection_name = self.parameterAsConnectionName(parameters, self.DATABASE, context)
            if not connection_name:
                raise QgsProcessingException(
                    self.tr('No connection specified'))
            try:
                md = QgsProviderRegistry.instance().providerMetadata('postgres')
                conn = md.createConnection(connection_name)
            except QgsProviderConnectionException:
                raise QgsProcessingException(
                    self.tr('Could not retrieve connection details for {}').format(connection_name))
            uri = conn.uri()
            connection_parts = QgsDataSourceUri(uri).connectionInfo(executing).split(' ')
            schema = self.parameterAsSchema(parameters, self.SCHEMA, context)
            connection_parts.append('active_schema={}'.format(schema or 'public'))
            pg_conn_string = 'PG:"{}"'.format(' '.join(connection_parts))
            # pg_conn_string = f'PG:"host={db_params["host"]} port={db_params["port"]} dbname={db_params["dbname"]} user={db_params["user"]} password={db_params["password"]} active_schema={db_params["schema"]}"'
            layernames = []
            for i in range(layer_count):
                layer = gdb.GetLayerByIndex(i)
                # 获取图层的几何类型
                geom_type = self.getLayerGeometryType(layer.GetGeomType())
                print(geom_type)
                layer_name = layer.GetName()
                t_layer_name = f'{uid}_{layer_name}'
                print(f"图层 {i}: {layer.GetName()}")
                command = f'ogr2ogr.exe -f "PostgreSQL" {pg_conn_string} "{gdb_folder}" {layer_name}  -overwrite  -nln {t_layer_name}'
                if geom_type:
                    geocolumn = self.parameterAsString(parameters, self.GEOCOLUMN, context)
                    tsrs = self.parameterAsCrs(parameters, self.T_SRS, context)
                    command = f'{command} -nlt {geom_type} -lco GEOMETRY_NAME={geocolumn} -t_srs {GdalUtils.gdal_crs_string(tsrs)}'
                print(command)
                # 使用 subprocess.Popen 在后台运行命令
                process = subprocess.Popen(command, shell=True, creationflags=subprocess.CREATE_NO_WINDOW)
                # 等待命令执行完毕
                process.wait()
                # subprocess.run(command, check=True)
                print(f"图层 {layer_name} 导入成功")
                layernames.append(t_layer_name)
                # 所属行政区划
                ssxzqh = getLastPrint()
                print("ssxzqh====" + ssxzqh)
                # 获取附件列表
                fileliststr = getFileListPrint()
                pgconn = PostgreSQL(schema='base')
                ogrLayer = t_layer_name
                pgconn.metadataStorage(parameters, ssxzqh, fileliststr, self.ywlxs[parameters.get("VECTOR_YWLX")],
                                       self.depts[parameters.get("VECTOR_GLBM")], ogrLayer,
                                       self.zymls[parameters.get("VECTOR_ZYML")])
                pgconn.close()
            return {
                "状态": f"GDB数据导入成功，共导入{len(layernames)}个图层，分别为：{','.join(layernames)}。"
            }
        else:
            return {
                "状态": f"{gdb_folder} 文件夹不包含任何图层"
            }

    def getConsoleCommands(self, parameters, context, feedback, executing=True):
        # print(parameters)
        # print(self.ywlxs)
        # connection_name = self.parameterAsConnectionName(parameters, self.DATABASE, context)
        # if not connection_name:
        #     raise QgsProcessingException(
        #         self.tr('No connection specified'))
        #
        # # resolve connection details to uri
        # try:
        #     md = QgsProviderRegistry.instance().providerMetadata('postgres')
        #     conn = md.createConnection(connection_name)
        # except QgsProviderConnectionException:
        #     raise QgsProcessingException(
        #         self.tr('Could not retrieve connection details for {}').format(connection_name))
        #
        # uri = conn.uri()
        # gdb_folder = self.parameterAsString(parameters, self.INPUTGDB, context)
        # ogrLayer, layername = self.getOgrCompatibleSource(self.INPUTGDB, parameters, context, feedback, executing)
        # shapeEncoding = self.parameterAsString(parameters, self.SHAPE_ENCODING, context)
        # ssrs = self.parameterAsCrs(parameters, self.S_SRS, context)
        # tsrs = self.parameterAsCrs(parameters, self.T_SRS, context)
        # asrs = self.parameterAsCrs(parameters, self.A_SRS, context)
        # table = self.parameterAsDatabaseTableName(parameters, self.TABLE, context)
        # schema = self.parameterAsSchema(parameters, self.SCHEMA, context)
        # pk = self.parameterAsString(parameters, self.PK, context)
        # pkstring = "-lco FID=" + pk
        # primary_key = self.parameterAsString(parameters, self.PRIMARY_KEY, context)
        # geocolumn = self.parameterAsString(parameters, self.GEOCOLUMN, context)
        # geocolumnstring = "-lco GEOMETRY_NAME=" + geocolumn
        # dim = self.DIMLIST[self.parameterAsEnum(parameters, self.DIM, context)]
        # dimstring = "-lco DIM=" + dim
        # simplify = self.parameterAsString(parameters, self.SIMPLIFY, context)
        # segmentize = self.parameterAsString(parameters, self.SEGMENTIZE, context)
        # spat = self.parameterAsExtent(parameters, self.SPAT, context)
        # clip = self.parameterAsBoolean(parameters, self.CLIP, context)
        # where = self.parameterAsString(parameters, self.WHERE, context)
        # wherestring = '-where "' + where + '"'
        # gt = self.parameterAsString(parameters, self.GT, context)
        # overwrite = self.parameterAsBoolean(parameters, self.OVERWRITE, context)
        # append = self.parameterAsBoolean(parameters, self.APPEND, context)
        # addfields = self.parameterAsBoolean(parameters, self.ADDFIELDS, context)
        # launder = self.parameterAsBoolean(parameters, self.LAUNDER, context)
        # launderstring = "-lco LAUNDER=NO"
        # index = self.parameterAsBoolean(parameters, self.INDEX, context)
        # indexstring = "-lco SPATIAL_INDEX=OFF"
        # skipfailures = self.parameterAsBoolean(parameters, self.SKIPFAILURES, context)
        # promotetomulti = self.parameterAsBoolean(parameters, self.PROMOTETOMULTI, context)
        # precision = self.parameterAsBoolean(parameters, self.PRECISION, context)
        # options = self.parameterAsString(parameters, self.OPTIONS, context)
        # table_alias = self.parameterAsString(parameters, self.TABLE_ALIAS, context)
        #
        # arguments = [
        #     '-progress',
        #     '--config PG_USE_COPY YES'
        # ]
        # if shapeEncoding:
        #     arguments.append('--config')
        #     arguments.append('SHAPE_ENCODING')
        #     arguments.append(shapeEncoding)
        # arguments.append('-f')
        # arguments.append('PostgreSQL')
        #
        # connection_parts = QgsDataSourceUri(uri).connectionInfo(executing).split(' ')
        # connection_parts.append('active_schema={}'.format(schema or 'public'))
        # arguments.append('PG:{}'.format(' '.join(connection_parts)))
        # # 检查数据库中该模式下是否存在相同的表
        # pgconn = PostgreSQL(schema='base')
        # resbool = pgconn.checkTableName(schemaname=schema,
        #                                 tablename=layername.lower() if len(table) == 0 else table.lower())
        # pgconn.close()
        # if resbool == False:
        #     QMessageBox.warning(None, "提示",
        #                         self.tr('数据库中已存在相同表名，请输入或修改表名！'))
        # arguments.append(dimstring)
        # arguments.append(ogrLayer)
        # arguments.append(layername)
        # if index:
        #     arguments.append(indexstring)
        # if launder:
        #     arguments.append(launderstring)
        # if append:
        #     arguments.append('-append')
        # if addfields:
        #     arguments.append('-addfields')
        # if overwrite:
        #     arguments.append('-overwrite')
        # if len(self.GEOMTYPE[self.parameterAsEnum(parameters, self.GTYPE, context)]) > 0:
        #     arguments.append('-nlt')
        #     arguments.append(self.GEOMTYPE[self.parameterAsEnum(parameters, self.GTYPE, context)])
        # if len(geocolumn) > 0:
        #     arguments.append(geocolumnstring)
        # if pk:
        #     arguments.append(pkstring)
        # elif primary_key:
        #     arguments.append("-lco FID=" + primary_key)
        # if len(table) == 0:
        #     # table = layername.lower()
        #     table = layername
        # if schema:
        #     table = f'{schema}.{table}'
        # # if len(table_alias) > 0:
        # #     arguments.append("-lco layer_name=" + table_alias)
        # arguments.append('-nln')
        # arguments.append(table)
        # if ssrs.isValid():
        #     arguments.append('-s_srs')
        #     arguments.append(GdalUtils.gdal_crs_string(ssrs))
        # if tsrs.isValid():
        #     arguments.append('-t_srs')
        #     arguments.append(GdalUtils.gdal_crs_string(tsrs))
        # if asrs.isValid():
        #     arguments.append('-a_srs')
        #     arguments.append(GdalUtils.gdal_crs_string(asrs))
        # if not spat.isNull():
        #     arguments.append('-spat')
        #     arguments.append(spat.xMinimum())
        #     arguments.append(spat.yMinimum())
        #     arguments.append(spat.xMaximum())
        #     arguments.append(spat.yMaximum())
        #     if clip:
        #         arguments.append('-clipsrc spat_extent')
        # if skipfailures:
        #     arguments.append('-skipfailures')
        # if where:
        #     arguments.append(wherestring)
        # if len(simplify) > 0:
        #     arguments.append('-simplify')
        #     arguments.append(simplify)
        # if len(segmentize) > 0:
        #     arguments.append('-segmentize')
        #     arguments.append(segmentize)
        # if len(gt) > 0:
        #     arguments.append('-gt')
        #     arguments.append(gt)
        # if promotetomulti:
        #     arguments.append('-nlt PROMOTE_TO_MULTI')
        # if precision is False:
        #     arguments.append('-lco PRECISION=NO')
        # if len(options) > 0:
        #     arguments.append(options)
        if isWindows():
            return ['cmd.exe', '/C ', 'ogr2ogr.exe']
        else:
            return ['ogr2ogr']

    def getTableName(self, table):
        sp = table.split(".")
        res = f'{sp[0]}."{sp[1]}"'
        print(res)
        return res

    def commandName(self):
        return "ogr2ogr"