GPKG/Spatialite: return fid and geom from selects

qgis · Jul 6, 2021 · 2c123dc · 2c123dc
1 parent 7b202fe
commit 2c123dc
Show file tree

Hide file tree

Showing 5 changed files with 162 additions and 18 deletions.
diff --git a/src/core/providers/ogr/qgsgeopackageproviderconnection.cpp b/src/core/providers/ogr/qgsgeopackageproviderconnection.cpp
@@ -13,6 +13,9 @@
  *   (at your option) any later version.                                   *
  *                                                                         *
  ***************************************************************************/
+
+#include <sqlite3.h>
+
 #include "qgsgeopackageproviderconnection.h"
 #include "qgsogrdbconnection.h"
 #include "qgssettings.h"
@@ -25,6 +28,7 @@
 
 #include <QTextCodec>
 
+
 QgsGeoPackageProviderConnection::QgsGeoPackageProviderConnection( const QString &name )
   : QgsAbstractDatabaseProviderConnection( name )
 {
@@ -278,8 +282,9 @@ QList<QgsGeoPackageProviderConnection::TableProperty> QgsGeoPackageProviderConne
       }
 
       QgsGeoPackageProviderConnection::TableProperty property;
-      property.setTableName( row.at( 0 ).toString() );
-      property.setPrimaryKeyColumns( { QStringLiteral( "fid" ) } );
+      const QString tableName { row.at( 0 ).toString() };
+      property.setTableName( tableName );
+      property.setPrimaryKeyColumns( { primaryKeyColumnName( tableName ) } );
       property.setGeometryColumnCount( 0 );
       static const QStringList aspatialTypes = { QStringLiteral( "attributes" ), QStringLiteral( "aspatial" ) };
       const QString dataType = row.at( 1 ).toString();
@@ -405,10 +410,25 @@ QgsAbstractDatabaseProviderConnection::QueryResult QgsGeoPackageProviderConnecti
       if ( fet.reset( OGR_L_GetNextFeature( ogrLayer ) ), fet )
       {
 
+        // pk column name
+        QString pkColumnName;
+
         QgsFields fields { QgsOgrUtils::readOgrFields( fet.get(), QTextCodec::codecForName( "UTF-8" ) ) };
 
-        // pk column name, hardcoded to "fid" (FIXME)
-        QString pkColumnName { QStringLiteral( "fid" ) };
+        // We try to guess the table name from the FROM clause
+        thread_local const QRegularExpression tableNameRegexp { QStringLiteral( R"re((?<=from|join)\s+(\w+)|"([^"]+)")re" ), QRegularExpression::PatternOption::CaseInsensitiveOption };
+        const auto match { tableNameRegexp.match( sql ) };
+        if ( match.hasMatch() )
+        {
+          pkColumnName = primaryKeyColumnName( match.captured( match.lastCapturedIndex() ) );
+        }
+
+        // default to "fid"
+        if ( pkColumnName.isEmpty() )
+        {
+          pkColumnName = QStringLiteral( "fid" );
+        }
+
         // geom column name
         QString geomColumnName;
 
@@ -424,15 +444,17 @@ QgsAbstractDatabaseProviderConnection::QueryResult QgsGeoPackageProviderConnecti
               geomColumnName = OGR_GFld_GetNameRef( geomFldDef );
             }
           }
-
         }
 
-        // May need to prepend PK and append geometries to the columns
-        thread_local const QRegularExpression pkRegExp { QStringLiteral( R"(^select\s+(\*|%1)[,\s+](.*)from)" ).arg( pkColumnName ),  QRegularExpression::PatternOption::CaseInsensitiveOption };
-        if ( pkRegExp.match( sql.trimmed() ).hasMatch() )
+        // May need to prepend PK and append geometry to the columns
+        if ( ! pkColumnName.isEmpty() )
         {
-          iterator->setPrimaryKeyColumnName( pkColumnName );
-          results.appendColumn( pkColumnName );
+          const QRegularExpression pkRegExp { QStringLiteral( R"(^select\s+(\*|%1)[,\s+](.*)from)" ).arg( pkColumnName ),  QRegularExpression::PatternOption::CaseInsensitiveOption };
+          if ( pkRegExp.match( sql.trimmed() ).hasMatch() )
+          {
+            iterator->setPrimaryKeyColumnName( pkColumnName );
+            results.appendColumn( pkColumnName );
+          }
         }
 
         // Add other fields
@@ -487,6 +509,47 @@ QgsAbstractDatabaseProviderConnection::QueryResult QgsGeoPackageProviderConnecti
   return QgsAbstractDatabaseProviderConnection::QueryResult();
 }
 
