Add multi search helper (#1186)

test/integration/helpers/msearch.test.js

@@ -0,0 +1,106 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+'use strict'
+
+const { createReadStream } = require('fs')
+const { join } = require('path')
+const split = require('split2')
+const { test, beforeEach, afterEach } = require('tap')
+const { waitCluster } = require('../../utils')
+const { Client, errors } = require('../../../')
+
+const INDEX = `test-helpers-${process.pid}`
+const client = new Client({
+  node: process.env.TEST_ES_SERVER || 'http://localhost:9200'
+})
+
+beforeEach(async () => {
+  await waitCluster(client)
+  await client.indices.create({ index: INDEX })
+  const stream = createReadStream(join(__dirname, '..', '..', 'fixtures', 'stackoverflow.ndjson'))
+  const result = await client.helpers.bulk({
+    datasource: stream.pipe(split()),
+    refreshOnCompletion: true,
+    onDocument (doc) {
+      return {
+        index: { _index: INDEX }
+      }
+    }
+  })
+  if (result.failed > 0) {
+    throw new Error('Failed bulk indexing docs')
+  }
+})
+
+afterEach(async () => {
+  await client.indices.delete({ index: INDEX }, { ignore: 404 })
+})
+
+test('Basic', t => {
+  t.plan(4)
+  const s = client.helpers.msearch({ operations: 1 })
+
+  s.search(
+    { index: INDEX },
+    { query: { match: { title: 'javascript' } } },
+    (err, result) => {
+      t.error(err)
+      t.strictEqual(result.body.hits.total.value, 106)
+    }
+  )
+
+  s.search(
+    { index: INDEX },
+    { query: { match: { title: 'ruby' } } },
+    (err, result) => {
+      t.error(err)
+      t.strictEqual(result.body.hits.total.value, 29)
+    }
+  )
+
+  t.teardown(() => s.stop())
+})
+
+test('Bad request', t => {
+  t.plan(3)
+  const s = client.helpers.msearch({ operations: 1 })
+
+  s.search(
+    { index: INDEX },
+    { query: { match: { title: 'javascript' } } },
+    (err, result) => {
+      t.error(err)
+      t.strictEqual(result.body.hits.total.value, 106)
+    }
+  )
+
+  s.search(
+    { index: INDEX },
+    { query: { foo: { title: 'ruby' } } },
+    (err, result) => {
+      t.true(err instanceof errors.ResponseError)
+    }
+  )
+
+  t.teardown(() => s.stop())
+})
+
+test('Send multiple request concurrently over the concurrency limit', t => {
+  t.plan(20)
+  const s = client.helpers.msearch({ operations: 1 })
+
+  for (let i = 0; i < 10; i++) {
+    s.search(
+      { index: INDEX },
+      { query: { match: { title: 'javascript' } } },
+      (err, result) => {
+        t.error(err)
+        t.strictEqual(result.body.hits.total.value, 106)
+      }
+    )
+  }
+
+  t.teardown(() => s.stop())
+})

test/types/helpers.test-d.ts

@@ -9,8 +9,10 @@ import {
   BulkStats,
   BulkHelperOptions,
   ScrollSearchResponse,
-  OnDropDocument
+  OnDropDocument,
+  MsearchHelper
 } from '../../lib/Helpers'
+import { ApiResponse, ApiError } from '../../lib/Transport'
 
 const client = new Client({
   node: 'http://localhost:9200'
@@ -429,4 +431,29 @@ expectError(
 
   expectType<Promise<Source[]>>(p)
   expectType<Source[]>(await p)
-}
+}
+
+/// .helpers.msearch
+
+const s = client.helpers.msearch({
+  operations: 20,
+  concurrency: 5,
+  retries: 5,
+  wait: 5000
+})
+
+expectType<MsearchHelper>(s)
+expectType<void>(s.stop())
+expectType<void>(s.stop(new Error('kaboom')))
+
+expectType<Promise<ApiResponse<Record<string, any>, unknown>>>(s.search({ index: 'foo'}, { query: {} }))
+expectType<Promise<ApiResponse<string, string>>>(s.search<string, Record<string, any>, string>({ index: 'foo'}, { query: {} }))
+
+expectType<void>(s.search({ index: 'foo'}, { query: {} }, (err, result) => {
+  expectType<ApiError>(err)
+  expectType<ApiResponse>(result)
+}))
+expectType<void>(s.search<string, Record<string, any>, string>({ index: 'foo'}, { query: {} }, (err, result) => {
+  expectType<ApiError>(err)
+  expectType<ApiResponse<string, string>>(result)
+}))

test/unit/helpers/msearch.test.js

@@ -0,0 +1,580 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+'use strict'
+
+const { test } = require('tap')
+const { Client, errors } = require('../../../')
+const { connection } = require('../../utils')
+
+test('Basic', async t => {
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1 })
+
+  const result = await s.search(
+    { index: 'test' },
+    { query: { match: { foo: 'bar' } } }
+  )
+
+  t.deepEqual(result.body, {
+    status: 200,
+    hits: {
+      hits: [
+        { _source: { one: 'one' } },
+        { _source: { two: 'two' } },
+        { _source: { three: 'three' } }
+      ]
+    }
+  })
+
+  t.deepEqual(result.documents, [
+    { one: 'one' },
+    { two: 'two' },
+    { three: 'three' }
+  ])
+
+  t.teardown(() => s.stop())
+})
+
+test('Multiple searches (inside async iterator)', t => {
+  t.plan(6)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }, {
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { four: 'four' } },
+                { _source: { five: 'five' } },
+                { _source: { six: 'six' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 2 })
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { one: 'one' } },
+          { _source: { two: 'two' } },
+          { _source: { three: 'three' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { one: 'one' },
+      { two: 'two' },
+      { three: 'three' }
+    ])
+  })
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { four: 'four' } },
+          { _source: { five: 'five' } },
+          { _source: { six: 'six' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { four: 'four' },
+      { five: 'five' },
+      { six: 'six' }
+    ])
+  })
+
+  t.teardown(() => s.stop())
+})
+
+test('Multiple searches (async iterator exits)', t => {
+  t.plan(6)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }, {
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { four: 'four' } },
+                { _source: { five: 'five' } },
+                { _source: { six: 'six' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { one: 'one' } },
+          { _source: { two: 'two' } },
+          { _source: { three: 'three' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { one: 'one' },
+      { two: 'two' },
+      { three: 'three' }
+    ])
+  })
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { four: 'four' } },
+          { _source: { five: 'five' } },
+          { _source: { six: 'six' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { four: 'four' },
+      { five: 'five' },
+      { six: 'six' }
+    ])
+  })
+
+  setImmediate(() => s.stop())
+})
+
+test('Stop a msearch processor (promises)', async t => {
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1 })
+
+  s.stop()
+
+  try {
+    await s.search(
+      { index: 'test' },
+      { query: { match: { foo: 'bar' } } }
+    )
+  } catch (err) {
+    t.strictEqual(err.message, 'The msearch processor has been stopped')
+  }
+
+  t.teardown(() => s.stop())
+})
+
+test('Stop a msearch processor (callbacks)', t => {
+  t.plan(1)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.stop()
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.strictEqual(err.message, 'The msearch processor has been stopped')
+  })
+})
+
+test('Bad header', t => {
+  t.plan(1)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.search(null, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.strictEqual(err.message, 'The header should be an object')
+  })
+
+  t.teardown(() => s.stop())
+})
+
+test('Bad body', t => {
+  t.plan(1)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.search({ index: 'test' }, null, (err, result) => {
+    t.strictEqual(err.message, 'The body should be an object')
+  })
+
+  t.teardown(() => s.stop())
+})
+
+test('Retry on 429', async t => {
+  let count = 0
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      if (count++ === 0) {
+        return {
+          body: {
+            responses: [{
+              status: 429,
+              error: {}
+            }]
+          }
+        }
+      } else {
+        return {
+          body: {
+            responses: [{
+              status: 200,
+              hits: {
+                hits: [
+                  { _source: { one: 'one' } },
+                  { _source: { two: 'two' } },
+                  { _source: { three: 'three' } }
+                ]
+              }
+            }]
+          }
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1, wait: 10 })
+
+  const result = await s.search(
+    { index: 'test' },
+    { query: { match: { foo: 'bar' } } }
+  )
+
+  t.deepEqual(result.body, {
+    status: 200,
+    hits: {
+      hits: [
+        { _source: { one: 'one' } },
+        { _source: { two: 'two' } },
+        { _source: { three: 'three' } }
+      ]
+    }
+  })
+
+  t.deepEqual(result.documents, [
+    { one: 'one' },
+    { two: 'two' },
+    { three: 'three' }
+  ])
+
+  t.teardown(() => s.stop())
+})
+
+test('Single search errors', async t => {
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 400,
+            error: { foo: 'bar' }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1 })
+
+  try {
+    await s.search(
+      { index: 'test' },
+      { query: { match: { foo: 'bar' } } }
+    )
+  } catch (err) {
+    t.true(err instanceof errors.ResponseError)
+  }
+
+  t.teardown(() => s.stop())
+})
+
+test('Entire msearch fails', t => {
+  t.plan(4)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        statusCode: 500,
+        body: {
+          status: 500,
+          error: { foo: 'bar' }
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1 })
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.true(err instanceof errors.ResponseError)
+    t.deepEqual(result.documents, [])
+  })
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.true(err instanceof errors.ResponseError)
+    t.deepEqual(result.documents, [])
+  })
+
+  t.teardown(() => s.stop())
+})
+
+test('Resolves the msearch helper', t => {
+  t.plan(1)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.stop()
+
+  s.then(
+    () => t.pass('called'),
+    e => t.fail('Should not fail')
+  )
+
+  s.catch(e => t.fail('Should not fail'))
+})
+
+test('Stop the msearch helper with an error', t => {
+  t.plan(3)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {}
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.stop(new Error('kaboom'))
+
+  s.then(
+    () => t.fail('Should fail'),
+    err => t.is(err.message, 'kaboom')
+  )
+
+  s.catch(err => t.is(err.message, 'kaboom'))
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.is(err.message, 'kaboom')
+  })
+})
+
+test('Multiple searches (concurrency = 1)', t => {
+  t.plan(6)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch({ operations: 1, concurrency: 1 })
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { one: 'one' } },
+          { _source: { two: 'two' } },
+          { _source: { three: 'three' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { one: 'one' },
+      { two: 'two' },
+      { three: 'three' }
+    ])
+  })
+
+  s.search({ index: 'test' }, { query: {} }, (err, result) => {
+    t.error(err)
+    t.deepEqual(result.body, {
+      status: 200,
+      hits: {
+        hits: [
+          { _source: { one: 'one' } },
+          { _source: { two: 'two' } },
+          { _source: { three: 'three' } }
+        ]
+      }
+    })
+
+    t.deepEqual(result.documents, [
+      { one: 'one' },
+      { two: 'two' },
+      { three: 'three' }
+    ])
+  })
+
+  t.teardown(() => s.stop())
+})