Add rowCount to connections API

+ fix memory leak in API fetcher + add progress bar to result widget
qgis · Jul 6, 2021 · 233a809 · 233a809
1 parent 357fa9c
commit 233a809
Show file tree

Hide file tree

Showing 21 changed files with 151 additions and 34 deletions.
diff --git a/python/core/auto_generated/providers/qgsabstractdatabaseproviderconnection.sip.in b/python/core/auto_generated/providers/qgsabstractdatabaseproviderconnection.sip.in
@@ -89,6 +89,13 @@ Returns the number of fetched rows.
 .. seealso:: :py:func:`rowCount`
 %End
 
+        qlonglong rowCount( ) const;
+%Docstring
+Returns the number of rows returned by a SELECT query or -1 if unknown
+
+.. seealso:: :py:func:`fetchedRowCount`
+%End
+
 
         // Python iterator
         QueryResult *__iter__();

diff --git a/src/core/providers/ogr/qgsgeopackageproviderconnection.cpp b/src/core/providers/ogr/qgsgeopackageproviderconnection.cpp
@@ -29,6 +29,8 @@
 #include <QTextCodec>
 
 
+///@cond PRIVATE
+
 QgsGeoPackageProviderConnection::QgsGeoPackageProviderConnection( const QString &name )
   : QgsAbstractDatabaseProviderConnection( name )
 {
@@ -616,6 +618,11 @@ bool QgsGeoPackageProviderResultIterator::hasNextRowPrivate() const
   return ! mNextRow.isEmpty();
 }
 
+qlonglong QgsGeoPackageProviderResultIterator::rowCountPrivate() const
+{
+  return  mRowCount;
+}
+
 void QgsGeoPackageProviderResultIterator::setFields( const QgsFields &fields )
 {
   mFields = fields;
@@ -1153,10 +1160,23 @@ QMap<QgsAbstractDatabaseProviderConnection::SqlKeywordCategory, QStringList> Qgs
   } );
 }
 
+QgsGeoPackageProviderResultIterator::QgsGeoPackageProviderResultIterator( gdal::ogr_datasource_unique_ptr hDS, OGRLayerH ogrLayer )
+  : mHDS( std::move( hDS ) )
+  , mOgrLayer( ogrLayer )
+{
+  if ( mOgrLayer )
+  {
+    // Do not scan the layer!
+    mRowCount = OGR_L_GetFeatureCount( mOgrLayer, false );
+  }
+}
+
 QgsGeoPackageProviderResultIterator::~QgsGeoPackageProviderResultIterator()
 {
   if ( mHDS )
   {
     GDALDatasetReleaseResultSet( mHDS.get(), mOgrLayer );
   }
 }
+
+///@endcond
diff --git a/src/core/providers/ogr/qgsgeopackageproviderconnection.h b/src/core/providers/ogr/qgsgeopackageproviderconnection.h
@@ -28,10 +28,7 @@
 struct QgsGeoPackageProviderResultIterator: public QgsAbstractDatabaseProviderConnection::QueryResult::QueryResultIterator
 {
 
-    QgsGeoPackageProviderResultIterator( gdal::ogr_datasource_unique_ptr hDS, OGRLayerH ogrLayer )
-      : mHDS( std::move( hDS ) )
-      , mOgrLayer( ogrLayer )
-    {}
+    QgsGeoPackageProviderResultIterator( gdal::ogr_datasource_unique_ptr hDS, OGRLayerH ogrLayer );
 
     ~QgsGeoPackageProviderResultIterator();
 
@@ -47,10 +44,12 @@ struct QgsGeoPackageProviderResultIterator: public QgsAbstractDatabaseProviderCo
     QVariantList mNextRow;
     QString mGeometryColumnName;
     QString mPrimaryKeyColumnName;
+    qlonglong mRowCount = -1;
 
     QVariantList nextRowPrivate() override;
-    QVariantList nextRowInternal();
     bool hasNextRowPrivate() const override;
+    qlonglong rowCountPrivate() const override;
+    QVariantList nextRowInternal();
 
 };
 

diff --git a/src/core/providers/qgsabstractdatabaseproviderconnection.cpp b/src/core/providers/qgsabstractdatabaseproviderconnection.cpp
@@ -1438,6 +1438,15 @@ qlonglong QgsAbstractDatabaseProviderConnection::QueryResult::fetchedRowCount()
   return mResultIterator->fetchedRowCount();
 }
 
