Add relation discovery for SpatiaLite

qgis · Sep 29, 2016 · ecba254 · ecba254
1 parent 31a1c23
commit ecba254
Show file tree

Hide file tree

Showing 3 changed files with 118 additions and 1 deletion.
diff --git a/src/providers/spatialite/qgsspatialiteprovider.cpp b/src/providers/spatialite/qgsspatialiteprovider.cpp
@@ -31,6 +31,7 @@ email                : a.furieri@lqt.it
 #include "qgsspatialitefeatureiterator.h"
 
 #include <qgsjsonutils.h>
+#include <qgsvectorlayer.h>
 
 #include <QMessageBox>
 #include <QFileInfo>
@@ -5264,6 +5265,80 @@ QgsAttributeList QgsSpatiaLiteProvider::pkAttributeIndexes() const
   return mPrimaryKeyAttrs;
 }
 
+QList<QgsVectorLayer*> QgsSpatiaLiteProvider::searchLayers( const QList<QgsVectorLayer*>& layers, const QString& connectionInfo, const QString& tableName )
+{
+  QList<QgsVectorLayer*> result;
+  Q_FOREACH ( QgsVectorLayer* layer, layers )
+  {
+    const QgsSpatiaLiteProvider* slProvider = qobject_cast<QgsSpatiaLiteProvider*>( layer->dataProvider() );
+    if ( slProvider && slProvider->mSqlitePath == connectionInfo && slProvider->mTableName == tableName )
+    {
+      result.append( layer );
+    }
+  }
+  return result;
+}
+
+
+QList<QgsRelation> QgsSpatiaLiteProvider::discoverRelations( const QgsVectorLayer* self, const QList<QgsVectorLayer*>& layers ) const
+{
+  QList<QgsRelation> output;
+  const QString sql = QString( "PRAGMA foreign_key_list(%1)" ).arg( QgsSpatiaLiteProvider::quotedIdentifier( mTableName ) );
+  char **results;
+  int rows;
+  int columns;
+  char *errMsg = nullptr;
+  int ret = sqlite3_get_table( mSqliteHandle, sql.toUtf8().constData(), &results, &rows, &columns, &errMsg );
+  if ( ret == SQLITE_OK )
+  {
+    int nbFound = 0;
+    for ( int row = 1; row <= rows; ++row )
+    {
+      const QString name = "fk_" + mTableName + "_" + QString::fromUtf8( results[row * columns + 0] );
+      const QString position = QString::fromUtf8( results[row * columns + 1] );
+      const QString refTable = QString::fromUtf8( results[row * columns + 2] );
+      const QString fkColumn = QString::fromUtf8( results[row * columns + 3] );
+      const QString refColumn = QString::fromUtf8( results[row * columns + 4] );
+      if ( position == "0" )
+      { // first reference field => try to find if we have layers for the referenced table
+        const QList<QgsVectorLayer*> foundLayers = searchLayers( layers, mSqlitePath, refTable );
+        Q_FOREACH ( const QgsVectorLayer* foundLayer, foundLayers )
+        {
+          QgsRelation relation;
+          relation.setRelationName( name );
+          relation.setReferencingLayer( self->id() );
+          relation.setReferencedLayer( foundLayer->id() );
+          relation.addFieldPair( fkColumn, refColumn );
+          relation.generateId();
+          if ( relation.isValid() )
+          {
+            output.append( relation );
+            ++nbFound;
+          }
+          else
+          {
+            QgsLogger::warning( "Invalid relation for " + name );
+          }
+        }
+      }
+      else
+      { // multi reference field => add the field pair to all the referenced layers found
+        for ( int i = 0; i < nbFound; ++i )
+        {
+          output[output.size() - 1 - i].addFieldPair( fkColumn, refColumn );
+        }
+      }
+    }
+    sqlite3_free_table( results );
+  }
+  else
+  {
+    QgsLogger::warning( QString( "SQLite error discovering relations: %1" ).arg( errMsg ) );
+    sqlite3_free( errMsg );
+  }
+  return output;
+}
+
 // ---------------------------------------------------------------------------
 
 QGISEXTERN bool saveStyle( const QString& uri, const QString& qmlStyle, const QString& sldStyle,

diff --git a/src/providers/spatialite/qgsspatialiteprovider.h b/src/providers/spatialite/qgsspatialiteprovider.h
@@ -209,6 +209,8 @@ class QgsSpatiaLiteProvider: public QgsVectorDataProvider
 
     void invalidateConnections( const QString& connection ) override;
 
+    QList<QgsRelation> discoverRelations( const QgsVectorLayer* self, const QList<QgsVectorLayer*>& layers ) const override;
+
     // static functions
     static void convertToGeosWKB( const unsigned char *blob, int blob_size,
                                   unsigned char **wkb, int *geom_size );
@@ -291,6 +293,11 @@ class QgsSpatiaLiteProvider: public QgsVectorDataProvider
     //! get SpatiaLite version string
     QString spatialiteVersion();
 
+    /**
+     * Search all the layers using the given table.
+     */
+    static QList<QgsVectorLayer*> searchLayers( const QList<QgsVectorLayer*>& layers, const QString& connectionInfo, const QString& tableName );
+
     QgsFields mAttributeFields;
 
     //! Flag indicating if the layer data source is a valid SpatiaLite layer

diff --git a/tests/src/python/test_provider_spatialite.py b/tests/src/python/test_provider_spatialite.py
@@ -19,7 +19,7 @@
 import shutil
 import tempfile
 
-from qgis.core import QgsVectorLayer, QgsPoint, QgsFeature, QgsGeometry
+from qgis.core import QgsVectorLayer, QgsPoint, QgsFeature, QgsGeometry, QgsProject, QgsMapLayerRegistry
 
 from qgis.testing import start_app, unittest
 from utilities import unitTestDataPath
@@ -123,6 +123,18 @@ def setUpClass(cls):
         sql += "VALUES (1, '[\"toto\",\"tutu\"]', '[1,-2,724562]', '[1.0, -232567.22]', GeomFromText('POLYGON((0 0,1 0,1 1,0 1,0 0))', 4326))"
         cur.execute(sql)
 
+        # 2 tables with relations
+        sql = "PRAGMA foreign_keys = ON;"
+        cur.execute(sql)
+        sql = "CREATE TABLE test_relation_a(artistid INTEGER PRIMARY KEY, artistname  TEXT);"
+        cur.execute(sql)
+        sql = "SELECT AddGeometryColumn('test_relation_a', 'Geometry', 4326, 'POLYGON', 'XY')"
+        cur.execute(sql)
+        sql = "CREATE TABLE test_relation_b(trackid INTEGER, trackname TEXT, trackartist INTEGER, FOREIGN KEY(trackartist) REFERENCES test_relation_a(artistid));"
+        cur.execute(sql)
+        sql = "SELECT AddGeometryColumn('test_relation_b', 'Geometry', 4326, 'POLYGON', 'XY')"
+        cur.execute(sql)
+
         cur.execute("COMMIT")
         con.close()
 
@@ -347,6 +359,29 @@ def test_arrays(self):
         self.assertEqual(read_back['ints'], new_f['ints'])
         self.assertEqual(read_back['reals'], new_f['reals'])
 
+    def test_discover_relation(self):
+        artist = QgsVectorLayer("dbname=%s table=test_relation_a (geometry)" % self.dbname, "test_relation_a", "spatialite")
+        self.assertTrue(artist.isValid())
+        track = QgsVectorLayer("dbname=%s table=test_relation_b (geometry)" % self.dbname, "test_relation_b", "spatialite")
+        self.assertTrue(track.isValid())
+        QgsMapLayerRegistry.instance().addMapLayer(artist)
+        QgsMapLayerRegistry.instance().addMapLayer(track)
+        try:
+            relMgr = QgsProject.instance().relationManager()
+            relations = relMgr.discoverRelations([], [artist, track])
+            relations = {r.name(): r for r in relations}
+            self.assertEqual({'fk_test_relation_b_0'}, set(relations.keys()))
+
+            a2t = relations['fk_test_relation_b_0']
+            self.assertTrue(a2t.isValid())
+            self.assertEqual('test_relation_b', a2t.referencingLayer().name())
+            self.assertEqual('test_relation_a', a2t.referencedLayer().name())
+            self.assertEqual([2], a2t.referencingFields())
+            self.assertEqual([0], a2t.referencedFields())
+        finally:
+            QgsMapLayerRegistry.instance().removeMapLayer(track.id())
+            QgsMapLayerRegistry.instance().removeMapLayer(artist.id())
+
     # This test would fail. It would require turning on WAL
     def XXXXXtestLocking(self):