PG: identify pk when exporting

When exporting a layer, if no PK information is passed in the destination URI, the first numeric field with NOT NULL and UNIQUE constraints is considered to be the primary key. Partial fix for #45286
qgis · Oct 4, 2021 · a0af938 · a0af938
1 parent 1006aa7
commit a0af938
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 1 deletion.
diff --git a/src/providers/postgres/qgspostgresprovider.cpp b/src/providers/postgres/qgspostgresprovider.cpp
@@ -4384,8 +4384,23 @@ Qgis::VectorExportResult QgsPostgresProvider::createEmptyLayer( const QString &u
   }
 
   // get the pk's name and type
+  // Try to find a PK candidate from numeric NOT NULL / UNIQUE columns
+  if ( primaryKey.isEmpty() )
+  {
+    for ( const auto &field : std::as_const( fields ) )
+    {
+      if ( field.isNumeric() &&
+           ( field.constraints().constraints() & QgsFieldConstraints::Constraint::ConstraintUnique ) &&
+           ( field.constraints().constraints() & QgsFieldConstraints::Constraint::ConstraintNotNull ) &&
+           ( field.constraints().constraints() & QgsFieldConstraints::ConstraintOrigin::ConstraintOriginProvider ) )
+      {
+        primaryKey = field.name();
+        break;
+      }
+    }
+  }
 
-  // if no pk name was passed, define the new pk field name
+  // if no pk name was passed or guessed, define the new pk field name
   if ( primaryKey.isEmpty() )
   {
     int index = 0;

diff --git a/tests/src/python/test_provider_postgres.py b/tests/src/python/test_provider_postgres.py
@@ -3387,6 +3387,43 @@ def testPkeyIntArray(self):
         self.assertTrue(feat.isValid())
         self.assertEqual(feat["name"], "test")
 
+    def testExportPkGuessLogic(self):
+        """Test that when creating an empty layer a NOT NULL UNIQUE numeric field is identified as a PK"""
+
+        md = QgsProviderRegistry.instance().providerMetadata("postgres")
+        conn = md.createConnection(self.dbconn, {})
+        conn.executeSql(
+            'DROP TABLE IF EXISTS qgis_test."testExportPkGuessLogic_source" CASCADE')
+        conn.executeSql(
+            'DROP TABLE IF EXISTS qgis_test."testExportPkGuessLogic_exported" CASCADE')
+        conn.executeSql(
+            """CREATE TABLE qgis_test."testExportPkGuessLogic_source" ( id bigint generated always as identity primary key,
+                    geom geometry(Point, 4326) check (st_isvalid(geom)),
+                    name text unique, author text not null)""")
+
+        source_layer = QgsVectorLayer(self.dbconn + ' sslmode=disable key=\'id\' srid=4326 type=POINT table="qgis_test"."testExportPkGuessLogic_source" (geom) sql=', 'testExportPkGuessLogic_source', 'postgres')
+
+        md = QgsProviderRegistry.instance().providerMetadata('postgres')
+        conn = md.createConnection(self.dbconn, {})
+        table = conn.table("qgis_test", "testExportPkGuessLogic_source")
+        self.assertEqual(table.primaryKeyColumns(), ['id'])
+
+        self.assertTrue(source_layer.isValid())
+
+        # Create the URI as the browser does (no PK information)
+        uri = self.dbconn + ' sslmode=disable srid=4326 type=POINT table="qgis_test"."testExportPkGuessLogic_exported" (geom) sql='
+
+        exporter = QgsVectorLayerExporter(uri, 'postgres', source_layer.fields(), source_layer.wkbType(), source_layer.crs(), True, {})
+        self.assertTrue(exporter.lastError() == '')
+
+        exported_layer = QgsVectorLayer(self.dbconn + ' sslmode=disable srid=4326 type=POINT table="qgis_test"."testExportPkGuessLogic_exported" (geom) sql=', 'testExportPkGuessLogic_exported', 'postgres')
+        self.assertTrue(exported_layer.isValid())
+
+        table = conn.table("qgis_test", "testExportPkGuessLogic_exported")
+        self.assertEqual(table.primaryKeyColumns(), ['id'])
+
+        self.assertEqual(exported_layer.fields().names(), ['id', 'name', 'author'])
+
 
 if __name__ == '__main__':
     unittest.main()