From 81a57dfec6c079e70a6dc75072d6bc03be15b8b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Thu, 9 Jan 2025 15:41:17 +0100
Subject: [PATCH] fix: fixed cli report command, adding --limit, and cleaned up
 query construction (fixes #88)

---
 aw_client/cli.py     | 15 ++++----
 aw_client/queries.py | 88 +++++++++++++++++++++++++-------------------
 2 files changed, 58 insertions(+), 45 deletions(-)

diff --git a/aw_client/cli.py b/aw_client/cli.py
index 6da4506..186fbd5 100755
--- a/aw_client/cli.py
+++ b/aw_client/cli.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+
 import json
 import logging
 import textwrap
@@ -138,6 +139,7 @@ def query(
 @click.option("--cache", is_flag=True)
 @click.option("--start", default=now - td1day, type=click.DateTime())
 @click.option("--stop", default=now + td1yr, type=click.DateTime())
+@click.option("--limit", default=10)
 @click.pass_obj
 def report(
     obj: _Context,
@@ -146,6 +148,7 @@ def report(
     start: datetime,
     stop: datetime,
     name: Optional[str] = None,
+    limit: int = 10,
 ):
     logger.info(f"Querying between {start} and {stop}")
     bid_window = f"aw-watcher-window_{hostname}"
@@ -182,11 +185,12 @@ def report(
         print_top(
             cat_events,
             lambda e: " > ".join(e.data["$category"]),
-            title="Top categories",
+            title="Categories",
+            n=limit,
         )
 
         title_events = _parse_events(period["window"]["title_events"])
-        print_top(title_events, lambda e: e.data["title"], title="Top titles")
+        print_top(title_events, lambda e: e.data["title"], title="Titles", n=limit)
 
         active_events = _parse_events(period["window"]["title_events"])
         print(
@@ -199,11 +203,8 @@ def _parse_events(events: List[dict]) -> List[Event]:
     return [Event(**event) for event in events]
 
 
-def print_top(events: List[Event], key=lambda e: e.data, title="Events"):
-    print(
-        title
-        + (f" (showing 10 out of {len(events)} events)" if len(events) > 10 else "")
-    )
+def print_top(events: List[Event], key=lambda e: e.data, title="Events", n=10):
+    print(f"Top {n} {title}" + (f" (out of {len(events)})" if len(events) > 10 else ""))
     print(
         tabulate(
             [
diff --git a/aw_client/queries.py b/aw_client/queries.py
index 814e6dc..615baa7 100644
--- a/aw_client/queries.py
+++ b/aw_client/queries.py
@@ -248,44 +248,56 @@ def fullDesktopQuery(
     params.bid_afk = escape_doublequote(params.bid_afk)
     params.bid_browsers = [escape_doublequote(bucket) for bucket in params.bid_browsers]
 
-    return (
-        f"""
-      {canonicalEvents(params)}
-      title_events = sort_by_duration(merge_events_by_keys(events, ["app", "title"]));
-      app_events   = sort_by_duration(merge_events_by_keys(title_events, ["app"]));
-      cat_events   = sort_by_duration(merge_events_by_keys(events, ["$category"]));
-      app_events  = limit_events(app_events, {default_limit});
-      title_events  = limit_events(title_events, {default_limit});
-      duration = sum_durations(events);
-      """  # Browser events are retrieved in canonicalQuery
-        + f"""
-      browser_events = split_url_events(browser_events);
-      browser_urls = merge_events_by_keys(browser_events, ["url"]);
-      browser_urls = sort_by_duration(browser_urls);
-      browser_urls = limit_events(browser_urls, {default_limit});
-      browser_domains = merge_events_by_keys(browser_events, ["$domain"]);
-      browser_domains = sort_by_duration(browser_domains);
-      browser_domains = limit_events(browser_domains, {default_limit});
-      browser_duration = sum_durations(browser_events);
-      """
-        + """
-      RETURN = {
-          "events": events,
-          "window": {
-              "app_events": app_events,
-              "title_events": title_events,
-              "cat_events": cat_events,
-              "active_events": not_afk,
-              "duration": duration
-          },
-          "browser": {
-              "domains": browser_domains,
-              "urls": browser_urls,
-              "duration": browser_duration
-          }
-      };
-      """
-    )
+    # Build the base query
+    query = f"""
+    {canonicalEvents(params)}
+    title_events = sort_by_duration(merge_events_by_keys(events, ["app", "title"]));
+    app_events   = sort_by_duration(merge_events_by_keys(title_events, ["app"]));
+    cat_events   = sort_by_duration(merge_events_by_keys(events, ["$category"]));
+    app_events  = limit_events(app_events, {default_limit});
+    title_events  = limit_events(title_events, {default_limit});
+    duration = sum_durations(events);
+    """
+
+    # Add browser-related query parts if browser buckets exist
+    if params.bid_browsers:
+        query += """
+        browser_events = split_url_events(browser_events);
+        browser_urls = merge_events_by_keys(browser_events, ["url"]);
+        browser_urls = sort_by_duration(browser_urls);
+        browser_urls = limit_events(browser_urls, {default_limit});
+        browser_domains = merge_events_by_keys(browser_events, ["$domain"]);
+        browser_domains = sort_by_duration(browser_domains);
+        browser_domains = limit_events(browser_domains, {default_limit});
+        browser_duration = sum_durations(browser_events);
+        """
+    else:
+        query += """
+        browser_events = [];
+        browser_urls = [];
+        browser_domains = [];
+        browser_duration = 0;
+        """
+
+    # Add the return statement
+    query += """
+        RETURN = {
+            "events": events,
+            "window": {
+                "app_events": app_events,
+                "title_events": title_events,
+                "cat_events": cat_events,
+                "active_events": not_afk,
+                "duration": duration
+            },
+            "browser": {
+                "domains": browser_domains,
+                "urls": browser_urls,
+                "duration": browser_duration
+            }
+        };
+    """
+    return query
 
 
 def test_fullDesktopQuery():