[FEATURE] Create attribute index support for spatialite provider

Allows creation of attribute indexes for spatialite layers
qgis · Dec 5, 2016 · bd8d212 · bd8d212
1 parent e927447
commit bd8d212
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 1 deletion.
diff --git a/src/providers/spatialite/qgsspatialiteprovider.cpp b/src/providers/spatialite/qgsspatialiteprovider.cpp
@@ -538,6 +538,7 @@ QgsSpatiaLiteProvider::QgsSpatiaLiteProvider( QString const &uri )
     mEnabledCapabilities |= QgsVectorDataProvider::ChangeAttributeValues;
     mEnabledCapabilities |= QgsVectorDataProvider::AddFeatures;
     mEnabledCapabilities |= QgsVectorDataProvider::AddAttributes;
+    mEnabledCapabilities |= QgsVectorDataProvider::CreateAttributeIndex;
   }
 
   alreadyDone = false;
@@ -4007,6 +4008,71 @@ bool QgsSpatiaLiteProvider::addFeatures( QgsFeatureList & flist )
   return ret == SQLITE_OK;
 }
 
+QString createIndexName( QString tableName, QString field )
+{
+  QRegularExpression safeExp( QStringLiteral( "[^a-zA-Z0-9]" ) );
+  tableName.replace( safeExp, QStringLiteral( "_" ) );
+  field.replace( safeExp, QStringLiteral( "_" ) );
+  return QStringLiteral( "%1_%2_idx" ).arg( tableName, field );
+}
+
+bool QgsSpatiaLiteProvider::createAttributeIndex( int field )
+{
+  char *errMsg = nullptr;
+  bool toCommit = false;
+
+  if ( field < 0 || field >= mAttributeFields.count() )
+    return false;
+
+  QString sql;
+  QString fieldName;
+
+  int ret = sqlite3_exec( mSqliteHandle, "BEGIN", nullptr, nullptr, &errMsg );
+  if ( ret != SQLITE_OK )
+  {
+    // some error occurred
+    goto abort;
+  }
+  toCommit = true;
+
+  fieldName = mAttributeFields.at( field ).name();
+
+  sql = QStringLiteral( "CREATE INDEX IF NOT EXISTS %1 ON \"%2\" (%3)" )
+        .arg( createIndexName( mTableName, fieldName ),
+              mTableName,
+              quotedIdentifier( fieldName ) );
+  ret = sqlite3_exec( mSqliteHandle, sql.toUtf8().constData(), nullptr, nullptr, &errMsg );
+  if ( ret != SQLITE_OK )
+  {
+    // some error occurred
+    goto abort;
+  }
+
+  ret = sqlite3_exec( mSqliteHandle, "COMMIT", nullptr, nullptr, &errMsg );
+  if ( ret != SQLITE_OK )
+  {
+    // some error occurred
+    goto abort;
+  }
+
+  return true;
+
+abort:
+  pushError( tr( "SQLite error: %2\nSQL: %1" ).arg( sql, errMsg ? errMsg : tr( "unknown cause" ) ) );
+  if ( errMsg )
+  {
+    sqlite3_free( errMsg );
+  }
+
+  if ( toCommit )
+  {
+    // ROLLBACK after some previous error
+    ( void )sqlite3_exec( mSqliteHandle, "ROLLBACK", nullptr, nullptr, nullptr );
+  }
+
+  return false;
+}
+
 bool QgsSpatiaLiteProvider::deleteFeatures( const QgsFeatureIds &id )
 {
   sqlite3_stmt *stmt = nullptr;

diff --git a/src/providers/spatialite/qgsspatialiteprovider.h b/src/providers/spatialite/qgsspatialiteprovider.h
@@ -170,6 +170,8 @@ class QgsSpatiaLiteProvider: public QgsVectorDataProvider
 
     QVariant defaultValue( int fieldId ) const override;
 
+    bool createAttributeIndex( int field ) override;
+
     /** The SpatiaLite provider does its own transforms so we return
      * true for the following three functions to indicate that transforms
      * should not be handled by the QgsCoordinateTransform object. See the

diff --git a/tests/src/python/test_provider_spatialite.py b/tests/src/python/test_provider_spatialite.py
@@ -19,7 +19,16 @@
 import shutil
 import tempfile
 
-from qgis.core import QgsVectorLayer, QgsPoint, QgsFeature, QgsGeometry, QgsProject, QgsMapLayerRegistry, QgsField, QgsFieldConstraints, QgsVectorLayerUtils
+from qgis.core import (QgsVectorLayer,
+                       QgsVectorDataProvider,
+                       QgsPoint,
+                       QgsFeature,
+                       QgsGeometry,
+                       QgsProject,
+                       QgsMapLayerRegistry,
+                       QgsField,
+                       QgsFieldConstraints,
+                       QgsVectorLayerUtils)
 
 from qgis.testing import start_app, unittest
 from utilities import unitTestDataPath
@@ -483,5 +492,40 @@ def testVectorLayerUtilsCreateFeatureWithProviderDefaultLiteral(self):
         f = QgsVectorLayerUtils.createFeature(vl, attributes={1: 'qgis is great', 0: 3})
         self.assertEqual(f.attributes(), [3, "qgis 'is good", 5, 12, None])
 
+    def testCreateAttributeIndex(self):
+        vl = QgsVectorLayer("dbname=%s table='test_defaults' key='id'" % self.dbname, "test_defaults", "spatialite")
+        self.assertTrue(vl.dataProvider().capabilities() & QgsVectorDataProvider.CreateAttributeIndex)
+        self.assertFalse(vl.dataProvider().createAttributeIndex(-1))
+        self.assertFalse(vl.dataProvider().createAttributeIndex(100))
+        self.assertTrue(vl.dataProvider().createAttributeIndex(1))
+
+        con = spatialite_connect(self.dbname, isolation_level=None)
+        cur = con.cursor()
+        rs = cur.execute("SELECT * FROM sqlite_master WHERE type='index' AND tbl_name='test_defaults'")
+        res = [row for row in rs]
+        self.assertEqual(len(res), 1)
+        index_name = res[0][1]
+        rs = cur.execute("PRAGMA index_info({})".format(index_name))
+        res = [row for row in rs]
+        self.assertEqual(len(res), 1)
+        self.assertEqual(res[0][2], 'name')
+
+        # second index
+        self.assertTrue(vl.dataProvider().createAttributeIndex(2))
+        rs = cur.execute("SELECT * FROM sqlite_master WHERE type='index' AND tbl_name='test_defaults'")
+        res = [row for row in rs]
+        self.assertEqual(len(res), 2)
+        indexed_columns = []
+        for row in res:
+            index_name = row[1]
+            rs = cur.execute("PRAGMA index_info({})".format(index_name))
+            res = [row for row in rs]
+            self.assertEqual(len(res), 1)
+            indexed_columns.append(res[0][2])
+
+        self.assertEqual(set(indexed_columns), set(['name', 'number']))
+        con.close()
+
+
 if __name__ == '__main__':
     unittest.main()