From 94a8f08926db4ea00db373bf84b0e79944a00c9a Mon Sep 17 00:00:00 2001
From: Felipe Elia <felipe.elia@get10up.com>
Date: Tue, 27 Feb 2024 16:11:51 -0300
Subject: [PATCH 1/2] Add aggregations back to the query

---
 includes/classes/Elasticsearch.php | 20 +++++++++++++++-----
 tests/php/indexables/TestPost.php  | 29 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/includes/classes/Elasticsearch.php b/includes/classes/Elasticsearch.php
index 69f8990173..cfc85c6aea 100644
--- a/includes/classes/Elasticsearch.php
+++ b/includes/classes/Elasticsearch.php
@@ -404,6 +404,12 @@ public function query( $index, $type, $query, $query_args, $query_object = null
 				 * @param  {array} $query_args Current WP Query arguments
 				 */
 				do_action( 'ep_retrieve_aggregations', $response['aggregations'], $query, '', $query_args );
+
+				if ( method_exists( $query_object, 'set' ) ) {
+					$query_object->set( 'ep_aggregations', $response['aggregations'] );
+				} else {
+					$query_object->query_vars['ep_aggregations'] = $response['aggregations'];
+				}
 			}
 
 			/**
@@ -456,11 +462,15 @@ public function query( $index, $type, $query, $query_args, $query_object = null
 			 *
 			 * @hook ep_es_query_results
 			 * @param {array} $results Results from Elasticsearch
-			 * @param  {response} $response Raw response from Elasticsearch
-			 * @param  {array} $query Raw Elasticsearch query
-			 * @param  {array} $query_args Query arguments
-			 * @param  {mixed} $query_object Could be WP_Query, WP_User_Query, etc.
-			 * @return  {array} New results
+			 *      @param {int}   $results.found_documents Total number of documents.
+			 *      @param {array} $results.documents       Array of documents.
+			 *      @param {array} $results.aggregations    Array of aggregations.
+			 *      @param {array} $results.suggest         Array of suggestions.
+			 * @param {response} $response Raw response from Elasticsearch
+			 * @param {array} $query Raw Elasticsearch query
+			 * @param {array} $query_args Query arguments
+			 * @param {mixed} $query_object Could be WP_Query, WP_User_Query, etc.
+			 * @return {array} New results
 			 */
 			return apply_filters(
 				'ep_es_query_results',
diff --git a/tests/php/indexables/TestPost.php b/tests/php/indexables/TestPost.php
index 0a5bf7f02f..ea3252754a 100644
--- a/tests/php/indexables/TestPost.php
+++ b/tests/php/indexables/TestPost.php
@@ -9146,4 +9146,33 @@ public function test_mapping_ep_stop_filter() {
 		$index_settings = $settings[ $index_name ]['settings'];
 		$this->assertSame( '_arabic_', $index_settings['index.analysis.filter.ep_stop.stopwords'] );
 	}
+
+	/**
+	 * Test if aggregations are set
+	 *
+	 * @since 5.1.0
+	 * @group post
+	 */
+	public function test_aggregations_return() {
+		$query = new \WP_Query(
+			[
+				'ep_integrate' => true,
+				'fields'       => 'ids',
+				'aggs'         => [
+					'name' => 'my_aggs',
+					'aggs' => [
+						'terms' => [
+							'size'  => 10000,
+							'field' => 'terms.category.slug',
+						],
+					],
+				],
+				'ep_custom_id' => 'my_query',
+			]
+		);
+
+		$this->assertTrue( $query->elasticsearch_success );
+		$this->assertArrayHasKey( 'ep_aggregations', $query->query_vars );
+		$this->assertArrayHasKey( 'my_aggs', $query->query_vars['ep_aggregations'] );
+	}
 }

From b3d8ab59de778eeea8ef8e9f4888e720d3960dee Mon Sep 17 00:00:00 2001
From: Felipe Elia <felipe.elia@get10up.com>
Date: Wed, 28 Feb 2024 10:00:00 -0300
Subject: [PATCH 2/2] Make sure $query_object is in fact an object (it could be
 null)

---
 includes/classes/Elasticsearch.php | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/includes/classes/Elasticsearch.php b/includes/classes/Elasticsearch.php
index cfc85c6aea..e66a23547c 100644
--- a/includes/classes/Elasticsearch.php
+++ b/includes/classes/Elasticsearch.php
@@ -405,10 +405,12 @@ public function query( $index, $type, $query, $query_args, $query_object = null
 				 */
 				do_action( 'ep_retrieve_aggregations', $response['aggregations'], $query, '', $query_args );
 
-				if ( method_exists( $query_object, 'set' ) ) {
-					$query_object->set( 'ep_aggregations', $response['aggregations'] );
-				} else {
-					$query_object->query_vars['ep_aggregations'] = $response['aggregations'];
+				if ( is_object( $query_object ) ) {
+					if ( method_exists( $query_object, 'set' ) ) {
+						$query_object->set( 'ep_aggregations', $response['aggregations'] );
+					} else {
+						$query_object->query_vars['ep_aggregations'] = $response['aggregations'];
+					}
 				}
 			}