Add DB connections factory to data items

(cherry picked from commit 5733970)
qgis · Feb 19, 2021 · 8e39816 · 8e39816
1 parent 7b3c4f3
commit 8e39816
Show file tree

Hide file tree

Showing 7 changed files with 181 additions and 38 deletions.
diff --git a/python/core/auto_generated/qgsdataitem.sip.in b/python/core/auto_generated/qgsdataitem.sip.in
@@ -405,6 +405,17 @@ Sets a custom sorting ``key`` for the item.
     void moveToThread( QThread *targetThread );
 %Docstring
 Move object and all its descendants to thread
+%End
+
+    virtual QgsAbstractDatabaseProviderConnection *databaseConnection() const /Factory/;
+%Docstring
+For data items that represent a DB connection or one of its children,
+this method returns a connection.
+All other data items will return NULL.
+
+Ownership of the returned objects is transferred to the caller.
+
+.. versionadded:: 3.16
 %End
 
   protected:
@@ -618,6 +629,7 @@ Returns icon for vector tile layer
 
 :return: the layer name
 %End
+
 };
 
 
@@ -663,6 +675,9 @@ Returns the standard browser data collection icon.
 .. seealso:: :py:func:`iconDir`
 %End
 
+    virtual QgsAbstractDatabaseProviderConnection *databaseConnection() const;
+
+
   protected:
 
     static QIcon openDirIcon();
@@ -678,6 +693,7 @@ Shared home directory icon.
 
 .. versionadded:: 3.4
 %End
+
 };
 
 
@@ -709,6 +725,8 @@ The optional ``providerKey`` string can be used to specify the key for the QgsDa
 
     ~QgsDatabaseSchemaItem();
 
+    virtual QgsAbstractDatabaseProviderConnection *databaseConnection() const;
+
 
     static QIcon iconDataCollection();
 %Docstring

