1# -*- coding: utf-8 -*-
2
3"""
4/***************************************************************************
5Name                 : DB Manager
6Description          : Database manager plugin for QGIS
7Date                 : May 23, 2011
8copyright            : (C) 2011 by Giuseppe Sucameli
9email                : brush.tyler@gmail.com
10
11 ***************************************************************************/
12
13/***************************************************************************
14 *                                                                         *
15 *   This program is free software; you can redistribute it and/or modify  *
16 *   it under the terms of the GNU General Public License as published by  *
17 *   the Free Software Foundation; either version 2 of the License, or     *
18 *   (at your option) any later version.                                   *
19 *                                                                         *
20 ***************************************************************************/
21"""
22from builtins import str
23
24# this will disable the dbplugin if the connector raise an ImportError
25from .connector import GPKGDBConnector
26
27from qgis.PyQt.QtCore import Qt, QFileInfo, QCoreApplication
28from qgis.PyQt.QtGui import QIcon
29from qgis.PyQt.QtWidgets import QApplication, QAction, QFileDialog
30from qgis.core import (
31    Qgis,
32    QgsApplication,
33    QgsDataSourceUri,
34    QgsSettings,
35    QgsProviderRegistry,
36)
37from qgis.gui import QgsMessageBar
38
39from ..plugin import DBPlugin, Database, Table, VectorTable, RasterTable, TableField, TableIndex, TableTrigger, \
40    InvalidDataException
41
42
43def classFactory():
44    return GPKGDBPlugin
45
46
47class GPKGDBPlugin(DBPlugin):
48
49    @classmethod
50    def icon(self):
51        return QgsApplication.getThemeIcon("/mGeoPackage.svg")
52
53    @classmethod
54    def typeName(self):
55        return 'gpkg'
56
57    @classmethod
58    def typeNameString(self):
59        return QCoreApplication.translate('db_manager', 'GeoPackage')
60
61    @classmethod
62    def providerName(self):
63        return 'ogr'
64
65    @classmethod
66    def connectionSettingsKey(self):
67        return 'providers/ogr/GPKG/connections'
68
69    def databasesFactory(self, connection, uri):
70        return GPKGDatabase(connection, uri)
71
72    def connect(self, parent=None):
73        conn_name = self.connectionName()
74
75        md = QgsProviderRegistry.instance().providerMetadata(self.providerName())
76        conn = md.findConnection(conn_name)
77
78        if conn is None:  # non-existent entry?
79            raise InvalidDataException(self.tr(u'There is no defined database connection "{0}".').format(conn_name))
80
81        uri = QgsDataSourceUri()
82        uri.setDatabase(conn.uri())
83        return self.connectToUri(uri)
84
85    @classmethod
86    def addConnection(self, conn_name, uri):
87        md = QgsProviderRegistry.instance().providerMetadata(self.providerName())
88        conn = md.createConnection(uri.database(), {})
89        md.saveConnection(conn, conn_name)
90        return True
91
92    @classmethod
93    def addConnectionActionSlot(self, item, action, parent, index):
94        QApplication.restoreOverrideCursor()
95        try:
96            filename, selected_filter = QFileDialog.getOpenFileName(parent,
97                                                                    parent.tr("Choose GeoPackage file"), None, "GeoPackage (*.gpkg)")
98            if not filename:
99                return
100        finally:
101            QApplication.setOverrideCursor(Qt.WaitCursor)
102
103        conn_name = QFileInfo(filename).fileName()
104        uri = QgsDataSourceUri()
105        uri.setDatabase(filename)
106        self.addConnection(conn_name, uri)
107        index.internalPointer().itemChanged()
108
109
110class GPKGDatabase(Database):
111
112    def __init__(self, connection, uri):
113        Database.__init__(self, connection, uri)
114
115    def connectorsFactory(self, uri):
116        return GPKGDBConnector(uri, self.connection())
117
118    def dataTablesFactory(self, row, db, schema=None):
119        return GPKGTable(row, db, schema)
120
121    def vectorTablesFactory(self, row, db, schema=None):
122        return GPKGVectorTable(row, db, schema)
123
124    def rasterTablesFactory(self, row, db, schema=None):
125        return GPKGRasterTable(row, db, schema)
126
127    def info(self):
128        from .info_model import GPKGDatabaseInfo
129
130        return GPKGDatabaseInfo(self)
131
132    def sqlResultModel(self, sql, parent):
133        from .data_model import GPKGSqlResultModel
134
135        return GPKGSqlResultModel(self, sql, parent)
136
137    def sqlResultModelAsync(self, sql, parent):
138        from .data_model import GPKGSqlResultModelAsync
139
140        return GPKGSqlResultModelAsync(self, sql, parent)
141
142    def registerDatabaseActions(self, mainWindow):
143        action = QAction(self.tr("Run &Vacuum"), self)
144        mainWindow.registerAction(action, self.tr("&Database"), self.runVacuumActionSlot)
145
146        Database.registerDatabaseActions(self, mainWindow)
147
148    def runVacuumActionSlot(self, item, action, parent):
149        QApplication.restoreOverrideCursor()
150        try:
151            if not isinstance(item, (DBPlugin, Table)) or item.database() is None:
152                parent.infoBar.pushMessage(self.tr("No database selected or you are not connected to it."),
153                                           Qgis.Info, parent.iface.messageTimeout())
154                return
155        finally:
156            QApplication.setOverrideCursor(Qt.WaitCursor)
157
158        self.runVacuum()
159
160    def runVacuum(self):
161        self.database().aboutToChange.emit()
162        self.database().connector.runVacuum()
163        self.database().refresh()
164
165    def runAction(self, action):
166        action = str(action)
167
168        if action.startswith("vacuum/"):
169            if action == "vacuum/run":
170                self.runVacuum()
171                return True
172
173        return Database.runAction(self, action)
174
175    def uniqueIdFunction(self):
176        return None
177
178    def toSqlLayer(self, sql, geomCol, uniqueCol, layerName="QueryLayer", layerType=None, avoidSelectById=False, filter=""):
179        from qgis.core import QgsVectorLayer
180
181        vl = QgsVectorLayer(self.uri().database() + '|subset=' + sql, layerName, 'ogr')
182        return vl
183
184    def supportsComment(self):
185        return False
186
187
188class GPKGTable(Table):
189
190    def __init__(self, row, db, schema=None):
191        """Constructs a GPKGTable
192
193        :param row: a three elements array with: [table_name, is_view, is_sys_table]
194        :type row: array [str, bool, bool]
195        :param db: database instance
196        :type db:
197        :param schema: schema name, defaults to None, ignored by GPKG
198        :type schema: str, optional
199        """
200
201        Table.__init__(self, db, None)
202        self.name, self.isView, self.isSysTable = row
203
204    def ogrUri(self):
205        ogrUri = u"%s|layername=%s" % (self.uri().database(), self.name)
206        return ogrUri
207
208    def mimeUri(self):
209        # QGIS has no provider to load Geopackage vectors, let's use OGR
210        return u"vector:ogr:%s:%s" % (self.name, self.ogrUri())
211
212    def toMapLayer(self, geometryType=None, crs=None):
213        from qgis.core import QgsVectorLayer
214
215        provider = "ogr"
216        uri = self.ogrUri()
217
218        if geometryType:
219            geom_mapping = {
220                'POINT': 'Point',
221                'LINESTRING': 'LineString',
222                'POLYGON': 'Polygon',
223            }
224            geometryType = geom_mapping[geometryType]
225            uri = "{}|geometrytype={}".format(uri, geometryType)
226
227        return QgsVectorLayer(uri, self.name, provider)
228
229    def tableFieldsFactory(self, row, table):
230        return GPKGTableField(row, table)
231
232    def tableIndexesFactory(self, row, table):
233        return GPKGTableIndex(row, table)
234
235    def tableTriggersFactory(self, row, table):
236        return GPKGTableTrigger(row, table)
237
238    def tableDataModel(self, parent):
239        from .data_model import GPKGTableDataModel
240
241        return GPKGTableDataModel(self, parent)
242
243
244class GPKGVectorTable(GPKGTable, VectorTable):
245
246    def __init__(self, row, db, schema=None):
247        GPKGTable.__init__(self, row[:-5], db, schema)
248        VectorTable.__init__(self, db, schema)
249        # GPKG does case-insensitive checks for table names, but the
250        # GPKG provider didn't do the same in QGIS < 1.9, so self.geomTableName
251        # stores the table name like stored in the geometry_columns table
252        self.geomTableName, self.geomColumn, self.geomType, self.geomDim, self.srid = row[-5:]
253        self.extent = self.database().connector.getTableExtent((self.schemaName(), self.name), self.geomColumn, force=False)
254
255    def uri(self):
256        uri = self.database().uri()
257        uri.setDataSource('', self.geomTableName, self.geomColumn)
258        return uri
259
260    def hasSpatialIndex(self, geom_column=None):
261        geom_column = geom_column if geom_column is not None else self.geomColumn
262        return self.database().connector.hasSpatialIndex((self.schemaName(), self.name), geom_column)
263
264    def createSpatialIndex(self, geom_column=None):
265        self.aboutToChange.emit()
266        ret = VectorTable.createSpatialIndex(self, geom_column)
267        if ret is not False:
268            self.database().refresh()
269        return ret
270
271    def deleteSpatialIndex(self, geom_column=None):
272        self.aboutToChange.emit()
273        ret = VectorTable.deleteSpatialIndex(self, geom_column)
274        if ret is not False:
275            self.database().refresh()
276        return ret
277
278    def refreshTableEstimatedExtent(self):
279        return
280
281    def refreshTableExtent(self):
282        prevExtent = self.extent
283        self.extent = self.database().connector.getTableExtent((self.schemaName(), self.name), self.geomColumn, force=True)
284        if self.extent != prevExtent:
285            self.refresh()
286
287    def runAction(self, action):
288        if GPKGTable.runAction(self, action):
289            return True
290        return VectorTable.runAction(self, action)
291
292
293class GPKGRasterTable(GPKGTable, RasterTable):
294
295    def __init__(self, row, db, schema=None):
296        GPKGTable.__init__(self, row[:-3], db, schema)
297        RasterTable.__init__(self, db, schema)
298        self.prefixName, self.geomColumn, self.srid = row[-3:]
299        self.geomType = 'RASTER'
300        self.extent = self.database().connector.getTableExtent((self.schemaName(), self.name), self.geomColumn)
301
302    def gpkgGdalUri(self):
303        gdalUri = u'GPKG:%s:%s' % (self.uri().database(), self.prefixName)
304        return gdalUri
305
306    def mimeUri(self):
307        # QGIS has no provider to load rasters, let's use GDAL
308        uri = u"raster:gdal:%s:%s" % (self.name, self.uri().database())
309        return uri
310
311    def toMapLayer(self, geometryType=None, crs=None):
312        from qgis.core import QgsRasterLayer, QgsContrastEnhancement
313
314        # QGIS has no provider to load rasters, let's use GDAL
315        uri = self.gpkgGdalUri()
316        rl = QgsRasterLayer(uri, self.name)
317        if rl.isValid():
318            rl.setContrastEnhancement(QgsContrastEnhancement.StretchToMinimumMaximum)
319        return rl
320
321
322class GPKGTableField(TableField):
323
324    def __init__(self, row, table):
325        TableField.__init__(self, table)
326        self.num, self.name, self.dataType, self.notNull, self.default, self.primaryKey = row
327        self.hasDefault = self.default
328
329
330class GPKGTableIndex(TableIndex):
331
332    def __init__(self, row, table):
333        TableIndex.__init__(self, table)
334        self.num, self.name, self.isUnique, self.columns = row
335
336
337class GPKGTableTrigger(TableTrigger):
338
339    def __init__(self, row, table):
340        TableTrigger.__init__(self, table)
341        self.name, self.function = row
342