From f280f335fc7c4c6401267e59c598be024be755e9 Mon Sep 17 00:00:00 2001
From: Andreas Braun <git@alcaeus.org>
Date: Mon, 10 Apr 2017 16:45:46 +0200
Subject: [PATCH] Format aggregation queries in data collector

---
 DataCollector/PrettyDataCollector.php         | 10 +++-
 .../DataCollector/PrettyDataCollectorTest.php | 50 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/DataCollector/PrettyDataCollector.php b/DataCollector/PrettyDataCollector.php
index ed53026a..abab5d46 100644
--- a/DataCollector/PrettyDataCollector.php
+++ b/DataCollector/PrettyDataCollector.php
@@ -84,7 +84,15 @@ public function collect(Request $request, Response $response, \Exception $except
                 }
 
                 // format the method call
-                if (isset($log['authenticate'])) {
+                if (isset($log['aggregate'])) {
+                    $query .= '.aggregate(' . $this->bsonEncode($log['pipeline']);
+
+                    if ($log['options']) {
+                        $query .= ', ' . $this->bsonEncode($log['options']);
+                    }
+
+                    $query .= ')';
+                } elseif (isset($log['authenticate'])) {
                     $query .= '.authenticate()';
                 } elseif (isset($log['batchInsert'])) {
                     if (1 === $log['num']) {
diff --git a/Tests/DataCollector/PrettyDataCollectorTest.php b/Tests/DataCollector/PrettyDataCollectorTest.php
index ac8bd12e..f9f99b65 100644
--- a/Tests/DataCollector/PrettyDataCollectorTest.php
+++ b/Tests/DataCollector/PrettyDataCollectorTest.php
@@ -312,4 +312,54 @@ public function testCollectLimitAndSort()
         $this->assertEquals(1, $collector->getQueryCount());
         $this->assertEquals($formatted, $collector->getQueries());
     }
+
+    public function testCollectAggregate()
+    {
+        $queries = [
+            [
+                'aggregate' => true,
+                'pipeline' => [
+                    [
+                        '$group' => [
+                            '_id' => '$verified',
+                            'count' => ['$sum' => 1],
+                        ],
+                    ],
+                ],
+                'options' => [],
+                'db' => 'foo',
+                'collection' => 'User',
+            ],
+            [
+                'aggregate' => true,
+                'pipeline' => [
+                    [
+                        '$group' => [
+                            '_id' => '$verified',
+                            'count' => ['$sum' => 1],
+                        ],
+                    ],
+                ],
+                'options' => ['group' => true],
+                'db' => 'foo',
+                'collection' => 'User',
+            ],
+        ];
+
+        $formatted = [
+            'use foo;',
+            'db.User.aggregate([ { "$group": { "_id": "$verified", "count": { "$sum": 1 } } } ]);',
+            'db.User.aggregate([ { "$group": { "_id": "$verified", "count": { "$sum": 1 } } } ], { "group": true });',
+        ];
+
+
+        $collector = new PrettyDataCollector();
+        foreach ($queries as $query) {
+            $collector->logQuery($query);
+        }
+        $collector->collect(new Request(), new Response());
+
+        $this->assertEquals(2, $collector->getQueryCount());
+        $this->assertEquals($formatted, $collector->getQueries());
+    }
 }