+QString QgsGeoPackageProviderConnection::primaryKeyColumnName( const QString &table ) const
+{
+  QString pkName;
+
+  sqlite3_database_unique_ptr sqliteHandle;
+  if ( SQLITE_OK == sqliteHandle.open_v2( uri(), SQLITE_OPEN_READONLY, nullptr ) )
+  {
+    char *errMsg;
+
+    const QString sql { QStringLiteral( "PRAGMA table_info(%1)" )
+                        .arg( QgsSqliteUtils::quotedString( table ) )};
+
+    std::vector<std::string> rows;
+    auto cb = [ ](
+                void *data /* Data provided in the 4th argument of sqlite3_exec() */,
+                int /* The number of columns in row */,
+                char **argv /* An array of strings representing fields in the row */,
+                char ** /* An array of strings representing column names */ ) -> int
+    {
+      if ( std::string( argv[5] ).compare( "1" ) == 0 )
+        static_cast<std::vector<std::string>*>( data )->push_back( argv[1] );
+      return 0;
+    };
+
+    // Columns 'cid', 'name', 'type', 'notnull', 'dflt_value', 'pk']
+    const int ret = sqlite3_exec( sqliteHandle.get(), sql.toUtf8(), cb, ( void * )&rows, &errMsg );
+
+    if ( errMsg )
+    {
+      sqlite3_free( errMsg );
+    }
+
+    if ( SQLITE_OK == ret && rows.size() > 0 )
+    {
+      pkName = QString::fromStdString( rows[0] );
+    }
+  }
+
+  return pkName;
+}
+
 QVariantList QgsGeoPackageProviderResultIterator::nextRowPrivate()
 {
   const QVariantList currentRow = mNextRow;
@@ -585,9 +648,12 @@ QgsFields QgsGeoPackageProviderConnection::fields( const QString &schema, const
   // Get fields from layer
   QgsFields fieldList;
 
-  // Prepend PK "fid" hardcoded (FIXME): there might be a way to get the PK name here
-  // but there is probably no way for the general execSql case.
-  fieldList.append( QgsField{ QStringLiteral( "fid" ), QVariant::LongLong } );
+  const QString pkname { primaryKeyColumnName( table ) };
+
+  if ( ! pkname.isEmpty() )
+  {
+    fieldList.append( QgsField{ pkname, QVariant::LongLong } );
+  }
 
   QgsVectorLayer::LayerOptions options { false, true };
   options.skipCrsValidation = true;

diff --git a/src/core/providers/ogr/qgsgeopackageproviderconnection.h b/src/core/providers/ogr/qgsgeopackageproviderconnection.h
@@ -88,6 +88,8 @@ class QgsGeoPackageProviderConnection : public QgsAbstractDatabaseProviderConnec
     void setDefaultCapabilities();
     //! Use GDAL to execute SQL
     QueryResult executeGdalSqlPrivate( const QString &sql, QgsFeedback *feedback = nullptr ) const;
+    //! Returns PK name for table
+    QString primaryKeyColumnName( const QString &table ) const;
 
 
 };

diff --git a/src/providers/spatialite/qgsspatialiteproviderconnection.cpp b/src/providers/spatialite/qgsspatialiteproviderconnection.cpp
@@ -476,28 +476,63 @@ QgsAbstractDatabaseProviderConnection::QueryResult QgsSpatiaLiteProviderConnecti
       if ( fet.reset( OGR_L_GetNextFeature( ogrLayer ) ), fet )
       {
 
-        QgsFields fields { QgsOgrUtils::readOgrFields( fet.get(), QTextCodec::codecForName( "UTF-8" ) ) };
-        iterator->setFields( fields );
+        const QgsFields fields { QgsOgrUtils::readOgrFields( fet.get(), QTextCodec::codecForName( "UTF-8" ) ) };
+
+        // geom column name
+        QString geomColumnName;
+
+        OGRFeatureDefnH featureDef = OGR_F_GetDefnRef( fet.get() );
+
+        if ( featureDef )
+        {
+          if ( OGR_F_GetGeomFieldCount( fet.get() ) > 0 )
+          {
+            OGRGeomFieldDefnH geomFldDef { OGR_F_GetGeomFieldDefnRef( fet.get(), 0 ) };
+            if ( geomFldDef )
+            {
+              geomColumnName = OGR_GFld_GetNameRef( geomFldDef );
+            }
+          }
+        }
 
+        // Add other fields
         for ( const auto &f : std::as_const( fields ) )
         {
           results.appendColumn( f.name() );
         }
+
+        // Append geom
+        if ( ! geomColumnName.isEmpty() )
+        {
+          results.appendColumn( geomColumnName );
+          iterator->setGeometryColumnName( geomColumnName );
+        }
+
+        iterator->setFields( fields );
       }
 
       // Check for errors
-      errCause = CPLGetLastErrorMsg( );
+      if ( CE_Failure == CPLGetLastErrorType() || CE_Fatal == CPLGetLastErrorType() )
+      {
+        errCause = CPLGetLastErrorMsg( );
+      }
 
       if ( ! errCause.isEmpty() )
       {
-        throw QgsProviderConnectionException( QObject::tr( "Error executing SQL %1: %2" ).arg( sql, errCause ) );
+        throw QgsProviderConnectionException( QObject::tr( "Error executing SQL statement %1: %2" ).arg( sql, errCause ) );
       }
 
       OGR_L_ResetReading( ogrLayer );
       iterator->nextRow();
       return results;
     }
