[processing] allow output directly on PostGIS tables

qgis · Nov 6, 2015 · 11b5092 · 11b5092
1 parent 22ace3d
commit 11b5092
Show file tree

Hide file tree

Showing 8 changed files with 256 additions and 21 deletions.
diff --git a/python/plugins/processing/core/outputs.py b/python/plugins/processing/core/outputs.py
@@ -306,5 +306,5 @@ def getVectorWriter(self, fields, geomType, crs, options=None):
 
         w = VectorWriter(self.value, self.encoding, fields, geomType,
                          crs, options)
-        self.memoryLayer = w.memLayer
+        self.layer = w.layer
         return w
diff --git a/python/plugins/processing/gui/OutputSelectionPanel.py b/python/plugins/processing/gui/OutputSelectionPanel.py
@@ -32,10 +32,12 @@
 from PyQt4.QtCore import QCoreApplication, QSettings
 from PyQt4.QtGui import QDialog, QMenu, QAction, QCursor, QFileDialog
 from qgis.gui import QgsEncodingFileDialog
+from qgis.core import *
 
 from processing.core.ProcessingConfig import ProcessingConfig
 from processing.core.outputs import OutputVector
 from processing.core.outputs import OutputDirectory
+from processing.gui.PostgisTableSelector import PostgisTableSelector
 
 pluginPath = os.path.split(os.path.dirname(__file__))[0]
 WIDGET, BASE = uic.loadUiType(
@@ -81,12 +83,41 @@ def selectOutput(self):
                     self.tr('Save to memory layer'), self.btnSelect)
                 actionSaveToMemory.triggered.connect(self.saveToMemory)
                 popupMenu.addAction(actionSaveToMemory)
+                actionSaveToPostGIS = QAction(
+                    self.tr('Save to PostGIS table...'), self.btnSelect)
+                actionSaveToPostGIS.triggered.connect(self.saveToPostGIS)
+                settings = QSettings()
+                settings.beginGroup('/PostgreSQL/connections/')
+                names = settings.childGroups()
+                settings.endGroup()
+                actionSaveToPostGIS.setEnabled(bool(names))
+                popupMenu.addAction(actionSaveToPostGIS)
 
             popupMenu.exec_(QCursor.pos())
 
     def saveToTemporaryFile(self):
         self.leText.setText('')
 
+    def saveToPostGIS(self):
+        dlg = PostgisTableSelector(self, self.output.name.lower())
+        dlg.exec_()
+        if dlg.connection:
+            settings = QSettings()
+            mySettings = '/PostgreSQL/connections/' + dlg.connection
+            dbname = settings.value(mySettings + '/database')
+            user = settings.value(mySettings + '/username')
+            host = settings.value(mySettings + '/host')
+            port = settings.value(mySettings + '/port')
+            password = settings.value(mySettings + '/password')
+            uri = QgsDataSourceURI()
+            uri.setConnection(host, str(port), dbname, user, password)
+            uri.setDataSource(dlg.schema, dlg.table, "the_geom")
+            connInfo = uri.connectionInfo()
+            (success, user, passwd ) = QgsCredentials.instance().get(connInfo, None, None)
+            if success:
+                QgsCredentials.instance().put(connInfo, user, passwd)
+            self.leText.setText("postgis:" + uri.uri())
+
     def saveToMemory(self):
         self.leText.setText('memory:')
 
@@ -124,10 +155,8 @@ def selectFile(self):
 
     def selectDirectory(self):
         lastDir = ''
-
-        dirName = QFileDialog.getExistingDirectory(self,
-                                                   self.tr('Select directory'), lastDir, QFileDialog.ShowDirsOnly)
-
+        dirName = QFileDialog.getExistingDirectory(self,self.tr('Select directory'),
+                                                   lastDir, QFileDialog.ShowDirsOnly)
         self.leText.setText(dirName)
 
     def getValue(self):
@@ -136,10 +165,14 @@ def getValue(self):
             value = None
         elif fileName.startswith('memory:'):
             value = fileName
+        elif fileName.startswith('postgis:'):
+            value = fileName
         elif not os.path.isabs(fileName):
             value = ProcessingConfig.getSetting(
                 ProcessingConfig.OUTPUT_FOLDER) + os.sep + fileName
         else:
             value = fileName
 
