Merge pull request #2423 from rldhont/processing_spatialite_output

[processing] allow output directly on Spatialite tables

python/plugins/processing/algs/qgis/spatialite_utils.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+
+"""
+***************************************************************************
+    spatialite_utils.py
+    ---------------------
+    Date                 : November 2015
+    Copyright            : (C) 2015 by René-Luc Dhont
+    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__ = 'René-Luc Dhont'
+__date__ = 'November 2015'
+__copyright__ = '(C) 2015, René-Luc Dhont'
+
+# This will get replaced with a git SHA1 when you do a git archive
+
+__revision__ = '$Format:%H$'
+
+from pyspatialite import dbapi2 as sqlite
+
+class DbError(Exception):
+
+    def __init__(self, message, query=None):
+        # Save error. funny that the variables are in utf-8
+        self.message = unicode(message, 'utf-8')
+        self.query = (unicode(query, 'utf-8') if query is not None else None)
+
+    def __str__(self):
+        return 'MESSAGE: %s\nQUERY: %s' % (self.message, self.query)
+
+class GeoDB:
+
+    def __init__(self, uri=None):
+        self.uri = uri
+        self.dbname = uri.database()
+
+        try:
+            self.con = sqlite.connect(self.con_info())
+
+        except (sqlite.InterfaceError, sqlite.OperationalError) as e:
+            raise DbError(e.message)
+
+        self.has_spatialite = self.check_spatialite()
+        if not self.has_spatialite:
+            self.has_spatialite = self.init_spatialite()
+
+    def con_info(self):
+        return unicode(self.dbname)
+
+    def init_spatialite(self):
+        # Get spatialite version
+        c = self.con.cursor()
+        try:
+            self._exec_sql(c, u'SELECT spatialite_version()')
+            rep = c.fetchall()
+            v = [int(a) for a in rep[0][0].split('.')]
+            vv = v[0] * 100000 + v[1] * 1000 + v[2] * 10
+
+            # Add spatialite support
+            if vv >= 401000:
+                # 4.1 and above
+                sql = "SELECT initspatialmetadata(1)"
+            else:
+                # Under 4.1
+                sql = "SELECT initspatialmetadata()"
+            self._exec_sql_and_commit(sql)
+        except:
+            return False
+        finally:
+            self.con.close()
+
+        try:
+            self.con = sqlite.connect(self.con_info())
+
+        except (sqlite.InterfaceError, sqlite.OperationalError) as e:
+            raise DbError(e.message)
+
+        return self.check_spatialite()
+
+    def check_spatialite(self):
+        try:
+            c = self.con.cursor()
+            self._exec_sql(c, u"SELECT CheckSpatialMetaData()")
+            v = c.fetchone()[0]
+            self.has_geometry_columns = v == 1 or v == 3
+            self.has_spatialite4 = v == 3
+        except Exception as e:
+            self.has_geometry_columns = False
+            self.has_spatialite4 = False
+
+        self.has_geometry_columns_access = self.has_geometry_columns
+        return self.has_geometry_columns
+
+    def _exec_sql(self, cursor, sql):
+        try:
+            cursor.execute(sql)
+        except (sqlite.Error, sqlite.ProgrammingError, sqlite.Warning, sqlite.InterfaceError, sqlite.OperationalError) as e:
+            raise DbError(e.message, sql)
+
+    def _exec_sql_and_commit(self, sql):
+        """Tries to execute and commit some action, on error it rolls
+        back the change.
+        """
+
+        try:
+            c = self.con.cursor()
+            self._exec_sql(c, sql)
+            self.con.commit()
+        except DbError as e:
+            self.con.rollback()
+            raise

python/plugins/processing/gui/OutputSelectionPanel.py

@@ -83,6 +83,10 @@ def selectOutput(self):
                     self.tr('Save to memory layer'), self.btnSelect)
                 actionSaveToMemory.triggered.connect(self.saveToMemory)
                 popupMenu.addAction(actionSaveToMemory)
+                actionSaveToSpatialite = QAction(
+                    self.tr('Save to Spatialite table...'), self.btnSelect)
+                actionSaveToSpatialite.triggered.connect(self.saveToSpatialite)
+                popupMenu.addAction(actionSaveToSpatialite)
                 actionSaveToPostGIS = QAction(
                     self.tr('Save to PostGIS table...'), self.btnSelect)
                 actionSaveToPostGIS.triggered.connect(self.saveToPostGIS)
@@ -118,6 +122,42 @@ def saveToPostGIS(self):
                 QgsCredentials.instance().put(connInfo, user, passwd)
             self.leText.setText("postgis:" + uri.uri())
 
+    def saveToSpatialite(self):
+        fileFilter = self.output.tr('Spatialite files(*.sqlite)', 'OutputFile')
+
+        settings = QSettings()
+        if settings.contains('/Processing/LastOutputPath'):
+            path = settings.value('/Processing/LastOutputPath')
+        else:
+            path = ProcessingConfig.getSetting(ProcessingConfig.OUTPUT_FOLDER)
+
+        encoding = settings.value('/Processing/encoding', 'System')
+        fileDialog = QgsEncodingFileDialog(
+            self, self.tr('Save Spatialite'), path, fileFilter, encoding)
+        fileDialog.setFileMode(QFileDialog.AnyFile)
+        fileDialog.setAcceptMode(QFileDialog.AcceptSave)
+        fileDialog.setConfirmOverwrite(False)
+
+        if fileDialog.exec_() == QDialog.Accepted:
+            files = fileDialog.selectedFiles()
+            encoding = unicode(fileDialog.encoding())
+            self.output.encoding = encoding
+            fileName = unicode(files[0])
+            selectedFileFilter = unicode(fileDialog.selectedNameFilter())
+            if not fileName.lower().endswith(
+                    tuple(re.findall("\*(\.[a-z]{1,10})", fileFilter))):
+                ext = re.search("\*(\.[a-z]{1,10})", selectedFileFilter)
+                if ext:
+                    fileName += ext.group(1)
+            settings.setValue('/Processing/LastOutputPath',
+                              os.path.dirname(fileName))
+            settings.setValue('/Processing/encoding', encoding)
+
+            uri = QgsDataSourceURI()
+            uri.setDatabase(fileName)
+            uri.setDataSource('', self.output.name.lower(), 'the_geom')
+            self.leText.setText("spatialite:" + uri.uri())
+
     def saveToMemory(self):
         self.leText.setText('memory:')
 
@@ -167,6 +207,8 @@ def getValue(self):
             value = fileName
         elif fileName.startswith('postgis:'):
             value = fileName
+        elif fileName.startswith('spatialite:'):
+            value = fileName
         elif not os.path.isabs(fileName):
             value = ProcessingConfig.getSetting(
                 ProcessingConfig.OUTPUT_FOLDER) + os.sep + fileName

python/plugins/processing/tools/vector.py

@@ -17,6 +17,7 @@
 ***************************************************************************
 """
 from processing.algs.qgis import postgis_utils
