Merge pull request #108 from axiomhq/islam/axm-22-tabular-result-form…

…at-support-for-axiom-node

feat: add tabular option for query()

e2e/nextjs-app/test/ingest.spec.ts

@@ -33,13 +33,20 @@ describe('Ingestion & query on different runtime', () => {
     await new Promise(r => setTimeout(r, 1000));
 
     // check dataset for ingested logs
-    const qResp = await axiom.query(`['${datasetName}'] | where ['test'] == "ingest_on_lambda"`, {
-      startTime,
-    });
-    expect(qResp.matches).toBeDefined();
-    expect(qResp.matches).toHaveLength(2);
-    expect(qResp.matches![0].data.foo).toEqual('bar');
-    expect(qResp.matches![1].data.bar).toEqual('baz');
+    const qResp = await axiom.query(
+      `['${datasetName}'] | where ['test'] == "ingest_on_lambda" | project _time, test, foo, bar`,
+      {
+        startTime,
+        format: 'tabular',
+      },
+    );
+    expect(qResp.status).toBeDefined();
+    expect(qResp.tables).toBeDefined();
+    expect(qResp.tables).toHaveLength(1);
+    expect(qResp.tables[0].columns).toHaveLength(4);
+    expect(qResp.tables[0].columns?.[1][0]).toEqual('ingest_on_lambda');
+    expect(qResp.tables[0].columns?.[2][0]).toEqual('bar');
+    expect(qResp.tables[0].columns?.[3][1]).toEqual('baz');
   });
 
   it('ingest on a edge function should succeed', async () => {
@@ -54,12 +61,18 @@ describe('Ingestion & query on different runtime', () => {
     await new Promise(resolve => setTimeout(resolve, 1000));
 
     // check dataset for ingested logs
-    const qResp = await axiom.query(`['${datasetName}'] | where ['test'] == "ingest_on_edge"`, {
-      startTime,
-    });
-    expect(qResp.matches).toBeDefined();
-    expect(qResp.matches).toHaveLength(2);
-    expect(qResp.matches![0].data.foo).toEqual('bar');
-    expect(qResp.matches![1].data.bar).toEqual('baz');
+    const qResp = await axiom.query(
+      `['${datasetName}'] | where ['test'] == "ingest_on_edge" | project _time, test, foo, bar`,
+      {
+        startTime,
+        format: 'tabular',
+      },
+    );
+    expect(qResp.status).toBeDefined();
+    expect(qResp.tables).toBeDefined();
+    expect(qResp.tables).toHaveLength(1);
+    expect(qResp.tables[0].columns).toHaveLength(4);
+    expect(qResp.tables[0].columns?.[2][0]).toEqual('bar');
+    expect(qResp.tables[0].columns?.[3][1]).toEqual('baz');
   });
 });

examples/js/src/query.ts

@@ -2,19 +2,23 @@
 // Processing Language (APL).
 import { Axiom } from '@axiomhq/js';
 
