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

"""
***************************************************************************
    ogr2ogrtopostgislist.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

from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QMessageBox

from qgis._core import QgsProcessingParameterDefinition
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

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


class Ogr2OgrToPostGisList(GdalAlgorithm):
    DATABASE = 'DATABASE'
    INPUT = 'INPUT'
    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(QgsProcessingParameterFeatureSource(self.INPUT,
                                                              self.tr('Input layer'),
                                                              types=[QgsProcessing.TypeVector]))
        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('行政区划字段'),
                                                      None,
                                                      self.INPUT,
                                                      QgsProcessingParameterField.Any,
                                                      optional=True))
        # 数据来源
        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))
        # self.addParameter(QgsProcessingParameterEnum(self.XLK,
        #                                              self.tr('下拉框'), options=self.XLKXZX,
        #                                              defaultValue=0))
        # 数据年份
        # self.addParameter(QgsProcessingParameterString(self.VECTOR_YEAR,
        #                                                self.tr('数据年份'), "",
        #                                                optional=False))

        # 行政区划
        # self.addParameter(QgsProcessingParameterString(self.VECTOR_XZQH,
        #                                                self.tr('行政区划'), "",
        #                                                optional=False))
        # 项目类型
        # self.addParameter(QgsProcessingParameterString(self.VECTOR_XMLX,
        #                                                self.tr('项目类型'), "",
        #                                                optional=False))

        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.SHAPE_ENCODING,
        #                                                self.tr('Shape encoding'), "", optional=True))
        #
        # self.addParameter(QgsProcessingParameterCrs(self.A_SRS,
        #                                             self.tr('Assign an output CRS'), defaultValue='', optional=True))
        # self.addParameter(QgsProcessingParameterCrs(self.S_SRS,
        #                                             self.tr('Override source CRS'), defaultValue='', optional=True))
        #
        # self.addParameter(QgsProcessingParameterEnum(self.GTYPE,
        #                                              self.tr('Output geometry type'), options=self.GEOMTYPE,
        #                                              defaultValue=0))
        #
        self.addParameter(QgsProcessingParameterString(self.PK,
                                                       self.tr('Primary key (new field)'), defaultValue='id',
                                                       optional=True))
        # self.addParameter(QgsProcessingParameterField(self.PRIMARY_KEY,
        #                                               self.tr(
        #                                                   'Primary key (existing field, used if the above option is left empty)'),
        #                                               parentLayerParameterName=self.INPUT, optional=True))
        self.addParameter(QgsProcessingParameterString(self.GEOCOLUMN,
                                                       self.tr('Geometry column name'), defaultValue='geom',
                                                       optional=True))
        # self.addParameter(QgsProcessingParameterEnum(self.DIM,
        #                                              self.tr('Vector dimensions'), options=self.DIMLIST,
        #                                              defaultValue=0))
        # self.addParameter(QgsProcessingParameterString(self.SIMPLIFY,
        #                                                self.tr('Distance tolerance for simplification'),
        #                                                defaultValue='', optional=True))
        # self.addParameter(QgsProcessingParameterString(self.SEGMENTIZE,
        #                                                self.tr('Maximum distance between 2 nodes (densification)'),
        #                                                defaultValue='', optional=True))
        # self.addParameter(QgsProcessingParameterExtent(self.SPAT,
        #                                                self.tr(
        #                                                    'Select features by extent (defined in input layer CRS)'),
        #                                                optional=True))
        # self.addParameter(QgsProcessingParameterBoolean(self.CLIP,
        #                                                 self.tr(
        #                                                     'Clip the input layer using the above (rectangle) extent'),
        #                                                 defaultValue=False))
        # self.addParameter(QgsProcessingParameterString(self.WHERE,
        #                                                self.tr(
        #                                                    'Select features using a SQL "WHERE" statement (Ex: column=\'value\')'),
        #                                                defaultValue='', optional=True))
        # self.addParameter(QgsProcessingParameterString(self.GT,
        #                                                self.tr('Group N features per transaction (Default: 20000)'),
        #                                                defaultValue='', optional=True))
        metadata = QgsProcessingParameterBoolean(self.Metadata_storage,
                                                 self.tr('元数据入库'), defaultValue=True)
        # metadata.setEnable(False)
        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 'importvectorintopostgisdatabaseavailableconnections'

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

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

    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 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()

        ogrLayer, layername = self.getOgrCompatibleSource(self.INPUT, parameters, context, feedback, executing)
        print(ogrLayer)
        print(layername)
        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',
                    GdalUtils.escapeAndJoin(arguments)]
        else:
            return ['ogr2ogr', GdalUtils.escapeAndJoin(arguments)]

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

    def commandName(self):
        return "ogr2ogr"