From aa8c013ebf70bede5ea30576a8a2a8f00421e171 Mon Sep 17 00:00:00 2001
From: Qiang Xue <qiang.xue@gmail.com>
Date: Tue, 23 Sep 2014 10:29:11 -0400
Subject: [PATCH] Fixes #3197: Using `ActiveQuery::indexBy()` may cause
 relational AR queries to generate incorrect relational results

---
 framework/CHANGELOG.md                       |  1 +
 framework/db/ActiveRelationTrait.php         | 35 ++++++++++++++------
 tests/unit/framework/db/ActiveRecordTest.php | 12 +++++++
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/framework/CHANGELOG.md b/framework/CHANGELOG.md
index 2d7f41b796..8488e29180 100644
--- a/framework/CHANGELOG.md
+++ b/framework/CHANGELOG.md
@@ -20,6 +20,7 @@ Yii Framework 2 Change Log
 - Bug #3153: Fixed the bug that using "between" operator to build a SQL query will cause a PHP notice (gonimar)
 - Bug #3184: Fixed the bug that client validation for string length comparison did not set error message correctly (Sergeygithub)
 - Bug #3194: Date formatter works only for timestamps in the year range 1970 to 2038 (kartik-v)
+- Bug #3197: Using `ActiveQuery::indexBy()` may cause relational AR queries to generate incorrect relational results (qiangxue)
 - Bug #3204: `yii\di\Container` did not handle the `$config` parameter well in case when it does not have a default value (qiangxue)
 - Bug #3216: Fixed the bug that `yii.activeForm.destroy()` did not remove `submit` event handlers (qiangxue)
 - Bug #3233: Ensure consistent behavior in ActiveRecord::afterSave() (cebe, qiangxue)
diff --git a/framework/db/ActiveRelationTrait.php b/framework/db/ActiveRelationTrait.php
index 2d484b2c5a..17267e7359 100644
--- a/framework/db/ActiveRelationTrait.php
+++ b/framework/db/ActiveRelationTrait.php
@@ -230,13 +230,23 @@ trait ActiveRelationTrait
 
             return [$model];
         } else {
+            // https://github.com/yiisoft/yii2/issues/3197
+            // delay indexing related models after buckets are built
+            $indexBy = $this->indexBy;
+            $this->indexBy = null;
             $models = $this->all();
+
             if (isset($viaModels, $viaQuery)) {
                 $buckets = $this->buildBuckets($models, $this->link, $viaModels, $viaQuery->link);
             } else {
                 $buckets = $this->buildBuckets($models, $this->link);
             }
 
+            $this->indexBy = $indexBy;
+            if ($this->indexBy !== null && $this->multiple) {
+                $buckets = $this->indexBuckets($buckets, $this->indexBy);
+            }
+
             $link = array_values(isset($viaQuery) ? $viaQuery->link : $this->link);
             foreach ($primaryModels as $i => $primaryModel) {
                 if ($this->multiple && count($link) == 1 && is_array($keys = $primaryModel[reset($link)])) {
@@ -358,22 +368,14 @@ trait ActiveRelationTrait
                 $key = $this->getModelKey($model, $linkKeys);
                 if (isset($map[$key])) {
                     foreach (array_keys($map[$key]) as $key2) {
-                        if ($this->indexBy !== null) {
-                            $buckets[$key2][$i] = $model;
-                        } else {
-                            $buckets[$key2][] = $model;
-                        }
+                        $buckets[$key2][] = $model;
                     }
                 }
             }
         } else {
             foreach ($models as $i => $model) {
                 $key = $this->getModelKey($model, $linkKeys);
-                if ($this->indexBy !== null) {
-                    $buckets[$key][$i] = $model;
-                } else {
-                    $buckets[$key][] = $model;
-                }
+                $buckets[$key][] = $model;
             }
         }
 
@@ -386,6 +388,19 @@ trait ActiveRelationTrait
         return $buckets;
     }
 
+    private function indexBuckets($buckets, $indexBy)
+    {
+        $result = [];
+        foreach ($buckets as $key => $models) {
+            $result[$key] = [];
+            foreach ($models as $model) {
+                $index = is_string($indexBy) ? $model[$indexBy] : call_user_func($indexBy, $model);
+                $result[$key][$index] = $model;
+            }
+        }
+        return $result;
+    }
+
     /**
      * @param array $attributes the attributes to prefix
      * @return array
diff --git a/tests/unit/framework/db/ActiveRecordTest.php b/tests/unit/framework/db/ActiveRecordTest.php
index 2c8e21abe3..f06c743e69 100644
--- a/tests/unit/framework/db/ActiveRecordTest.php
+++ b/tests/unit/framework/db/ActiveRecordTest.php
@@ -630,5 +630,17 @@ class ActiveRecordTest extends DatabaseTestCase
         $this->assertEquals(3, $category->getItems()->count());
         $this->assertEquals(1, $category->getLimitedItems()->count());
         $this->assertEquals(1, $category->getLimitedItems()->distinct(true)->count());
+
+        // https://github.com/yiisoft/yii2/issues/3197
+        $orders = Order::find()->with('orderItems')->orderBy('id')->all();
+        $this->assertEquals(3, count($orders));
+        $this->assertEquals(2, count($orders[0]->orderItems));
+        $this->assertEquals(3, count($orders[1]->orderItems));
+        $this->assertEquals(1, count($orders[2]->orderItems));
+        $orders = Order::find()->with(['orderItems' => function ($q) { $q->indexBy('item_id'); }])->orderBy('id')->all();
+        $this->assertEquals(3, count($orders));
+        $this->assertEquals(2, count($orders[0]->orderItems));
+        $this->assertEquals(3, count($orders[1]->orderItems));
+        $this->assertEquals(1, count($orders[2]->orderItems));
     }
 }