-const axiom = new Axiom({ token: process.env.AXIOM_TOKEN || ''});
+const axiom = new Axiom({ token: process.env.AXIOM_TOKEN || '', url: process.env.AXIOM_URL || '' });
 
 async function query() {
-  const aplQuery = "['my-dataset']";
+  const aplQuery = "['new-lambda-test']";
 
-  const res = await axiom.query(aplQuery);
-  if (!res.matches || res.matches.length === 0) {
-    console.warn('no matches found');
+  const res = await axiom.query(aplQuery, {
+    startTime: '2023-10-23T15:46:25.089482+02:00',
+    format: 'tabular',
+  });
+
+  if (!res.tables || res.tables.length === 0) {
+    console.warn('no tables found');
     return;
   }
 
-  for (let matched of res.matches) {
-    console.log(matched.data);
+  for (let table of res.tables) {
+    console.table(table);
   }
 }
 

integration/src/client.spec.ts

@@ -6,7 +6,11 @@ const datasetSuffix = process.env.AXIOM_DATASET_SUFFIX || 'local';
 
 describe('Axiom', () => {
   const datasetName = `test-axiom-js-client-${datasetSuffix}`;
-  const axiom = new AxiomWithoutBatching({ token: process.env.AXIOM_TOKEN || '', url: process.env.AXIOM_URL, orgId: process.env.AXIOM_ORG_ID });
+  const axiom = new AxiomWithoutBatching({
+    token: process.env.AXIOM_TOKEN || '',
+    url: process.env.AXIOM_URL,
+    orgId: process.env.AXIOM_ORG_ID,
+  });
 
   beforeAll(async () => {
     await axiom.datasets.create({
@@ -112,4 +116,29 @@ baz`,
       expect(result.matches?.length).toEqual(11);
     });
   });
+
+  describe('apl tabular query', () => {
+    it('returns a valid response', async () => {
+      const status = await axiom.ingest(datasetName, { test: 'apl' });
+      expect(status.ingested).toEqual(1);
+
+      // wait 1 sec for ingestion to finish
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+
+      const result = await axiom.query("['" + datasetName + "'] | where ['test'] == 'apl' | project _time, ['test']", {
+        format: 'tabular',
+      });
+      expect(result.status.rowsMatched).toEqual(1);
+      expect(result.tables?.length).toEqual(1);
+      expect(result.tables[0].columns?.length).toEqual(2); // _time and test
+      expect(result.tables[0].columns?.[0]).toBeDefined();
+      expect(result.tables[0].columns?.[1]).toBeDefined();
+      expect(result.tables[0].columns?.[1].length).toEqual(1); // only one row
+      expect(Array.from(result.tables[0].events()).length).toEqual(result.tables[0].columns?.[0].length);
+      expect(Object.keys(Array.from(result.tables[0].events())[0])).toEqual(result.tables[0].fields.map((f) => f.name));
+      expect(Array.from(result.tables[0].events())[0][result.tables[0].fields[0].name]).toEqual(
+        result.tables[0].columns?.[0][0],
+      );
+    });
+  });
 });

packages/js/src/client.ts

@@ -95,33 +95,75 @@ class BaseClient extends HTTPClient {
    *
    * @param apl - the apl query
    * @param options - optional query options
-   * @returns result of the query, check: {@link QueryResult}
+   * @returns result of the query depending on the format in options, check: {@link QueryResult} and {@link TabularQueryResult}
    *
    * @example
    * ```
    * await axiom.query("['dataset'] | count");
    * ```
    *
    */
-  query = (apl: string, options?: QueryOptions): Promise<QueryResult> => {
+  query = <
+    TOptions extends QueryOptions,
+    TResult = TOptions['format'] extends 'tabular' ? Promise<TabularQueryResult> : Promise<QueryResult>,
+  >(
+    apl: string,
+    options?: TOptions,
+  ): Promise<TResult> => {
     const req: Query = { apl: apl };
     if (options?.startTime) {
       req.startTime = options?.startTime;
     }
     if (options?.endTime) {
       req.endTime = options?.endTime;
     }
-    return this.client.post<QueryResult>(
-      this.localPath + '/datasets/_apl',
-      {
-        body: JSON.stringify(req),
-      },
-      {
-        'streaming-duration': options?.streamingDuration as string,
-        nocache: options?.noCache as boolean,
-        format: 'legacy',
-      },
-    );
+
+    return this.client
+      .post<TOptions['format'] extends 'tabular' ? RawTabularQueryResult : QueryResult>(
+        this.localPath + '/datasets/_apl',
+        {
+          body: JSON.stringify(req),
+        },
+        {
+          'streaming-duration': options?.streamingDuration as string,
+          nocache: options?.noCache as boolean,
+          format: options?.format ?? 'legacy',
+        },
+      )
+      .then((res) => {
+        if (options?.format !== 'tabular') {
+          return res;
+        }
+
+        const result = res as RawTabularQueryResult;
+        return {
+          ...res,
+          tables: result.tables.map((t) => {
+            return {
+              ...t,
+              events: function* () {
+                let iteration = 0;
+                if (!this.columns) {
+                  return;
+                }
+
+                while (iteration <= this.columns[0].length) {
+                  const value = Object.fromEntries(
+                    this.fields.map((field, fieldIdx) => [field.name, this.columns![fieldIdx][iteration]]),
+                  );
+
+                  if (iteration >= this.columns[0].length) {
+                    return value;
+                  }
+
+                  yield value;
+                  iteration++;
+                }
+              },
+            };
+          }),
+        };
+      }) as Promise<TResult>;
   };
 
   /**
@@ -137,7 +179,13 @@ class BaseClient extends HTTPClient {
    * await axiom.aplQuery("['dataset'] | count");
    * ```
    */
-  aplQuery = (apl: string, options?: QueryOptions): Promise<QueryResult> => this.query(apl, options);
+  aplQuery = <
+    TOptions extends QueryOptions,
+    TResult = TOptions['format'] extends 'tabular' ? Promise<TabularQueryResult> : Promise<QueryResult>,
+  >(
+    apl: string,
+    options?: TOptions,
+  ): Promise<TResult> => this.query(apl, options);
 }
 
 /**
@@ -313,6 +361,7 @@ export interface QueryOptionsBase {
 export interface QueryOptions extends QueryOptionsBase {
   startTime?: string;
   endTime?: string;
+  format?: 'legacy' | 'tabular';
 }
 
 export interface QueryLegacy {
@@ -337,6 +386,12 @@ export interface Aggregation {
   op: AggregationOp;
 }
 
+export interface TabularAggregation {
+  name: AggregationOp;
+  args: any[];
+  fields: string[];
+}
+
 export enum AggregationOp {
   Count = 'count',
   Distinct = 'distinct',
@@ -422,6 +477,49 @@ export interface QueryResult {
   status: Status;
 }
 
+export interface RawTabularQueryResult {
+  datasetNames: string[];
+  fieldsMetaMap: Record<
+    string,
+    Array<{ description: string; hidden: boolean; name: string; type: string; unit: string }>
+  >;
+  format: string;
+  status: Status;
+  tables: Array<RawAPLResultTable>;
+}
+
+export interface TabularQueryResult extends RawTabularQueryResult {
+  tables: Array<APLResultTable>;
+}
+
+export interface RawAPLResultTable {
+  name: string;
+  sources: Array<{ name: string }>;
+  fields: Array<{ name: string; type: string; agg?: TabularAggregation }>;
+  order: Array<{
+    name: string;
+    desc: boolean;
+  }>;
+  groups: Array<{ name: string }>;
+  range?: {
+    field: string;
+    start: string;
+    end: string;
+  };
+  buckets?: { field: string; size: any };
+  columns?: Array<Array<any>>;
+}
+
+export interface APLResultTable extends RawAPLResultTable {
+  /**
+   * Returns an iterable that yields each row of the table as a record,
+   * where the keys are the field names and the values are the values in the columns.
+   *
+   * @returns {Generator<Record<string, any>, undefined, unknown>}
+   */
+  events: () => Generator<Record<string, any>, undefined, unknown>;
+}
+
 export interface Timeseries {
   series?: Array<Interval>;
   totals?: Array<EntryGroup>;

packages/js/src/index.ts

@@ -1,4 +1,4 @@
-export { AxiomWithoutBatching, Axiom, ContentType, ContentEncoding, IngestOptions, IngestStatus, IngestFailure, QueryOptionsBase, QueryOptions, QueryLegacy, Aggregation, AggregationOp, Filter, FilterOp, Order, Projection, VirtualColumn, QueryLegacyResult, QueryResult, Timeseries, Interval, EntryGroup, EntryGroupAgg, Entry, Status, Message, Query } from './client.js';
+export { AxiomWithoutBatching, Axiom, ContentType, ContentEncoding, IngestOptions, IngestStatus, IngestFailure, QueryOptionsBase, QueryOptions, QueryLegacy, Aggregation, AggregationOp, Filter, FilterOp, Order, Projection, VirtualColumn, QueryResult, QueryLegacyResult, TabularQueryResult, RawTabularQueryResult, Timeseries, Interval, EntryGroup, EntryGroupAgg, Entry, Status, Message, Query } from './client.js';
 export { ClientOptions } from './httpClient.js';
 export { datasets } from './datasets.js';
 export { annotations } from './annotations.js';