Notify error on SQL layer creation

qgis · Jul 6, 2021 · 94056c8 · 94056c8
1 parent 9adbdfa
commit 94056c8
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 6 deletions.
diff --git a/src/app/browser/qgsinbuiltdataitemproviders.cpp b/src/app/browser/qgsinbuiltdataitemproviders.cpp
@@ -1066,7 +1066,7 @@ void QgsDatabaseItemGuiProvider::populateContextMenu( QgsDataItem *item, QMenu *
 
         QAction *sqlAction = new QAction( QObject::tr( "Run SQL command…" ), menu );
 
-        QObject::connect( sqlAction, &QAction::triggered, collectionItem, [ collectionItem ]
+        QObject::connect( sqlAction, &QAction::triggered, collectionItem, [ collectionItem, context ]
         {
           std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn2( collectionItem->databaseConnection() );
           // This should never happen but let's play safe
@@ -1083,10 +1083,19 @@ void QgsDatabaseItemGuiProvider::populateContextMenu( QgsDataItem *item, QMenu *
           widget->layout()->setMargin( 0 );
           dialog.layout()->addWidget( widget );
 
-          connect( widget, &QgsQueryResultWidget::createSqlVectorLayer, widget, [collectionItem]( const QString &, const QString &, const QgsAbstractDatabaseProviderConnection::SqlVectorLayerOptions & options )
+          connect( widget, &QgsQueryResultWidget::createSqlVectorLayer, widget, [ collectionItem, context ]( const QString &, const QString &, const QgsAbstractDatabaseProviderConnection::SqlVectorLayerOptions & options )
           {
             std::unique_ptr<QgsAbstractDatabaseProviderConnection> conn3( collectionItem->databaseConnection() );
-            conn3->createSqlVectorLayer( options );
+            try
+            {
+              QgsMapLayer *sqlLayer { conn3->createSqlVectorLayer( options ) };
+              QgsProject::instance()->addMapLayers( { sqlLayer } );
+            }
+            catch ( QgsProviderConnectionException &ex )
+            {
+              notify( QObject::tr( "New SQL Layer Creation Error" ), QObject::tr( "Error creating new SQL layer: %1" ).arg( ex.what() ), context, Qgis::MessageLevel::Critical );
+            }
+
           } );
           dialog.exec();
         } );

diff --git a/src/providers/postgres/qgspostgresproviderconnection.cpp b/src/providers/postgres/qgspostgresproviderconnection.cpp
@@ -754,7 +754,7 @@ QgsVectorLayer *QgsPostgresProviderConnection::createSqlVectorLayer( const SqlVe
   }
 
   QgsDataSourceUri tUri( uri( ) );
-  tUri.setKeyColumn( QStringLiteral( "_uid_" ) );
+
   tUri.setSql( options.filter );
 
   if ( ! options.primaryKeyColumns.isEmpty() )
@@ -764,10 +764,27 @@ QgsVectorLayer *QgsPostgresProviderConnection::createSqlVectorLayer( const SqlVe
   }
   else
   {
-    tUri.setTable( QStringLiteral( "(SELECT row_number() over () AS _uid_, * FROM (%1\n) AS _subq_1_\n)" ).arg( options.sql ) );
+    int pkId { 0 };
+    while ( options.sql.contains( QStringLiteral( "_uid%1_" ).arg( pkId ), Qt::CaseSensitivity::CaseInsensitive ) )
+    {
+      pkId ++;
+    }
+    tUri.setKeyColumn( QStringLiteral( "_uid%1_" ).arg( pkId ) );
+
+    int sqlId { 0 };
+    while ( options.sql.contains( QStringLiteral( "_subq_%s_" ).arg( sqlId ), Qt::CaseSensitivity::CaseInsensitive ) )
+    {
+      sqlId ++;
+    }
+    tUri.setTable( QStringLiteral( "(SELECT row_number() over () AS _uid%1_, * FROM (%2\n) AS _subq_%3_\n)" ).arg( QString::number( pkId ), options.sql, QString::number( sqlId ) ) );
+  }
+
+  if ( ! options.geometryColumn.isEmpty() )
+  {
+    tUri.setGeometryColumn( options.geometryColumn );
   }
 
-  return new QgsVectorLayer{ tUri.uri(), options.layerName.isEmpty() ? QStringLiteral( "sql_layer" ) : options.layerName, mProviderKey };
+  return new QgsVectorLayer{ tUri.uri(), options.layerName.isEmpty() ? QStringLiteral( "QueryLayer" ) : options.layerName, mProviderKey };
 }
 
 QgsFields QgsPostgresProviderConnection::fields( const QString &schema, const QString &tableName ) const

diff --git a/tests/src/python/test_qgsproviderconnection_postgres.py b/tests/src/python/test_qgsproviderconnection_postgres.py
@@ -449,6 +449,118 @@ def test_table_scan(self):
         self.assertTrue(vl.isValid())
         self.assertEqual(vl.featureCount(), 10)
 
+    def test_create_vector_layer(self):
+        """Test query layers"""
+
+        md = QgsProviderRegistry.instance().providerMetadata('postgres')
+        conn = md.createConnection(self.uri, {})
+
+        sql = """
+        DROP TABLE IF EXISTS qgis_test.query_layer1;
+        CREATE TABLE qgis_test.query_layer1 (
+            id SERIAL PRIMARY KEY,
+            geom geometry(POINT,4326)
+        );
+        INSERT INTO qgis_test.query_layer1 (id, geom) VALUES (221, ST_GeomFromText('point(9 45)', 4326));
+        INSERT INTO qgis_test.query_layer1 (id, geom) VALUES (201, ST_GeomFromText('point(9.5 45.5)', 4326));
+        """
+
+        conn.executeSql(sql)
+
+        options = QgsAbstractDatabaseProviderConnection.SqlVectorLayerOptions()
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 WHERE id < 200 LIMIT 2'
+        options.primaryKeyColumns = ['id']
+        options.geometryColumn = 'geom'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 0)
+
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 WHERE id > 200 LIMIT 2'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 2)
+
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 WHERE id > 210 LIMIT 2'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 LIMIT 2'
+        options.filter = 'id > 210'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
+        # Wrong calls
+        options.primaryKeyColumns = ['DOES_NOT_EXIST']
+        vl = conn.createSqlVectorLayer(options)
+        self.assertFalse(vl.isValid())
+
+        options.primaryKeyColumns = ['id']
+        options.geometryColumn = 'DOES_NOT_EXIST'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertFalse(vl.isValid())
+
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 WHERE id > 210 LIMIT 2'
+        options.primaryKeyColumns = []
+        options.geometryColumn = ''
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
+        # No geometry and no PK, aspatial layer
+        options.sql = 'SELECT id, geom FROM qgis_test.query_layer1 WHERE id > 210 LIMIT 2'
+        options.primaryKeyColumns = []
+        options.geometryColumn = ''
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertNotEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
+        # Composite keys
+        sql = """
+        DROP TABLE IF EXISTS qgis_test.query_layer2;
+        CREATE TABLE qgis_test.query_layer2 (
+            id SERIAL,
+            id2 SERIAL,
+            geom geometry(POINT,4326),
+            PRIMARY KEY(id, id2)
+        );
+        INSERT INTO qgis_test.query_layer2 (id, id2, geom) VALUES (101, 101, ST_GeomFromText('point(9 45)', 4326));
+        INSERT INTO qgis_test.query_layer2 (id, id2, geom) VALUES (201, 201, ST_GeomFromText('point(9.5 45.5)', 4326));
+        """
+
+        conn.executeSql(sql)
+
+        options = QgsAbstractDatabaseProviderConnection.SqlVectorLayerOptions()
+        options.sql = 'SELECT id, id2, geom FROM qgis_test.query_layer2 ORDER BY id ASC LIMIT 1'
+        options.primaryKeyColumns = ['id', 'id2']
+        options.geometryColumn = 'geom'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
+        # No PKs
+        options.primaryKeyColumns = []
+        options.geometryColumn = 'geom'
+        vl = conn.createSqlVectorLayer(options)
+        self.assertTrue(vl.isValid())
+        self.assertEqual(vl.geometryType(), QgsWkbTypes.PointGeometry)
+        features = [f for f in vl.getFeatures()]
+        self.assertEqual(len(features), 1)
+
 
 if __name__ == '__main__':
     unittest.main()