From 444a0902e04006357b18834e0004fea857a23f0d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20S=CC=8Cra=CC=81mek?= <sramek.jn@gmail.com>
Date: Tue, 7 Apr 2020 12:19:35 +0200
Subject: [PATCH] Add id keys into obatained collection properties

---
 src/LeanMapper/Entity.php | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/LeanMapper/Entity.php b/src/LeanMapper/Entity.php
index 78f96f6..0ec72be 100644
--- a/src/LeanMapper/Entity.php
+++ b/src/LeanMapper/Entity.php
@@ -941,6 +941,7 @@ private function getHasManyValue(
         Filtering $relTableFiltering = null
     ) {
         $targetTable = $relationship->getTargetTable();
+        $primaryKey = $this->mapper->getPrimaryKey($targetTable);
         $columnReferencingTargetTable = $relationship->getColumnReferencingTargetTable();
         $rows = $this->row->referencing(
             $relationship->getRelationshipTable(),
@@ -956,7 +957,7 @@ private function getHasManyValue(
                 $entity = $this->entityFactory->createEntity($entityClass, $valueRow);
                 $this->checkConsistency($property, $entityClass, $entity);
                 $entity->makeAlive($this->entityFactory);
-                $value[] = $entity;
+                $value[$entity->get($primaryKey)] = $entity;
             }
         }
         return $this->entityFactory->createCollection($value);
@@ -1008,6 +1009,7 @@ private function getBelongsToOneValue(Property $property, Relationship\BelongsTo
     private function getBelongsToManyValue(Property $property, Relationship\BelongsToMany $relationship, Filtering $filtering = null)
     {
         $targetTable = $relationship->getTargetTable();
+        $primaryKey = $this->mapper->getPrimaryKey($targetTable);
         $rows = $this->row->referencing($targetTable, $relationship->getColumnReferencingSourceTable(), $filtering, $relationship->getStrategy());
         $value = [];
         foreach ($rows as $row) {
@@ -1015,7 +1017,7 @@ private function getBelongsToManyValue(Property $property, Relationship\BelongsT
             $entity = $this->entityFactory->createEntity($entityClass, $row);
             $this->checkConsistency($property, $entityClass, $entity);
             $entity->makeAlive($this->entityFactory);
-            $value[] = $entity;
+            $value[$entity->get($primaryKey)] = $entity;
         }
         return $this->entityFactory->createCollection($value);
     }