Fix use of legacy format. And do some code QOL

Co-authored-by: Silas Scheele-Simonsen <[email protected]>
Signed-off-by: Heinz Gies <[email protected]>

examples/cli.rs

@@ -84,7 +84,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 println!("{:?}", ingest_status);
             }
             Datasets::Query { apl } => {
-                let result = client.query(apl, None).await?;
+                let result = client.query(&apl, None).await?;
                 for table in result.tables {
                     println!("{}:", table.name());
 

src/client.rs

@@ -107,37 +107,14 @@ impl Client {
     /// Executes the given query specified using the Axiom Processing Language (APL).
     /// To learn more about APL, see the APL documentation at https://www.axiom.co/docs/apl/introduction.
     #[instrument(skip(self, opts))]
-    pub async fn query<S, O>(&self, apl: S, opts: O) -> Result<QueryResult>
+    pub async fn query<S, O>(&self, apl: &S, opts: O) -> Result<QueryResult>
     where
-        S: Into<String> + FmtDebug,
+        S: ToString + FmtDebug + ?Sized,
         O: Into<Option<QueryOptions>>,
     {
-        let (req, query_params) = match opts.into() {
-            Some(opts) => {
-                let req = Query {
-                    apl: apl.into(),
-                    start_time: opts.start_time,
-                    end_time: opts.end_time,
-                    cursor: opts.cursor,
-                    include_cursor: opts.include_cursor,
-                };
-
-                let query_params = QueryParams {
-                    no_cache: opts.no_cache,
-                    save: opts.save,
-                    format: opts.format,
-                };
-
-                (req, query_params)
-            }
-            None => (
-                Query {
-                    apl: apl.into(),
-                    ..Default::default()
-                },
-                QueryParams::default(),
-            ),
-        };
+        let opts: QueryOptions = opts.into().unwrap_or_default();
+        let query_params = QueryParams::from(&opts);
+        let req = Query::new(apl, opts);
 
         let query_params = serde_qs::to_string(&query_params)?;
         let path = format!("/v1/datasets/_apl?{query_params}");

src/datasets/model.rs

@@ -260,6 +260,19 @@ pub struct Query {
     pub include_cursor: bool,
 }
 
+impl Query {
+    /// Creates a new query with the given APL and options.
+    pub fn new<S: ToString + ?Sized>(apl: &S, opts: QueryOptions) -> Self {
+        Self {
+            apl: apl.to_string(),
+            start_time: opts.start_time,
+            end_time: opts.end_time,
+            cursor: opts.cursor,
+            include_cursor: opts.include_cursor,
+        }
+    }
+}
+
 // QueryParams is the part of `QueryOptions` that is added to the request url.
 #[derive(Serialize, Debug, Default)]
 pub(crate) struct QueryParams {
@@ -268,10 +281,25 @@ pub(crate) struct QueryParams {
     #[serde(rename = "saveAsKind")]
     pub save: bool,
     pub format: AplResultFormat,
+    #[serde(rename = "includeCursorField")]
+    pub include_cursor_field: bool,
+}
+
+impl From<&QueryOptions> for QueryParams {
+    fn from(options: &QueryOptions) -> Self {
+        Self {
+            no_cache: options.no_cache,
+            save: options.save,
+            format: options.format,
+            include_cursor_field: options.include_cursor_field,
+        }
+    }
 }
 
+// This is a configuration that just happens to have many flags.
+#[allow(clippy::struct_excessive_bools)]
 /// The optional parameters to APL query methods.
-#[derive(Debug)]
+#[derive(Debug, Default, Serialize, Clone)]
 pub struct QueryOptions {
     /// The start time of the query.
     pub start_time: Option<DateTime<Utc>>,
@@ -282,7 +310,6 @@ pub struct QueryOptions {
     /// Specifies whether the event that matches the cursor should be
     /// included in the result.
     pub include_cursor: bool,
-
     /// Omits the query cache.
     pub no_cache: bool,
     /// Save the query on the server, if set to `true`. The ID of the saved query
@@ -294,29 +321,15 @@ pub struct QueryOptions {
     pub save: bool,
     /// Format specifies the format of the APL query. Defaults to Legacy.
     pub format: AplResultFormat,
-}
-
-impl Default for QueryOptions {
-    fn default() -> Self {
-        QueryOptions {
-            start_time: None,
-            end_time: None,
-            cursor: None,
-            include_cursor: false,
-            no_cache: false,
-            save: false,
-            format: AplResultFormat::Legacy,
-        }
-    }
+    /// Requests the cursor to be included in the response
+    pub include_cursor_field: bool,
 }
 
 /// The result format of an APL query.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Default)]
 #[non_exhaustive]
 #[serde(rename_all = "camelCase")]
 pub enum AplResultFormat {
-    /// Legacy result format
-    Legacy,
     /// Tabular result format
     #[default]
     Tabular,

tests/cursor.rs

@@ -107,11 +107,12 @@ async fn test_cursor_impl(ctx: &mut Context) {
     let apl_query_result = ctx
         .client
         .query(
-            format!("['{}'] | sort by _time desc", ctx.dataset.name),
+            &format!("['{}'] | sort by _time desc", ctx.dataset.name),
             QueryOptions {
                 start_time: Some(start_time),
                 end_time: Some(end_time),
                 save: true,
+                include_cursor_field: true,
                 ..Default::default()
             },
         )
@@ -124,12 +125,12 @@ async fn test_cursor_impl(ctx: &mut Context) {
 
     let row = table.get_row(500).unwrap();
 
-    let mid_row_id = &row.get_field("row_id").unwrap();
+    let mid_row_id = &row.get_field("_cursor").expect("column _cursor not found");
 
     let apl_query_result = ctx
         .client
         .query(
-            format!("['{}'] | sort by _time desc", ctx.dataset.name),
+            &format!("['{}'] | sort by _time desc", ctx.dataset.name),
             QueryOptions {
                 start_time: Some(start_time),
                 end_time: Some(end_time),

tests/datasets.rs

@@ -169,7 +169,7 @@ async fn test_datasets_impl(ctx: &mut Context) -> Result<(), Box<dyn std::error:
     let apl_query_result = ctx
         .client
         .query(
-            format!("['{}']", ctx.dataset.name),
+            &format!("['{}']", ctx.dataset.name),
             QueryOptions {
                 save: true,
                 ..Default::default()