Merge pull request #5787 from nyalldawson/index

Fix OGR provider cannot create attribute or spatial indexes for GeoPackage/SQLite layers
qgis · Dec 2, 2017 · b8b8c1a · b8b8c1a
2 parents 31c79da + aaa18e0
commit b8b8c1a
Show file tree

Hide file tree

Showing 3 changed files with 154 additions and 19 deletions.
diff --git a/src/providers/ogr/qgsogrprovider.cpp b/src/providers/ogr/qgsogrprovider.cpp
@@ -1966,36 +1966,74 @@ bool QgsOgrProvider::createSpatialIndex()
   if ( !doInitialActionsForEdition() )
     return false;
 
-  if ( mGDALDriverName != QLatin1String( "ESRI Shapefile" ) )
-    return false;
-
   QByteArray layerName = mOgrOrigLayer->name();
+  if ( mGDALDriverName == QLatin1String( "ESRI Shapefile" ) )
+  {
+    QByteArray sql = QByteArray( "CREATE SPATIAL INDEX ON " ) + quotedIdentifier( layerName );  // quote the layer name so spaces are handled
+    QgsDebugMsg( QString( "SQL: %1" ).arg( QString::fromUtf8( sql ) ) );
+    mOgrOrigLayer->ExecuteSQLNoReturn( sql );
 
-  QByteArray sql = QByteArray( "CREATE SPATIAL INDEX ON " ) + quotedIdentifier( layerName );  // quote the layer name so spaces are handled
-  QgsDebugMsg( QString( "SQL: %1" ).arg( QString::fromUtf8( sql ) ) );
-  mOgrOrigLayer->ExecuteSQLNoReturn( sql );
+    QFileInfo fi( mFilePath );     // to get the base name
+    //find out, if the .qix file is there
+    return QFileInfo::exists( fi.path().append( '/' ).append( fi.completeBaseName() ).append( ".qix" ) );
+  }
+  else if ( mGDALDriverName == QLatin1String( "GPKG" ) ||
+            mGDALDriverName == QLatin1String( "SQLite" ) )
+  {
+    QMutex *mutex = nullptr;
+    OGRLayerH layer = mOgrOrigLayer->getHandleAndMutex( mutex );
+    QByteArray sql = QByteArray( "SELECT CreateSpatialIndex(" + quotedIdentifier( layerName ) + ","
+                                 + quotedIdentifier( OGR_L_GetGeometryColumn( layer ) ) + ") " ); // quote the layer name so spaces are handled
+    mOgrOrigLayer->ExecuteSQLNoReturn( sql );
+    return true;
+  }
+  return false;
+}
 
