[9.4] Improve server error handling in Apache Arrow helper (#3276) (#3277)

Co-authored-by: Josh Mock <joshua.mock@elastic.co>

Author: github-actions[bot]

src/helpers.ts

@@ -1024,13 +1024,34 @@ export default class Helpers {
         if (metaHeader !== null) {
           reqOptions.headers = reqOptions.headers ?? {}
           reqOptions.headers['x-elastic-client-meta'] = `${metaHeader as string},h=qa`
-          reqOptions.asStream = true
         }
 
         params.format = 'arrow'
+        reqOptions.asStream = true
+        reqOptions.meta = true
+
+        const result = await client.esql.query(params, reqOptions) as unknown as TransportResult<Readable, unknown>
+        if (result.statusCode >= 400) {
+          const chunks: Buffer[] = []
+          if (Buffer.isBuffer(result.body)) {
+            chunks.push(result.body)
+          } else {
+            for await (const chunk of result.body) {
+              chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk))
+            }
+          }
+          const body = Buffer.concat(chunks)
+          let parsedBody: any
+          try {
+            parsedBody = JSON.parse(body.toString())
+          } catch {
+            parsedBody = body.toString()
+          }
+          result.body = parsedBody
+          throw new ResponseError(result as any)
+        }
 
-        const response = await client.esql.query(params, reqOptions) as unknown as Readable
-        return await AsyncRecordBatchStreamReader.from(response)
+        return await AsyncRecordBatchStreamReader.from(result.body)
       }
     }
 

test/unit/helpers/bulk.test.ts

@@ -904,6 +904,32 @@ test('bulk index', t => {
       t.ok(called, 'async onFlush should be awaited')
     })
 
+    t.test('onFlush error causes bulk to reject', async t => {
+      const MockConnection = connection.buildMockConnection({
+        onRequest () {
+          return { body: { errors: false, items: [{}] } }
+        }
+      })
+
+      const client = new Client({
+        node: 'http://localhost:9200',
+        Connection: MockConnection
+      })
+
+      await t.rejects(
+        client.helpers.bulk({
+          datasource: dataset.slice(),
+          onDocument (doc) {
+            return { index: { _index: 'test' } }
+          },
+          onFlush () {
+            return Promise.reject(new Error('onFlush failed'))
+          }
+        }),
+        { message: 'onFlush failed' }
+      )
+    })
+
     t.end()
   })
 

test/unit/helpers/esql.test.ts

@@ -6,7 +6,7 @@
 import { test } from 'tap'
 import * as arrow from 'apache-arrow'
 import { connection } from '../../utils'
-import { Client } from '../../../'
+import { Client, errors } from '../../../'
 
 test('ES|QL helper', t => {
   test('toRecords', t => {
@@ -312,6 +312,77 @@ test('ES|QL helper', t => {
       t.end()
     })
 
+    t.test('Throws ResponseError when server returns a JSON error', async t => {
+      const errorBody = {
+        error: {
+          root_cause: [{
+            type: 'illegal_argument_exception',
+            reason: 'ES|QL type [date_range] is not supported by the Arrow format'
+          }],
+          type: 'illegal_argument_exception',
+          reason: 'ES|QL type [date_range] is not supported by the Arrow format'
+        },
+        status: 400
+      }
+
+      const MockConnection = connection.buildMockConnection({
+        onRequest (_params) {
+          return {
+            body: JSON.stringify(errorBody),
+            statusCode: 400,
+            headers: {
+              'content-type': 'application/vnd.elasticsearch+json;compatible-with=9'
+            }
+          }
+        }
+      })
+
+      const client = new Client({
+        node: 'http://localhost:9200',
+        Connection: MockConnection
+      })
+
+      try {
+        await client.helpers.esql({ query: 'FROM test-range | LIMIT 1' }).toArrowReader()
+        t.fail('Should throw')
+      } catch (err: any) {
+        t.ok(err instanceof errors.ResponseError)
+        t.equal(err.statusCode, 400)
+        t.equal(err.message, 'illegal_argument_exception\n\tRoot causes:\n\t\tillegal_argument_exception: ES|QL type [date_range] is not supported by the Arrow format')
+      }
+
+      t.end()
+    })
+
+    t.test('Throws ResponseError with non-JSON error body', async t => {
+      const MockConnection = connection.buildMockConnection({
+        onRequest (_params) {
+          return {
+            body: 'internal server error',
+            statusCode: 500,
+            headers: {
+              'content-type': 'text/plain'
+            }
+          }
+        }
+      })
+
+      const client = new Client({
+        node: 'http://localhost:9200',
+        Connection: MockConnection
+      })
+
+      try {
+        await client.helpers.esql({ query: 'FROM test | LIMIT 1' }).toArrowReader()
+        t.fail('Should throw')
+      } catch (err: any) {
+        t.ok(err instanceof errors.ResponseError)
+        t.equal(err.statusCode, 500)
+      }
+
+      t.end()
+    })
+
     t.end()
   })
 

test/unit/helpers/msearch.test.ts

@@ -758,3 +758,34 @@ test('Should use req options', async t => {
 
   t.teardown(() => m.stop())
 })
+
+test('documents returns empty array when hits is absent', async t => {
+  const MockConnection = connection.buildMockConnection({
+    onRequest (_params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            aggregations: { count: { value: 5 } }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const m = client.helpers.msearch({ operations: 1 })
+
+  const result = await m.search(
+    { index: 'test' },
+    { query: { match_all: {} } }
+  )
+
+  t.same(result.documents, [])
+
+  t.teardown(() => m.stop())
+})