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

strk · strk · commit ada9348e2b8d · 2016-10-14T10:43:41.000+02:00
Fixes #15226 (drag & drop of postgresql views) Includes test
diff --git a/src/providers/postgres/qgspostgresprovider.cpp b/src/providers/postgres/qgspostgresprovider.cpp
@@ -1360,56 +1360,73 @@ bool QgsPostgresProvider::determinePrimaryKey()
   return mValid;
 }
 
-void QgsPostgresProvider::determinePrimaryKeyFromUriKeyColumn()
+/* static */
+QStringList QgsPostgresProvider::parseUriKey( const QString& key )
 {
-  QString primaryKey = mUri.keyColumn();
-  mPrimaryKeyType = pktUnknown;
+  if ( key.isEmpty() ) return QStringList();
 
-  if ( !primaryKey.isEmpty() )
-  {
-    QStringList cols;
+  QStringList cols;
 
-    // remove quotes from key list
-    if ( primaryKey.startsWith( '"' ) && primaryKey.endsWith( '"' ) )
+  // remove quotes from key list
+  if ( key.startsWith( '"' ) && key.endsWith( '"' ) )
+  {
+    int i = 1;
+    QString col;
+    while ( i < key.size() )
     {
-      int i = 1;
-      QString col;
-      while ( i < primaryKey.size() )
+      if ( key[i] == '"' )
       {
-        if ( primaryKey[i] == '"' )
+        if ( i + 1 < key.size() && key[i+1] == '"' )
         {
-          if ( i + 1 < primaryKey.size() && primaryKey[i+1] == '"' )
-          {
-            i++;
-          }
-          else
-          {
-            cols << col;
-            col = "";
+          i++;
+        }
+        else
+        {
+          cols << col;
+          col = "";
 
-            if ( ++i == primaryKey.size() )
-              break;
+          if ( ++i == key.size() )
+            break;
 
-            Q_ASSERT( primaryKey[i] == ',' );
-            i++;
-            Q_ASSERT( primaryKey[i] == '"' );
-            i++;
-            col = "";
-            continue;
-          }
+          Q_ASSERT( key[i] == ',' );
+          i++;
+          Q_ASSERT( key[i] == '"' );
+          i++;
+          col = "";
+          continue;
         }
-
-        col += primaryKey[i++];
       }
+
+      col += key[i++];
     }
-    else if ( primaryKey.contains( ',' ) )
-    {
-      cols = primaryKey.split( ',' );
-    }
-    else
+  }
+  else if ( key.contains( ',' ) )
+  {
+    cols = key.split( ',' );
+  }
+  else
+  {
+    cols << key;
+  }
+
+  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 )
     {
-      cols << primaryKey;
-      primaryKey = quotedIdentifier( primaryKey );
+      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() )
   {
@@ -3503,45 +3523,39 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
   }
   else
   {
-    // search for the passed field
-    for ( int fldIdx = 0; fldIdx < fields.count(); ++fldIdx )
+    pkList = parseUriKey( primaryKey );
+    Q_FOREACH ( const QString& col, pkList )
     {
-      if ( fields.at( fldIdx ).name() == primaryKey )
+      // search for the passed field
+      QString type;
+      for ( int fldIdx = 0; fldIdx < fields.count(); ++fldIdx )
       {
-        // found, get the field type
-        QgsField fld = fields.at( fldIdx );
-        if ( convertField( fld, options ) )
+        if ( fields[fldIdx].name() == col )
         {
-          primaryKeyType = fld.typeName();
+          // found, get the field type
+          QgsField fld = fields[fldIdx];
+          if ( convertField( fld, options ) )
+          {
+            type = fld.typeName();
+            break;
+          }
         }
       }
-    }
-  }
-
-  // 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,
-    // use the equivalent autoincremental type (serialN)
-    if ( primaryKeyType == "int2" || primaryKeyType == "int4" )
-    {
-      primaryKeyType = "serial";
-    }
-    else if ( primaryKeyType == "int8" )
-    {
-      primaryKeyType = "serial8";
+      if ( type.isEmpty() ) type = "serial";
+      else
+      {
+        // if the pk field's type is one of the postgres integer types,
+        // use the equivalent autoincremental type (serialN)
+        if ( primaryKeyType == "int2" || primaryKeyType == "int4" )
+        {
+          primaryKeyType = "serial";
+        }
+        else if ( primaryKeyType == "int8" )
+        {
+          primaryKeyType = "serial8";
+        }
+      }
+      pkType << type;
     }
   }
 
@@ -3577,17 +3591,32 @@ QgsVectorLayerImport::ImportError QgsPostgresProvider::createEmptyLayer( const Q
         throw PGException( result );
     }
 
-    if ( options && options->value( "lowercaseFieldNames", false ).toBool() )
+    sql = QString( "CREATE TABLE %1(" ) .arg( schemaTableName );
+    QString pk;
+    for ( int i = 0; i < pkList.size(); ++i )
     {
-      //convert primary key name to lowercase
-      //this must happen after determining the field type of the primary key
-      primaryKey = primaryKey.toLower();
-    }
+      QString col = pkList[i];
+      const QString& type = pkType[i];
 
-    sql = QString( "CREATE TABLE %1(%2 %3 PRIMARY KEY)" )
-          .arg( schemaTableName,
-                quotedIdentifier( primaryKey ),
-                primaryKeyType );
+      if ( options && options->value( "lowercaseFieldNames", false ).toBool() )
+      {
+        col = col.toLower();
+      }
+      else
+      {
+        col = quotedIdentifier( col ); // no need to quote lowercase field
+      }
+
+      if ( i )
+      {
+        pk  += ",";
+        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 )
+      {
+        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, 0 );
+        oldToNewAttrIdxMap->insert( fldIdx, pkIdx );
         continue;
       }
 
diff --git a/src/providers/postgres/qgspostgresprovider.h b/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();
 
diff --git a/tests/src/python/test_provider_postgres.py b/tests/src/python/test_provider_postgres.py
@@ -461,6 +461,37 @@ def testNumericPrecision(self):
         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):
 