+qlonglong QgsAbstractDatabaseProviderConnection::QueryResult::rowCount() const
+{
+  if ( ! mResultIterator )
+  {
+    return -1;
+  }
+  return mResultIterator->rowCount();
+}
+
 
 bool QgsAbstractDatabaseProviderConnection::QueryResult::hasNextRow() const
 {
@@ -1482,5 +1491,11 @@ qlonglong QgsAbstractDatabaseProviderConnection::QueryResult::QueryResultIterato
   return mFetchedRowCount;
 }
 
+qlonglong QgsAbstractDatabaseProviderConnection::QueryResult::QueryResultIterator::rowCount()
+{
+  QMutexLocker lock( &mMutex );
+  return rowCountPrivate();
+}
+
 
 ///@endcond private
diff --git a/src/core/providers/qgsabstractdatabaseproviderconnection.h b/src/core/providers/qgsabstractdatabaseproviderconnection.h
@@ -119,6 +119,13 @@ class CORE_EXPORT QgsAbstractDatabaseProviderConnection : public QgsAbstractProv
          */
         qlonglong fetchedRowCount( ) const;
 
+        /**
+         * Returns the number of rows returned by a SELECT query or -1 if unknown
+         *
+         * \see fetchedRowCount()
+         */
+        qlonglong rowCount( ) const;
+
 
 #ifdef SIP_RUN
         // Python iterator
@@ -154,15 +161,26 @@ class CORE_EXPORT QgsAbstractDatabaseProviderConnection : public QgsAbstractProv
          */
         struct CORE_EXPORT QueryResultIterator SIP_SKIP
         {
+            //! Returns the next result row
             QVariantList nextRow();
+
+            //! Returns TRUE if there is another row to fetch
             bool hasNextRow() const;
+
+            //! Returns the number of actually fetched rows
             qlonglong fetchedRowCount();
+
+            //! Returns the total number of rows returned by a SELECT query or -1 if this is not known.
+            qlonglong rowCount();
+
             virtual ~QueryResultIterator() = default;
 
           private:
 
             virtual QVariantList nextRowPrivate() = 0;
             virtual bool hasNextRowPrivate() const = 0;
+            virtual qlonglong rowCountPrivate() const = 0;
+
             mutable qlonglong mFetchedRowCount = 0;
             mutable QMutex mMutex;
 

diff --git a/src/core/qgsqueryresultmodel.cpp b/src/core/qgsqueryresultmodel.cpp
@@ -21,12 +21,12 @@ QgsQueryResultModel::QgsQueryResultModel( const QgsAbstractDatabaseProviderConne
   , mColumns( queryResult.columns() )
 {
   qRegisterMetaType< QList<QList<QVariant>>>( "QList<QList<QVariant>>" );
-  mWorker = new QgsQueryResultFetcher( &mQueryResult );
+  mWorker = std::make_unique<QgsQueryResultFetcher>( &mQueryResult );
   mWorker->moveToThread( &mWorkerThread );
-  connect( &mWorkerThread, &QThread::started, mWorker, &QgsQueryResultFetcher::fetchRows );
-  connect( mWorker, &QgsQueryResultFetcher::rowsReady, this, &QgsQueryResultModel::rowsReady );
-  // Forward signal
-  connect( mWorker, &QgsQueryResultFetcher::fetchingComplete, this, &QgsQueryResultModel::fetchingComplete );
+  connect( &mWorkerThread, &QThread::started, mWorker.get(), &QgsQueryResultFetcher::fetchRows );
+  // Forward signals to the model
+  connect( mWorker.get(), &QgsQueryResultFetcher::rowsReady, this, &QgsQueryResultModel::rowsReady );
+  connect( mWorker.get(), &QgsQueryResultFetcher::fetchingComplete, this, &QgsQueryResultModel::fetchingComplete );
   mWorkerThread.start();
 }
 
@@ -57,7 +57,6 @@ QgsQueryResultModel::~QgsQueryResultModel()
     mWorker->stopFetching();
     mWorkerThread.quit();
     mWorkerThread.wait();
-    mWorker->deleteLater();
   }
   else
   {

diff --git a/src/core/qgsqueryresultmodel.h b/src/core/qgsqueryresultmodel.h
@@ -122,7 +122,7 @@ class CORE_EXPORT QgsQueryResultModel : public QAbstractTableModel
     QgsAbstractDatabaseProviderConnection::QueryResult mQueryResult;
     QStringList mColumns;
     QThread mWorkerThread;
-    QgsQueryResultFetcher *mWorker = nullptr;
+    std::unique_ptr<QgsQueryResultFetcher> mWorker;
     QList<QVariantList> mRows;
 
 };

diff --git a/src/gui/qgsqueryresultwidget.cpp b/src/gui/qgsqueryresultwidget.cpp
@@ -219,7 +219,6 @@ void QgsQueryResultWidget::cancelApiFetcher()
     mApiFetcher->stopFetching();
     mApiFetcherWorkerThread.quit();
     mApiFetcherWorkerThread.wait();
-    mApiFetcherWorkerThread.deleteLater();
   }
 }
 
