diff --git a/src/Codeception/Module/Db.php b/src/Codeception/Module/Db.php
index 9414278b..246ddc44 100644
--- a/src/Codeception/Module/Db.php
+++ b/src/Codeception/Module/Db.php
@@ -769,17 +769,13 @@ private function addInsertedRow($table, array $row, $id)
         $primaryKey = $this->_getDriver()->getPrimaryKey($table);
         $primary = [];
         if ($primaryKey) {
-            if ($id && count($primaryKey) === 1) {
-                $primary [$primaryKey[0]] = $id;
-            } else {
-                foreach ($primaryKey as $column) {
-                    if (isset($row[$column])) {
-                        $primary[$column] = $row[$column];
-                    } else {
-                        throw new \InvalidArgumentException(
-                            'Primary key field ' . $column . ' is not set for table ' . $table
-                        );
-                    }
+            foreach ($primaryKey as $column) {
+                if (isset($row[$column]) || $id) {
+                    $primary[$column] = isset($row[$column]) ? $row[$column] : $id;
+                } else {
+                    throw new \InvalidArgumentException(
+                        'Primary key field ' . $column . ' is not set for table ' . $table
+                    );
                 }
             }
         } else {