+
+
         return value
diff --git a/python/plugins/processing/gui/PostgisTableSelector.py b/python/plugins/processing/gui/PostgisTableSelector.py
@@ -0,0 +1,92 @@
+import os
+from PyQt4 import uic, QtCore, QtGui
+from processing.algs.qgis.postgis_utils import GeoDB
+from qgis.core import *
+from PyQt4.QtGui import QMessageBox
+
+pluginPath = os.path.split(os.path.dirname(__file__))[0]
+WIDGET, BASE = uic.loadUiType(
+    os.path.join(pluginPath, 'ui', 'DlgPostgisTableSelector.ui'))
+
+
+class PostgisTableSelector(BASE, WIDGET):
+
+    def __init__(self, parent, tablename):
+        super(PostgisTableSelector, self).__init__(parent)
+        self.connection = None
+        self.table = None
+        self.schema = None
+        self.setupUi(self)
+        settings = QtCore.QSettings()
+        settings.beginGroup('/PostgreSQL/connections/')
+        names = settings.childGroups()
+        settings.endGroup()
+        for n in names:
+            item = ConnectionItem(n)
+            self.treeConnections.addTopLevelItem(item)
+
+        def itemExpanded(item):
+            try:
+                item.populateSchemas()
+            except:
+                pass
+
+        self.treeConnections.itemExpanded.connect(itemExpanded)
+
+        self.textTableName.setText(tablename)
+
+        self.buttonBox.accepted.connect(self.okPressed)
+        self.buttonBox.rejected.connect(self.cancelPressed)
+
+    def cancelPressed(self):
+        self.close()
+
+    def okPressed(self):
+        if self.textTableName.text().strip() == "":
+            self.textTableName.setStyleSheet("QLineEdit{background: yellow}")
+            return
+        item = self.treeConnections.currentItem()
+        if isinstance(item, ConnectionItem):
+            QMessageBox.warning(self, "Wrong selection", "Select a schema item in the tree")
+            return
+        self.schema = item.text(0)
+        self.table = self.textTableName.text().strip()
+        self.connection = item.parent().text(0)
+        self.close()
+
+class ConnectionItem(QtGui.QTreeWidgetItem):
+
+    connIcon = QtGui.QIcon(os.path.dirname(__file__) + '/../images/postgis.png')
+    schemaIcon = QtGui.QIcon(os.path.dirname(__file__) + '/../images/namespace.png')
+
+    def __init__(self, connection):
+        QtGui.QTreeWidgetItem.__init__(self)
+        self.setChildIndicatorPolicy(QtGui.QTreeWidgetItem.ShowIndicator)
+        self.connection = connection
+        self.setText(0, connection)
+        self.setIcon(0, self.connIcon)
+
+    def populateSchemas(self):
+        if self.childCount() != 0:
+            return
+        settings = QtCore.QSettings()
+        connSettings = '/PostgreSQL/connections/' + self.connection
+        database = settings.value(connSettings + '/database')
+        user = settings.value(connSettings + '/username')
+        host = settings.value(connSettings + '/host')
+        port = settings.value(connSettings + '/port')
+        passwd = settings.value(connSettings + '/password')
+        uri = QgsDataSourceURI()
+        uri.setConnection(host, str(port), database, user, passwd)
+        connInfo = uri.connectionInfo()
+        (success, user, passwd ) = QgsCredentials.instance().get(connInfo, None, None)
+        if success:
+            QgsCredentials.instance().put(connInfo, user, passwd)
+            geodb = GeoDB(host, int(port), database, user, passwd)
+            schemas = geodb.list_schemas()
+            for oid, name, owner, perms in schemas:
+                item = QtGui.QTreeWidgetItem()
+                item.setText(0, name)
+                item.setIcon(0, self.schemaIcon)
+                self.addChild(item)
+
diff --git a/python/plugins/processing/gui/Postprocessing.py b/python/plugins/processing/gui/Postprocessing.py
@@ -59,9 +59,8 @@ def handleAlgorithmResults(alg, progress=None, showResults=True):
             continue
         if isinstance(out, (OutputRaster, OutputVector, OutputTable)):
             try:
