From aeb841393ca9dfd9449b5cb2b04bbc285d8dc69b Mon Sep 17 00:00:00 2001
From: Hung Neo <hungnq1989@gmail.com>
Date: Fri, 19 Jan 2018 07:47:09 +0200
Subject: [PATCH] Support for setting _source

---
 src/Search/Search.php       | 27 +++++++++++++++++++++++++++
 tests/Search/SearchTest.php | 23 +++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/Search/Search.php b/src/Search/Search.php
index 6752fa4..14a8097 100644
--- a/src/Search/Search.php
+++ b/src/Search/Search.php
@@ -21,6 +21,11 @@ class Search
      */
     private $query;
 
+    /**
+     * @var array
+     */
+    private $source;
+
     /**
      * @var Sort
      */
@@ -135,6 +140,24 @@ public function getSort()
         return $this->sort;
     }
 
+    /**
+     * @param array $source
+     *
+     * @return $this
+     */
+    public function setSource(array $source)
+    {
+        $this->source = $source;
+        return $this;
+    }
+
+    /**
+     * @return array
+     */
+    public function getSource()
+    {
+        return $this->source;
+    }
 
     /**
      * @return AggregationCollection
@@ -253,6 +276,10 @@ public function buildBody()
             }
         }
 
+        if (!empty($this->getSource())) {
+            $body['_source'] = $this->getSource();
+        }
+
         $aggregations = $this->getAggregations();
         if ($aggregations->count() > 0) {
             $body['aggs'] = $aggregations->toArray();
diff --git a/tests/Search/SearchTest.php b/tests/Search/SearchTest.php
index 3f3051f..f01f35c 100644
--- a/tests/Search/SearchTest.php
+++ b/tests/Search/SearchTest.php
@@ -86,6 +86,9 @@ public function testSetterGetter()
         $this->search->setSize(100);
         $this->assertEquals(100, $this->search->getSize());
 
+        $this->search->setSource(['*']);
+        $this->assertEquals(['*'], $this->search->getSource());
+
         $this->search->setSort($this->sort);
         $this->assertInstanceOf(Sort::class, $this->search->getSort());
 
@@ -173,6 +176,26 @@ public function testToArray()
         ], $this->search->buildBody());
     }
 
+    public function testToArrayWithSource()
+    {
+        $this->search = $this->service->createSearch();
+        $this->search->setPage(1);
+        $this->search->setSize(100);
+        $this->search->setSource(['id', 'title', 'description']);
+
+        $this->assertEquals([
+            'query' => ['match_all' => new \stdClass()],
+            'size' => 100,
+            'from' => 0,
+            '_source' => [
+                'id',
+                'title',
+                'description'
+            ],
+        ], $this->search->buildBody());
+    }
+
+
     /**
      * Test adding an array of aggregation to Search
      */