diff --git a/src/app/browser/qgsinbuiltdataitemproviders.cpp b/src/app/browser/qgsinbuiltdataitemproviders.cpp
@@ -878,57 +878,22 @@ void QgsDatabaseItemGuiProvider::populateContextMenu( QgsDataItem *item, QMenu *
 {
   Q_UNUSED( selectedItems )
 
-  // Small helper to get a connection from either a connection name (if it's stored) or a data item path
-  auto getConnectionFromDataItem = [ ]( QgsDataCollectionItem * collectionItem, bool isSchema ) -> QgsAbstractDatabaseProviderConnection *
-  {
-    // This is super messy: we need the QgsDataProvider key and NOT the QgsDataItemProvider key!
-    const QString dataProviderKey { QgsApplication::dataItemProviderRegistry()->dataProviderKey( collectionItem->providerKey() ) };
-    QgsProviderMetadata *md { QgsProviderRegistry::instance()->providerMetadata( dataProviderKey ) };
-    if ( ! md )
-      return nullptr;
-
-    const QString connectionName { isSchema ? collectionItem->parent()->name() : collectionItem->name() };
-    QString connectionPath { isSchema ? collectionItem->parent()->path() : collectionItem->path() };
-    // The path is normally prefixed with the data item provider key (which it NOT necessarily the data provider name)
-    // such as "gpkg://<path_to_gpkg>": remove it.
-    connectionPath.remove( QRegularExpression( QStringLiteral( R"re(^[aZ]{2,}://)re" ) ) );
-    try
-    {
-      // If we are coming here from the browser, there is no stored connection
-      if ( md->findConnection( connectionName ) )
-      {
-        return static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( connectionName ) );
-      }
-      else
-      {
-        return static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( connectionPath, {} ) );
-      }
-    }
-    catch ( QgsProviderConnectionException &ex )
-    {
-      // This is expected and it is not an error in case the provider does not implement the connections API
-    }
-    return nullptr;
-  };
-
   // Add create new table for collection items but not not if it is a root item
   if ( ! qobject_cast<QgsConnectionsRootItem *>( item ) )
   {
     if ( QgsDataCollectionItem * collectionItem { qobject_cast<QgsDataCollectionItem *>( item ) } )
     {
-      // Note: we could have used layerCollection() but casting to QgsDatabaseSchemaItem is more explicit
-      const bool isSchema { qobject_cast<QgsDatabaseSchemaItem *>( collectionItem ) != nullptr };
-      std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn( getConnectionFromDataItem( collectionItem, isSchema ) );
+      std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn( item->databaseConnection() );
 
       if ( conn && conn->capabilities().testFlag( QgsAbstractDatabaseProviderConnection::Capability::CreateVectorTable ) )
       {
 
         QAction *newTableAction = new QAction( QObject::tr( "New Table…" ), menu );
 
-        QObject::connect( newTableAction, &QAction::triggered, collectionItem, [ &getConnectionFromDataItem, isSchema, collectionItem, context]
+        QObject::connect( newTableAction, &QAction::triggered, collectionItem, [ collectionItem, context]
         {
 
-          std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn2( getConnectionFromDataItem( collectionItem, isSchema ) );
+          std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn2( collectionItem->databaseConnection() );
           // This should never happen but let's play safe
           if ( ! conn2 )
           {
@@ -939,6 +904,8 @@ void QgsDatabaseItemGuiProvider::populateContextMenu( QgsDataItem *item, QMenu *
           QgsNewVectorTableDialog dlg { conn2.get(), nullptr };
           dlg.setCrs( QgsProject::instance()->defaultCrsForNewLayers() );
 
+          const bool isSchema { qobject_cast<QgsDatabaseSchemaItem *>( collectionItem ) != nullptr };
+
           if ( isSchema )
           {
             dlg.setSchemaName( collectionItem->name() );

diff --git a/src/core/providers/ogr/qgsogrdataitems.cpp b/src/core/providers/ogr/qgsogrdataitems.cpp
@@ -461,6 +461,60 @@ QgsMimeDataUtils::Uri QgsOgrDataCollectionItem::mimeUri() const
   return u;
 }
 
+QgsAbstractDatabaseProviderConnection *QgsOgrDataCollectionItem::databaseConnection() const
+{
+
+  QgsAbstractDatabaseProviderConnection *conn { QgsDataCollectionItem::databaseConnection() };
+
+  // There is a chance that this is a spatialite file, but spatialite is not handled by OGR and
+  // it's not even in core.
+  if ( ! conn )
+  {
+
+    // test that file is valid with OGR
+    if ( OGRGetDriverCount() == 0 )
+    {
+      OGRRegisterAll();
+    }
+    // do not print errors, but write to debug
+    CPLPushErrorHandler( CPLQuietErrorHandler );
+    CPLErrorReset();
+    gdal::dataset_unique_ptr hDS( GDALOpenEx( path().toUtf8().constData(), GDAL_OF_VECTOR, nullptr, nullptr, nullptr ) );
+    CPLPopErrorHandler();
+
+    if ( ! hDS )
+    {
+      QgsDebugMsgLevel( QStringLiteral( "GDALOpen error # %1 : %2 on %3" ).arg( CPLGetLastErrorNo() ).arg( CPLGetLastErrorMsg() ).arg( path() ), 2 );
+      return nullptr;
+    }
+
+    GDALDriverH hDriver = GDALGetDatasetDriver( hDS.get() );
+    QString driverName = GDALGetDriverShortName( hDriver );
+
+    if ( driverName == QLatin1String( "SQLite" ) )
+    {
+      QgsProviderMetadata *md { QgsProviderRegistry::instance()->providerMetadata( QStringLiteral( "spatialite" ) ) };
+      if ( md )
+      {
+        QgsDataSourceUri uri;
+        uri.setDatabase( path( ) );
+        conn = static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( uri.uri(), {} ) );
+      }
+    }
+  }
+  return conn;
+}
+
+void QgsOgrDataCollectionItem::setDriverName( const QString &driverName )
+{
+  mDriverName = driverName;
+}
+
+QString QgsOgrDataCollectionItem::driverName() const
+{
+  return mDriverName;
+}
+
 // ---------------------------------------------------------------------------
 
 QString QgsOgrDataItemProvider::name()
@@ -720,6 +774,7 @@ QgsDataItem *QgsOgrDataItemProvider::createDataItem( const QString &pathIn, QgsD
   else if ( numLayers > 1 || sOgrSupportedDbDriverNames.contains( driverName ) )
   {
     item = new QgsOgrDataCollectionItem( parentItem, name, path );
+    qobject_cast<QgsOgrDataCollectionItem *>( item )->setDriverName( driverName );
   }
   else
   {

diff --git a/src/core/providers/ogr/qgsogrdataitems.h b/src/core/providers/ogr/qgsogrdataitems.h
@@ -132,6 +132,24 @@ class CORE_EXPORT QgsOgrDataCollectionItem final: public QgsDataCollectionItem
 
     bool hasDragEnabled() const override;
     QgsMimeDataUtils::Uri mimeUri() const override;
+
+    QgsAbstractDatabaseProviderConnection *databaseConnection() const override;
+
+    /**
+     * Sets GDAL driver name.
+     * \since QGIS 3.18
+     */
+    void setDriverName( const QString &driverName );
+
+    /**
+     * Returns GDAL driver name.
+     * \since QGIS 3.18
+     */
+    QString driverName() const;
+
+  private:
+
+    QString mDriverName;
 };
 
 //! Provider for OGR root data item

diff --git a/src/core/qgsdataitem.cpp b/src/core/qgsdataitem.cpp
@@ -106,6 +106,44 @@ QIcon QgsDataCollectionItem::homeDirIcon()
   return QgsApplication::getThemeIcon( QStringLiteral( "mIconFolderHome.svg" ) );
 }
 
+QgsAbstractDatabaseProviderConnection *QgsDataCollectionItem::databaseConnection() const
+{
+  const QString dataProviderKey { QgsApplication::dataItemProviderRegistry()->dataProviderKey( providerKey() ) };
+  QgsProviderMetadata *md { QgsProviderRegistry::instance()->providerMetadata( dataProviderKey ) };
+
+  if ( ! md )
+  {
+    return nullptr;
+  }
+
+  const QString connectionName { name() };
+
+  try
+  {
+    // First try to retrieve the connection by name if this is a stored connection
+    if ( md->findConnection( connectionName ) )
+    {
+      return static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( connectionName ) );
+    }
+
+    // If that fails, try to create a connection from the path, in case this is a
+    // filesystem-based DB (gpkg or spatialite)
+    // The name is useless, we need to get the file path from the data item path
+    const QString databaseFilePath { path().remove( QRegularExpression( R"re([\aZ]{2,}://)re" ) ) };
+
+    if ( QFile::exists( databaseFilePath ) )
+    {
+      return static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( databaseFilePath, {} ) );
+    }
+  }
+  catch ( QgsProviderConnectionException &ex )
+  {
+    // This is expected and it is not an error in case the provider does not implement
+    // the connections API
+  }
+  return nullptr;
+}
+
 QIcon QgsDataCollectionItem::iconDir()
 {
   return QgsApplication::getThemeIcon( QStringLiteral( "/mIconFolder.svg" ) );
@@ -398,6 +436,11 @@ void QgsDataItem::moveToThread( QThread *targetThread )
   QObject::moveToThread( targetThread );
 }
 
+QgsAbstractDatabaseProviderConnection *QgsDataItem::databaseConnection() const
+{
+  return nullptr;
+}
+
 QIcon QgsDataItem::icon()
 {
   if ( state() == Populating && sPopulatingIcon )
@@ -1812,6 +1855,27 @@ QIcon QgsDatabaseSchemaItem::iconDataCollection()
   return QgsApplication::getThemeIcon( QStringLiteral( "/mIconDbSchema.svg" ) );
 }
 
+QgsAbstractDatabaseProviderConnection *QgsDatabaseSchemaItem::databaseConnection() const
+{
+  const QString dataProviderKey { QgsApplication::dataItemProviderRegistry()->dataProviderKey( providerKey() ) };
+  QgsProviderMetadata *md { QgsProviderRegistry::instance()->providerMetadata( dataProviderKey ) };
+  if ( ! md )
+  {
+    return nullptr;
+  }
+  const QString connectionName { parent()->name() };
+  try
+  {
+    return static_cast<QgsAbstractDatabaseProviderConnection *>( md->createConnection( connectionName ) );
+  }
+  catch ( QgsProviderConnectionException &ex )
+  {
+    // This is expected and it is not an error in case the provider does not implement
+    // the connections API
+  }
+  return nullptr;
+}
+
 
 QgsConnectionsRootItem::QgsConnectionsRootItem( QgsDataItem *parent, const QString &name, const QString &path, const QString &providerKey )
   : QgsDataCollectionItem( parent, name, path, providerKey )

diff --git a/src/core/qgsdataitem.h b/src/core/qgsdataitem.h
@@ -416,6 +416,17 @@ class CORE_EXPORT QgsDataItem : public QObject
     //! Move object and all its descendants to thread
     void moveToThread( QThread *targetThread );
 
+    /**
+     * For data items that represent a DB connection or one of its children,
+     * this method returns a connection.
+     * All other data items will return NULL.
+     *
+     * Ownership of the returned objects is transferred to the caller.
+     *
+     * \since QGIS 3.16
+     */
+    virtual QgsAbstractDatabaseProviderConnection *databaseConnection() const SIP_FACTORY;
+
   protected:
     virtual void populate( const QVector<QgsDataItem *> &children );
 
@@ -640,6 +651,7 @@ class CORE_EXPORT QgsLayerItem : public QgsDataItem
 
     //! \returns the layer name
     virtual QString layerName() const { return name(); }
+
 };
 
 
@@ -681,6 +693,8 @@ class CORE_EXPORT QgsDataCollectionItem : public QgsDataItem
      */
     static QIcon iconDataCollection();
 
+    QgsAbstractDatabaseProviderConnection *databaseConnection() const override;
+
   protected:
 
     /**
@@ -694,6 +708,7 @@ class CORE_EXPORT QgsDataCollectionItem : public QgsDataItem
      * \since QGIS 3.4
      */
     static QIcon homeDirIcon();
+
 };
 
 
@@ -722,6 +737,7 @@ class CORE_EXPORT QgsDatabaseSchemaItem : public QgsDataCollectionItem
 
     ~QgsDatabaseSchemaItem() override;
 
+    QgsAbstractDatabaseProviderConnection *databaseConnection() const override;
 
     /**
      * Returns the standard browser data collection icon.

diff --git a/tests/src/python/test_qgsdataitem.py b/tests/src/python/test_qgsdataitem.py
@@ -76,6 +76,11 @@ def testPythonCreateChildrenCalledFromCplusplus(self):
             self.assertTrue(tabSetDestroyedFlag[0])
             tabSetDestroyedFlag[0] = False
 
+    def test_databaseConnection(self):
+
+        from IPython import embed
+        embed(using=False)
+
 
 if __name__ == '__main__':
     unittest.main()