Merge pull request #136 from axiomhq/arne/add-tabular

feat: Add support for tabular format

src/axiom_py/client.py

@@ -78,6 +78,7 @@ class AplResultFormat(Enum):
     """The result format of an APL query."""
 
     Legacy = "legacy"
+    Tabular = "tabular"
 
 
 class ContentType(Enum):

src/axiom_py/query/__init__.py

@@ -3,7 +3,6 @@
 from .filter import FilterOperation, BaseFilter, Filter
 from .aggregation import AggregationOperation, Aggregation
 from .result import (
-    MessageCode,
     MessagePriority,
     Message,
     QueryStatus,
@@ -28,7 +27,6 @@
     Filter,
     AggregationOperation,
     Aggregation,
-    MessageCode,
     MessagePriority,
     Message,
     QueryStatus,

src/axiom_py/query/result.py

@@ -1,21 +1,11 @@
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import List, Dict, Optional
+from typing import List, Dict, Optional, Union
 from enum import Enum
 
 from .query import QueryLegacy
 
 
-class MessageCode(Enum):
-    """Message codes represents the code associated with the query."""
-
-    UNKNOWN_MESSAGE_CODE = ""
-    VIRTUAL_FIELD_FINALIZE_ERROR = "virtual_field_finalize_error"
-    MISSING_COLUMN = "missing_column"
-    LICENSE_LIMIT_FOR_QUERY_WARNING = "license_limit_for_query_warning"
-    DEFAULT_LIMIT_WARNING = "default_limit_warning"
-
-
 class MessagePriority(Enum):
     """Message priorities represents the priority of a message associated with a query."""
 
@@ -36,7 +26,7 @@ class Message:
     # describes how often a message of this type was raised by the query.
     count: int
     # code of the message.
-    code: MessageCode
+    code: str
     # a human readable text representation of the message.
     msg: str
 
@@ -151,17 +141,122 @@ class QueryLegacyResult:
     savedQueryID: Optional[str] = field(default=None)
 
 
+@dataclass
+class Source:
+    name: str
+
+
+@dataclass
+class Order:
+    desc: bool
+    field: str
+
+
+@dataclass
+class Group:
+    name: str
+
+
+@dataclass
+class Range:
+    # Start is the starting time the query is limited by.
+    start: datetime
+    # End is the ending time the query is limited by.
+    end: datetime
+    # Field specifies the field name on which the query range was restricted.
+    field: str
+
+
+@dataclass
+class Aggregation:
+    # Args specifies any non-field arguments for the aggregation.
+    args: Optional[List[object]]
+    # Fields specifies the names of the fields this aggregation is computed on.
+    fields: Optional[List[str]]
+    # Name is the system name of the aggregation.
+    name: str
+
+
+@dataclass
+class Field:
+    name: str
+    type: str
+    agg: Optional[Aggregation]
+
+
+@dataclass
+class Bucket:
+    # Field specifies the field used to create buckets on.
+    field: str
+    # An integer or float representing the fixed bucket size.
+    size: Union[int, float]
+
+
+@dataclass
+class Table:
+    buckets: Optional[Bucket]
+    # Columns contain a series of arrays with the raw result data.
+    # The columns here line up with the fields in the Fields array.
+    columns: Optional[List[List[object]]]
+    # Fields contain information about the fields included in these results.
+    # The order of the fields match up with the order of the data in Columns.
+    fields: List[Field]
+    # Groups specifies which grouping operations has been performed on the
+    # results.
+    groups: List[Group]
+    # Name is the name assigned to this table. Defaults to "0".
+    # The name "_totals" is reserved for system use.
+    name: str
+    # Order echoes the ordering clauses that was used to sort the results.
+    order: List[Order]
+    range: Optional[Range]
+    # Sources contain the names of the datasets that contributed data to these
+    # results.
+    sources: List[Source]
+
+    def events(self):
+        return ColumnIterator(self)
+
+
+class ColumnIterator:
+    table: Table
+    i: int = 0
+
+    def __init__(self, table: Table):
+        self.table = table
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if (
+            self.table.columns is None
+            or len(self.table.columns) == 0
+            or self.i >= len(self.table.columns[0])
+        ):
+            raise StopIteration
+
+        event = {}
+        for j, f in enumerate(self.table.fields):
+            event[f.name] = self.table.columns[j][self.i]
+
+        self.i += 1
+        return event
+
+
 @dataclass
 class QueryResult:
     """Result is the result of apl query."""
 
-    request: QueryLegacy
+    request: Optional[QueryLegacy]
     # Status of the apl query result.
     status: QueryStatus
     # Matches are the events that matched the apl query.
-    matches: List[Entry]
+    matches: Optional[List[Entry]]
     # Buckets are the time series buckets.
-    buckets: Timeseries
+    buckets: Optional[Timeseries]
+    # Tables is populated in tabular queries.
+    tables: Optional[List[Table]]
     # Dataset names are the datasets that were used in the apl query.
     dataset_names: List[str] = field(default_factory=lambda: [])
     # savedQueryID is the ID of the apl query that generated this result when it

src/axiom_py/util.py

@@ -7,7 +7,7 @@
 
 from .query import QueryKind
 from .query.aggregation import AggregationOperation
-from .query.result import MessageCode, MessagePriority
+from .query.result import MessagePriority
 from .query.filter import FilterOperation
 
 
@@ -52,7 +52,6 @@ def from_dict(data_class: Type[T], data) -> T:
             datetime: _convert_string_to_datetime,
             AggregationOperation: AggregationOperation,
             FilterOperation: FilterOperation,
-            MessageCode: MessageCode,
             MessagePriority: MessagePriority,
             timedelta: _convert_string_to_timedelta,
         }

tests/test_client.py

@@ -174,6 +174,23 @@ def test_step005_apl_query(self):
 
         self.assertEqual(len(qr.matches), len(self.events))
 
+    def test_step005_apl_query_tabular(self):
+        """Test apl query (tabular)"""
+        # query the events we ingested in step2
+        startTime = datetime.utcnow() - timedelta(minutes=2)
+        endTime = datetime.utcnow()
+
+        apl = "['%s']" % self.dataset_name
+        opts = AplOptions(
+            start_time=startTime,
+            end_time=endTime,
+            format=AplResultFormat.Tabular,
+        )
+        qr = self.client.query(apl, opts)
+
+        events = list(qr.tables[0].events())
+        self.assertEqual(len(events), len(self.events))
+
     def test_step005_wrong_query_kind(self):
         """Test wrong query kind"""
         startTime = datetime.utcnow() - timedelta(minutes=2)