-    errCause = CPLGetLastErrorMsg( );
+
+    // Check for errors
+    if ( CE_Failure == CPLGetLastErrorType() || CE_Fatal == CPLGetLastErrorType() )
+    {
+      errCause = CPLGetLastErrorMsg( );
+    }
+
   }
   else
   {
@@ -518,6 +553,7 @@ void QgsSpatialiteProviderResultIterator::setFields( const QgsFields &fields )
   mFields = fields;
 }
 
+
 QgsSpatialiteProviderResultIterator::~QgsSpatialiteProviderResultIterator()
 {
   if ( mHDS )
@@ -549,6 +585,13 @@ QVariantList QgsSpatialiteProviderResultIterator::nextRowInternal()
         {
           row.push_back( attribute );
         }
+
+        // Geom goes last
+        if ( ! mGeometryColumnName.isEmpty( ) )
+        {
+          row.push_back( f.geometry().asWkb() );
+        }
+
       }
       else // Fallback to strings
       {
@@ -568,6 +611,13 @@ QVariantList QgsSpatialiteProviderResultIterator::nextRowInternal()
   return row;
 }
 
+
+void QgsSpatialiteProviderResultIterator::setGeometryColumnName( const QString &geometryColumnName )
+{
+  mGeometryColumnName = geometryColumnName;
+}
+
+
 bool QgsSpatialiteProviderResultIterator::hasNextRowPrivate() const
 {
   return ! mNextRow.isEmpty();

diff --git a/src/providers/spatialite/qgsspatialiteproviderconnection.h b/src/providers/spatialite/qgsspatialiteproviderconnection.h
@@ -34,12 +34,15 @@ struct QgsSpatialiteProviderResultIterator: public QgsAbstractDatabaseProviderCo
 
     void setFields( const QgsFields &fields );
 
+    void setGeometryColumnName( const QString &geometryColumnName );
+
   private:
 
     gdal::ogr_datasource_unique_ptr mHDS;
     OGRLayerH mOgrLayer;
     QgsFields mFields;
     QVariantList mNextRow;
+    QString mGeometryColumnName;
 
     QVariantList nextRowPrivate() override;
     bool hasNextRowPrivate() const override;

diff --git a/tests/src/python/test_qgsproviderconnection_spatialite.py b/tests/src/python/test_qgsproviderconnection_spatialite.py
@@ -157,6 +157,29 @@ def test_create_vector_layer(self):
         self.assertEqual(len(features), 1)
         self.assertEqual(features[0].attributes(), [8, 'Sülfeld'])
 
+    def test_execute_sql_pk_geoms(self):
+        """OGR hides fid and geom from attributes, check if we can still get them"""
+
+        md = QgsProviderRegistry.instance().providerMetadata('spatialite')
+        conn = md.createConnection(self.uri, {})
+
+        # Check errors
+        with self.assertRaises(QgsProviderConnectionException):
+            sql = 'SELECT not_exists, name, geom FROM cdb_lines WHERE name LIKE \'S%\' LIMIT 2'
+            results = conn.executeSql(sql)
+
+        sql = 'SELECT fid, name, geom FROM cdb_lines WHERE name LIKE \'S%\' LIMIT 2'
+        results = conn.executeSql(sql)
+        self.assertEqual(results[0][:2], [8, 'Sülfeld'])
+        self.assertEqual(results[1][:2], [16, 'Steimker Berg'])
+        self.assertEqual(results[0][2][:20], b'\x01\x03\x00\x00\x00\x01\x00\x00\x00/\x00\x00\x00\xf6\x88\x16Y\xad\xb2"')
+        self.assertEqual(results[1][2][:20], b'\x01\x03\x00\x00\x00\x01\x00\x00\x00F\x00\x00\x00 \xc1\x9f\xda\xb4\xfb"')
+
+        sql = 'SELECT name, st_astext(geom) FROM cdb_lines WHERE name LIKE \'S%\' LIMIT 2'
+        results = conn.executeSql(sql)
+        self.assertEqual(results[0], ['Sülfeld',
+                                      'POLYGON((612694.674 5807839.658, 612668.715 5808176.815, 612547.354 5808414.452, 612509.527 5808425.73, 612522.932 5808473.02, 612407.901 5808519.082, 612505.836 5808632.763, 612463.449 5808781.115, 612433.57 5808819.061, 612422.685 5808980.281999, 612473.423 5808995.424999, 612333.856 5809647.731, 612307.316 5809781.446, 612267.099 5809852.803, 612308.221 5810040.995, 613920.397 5811079.478, 613947.16 5811129.3, 614022.726 5811154.456, 614058.436 5811260.36, 614194.037 5811331.972, 614307.176 5811360.06, 614343.842 5811323.238, 614443.449 5811363.03, 614526.199 5811059.031, 614417.83 5811057.603, 614787.296 5809648.422, 614772.062 5809583.246, 614981.93 5809245.35, 614811.885 5809138.271, 615063.452 5809100.954, 615215.476 5809029.413, 615469.441 5808883.282, 615569.846 5808829.522, 615577.239 5808806.242, 615392.964 5808736.873, 615306.34 5808662.171, 615335.445 5808290.588, 615312.192 5808290.397, 614890.582 5808077.956, 615018.854 5807799.895, 614837.326 5807688.363, 614435.698 5807646.847, 614126.351 5807661.841, 613555.813 5807814.801, 612826.66 5807964.828, 612830.113 5807856.315, 612694.674 5807839.658))'])
+
 
 if __name__ == '__main__':
     unittest.main()