+from processing.algs.qgis import spatialite_utils
 
 __author__ = 'Victor Olaya'
 __date__ = 'February 2013'
@@ -69,6 +70,13 @@
     QVariant.Bool: "BOOLEAN"
 }
 
+TYPE_MAP_SPATIALITE_LAYER = {
+    QVariant.String: "VARCHAR",
+    QVariant.Double: "REAL",
+    QVariant.Int: "INTEGER",
+    QVariant.Bool: "INTEGER"
+}
+
 
 def features(layer):
     """This returns an iterator over features in a vector layer,
@@ -423,6 +431,7 @@ class VectorWriter:
 
     MEMORY_LAYER_PREFIX = 'memory:'
     POSTGIS_LAYER_PREFIX = 'postgis:'
+    SPATIALITE_LAYER_PREFIX = 'spatialite:'
 
     def __init__(self, destination, encoding, fields, geometryType,
                  crs, options=None):
@@ -485,7 +494,37 @@ def _runSQL(sql):
                                     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.layer = QgsVectorLayer(uri.uri(), uri.table(), "spatialite")
+            self.writer = self.layer.dataProvider()
+        elif self.destination.startswith(self.SPATIALITE_LAYER_PREFIX):
+            self.isNotFileBased = True
+            uri = QgsDataSourceURI(self.destination[len(self.SPATIALITE_LAYER_PREFIX):])
+            print uri.uri()
+            try:
+                db = spatialite_utils.GeoDB(uri=uri)
+            except spatialite_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 spatialite_utils.DbError as e:
+                    raise GeoAlgorithmExecutionException(
+                        'Error creating output Spatialite table:\n%s' % e.message)
+
+            fields = [_toQgsField(f) for f in fields]
+            fieldsdesc = ",".join('%s %s' % (f.name(),
+                                            TYPE_MAP_SPATIALITE_LAYER.get(f.type(), "VARCHAR"))
+                                  for f in fields)
+
+            _runSQL("DROP TABLE IF EXISTS %s" % uri.table().lower())
+            _runSQL("CREATE TABLE %s (%s)" % (uri.table().lower(), fieldsdesc))
+            _runSQL("SELECT AddGeometryColumn('{table}', 'the_geom', {srid}, '{typmod}', 2)".format(
+                                    table=uri.table().lower(), srid=crs.authid().split(":")[-1],
+                                    typmod=GEOM_TYPE_MAP[geometryType].upper()))
+
+            self.layer = QgsVectorLayer(uri.uri(), uri.table(), "spatialite")
             self.writer = self.layer.dataProvider()
         else:
             formats = QgsVectorFileWriter.supportedFiltersAndFormats()