-  QFileInfo fi( mFilePath );     // to get the base name
-  //find out, if the .qix file is there
-  QFile indexfile( fi.path().append( '/' ).append( fi.completeBaseName() ).append( ".qix" ) );
-  return indexfile.exists();
+QString createIndexName( QString tableName, QString field )
+{
+  QRegularExpression safeExp( QStringLiteral( "[^a-zA-Z0-9]" ) );
+  tableName.replace( safeExp, QStringLiteral( "_" ) );
+  field.replace( safeExp, QStringLiteral( "_" ) );
+  return tableName + "_" + field + "_idx";
 }
 
 bool QgsOgrProvider::createAttributeIndex( int field )
 {
+  if ( field < 0 || field >= mAttributeFields.count() )
+    return false;
+
   if ( !doInitialActionsForEdition() )
     return false;
 
   QByteArray quotedLayerName = quotedIdentifier( mOgrOrigLayer->name() );
-  QByteArray dropSql = "DROP INDEX ON " + quotedLayerName;
-  mOgrOrigLayer->ExecuteSQLNoReturn( dropSql );
-  QByteArray createSql = "CREATE INDEX ON " + quotedLayerName + " USING " + textEncoding()->fromUnicode( fields().at( field ).name() );
-  mOgrOrigLayer->ExecuteSQLNoReturn( createSql );
+  if ( mGDALDriverName == QLatin1String( "GPKG" ) ||
+       mGDALDriverName == QLatin1String( "SQLite" ) )
+  {
+    if ( field == 0 && mFirstFieldIsFid )
+    {
+      // already an index on this field, no need to re-created
+      return false;
+    }
+
+    QString indexName = createIndexName( mOgrOrigLayer->name(), fields().at( field ).name() );
+    QByteArray createSql = "CREATE INDEX IF NOT EXISTS " + textEncoding()->fromUnicode( indexName ) + " ON " + quotedLayerName + " (" + textEncoding()->fromUnicode( fields().at( field ).name() ) + ")";
+    mOgrOrigLayer->ExecuteSQLNoReturn( createSql );
+    return true;
+  }
+  else
+  {
+    QByteArray dropSql = "DROP INDEX ON " + quotedLayerName;
+    mOgrOrigLayer->ExecuteSQLNoReturn( dropSql );
+    QByteArray createSql = "CREATE INDEX ON " + quotedLayerName + " USING " + textEncoding()->fromUnicode( fields().at( field ).name() );
+    mOgrOrigLayer->ExecuteSQLNoReturn( createSql );
 
-  QFileInfo fi( mFilePath );     // to get the base name
-  //find out, if the .idm file is there
-  QFile indexfile( fi.path().append( '/' ).append( fi.completeBaseName() ).append( ".idm" ) );
-  return indexfile.exists();
+    QFileInfo fi( mFilePath );     // to get the base name
+    //find out, if the .idm/.ind file is there
+    QString idmFile( fi.path().append( '/' ).append( fi.completeBaseName() ).append( ".idm" ) );
+    QString indFile( fi.path().append( '/' ).append( fi.completeBaseName() ).append( ".ind" ) );
+    return QFile::exists( idmFile ) || QFile::exists( indFile );
+  }
 }
 
 bool QgsOgrProvider::deleteFeatures( const QgsFeatureIds &id )
@@ -2241,6 +2279,12 @@ void QgsOgrProvider::computeCapabilities()
         ability &= ~( AddAttributes | DeleteFeatures );
       }
     }
+    else if ( mGDALDriverName == QLatin1String( "GPKG" ) ||
+              mGDALDriverName == QLatin1String( "SQLite" ) )
+    {
+      ability |= CreateSpatialIndex;
+      ability |= CreateAttributeIndex;
+    }
 
     /* Curve geometries are available in some drivers starting with GDAL 2.0 */
     if ( mOgrLayer->TestCapability( "CurveGeometries" ) )

diff --git a/tests/src/python/test_provider_ogr_gpkg.py b/tests/src/python/test_provider_ogr_gpkg.py
@@ -34,9 +34,11 @@
                        QgsPointXY,
                        QgsProject,
                        QgsWkbTypes,
-                       QgsDataProvider)
+                       QgsDataProvider,
+                       QgsVectorDataProvider)
 from qgis.PyQt.QtCore import QCoreApplication, QVariant
 from qgis.testing import start_app, unittest
+from qgis.utils import spatialite_connect
 
 
 def GDAL_COMPUTE_VERSION(maj, min, rev):
@@ -771,6 +773,67 @@ def test_SplitFeature(self):
         self.assertEqual([f for f in layer.getFeatures()][0].geometry().asWkt(), 'Polygon ((0.5 0, 0.5 1, 1 1, 1 0, 0.5 0))')
         self.assertEqual([f for f in layer.getFeatures()][1].geometry().asWkt(), 'Polygon ((0.5 1, 0.5 0, 0 0, 0 1, 0.5 1))')
 
