Merge pull request #2210 from 10up/add/get_documents

Added Elasticsearch::get_documents and Indexable::multi_get
10up · May 28, 2021 · de9d2ae · de9d2ae
2 parents 5f76f26 + 8fe3a14
commit de9d2ae
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 0 deletions.
diff --git a/includes/classes/Elasticsearch.php b/includes/classes/Elasticsearch.php
@@ -650,6 +650,67 @@ public function delete_network_alias( $alias ) {
 		return false;
 	}
 
+	/**
+	 * Get multiple documents from Elasticsearch given an array of ids
+	 *
+	 * @param  string $index Index name.
+	 * @param  string $type Index type. Previously this was used for index type. Now it's just passed to hooks for legacy reasons.
+	 * @param  array  $document_ids Array of document ids to get.
+	 * @since  3.6.0
+	 * @return boolean|array
+	 */
+	public function get_documents( $index, $type, $document_ids ) {
+		if ( version_compare( $this->get_elasticsearch_version(), '7.0', '<' ) ) {
+			$path = apply_filters( 'ep_index_' . $type . '_request_path', $index . '/' . $type . '/_mget', $document_ids, $type );
+		} else {
+			$path = apply_filters( 'ep_index_' . $type . '_request_path', $index . '/_doc/_mget', $document_ids, $type );
+		}
+
+		$request_args = [
+			'method' => 'POST',
+			'body'   => wp_json_encode(
+				array(
+					'ids' => $document_ids,
+				)
+			),
+		];
+
+		$request = $this->remote_request( $path, $request_args, [], 'post' );
+
+		if ( is_wp_error( $request ) ) {
+			return false;
+		}
+
+		$response_body = wp_remote_retrieve_body( $request );
+
+		$response = json_decode( $response_body, true );
+
+		$docs = [];
+
+		if ( isset( $response['docs'] ) && is_array( $response['docs'] ) ) {
+			foreach ( $response['docs'] as $doc ) {
+				if ( ! empty( $doc['exists'] ) || ! empty( $doc['found'] ) ) {
+					$docs[ $doc['_id'] ] = $doc['_source'];
+				}
+			}
+		}
+
+		/**
+		 * Filter documents found by Elasticsearch through the /_mget endpoint.
+		 *
+		 * @hook ep_get_documents
+		 * @since 3.6.0
+		 * @param {array} $docs Documents found indexed by ID
+		 * @param  {string} $index Index name
+		 * @param  {string} $type Index type
+		 * @param  {array} $document_ids Array of document ids
+		 * @return  {array} Documents to be returned
+		 */
+		$docs = apply_filters( 'ep_get_documents', $docs, $index, $type, $document_ids );
+
+		return $docs;
+	}
+
 	/**
 	 * Create the network alias.
 	 *

diff --git a/includes/classes/Indexable.php b/includes/classes/Indexable.php
@@ -208,6 +208,17 @@ public function get( $object_id ) {
 		return Elasticsearch::factory()->get_document( $this->get_index_name(), $this->slug, $object_id );
 	}
 
+	/**
+	 * Get objects within the indexable
+	 *
+	 * @param  int $object_ids Array of object ids to get.
+	 * @since  3.6.0
+	 * @return boolean|array
+	 */
+	public function multi_get( $object_ids ) {
+		return Elasticsearch::factory()->get_documents( $this->get_index_name(), $this->slug, $object_ids );
+	}
+
 	/**
 	 * Delete an index within the indexable
 	 *

diff --git a/tests/php/TestElasticsearch.php b/tests/php/TestElasticsearch.php
@@ -51,4 +51,48 @@ public function testGetClusterStatus() {
 
 		}
 	}
+
+	/**
+	 * Test get documents
+	 *
+	 * @since 3.6.0
+	 * @group elasticsearch
+	 */
+	public function testGetDocuments() {
+
+		$post_ids = array();
+		$post_ids[] = Functions\create_and_sync_post();
+		$post_ids[] = Functions\create_and_sync_post();
+
+		ElasticPress\Elasticsearch::factory()->refresh_indices();
+
+		$index_name = ElasticPress\Indexables::factory()->get( 'post' )->get_index_name();
+
+		$documents = ElasticPress\Elasticsearch::factory()->get_documents( $index_name, 'post', $post_ids );
+
+		$this->assertIsArray( $documents );
+		$this->assertEquals( 2, count( $documents ) );
+		$this->assertArrayHasKey( $post_ids[0], $documents );
+		$this->assertArrayHasKey( $post_ids[1], $documents );
+
+		$post_ids[] = 99999999; // Adding an id that doesn't exist
+
+		$documents = ElasticPress\Elasticsearch::factory()->get_documents( $index_name, 'post', $post_ids );
+
+		$this->assertIsArray( $documents );
+		$this->assertEquals( 2, count( $documents ) );
+		$this->assertArrayHasKey( $post_ids[0], $documents );
+		$this->assertArrayHasKey( $post_ids[1], $documents );
+
+		// Trying to get a document that doesn't exist
+		$documents = ElasticPress\Elasticsearch::factory()->get_documents( $index_name, 'post', [ 99999999 ] );
+
+		$this->assertIsArray( $documents );
+		$this->assertEmpty( $documents );
+
+		$documents = ElasticPress\Elasticsearch::factory()->get_documents( $index_name, 'post', []  );
+
+		$this->assertIsArray( $documents );
+		$this->assertEmpty( $documents );
+	}
 }