@@ -250,10 +249,19 @@ void QgsQueryResultWidget::startFetching()
           mQueryResultsTableView->show();
           updateButtons();
           updateSqlLayerColumns( );
+          mActualRowCount = result.rowCount();
+          if ( mActualRowCount != -1 )
+          {
+            mProgressBar->setRange( 0, mActualRowCount );
+          }
         }
         mStatusLabel->setText( tr( "Fetched rows: %1 %2" )
                                .arg( mModel->rowCount( mModel->index( -1, -1 ) ) )
                                .arg( mWasCanceled ? tr( "(stopped)" ) : QString() ) );
+        if ( mActualRowCount != -1 )
+        {
+          mProgressBar->setValue( mModel->rowCount( mModel->index( -1, -1 ) ) );
+        }
       } );
 
       mQueryResultsTableView->setModel( mModel.get() );
@@ -356,6 +364,11 @@ void QgsQueryResultWidget::setConnection( QgsAbstractDatabaseProviderConnection
       mApiFetcher = nullptr;
     } );
     connect( mApiFetcher, &QgsConnectionsApiFetcher::tokensReady, this, &QgsQueryResultWidget::tokensReady );
+    connect( mApiFetcher, &QgsConnectionsApiFetcher::fetchingFinished, &mApiFetcherWorkerThread, [ =  ]
+    {
+      mApiFetcherWorkerThread.quit();
+      mApiFetcherWorkerThread.wait();
+    } );
     mApiFetcherWorkerThread.start();
   }
 

diff --git a/src/gui/qgsqueryresultwidget.h b/src/gui/qgsqueryresultwidget.h
@@ -40,9 +40,9 @@ class GUI_EXPORT QgsConnectionsApiFetcher: public QObject
 
   public:
 
-    //! Constructs a result fetcher from \a queryResult.
-    QgsConnectionsApiFetcher( const QgsAbstractDatabaseProviderConnection *conn )
-      : mConnection( conn )
+    //! Constructs a result fetcher from \a connection.
+    QgsConnectionsApiFetcher( const QgsAbstractDatabaseProviderConnection *connection )
+      : mConnection( connection )
     {}
 
     //! Start fetching
@@ -155,6 +155,7 @@ class GUI_EXPORT QgsQueryResultWidget: public QWidget, private Ui::QgsQueryResul
     bool mFirstRowFetched = false;
     QFutureWatcher<QgsAbstractDatabaseProviderConnection::QueryResult> mQueryResultWatcher;
     QString mSqlErrorMessage;