-                if out.value.startswith('memory:'):
-                    layer = out.memoryLayer
-                    QgsMapLayerRegistry.instance().addMapLayers([layer])
+                if out.layer is not None:
+                    QgsMapLayerRegistry.instance().addMapLayers([out.layer])
                 else:
                     if ProcessingConfig.getSetting(
                             ProcessingConfig.USE_FILENAME_AS_LAYER_NAME):

diff --git a/python/plugins/processing/images/namespace.png b/python/plugins/processing/images/namespace.png
diff --git a/python/plugins/processing/images/postgis.png b/python/plugins/processing/images/postgis.png
diff --git a/python/plugins/processing/tools/vector.py b/python/plugins/processing/tools/vector.py
@@ -16,6 +16,7 @@
 *                                                                         *
 ***************************************************************************
 """
+from processing.algs.qgis import postgis_utils
 
 __author__ = 'Victor Olaya'
 __date__ = 'February 2013'
@@ -33,6 +34,9 @@
 from PyQt4.QtCore import QVariant, QSettings
 from qgis.core import QGis, QgsFields, QgsField, QgsGeometry, QgsRectangle, QgsSpatialIndex, QgsMapLayerRegistry, QgsMapLayer, QgsVectorLayer, QgsVectorFileWriter, QgsDistanceArea
 from processing.core.ProcessingConfig import ProcessingConfig
+from PyQt4 import QtSql
+from processing.core.GeoAlgorithmExecutionException import GeoAlgorithmExecutionException
+from qgis.core import *
 
 
 GEOM_TYPE_MAP = {
@@ -58,6 +62,13 @@
     QVariant.Int: "integer"
 }
 
+TYPE_MAP_POSTGIS_LAYER = {
+    QVariant.String: "VARCHAR",
+    QVariant.Double: "REAL",
+    QVariant.Int: "INTEGER",
+    QVariant.Bool: "BOOLEAN"
+}
+
 
 def features(layer):
     """This returns an iterator over features in a vector layer,
@@ -411,20 +422,21 @@ def bufferedBoundingBox(bbox, buffer_size):
 class VectorWriter:
 
     MEMORY_LAYER_PREFIX = 'memory:'
+    POSTGIS_LAYER_PREFIX = 'postgis:'
 
