Fix PostgreSQL import of layers with multi-column or quoted-column keys

Fixes #15226 (drag & drop of postgresql views)
Includes test

src/providers/postgres/qgspostgresprovider.cpp

@@ -1360,25 +1360,23 @@ bool QgsPostgresProvider::determinePrimaryKey()
   return mValid;
 }
 
-void QgsPostgresProvider::determinePrimaryKeyFromUriKeyColumn()
+/* static */
+QStringList QgsPostgresProvider::parseUriKey( const QString& key )
 {
-  QString primaryKey = mUri.keyColumn();
+  if ( key.isEmpty() ) return QStringList();
-  mPrimaryKeyType = pktUnknown;
 
-  if ( !primaryKey.isEmpty() )
-  {
   QStringList cols;
 
   // remove quotes from key list
-    if ( primaryKey.startsWith( '"' ) && primaryKey.endsWith( '"' ) )
+  if ( key.startsWith( '"' ) && key.endsWith( '"' ) )
   {
     int i = 1;
     QString col;
-      while ( i < primaryKey.size() )
+    while ( i < key.size() )
     {
-        if ( primaryKey[i] == '"' )
+      if ( key[i] == '"' )
       {
-          if ( i + 1 < primaryKey.size() && primaryKey[i+1] == '"' )
+        if ( i + 1 < key.size() && key[i+1] == '"' )
         {
           i++;
         }
@@ -1387,29 +1385,48 @@ void QgsPostgresProvider::determinePrimaryKeyFromUriKeyColumn()
           cols << col;
           col = "";
 
-            if ( ++i == primaryKey.size() )
+          if ( ++i == key.size() )
             break;
 
-            Q_ASSERT( primaryKey[i] == ',' );
+          Q_ASSERT( key[i] == ',' );
           i++;
-            Q_ASSERT( primaryKey[i] == '"' );
+          Q_ASSERT( key[i] == '"' );
           i++;
           col = "";
           continue;
         }
       }
 
-        col += primaryKey[i++];
+      col += key[i++];
     }
   }
-    else if ( primaryKey.contains( ',' ) )
+  else if ( key.contains( ',' ) )
   {
-      cols = primaryKey.split( ',' );
+    cols = key.split( ',' );
   }
   else
   {
-      cols << primaryKey;
+    cols << key;
-      primaryKey = quotedIdentifier( primaryKey );
+  }
+
+  return cols;
+}
+
+void QgsPostgresProvider::determinePrimaryKeyFromUriKeyColumn()
+{
+  QString primaryKey = mUri.keyColumn();
+  mPrimaryKeyType = pktUnknown;
+
+  if ( !primaryKey.isEmpty() )
+  {
+    QStringList cols = parseUriKey( primaryKey );
+
+    primaryKey = "";
+    QString del = "";
+    Q_FOREACH ( const QString& col, cols )
+    {
+      primaryKey += del + quotedIdentifier( col );
+      del = ",";
     }
 
     Q_FOREACH ( const QString& col, cols )
@@ -3463,6 +3480,9 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
   QString primaryKey = dsUri.keyColumn();
   QString primaryKeyType;
 
+  QStringList pkList;
+  QStringList pkType;
+
   QString schemaTableName = "";
   if ( !schemaName.isEmpty() )
   {
@@ -3502,35 +3522,26 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
     }
   }
   else
+  {
+    pkList = parseUriKey( primaryKey );
+    Q_FOREACH ( const QString& col, pkList )
     {
       // search for the passed field
+      QString type;
       for ( int fldIdx = 0; fldIdx < fields.count(); ++fldIdx )
       {
-      if ( fields.at( fldIdx ).name() == primaryKey )
+        if ( fields[fldIdx].name() == col )
         {
           // found, get the field type
-        QgsField fld = fields.at( fldIdx );
+          QgsField fld = fields[fldIdx];
           if ( convertField( fld, options ) )
           {
-          primaryKeyType = fld.typeName();
+            type = fld.typeName();
+            break;
           }
         }
       }
-  }
+      if ( type.isEmpty() ) type = "serial";
-
-  // if the pk field doesn't exist yet, create a serial pk field
-  // as it's autoincremental
-  if ( primaryKeyType.isEmpty() )
-  {
-    primaryKeyType = "serial";
-#if 0
-    // TODO: check the feature count to choose if create a serial8 pk field
-    if ( layer->featureCount() > 0xffffffff )
-    {
-      primaryKeyType = "serial8";
-    }
-#endif
-  }
       else
       {
         // if the pk field's type is one of the postgres integer types,
@@ -3544,6 +3555,9 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
           primaryKeyType = "serial8";
         }
       }
+      pkType << type;
+    }
+  }
 
   try
   {
@@ -3577,17 +3591,32 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
         throw PGException( result );
     }
 
