Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,8 @@
"zx": "^7.2.2"
},
"dependencies": {
"@elastic/transport": "^8.8.1",
"@elastic/transport": "^8.9.0",
"@apache-arrow/esnext-cjs": "^17.0.0",
"tslib": "^2.4.0"
},
"tap": {
Expand Down
27 changes: 22 additions & 5 deletions src/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import assert from 'node:assert'
import * as timersPromises from 'node:timers/promises'
import { Readable } from 'node:stream'
import { errors, TransportResult, TransportRequestOptions, TransportRequestOptionsWithMeta } from '@elastic/transport'
import { Table, TypeMap, tableFromIPC } from '@apache-arrow/esnext-cjs'
import Client from './client'
import * as T from './api/types'

Expand Down Expand Up @@ -155,6 +156,7 @@ export interface EsqlResponse {

export interface EsqlHelper {
toRecords: <TDocument>() => Promise<EsqlToRecords<TDocument>>
toArrow: () => Promise<Table<TypeMap>>
}

export interface EsqlToRecords<TDocument> {
Expand Down Expand Up @@ -965,11 +967,6 @@ export default class Helpers {
* @returns {object} EsqlHelper instance
*/
esql (params: T.EsqlQueryRequest, reqOptions: TransportRequestOptions = {}): EsqlHelper {
if (this[kMetaHeader] !== null) {
reqOptions.headers = reqOptions.headers ?? {}
reqOptions.headers['x-elastic-client-meta'] = `${this[kMetaHeader] as string},h=qo`
}

const client = this[kClient]

function toRecords<TDocument> (response: EsqlResponse): TDocument[] {
Expand All @@ -985,17 +982,37 @@ export default class Helpers {
})
}

const metaHeader = this[kMetaHeader]

const helper: EsqlHelper = {
/**
* Pivots ES|QL query results into an array of row objects, rather than the default format where each row is an array of values.
*/
async toRecords<TDocument>(): Promise<EsqlToRecords<TDocument>> {
if (metaHeader !== null) {
reqOptions.headers = reqOptions.headers ?? {}
reqOptions.headers['x-elastic-client-meta'] = `${metaHeader as string},h=qo`
}

params.format = 'json'
params.columnar = false
// @ts-expect-error it's typed as ArrayBuffer but we know it will be JSON
const response: EsqlResponse = await client.esql.query(params, reqOptions)
const records: TDocument[] = toRecords(response)
const { columns } = response
return { records, columns }
},

async toArrow (): Promise<Table<TypeMap>> {
if (metaHeader !== null) {
reqOptions.headers = reqOptions.headers ?? {}
reqOptions.headers['x-elastic-client-meta'] = `${metaHeader as string},h=qa`
}

params.format = 'arrow'

const response = await client.esql.query(params, reqOptions)
return tableFromIPC(response)
}
}

Expand Down
62 changes: 62 additions & 0 deletions test/unit/helpers/esql.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
*/

import { test } from 'tap'
import { Table } from '@apache-arrow/esnext-cjs'
import { connection } from '../../utils'
import { Client } from '../../../'

Expand Down Expand Up @@ -109,5 +110,66 @@ test('ES|QL helper', t => {

t.end()
})

test('toArrow', t => {
t.test('Parses a binary response into an Arrow table', async t => {
const binaryContent = '/////zABAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAIAAAB8AAAABAAAAJ7///8UAAAARAAAAEQAAAAAAAoBRAAAAAEAAAAEAAAAjP///wgAAAAQAAAABAAAAGRhdGUAAAAADAAAAGVsYXN0aWM6dHlwZQAAAAAAAAAAgv///wAAAQAEAAAAZGF0ZQAAEgAYABQAEwASAAwAAAAIAAQAEgAAABQAAABMAAAAVAAAAAAAAwFUAAAAAQAAAAwAAAAIAAwACAAEAAgAAAAIAAAAEAAAAAYAAABkb3VibGUAAAwAAABlbGFzdGljOnR5cGUAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAGAAAAYW1vdW50AAAAAAAA/////7gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAABgAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAABYAAAABQAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAACgAAAAAAAAAMAAAAAAAAAABAAAAAAAAADgAAAAAAAAAKAAAAAAAAAAAAAAAAgAAAAUAAAAAAAAAAAAAAAAAAAAFAAAAAAAAAAAAAAAAAAAAHwAAAAAAAAAAAACgmZkTQAAAAGBmZiBAAAAAAAAAL0AAAADAzMwjQAAAAMDMzCtAHwAAAAAAAADV6yywkgEAANWPBquSAQAA1TPgpZIBAADV17mgkgEAANV7k5uSAQAA/////wAAAAA='

const MockConnection = connection.buildMockConnection({
onRequest (_params) {
return {
body: Buffer.from(binaryContent, 'base64'),
statusCode: 200,
headers: {
'content-type': 'application/vnd.elasticsearch+arrow+stream'
}
}
}
})

const client = new Client({
node: 'http://localhost:9200',
Connection: MockConnection
})

const result = await client.helpers.esql({ query: 'FROM sample_data' }).toArrow()
t.ok(result instanceof Table)

const table = [...result]
t.same(table[0], [
["amount", 4.900000095367432],
["date", 1729532586965],
])
t.end()
})

t.test('ESQL helper uses correct x-elastic-client-meta helper value', async t => {
const binaryContent = '/////zABAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAIAAAB8AAAABAAAAJ7///8UAAAARAAAAEQAAAAAAAoBRAAAAAEAAAAEAAAAjP///wgAAAAQAAAABAAAAGRhdGUAAAAADAAAAGVsYXN0aWM6dHlwZQAAAAAAAAAAgv///wAAAQAEAAAAZGF0ZQAAEgAYABQAEwASAAwAAAAIAAQAEgAAABQAAABMAAAAVAAAAAAAAwFUAAAAAQAAAAwAAAAIAAwACAAEAAgAAAAIAAAAEAAAAAYAAABkb3VibGUAAAwAAABlbGFzdGljOnR5cGUAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAGAAAAYW1vdW50AAAAAAAA/////7gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAABgAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAABYAAAABQAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAACgAAAAAAAAAMAAAAAAAAAABAAAAAAAAADgAAAAAAAAAKAAAAAAAAAAAAAAAAgAAAAUAAAAAAAAAAAAAAAAAAAAFAAAAAAAAAAAAAAAAAAAAHwAAAAAAAAAAAACgmZkTQAAAAGBmZiBAAAAAAAAAL0AAAADAzMwjQAAAAMDMzCtAHwAAAAAAAADV6yywkgEAANWPBquSAQAA1TPgpZIBAADV17mgkgEAANV7k5uSAQAA/////wAAAAA='

const MockConnection = connection.buildMockConnection({
onRequest (params) {
const header = params.headers?.['x-elastic-client-meta'] ?? ''
t.ok(header.includes('h=qa'), `Client meta header does not include ESQL helper value: ${header}`)
return {
body: Buffer.from(binaryContent, 'base64'),
statusCode: 200,
headers: {
'content-type': 'application/vnd.elasticsearch+arrow+stream'
}
}
}
})

const client = new Client({
node: 'http://localhost:9200',
Connection: MockConnection
})

await client.helpers.esql({ query: 'FROM sample_data' }).toArrow()
t.end()
})

t.end()
})
t.end()
})
5 changes: 3 additions & 2 deletions tsconfig.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"compilerOptions": {
"target": "es2019",
"target": "ES2019",
"module": "commonjs",
"moduleResolution": "node",
"declaration": true,
Expand All @@ -21,7 +21,8 @@
"importHelpers": true,
"outDir": "lib",
"lib": [
"esnext"
"ES2019",
"dom"
]
},
"formatCodeOptions": {
Expand Down
Loading