+    def testCreateAttributeIndex(self):
+        tmpfile = os.path.join(self.basetestpath, 'testGeopackageAttributeIndex.gpkg')
+        ds = ogr.GetDriverByName('GPKG').CreateDataSource(tmpfile)
+        lyr = ds.CreateLayer('test', geom_type=ogr.wkbPolygon)
+        lyr.CreateField(ogr.FieldDefn('str_field', ogr.OFTString))
+        lyr.CreateField(ogr.FieldDefn('str_field2', ogr.OFTString))
+        f = None
+        ds = None
+
+        vl = QgsVectorLayer(u'{}'.format(tmpfile) + "|layername=" + "test", 'test', u'ogr')
+        self.assertTrue(vl.isValid())
+        self.assertTrue(vl.dataProvider().capabilities() & QgsVectorDataProvider.CreateAttributeIndex)
+        self.assertFalse(vl.dataProvider().createAttributeIndex(-1))
+        self.assertFalse(vl.dataProvider().createAttributeIndex(100))
+
+        # should not be allowed - there's already a index on the primary key
+        self.assertFalse(vl.dataProvider().createAttributeIndex(0))
+
+        self.assertTrue(vl.dataProvider().createAttributeIndex(1))
+
+        con = spatialite_connect(tmpfile, isolation_level=None)
+        cur = con.cursor()
+        rs = cur.execute("SELECT * FROM sqlite_master WHERE type='index' AND tbl_name='test'")
+        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], 'str_field')
+
+        # second index
+        self.assertTrue(vl.dataProvider().createAttributeIndex(2))
+        rs = cur.execute("SELECT * FROM sqlite_master WHERE type='index' AND tbl_name='test'")
+        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.assertCountEqual(indexed_columns, ['str_field', 'str_field2'])
+        con.close()
+
+    def testCreateSpatialIndex(self):
+        tmpfile = os.path.join(self.basetestpath, 'testGeopackageSpatialIndex.gpkg')
+        ds = ogr.GetDriverByName('GPKG').CreateDataSource(tmpfile)
+        lyr = ds.CreateLayer('test', geom_type=ogr.wkbPolygon, options=['SPATIAL_INDEX=NO'])
+        lyr.CreateField(ogr.FieldDefn('str_field', ogr.OFTString))
+        lyr.CreateField(ogr.FieldDefn('str_field2', ogr.OFTString))
+        f = None
+        ds = None
+
+        vl = QgsVectorLayer(u'{}'.format(tmpfile) + "|layername=" + "test", 'test', u'ogr')
+        self.assertTrue(vl.isValid())
+        self.assertTrue(vl.dataProvider().capabilities() & QgsVectorDataProvider.CreateSpatialIndex)
+        self.assertTrue(vl.dataProvider().createSpatialIndex())
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/src/python/test_provider_shapefile.py b/tests/src/python/test_provider_shapefile.py
@@ -584,6 +584,34 @@ def testOpenWithFilter(self):
             # force close of data provider
             vl.setDataSource('', 'test', 'ogr')
 
+    def testCreateAttributeIndex(self):
+        tmpdir = tempfile.mkdtemp()
+        self.dirs_to_cleanup.append(tmpdir)
+        srcpath = os.path.join(TEST_DATA_DIR, 'provider')
+        for file in glob.glob(os.path.join(srcpath, 'shapefile.*')):
+            shutil.copy(os.path.join(srcpath, file), tmpdir)
+        datasource = os.path.join(tmpdir, 'shapefile.shp')
+
+        vl = QgsVectorLayer('{}|layerid=0'.format(datasource), 'test', 'ogr')
+        self.assertTrue(vl.isValid())
+        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))
+
+    def testCreateSpatialIndex(self):
+        tmpdir = tempfile.mkdtemp()
+        self.dirs_to_cleanup.append(tmpdir)
+        srcpath = os.path.join(TEST_DATA_DIR, 'provider')
+        for file in glob.glob(os.path.join(srcpath, 'shapefile.*')):
+            shutil.copy(os.path.join(srcpath, file), tmpdir)
+        datasource = os.path.join(tmpdir, 'shapefile.shp')
+
+        vl = QgsVectorLayer('{}|layerid=0'.format(datasource), 'test', 'ogr')
+        self.assertTrue(vl.isValid())
+        self.assertTrue(vl.dataProvider().capabilities() & QgsVectorDataProvider.CreateSpatialIndex)
+        self.assertTrue(vl.dataProvider().createSpatialIndex())
+
 
 if __name__ == '__main__':
     unittest.main()