-    def __init__(self, fileName, encoding, fields, geometryType,
+    def __init__(self, destination, encoding, fields, geometryType,
                  crs, options=None):
-        self.fileName = fileName
-        self.isMemory = False
-        self.memLayer = None
+        self.destination = destination
+        self.isNotFileBased = False
+        self.layer = None
         self.writer = None
 
         if encoding is None:
             settings = QSettings()
             encoding = settings.value('/Processing/encoding', 'System', type=str)
 
-        if self.fileName.startswith(self.MEMORY_LAYER_PREFIX):
-            self.isMemory = True
+        if self.destination.startswith(self.MEMORY_LAYER_PREFIX):
+            self.isNotFileBased = True
 
             uri = GEOM_TYPE_MAP[geometryType] + "?uuid=" + unicode(uuid.uuid4())
             if crs.isValid():
@@ -437,8 +449,44 @@ def __init__(self, fileName, encoding, fields, geometryType,
             if fieldsdesc:
                 uri += '&' + '&'.join(fieldsdesc)
 
-            self.memLayer = QgsVectorLayer(uri, self.fileName, 'memory')
-            self.writer = self.memLayer.dataProvider()
+            self.layer = QgsVectorLayer(uri, self.destination, 'memory')
+            self.writer = self.layer.dataProvider()
+        elif self.destination.startswith(self.POSTGIS_LAYER_PREFIX):
+            self.isNotFileBased = True
+            uri = QgsDataSourceURI(self.destination[len(self.POSTGIS_LAYER_PREFIX):])
+            connInfo = uri.connectionInfo()
+            (success, user, passwd ) = QgsCredentials.instance().get(connInfo, None, None)
+            if success:
+                QgsCredentials.instance().put(connInfo, user, passwd)
+            else:
+                raise GeoAlgorithmExecutionException("Couldn't connect to database")
+            print uri.uri()
+            try:
+                db = postgis_utils.GeoDB(host=uri.host(), port=int(uri.port()),
+                    dbname=uri.database(), user=user, passwd=passwd)
+            except postgis_utils.DbError as e:
+                raise GeoAlgorithmExecutionException(
+                    "Couldn't connect to database:\n%s" % e.message)
+
+            def _runSQL(sql):
+                try:
+                    db._exec_sql_and_commit(unicode(sql))
+                except postgis_utils.DbError as e:
+                    raise GeoAlgorithmExecutionException(
+                        'Error creating output PostGIS table:\n%s' % e.message)
+
+            fields = [_toQgsField(f) for f in fields]
+            fieldsdesc = ",".join('%s %s' % (f.name(),
+                                            TYPE_MAP_POSTGIS_LAYER.get(f.type(), "VARCHAR"))
+                                  for f in fields)
+
+            _runSQL("CREATE TABLE %s.%s (%s)" % (uri.schema(), uri.table().lower(), fieldsdesc))
+            _runSQL("SELECT AddGeometryColumn('{schema}', '{table}', 'the_geom', {srid}, '{typmod}', 2)".format(
+                                    table=uri.table().lower(), schema=uri.schema(), srid=crs.authid().split(":")[-1],
+                                    typmod=GEOM_TYPE_MAP[geometryType].upper()))
+
+            self.layer = QgsVectorLayer(uri.uri(), uri.table(), "postgres")
+            self.writer = self.layer.dataProvider()
         else:
             formats = QgsVectorFileWriter.supportedFiltersAndFormats()
             OGRCodes = {}
@@ -448,21 +496,20 @@ def __init__(self, fileName, encoding, fields, geometryType,
                 extension = extension[:extension.find(' ')]
                 OGRCodes[extension] = value
 
-            extension = self.fileName[self.fileName.rfind('.') + 1:]
+            extension = self.destination[self.destination.rfind('.') + 1:]
             if extension not in OGRCodes:
                 extension = 'shp'
-                self.filename = self.filename + 'shp'
+                self.destination = self.destination + '.shp'
 
             qgsfields = QgsFields()
             for field in fields:
                 qgsfields.append(_toQgsField(field))
 
-            self.writer = QgsVectorFileWriter(
-                self.fileName, encoding,
+            self.writer = QgsVectorFileWriter(self.destination, encoding,
                 qgsfields, geometryType, crs, OGRCodes[extension])
 
     def addFeature(self, feature):
-        if self.isMemory:
+        if self.isNotFileBased:
             self.writer.addFeatures([feature])
         else:
             self.writer.addFeature(feature)

diff --git a/python/plugins/processing/ui/DlgPostgisTableSelector.ui b/python/plugins/processing/ui/DlgPostgisTableSelector.ui
@@ -0,0 +1,64 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<ui version="4.0">
+ <class>Dialog</class>
+ <widget class="QDialog" name="Dialog">
+  <property name="geometry">
+   <rect>
+    <x>0</x>
+    <y>0</y>
+    <width>464</width>
+    <height>395</height>
+   </rect>
+  </property>
+  <property name="windowTitle">
+   <string> output table</string>
+  </property>
+  <layout class="QVBoxLayout" name="verticalLayout">
+   <item>
+    <widget class="QLabel" name="label_3">
+     <property name="text">
+      <string>Select connection and schema</string>
+     </property>
+    </widget>
+   </item>
+   <item>
+    <widget class="QTreeWidget" name="treeConnections">
+     <attribute name="headerVisible">
+      <bool>false</bool>
+     </attribute>
+     <column>
+      <property name="text">
+       <string notr="true">1</string>
+      </property>
+     </column>
+    </widget>
+   </item>
+   <item>
+    <layout class="QHBoxLayout" name="horizontalLayout">
+     <item>
+      <widget class="QLabel" name="label_2">
+       <property name="text">
+        <string>Table name</string>
+       </property>
+      </widget>
+     </item>
+     <item>
+      <widget class="QLineEdit" name="textTableName"/>
+     </item>
+    </layout>
+   </item>
+   <item>
+    <widget class="QDialogButtonBox" name="buttonBox">
+     <property name="orientation">
+      <enum>Qt::Horizontal</enum>
+     </property>
+     <property name="standardButtons">
+      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>
+     </property>
+    </widget>
+   </item>
+  </layout>
+ </widget>
+ <resources/>
+ <connections/>
+</ui>