+    qlonglong mActualRowCount = -1;
 
     /**
      * Updates buttons status.

diff --git a/src/providers/hana/qgshanaproviderconnection.cpp b/src/providers/hana/qgshanaproviderconnection.cpp
@@ -52,6 +52,12 @@ bool QgsHanaProviderResultIterator::hasNextRowPrivate() const
   return mNextRow;
 }
 
+qlonglong QgsHanaProviderResultIterator::rowCountPrivate() const
+{
+  // TODO: hana team, this is for you.
+  return -1;
+}
+
 QgsHanaProviderConnection::QgsHanaProviderConnection( const QString &name )
   : QgsAbstractDatabaseProviderConnection( name )
 {

diff --git a/src/providers/hana/qgshanaproviderconnection.h b/src/providers/hana/qgshanaproviderconnection.h
@@ -25,8 +25,9 @@ struct QgsHanaEmptyProviderResultIterator: public QgsAbstractDatabaseProviderCon
 {
     // QueryResultIterator interface
   private:
-    QVariantList nextRowPrivate();
-    bool hasNextRowPrivate() const;
+    QVariantList nextRowPrivate() override;
+    bool hasNextRowPrivate() const override;
+    qlonglong rowCountPrivate() const override { return 0; };
 };
 
 struct QgsHanaProviderResultIterator: public QgsAbstractDatabaseProviderConnection::QueryResult::QueryResultIterator
@@ -40,8 +41,9 @@ struct QgsHanaProviderResultIterator: public QgsAbstractDatabaseProviderConnecti
 
     // QueryResultIterator interface
   private:
-    QVariantList nextRowPrivate();
-    bool hasNextRowPrivate() const;
+    QVariantList nextRowPrivate() override;
+    bool hasNextRowPrivate() const override;
+    qlonglong rowCountPrivate() const override;
 };
 
 class QgsHanaConnectionRef;

diff --git a/src/providers/mssql/qgsmssqlproviderconnection.cpp b/src/providers/mssql/qgsmssqlproviderconnection.cpp
@@ -331,6 +331,11 @@ QVariantList QgssMssqlProviderResultIterator::nextRowInternal()
   return row;
 }
 
+qlonglong QgssMssqlProviderResultIterator::rowCountPrivate() const
+{
+  return mQuery.size();
+}
+
 
 QList<QgsMssqlProviderConnection::TableProperty> QgsMssqlProviderConnection::tables( const QString &schema, const TableFlags &flags ) const
 {

diff --git a/src/providers/mssql/qgsmssqlproviderconnection.h b/src/providers/mssql/qgsmssqlproviderconnection.h
@@ -39,6 +39,9 @@ struct QgssMssqlProviderResultIterator: public QgsAbstractDatabaseProviderConnec
 
     QVariantList nextRowInternal();
 
+
+    // QueryResultIterator interface
+    qlonglong rowCountPrivate() const override;
 };
 
 class QgsMssqlProviderConnection : public QgsAbstractDatabaseProviderConnection

diff --git a/src/providers/oracle/qgsoracleproviderconnection.cpp b/src/providers/oracle/qgsoracleproviderconnection.cpp
@@ -1407,6 +1407,11 @@ QVariantList QgsOracleProviderResultIterator::nextRowInternal()
   return row;
 }
 
+qlonglong QgsOracleProviderResultIterator::rowCountPrivate() const
+{
+  return mQuery.size();
+}
+
 void QgsOracleProviderConnection::createVectorTable( const QString &schema,
     const QString &name,
     const QgsFields &fields,

diff --git a/src/providers/oracle/qgsoracleproviderconnection.h b/src/providers/oracle/qgsoracleproviderconnection.h
@@ -37,7 +37,7 @@ struct QgsOracleProviderResultIterator: public QgsAbstractDatabaseProviderConnec
     QVariantList mNextRow;
 
     QVariantList nextRowInternal();
-
+    qlonglong rowCountPrivate() const override;
 };
 
 class QgsOracleProviderConnection : public QgsAbstractDatabaseProviderConnection

diff --git a/src/providers/postgres/qgspostgresproviderconnection.cpp b/src/providers/postgres/qgspostgresproviderconnection.cpp
@@ -427,6 +427,11 @@ bool QgsPostgresProviderResultIterator::hasNextRowPrivate() const
   return result && mRowIndex < result->PQntuples();
 }
 
+qlonglong QgsPostgresProviderResultIterator::rowCountPrivate() const
+{
+  return result ? result->PQntuples() : -1;
+}
+
 
 void QgsPostgresProviderConnection::vacuum( const QString &schema, const QString &name ) const
 {

diff --git a/src/providers/postgres/qgspostgresproviderconnection.h b/src/providers/postgres/qgspostgresproviderconnection.h
@@ -32,11 +32,11 @@ struct QgsPostgresProviderResultIterator: public QgsAbstractDatabaseProviderConn
 
     QVariantList nextRowPrivate() override;
     bool hasNextRowPrivate() const override;
+    qlonglong rowCountPrivate() const override;
 
     bool mResolveTypes = true;
     std::shared_ptr<QgsPoolPostgresConn> mConn;
     qlonglong mRowIndex = 0;
-
 };
 
 class QgsPostgresProviderConnection : public QgsAbstractDatabaseProviderConnection

diff --git a/src/providers/spatialite/qgsspatialiteproviderconnection.cpp b/src/providers/spatialite/qgsspatialiteproviderconnection.cpp
@@ -554,6 +554,17 @@ void QgsSpatialiteProviderResultIterator::setFields( const QgsFields &fields )
 }
 
 
+QgsSpatialiteProviderResultIterator::QgsSpatialiteProviderResultIterator( gdal::ogr_datasource_unique_ptr hDS, OGRLayerH ogrLayer )
+  : mHDS( std::move( hDS ) )
+  , mOgrLayer( ogrLayer )
+{
+  if ( mOgrLayer )
+  {
+    // Do not scan the layer!
+    mRowCount = OGR_L_GetFeatureCount( mOgrLayer, false );
+  }
+}
+
 QgsSpatialiteProviderResultIterator::~QgsSpatialiteProviderResultIterator()
 {
   if ( mHDS )
@@ -611,6 +622,11 @@ QVariantList QgsSpatialiteProviderResultIterator::nextRowInternal()
   return row;
 }
 
+qlonglong QgsSpatialiteProviderResultIterator::rowCountPrivate() const
+{
+  return mRowCount;
+}
+
 
 void QgsSpatialiteProviderResultIterator::setGeometryColumnName( const QString &geometryColumnName )
 {