+    sql = QString( "CREATE TABLE %1(" ) .arg( schemaTableName );
+    QString pk;
+    for ( int i = 0; i < pkList.size(); ++i )
+    {
+      QString col = pkList[i];
+      const QString& type = pkType[i];
+
       if ( options && options->value( "lowercaseFieldNames", false ).toBool() )
       {
-      //convert primary key name to lowercase
+        col = col.toLower();
-      //this must happen after determining the field type of the primary key
+      }
-      primaryKey = primaryKey.toLower();
+      else
+      {
+        col = quotedIdentifier( col ); // no need to quote lowercase field
       }
 
-    sql = QString( "CREATE TABLE %1(%2 %3 PRIMARY KEY)" )
+      if ( i )
-          .arg( schemaTableName,
+      {
-                quotedIdentifier( primaryKey ),
+        pk  += ",";
-                primaryKeyType );
+        sql += ",";
+      }
+
+      pk += col;
+      sql += col + " " + type;
+    }
+    sql += QString( ", PRIMARY KEY (%1) )" ) .arg( pk );
 
     result = conn->PQexec( sql );
     if ( result.PQresultStatus() != PGRES_COMMAND_OK )
@@ -3678,9 +3707,25 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
         fld.setName( fld.name().toLower() );
       }
 
-      if ( fld.name() == primaryKey )
+      int pkIdx = -1;
+      for ( int i = 0; i < pkList.size(); ++i )
       {
-        oldToNewAttrIdxMap->insert( fldIdx, 0 );
+        QString col = pkList[i];
+        if ( options && options->value( "lowercaseFieldNames", false ).toBool() )
+        {
+          //convert field name to lowercase (TODO: avoid doing this
+          //over and over)
+          col =  col.toLower();
+        }
+        if ( fld.name() == col )
+        {
+          pkIdx = i;
+          break;
+        }
+      }
+      if ( pkIdx >= 0 )
+      {
+        oldToNewAttrIdxMap->insert( fldIdx, pkIdx );
         continue;
       }
 

src/providers/postgres/qgspostgresprovider.h

@@ -114,6 +114,11 @@ class QgsPostgresProvider : public QgsVectorDataProvider
      */
     static QString endianString();
 
+    /**
+     * Returns a list of unquoted column names from an uri key
+     */
+    static QStringList parseUriKey( const QString& key );
+
     /**
      * Changes the stored extent for this layer to the supplied extent.
      * For example, this is called when the extent worker thread has a result.
@@ -126,11 +131,16 @@ class QgsPostgresProvider : public QgsVectorDataProvider
      */
     virtual void updateExtents() override;
 
-    /** Determine the fields making up the primary key
+    /**
+     * Determine the fields making up the primary key
      */
     bool determinePrimaryKey();
 
-    /** Determine the fields making up the primary key from the uri attribute keyColumn
+    /**
+     * Determine the fields making up the primary key from the uri attribute keyColumn
+     *
+     * Fills mPrimaryKeyType and mPrimaryKeyAttrs
+     * from mUri
      */
     void determinePrimaryKeyFromUriKeyColumn();
 

tests/src/python/test_provider_postgres.py

@@ -461,6 +461,37 @@ class TestPyQgsPostgresProvider(unittest.TestCase, ProviderTestCase):
         self.assertEqual(f['f2'], 123.456)
         self.assertEqual(f['f3'], '12345678.90123456789')
 
+    # See http://hub.qgis.org/issues/15226
+    def testImportKey(self):
+        uri = 'point?field=f1:int'
+        uri += '&field=F2:double(6,4)'
+        uri += '&field=f3:string(20)'
+        lyr = QgsVectorLayer(uri, "x", "memory")
+        self.assertTrue(lyr.isValid())
+
+        def testKey(lyr, key, kfnames):
+            self.execSQLCommand('DROP TABLE IF EXISTS qgis_test.import_test')
+            uri = '%s table="qgis_test"."import_test" (g) key=\'%s\'' % (self.dbconn, key)
+            err = QgsVectorLayerImport.importLayer(lyr, uri, "postgres", lyr.crs())
+            self.assertEqual(err[0], QgsVectorLayerImport.NoError,
+                             'unexpected import error {0}'.format(err))
+            olyr = QgsVectorLayer(uri, "y", "postgres")
+            self.assertTrue(olyr.isValid())
+            flds = lyr.fields()
+            oflds = olyr.fields()
+            self.assertEquals(oflds.size(), flds.size())
+            for i in range(0, oflds.size()):
+                self.assertEqual(oflds[i].name(), flds[i].name())
+            pks = olyr.pkAttributeList()
+            self.assertEquals(len(pks), len(kfnames))
+            for i in range(0, len(kfnames)):
+                self.assertEqual(oflds[pks[i]].name(), kfnames[i])
+
+        testKey(lyr, 'f1', ['f1'])
+        testKey(lyr, '"f1"', ['f1'])
+        testKey(lyr, '"f1","F2"', ['f1', 'F2'])
+        testKey(lyr, '"f1","F2","f3"', ['f1', 'F2', 'f3'])
+
 
 class TestPyQgsPostgresProviderCompoundKey(unittest.TestCase, ProviderTestCase):