fix: remove unnecessary functions in query class Keep feature parity

with aw-client

Author: 0xbrayo

aw-client-rust/src/classes.rs

@@ -2,13 +2,9 @@
 //!
 //! Taken from default classes in aw-webui
 
-use log::warn;
-use rand::Rng;
 use serde::{Deserialize, Serialize};
 use serde_json;
 
-use super::blocking::AwClient as ActivityWatchClient;
-
 pub type CategoryId = Vec<String>;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -146,61 +142,3 @@ pub fn default_classes() -> Vec<(CategoryId, CategorySpec)> {
         ),
     ]
 }
-
-/// Get classes from server-side settings using default localhost:5600.
-/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
-pub fn get_classes() -> Vec<(CategoryId, CategorySpec)> {
-    get_classes_from_server("localhost", 5600)
-}
-
-/// Get classes from server-side settings with custom host and port.
-/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
-pub fn get_classes_from_server(host: &str, port: u16) -> Vec<(CategoryId, CategorySpec)> {
-    let mut rng = rand::rng();
-    let random_int = rng.random_range(0..10001);
-    let client_id = format!("get-setting-{}", random_int);
-
-    // Create a client with a random ID, similar to the Python implementation
-    let awc = match ActivityWatchClient::new(host, port, &client_id) {
-        Ok(client) => client,
-        Err(_) => {
-            warn!(
-                "Failed to create ActivityWatch client for {}:{}, using default classes",
-                host, port
-            );
-            return default_classes();
-        }
-    };
-
-    awc.get_setting("classes")
-        .map(|setting_value| {
-            // Try to deserialize the setting into Vec<ClassSetting>
-            if setting_value.is_null() {
-                return default_classes();
-            }
-
-            let class_settings: Vec<ClassSetting> = match serde_json::from_value(setting_value) {
-                Ok(classes) => classes,
-                Err(e) => {
-                    warn!(
-                        "Failed to deserialize classes setting: {}, using default classes",
-                        e
-                    );
-                    return default_classes();
-                }
-            };
-
-            // Convert ClassSetting to (CategoryId, CategorySpec) format
-            class_settings
-                .into_iter()
-                .map(|class| (class.name, class.rule))
-                .collect()
-        })
-        .unwrap_or_else(|_| {
-            warn!(
-                "Failed to get classes from server {}:{}, using default classes as fallback",
-                host, port
-            );
-            default_classes()
-        })
-}

aw-client-rust/src/lib.rs

@@ -36,7 +36,7 @@ impl std::fmt::Debug for AwClient {
 }
 
 fn get_hostname() -> String {
-    return gethostname::gethostname().to_string_lossy().to_string();
+    gethostname::gethostname().to_string_lossy().to_string()
 }
 
 impl AwClient {
@@ -123,8 +123,6 @@ impl AwClient {
             .map(|(start, stop)| format!("{}/{}", start, stop))
             .collect();
 
-        let query_lines: Vec<&str> = query.split('\n').collect();
-
         // Result is a sequence, one element per timeperiod
         self.client
             .post(url)

aw-client-rust/src/queries.rs

@@ -78,8 +78,6 @@ pub static BROWSER_APPNAMES: phf::Map<&'static str, &'static [&'static str]> = p
     "vivaldi" => &["Vivaldi-stable", "Vivaldi-snapshot", "vivaldi.exe"],
 };
 
-pub const DEFAULT_LIMIT: u32 = 100;
-
 /// Type alias for categorization classes
 pub type ClassRule = (CategoryId, CategorySpec);
 
@@ -135,33 +133,6 @@ impl QueryParams {
             QueryParams::Android(params) => build_android_canonical_events(params),
         }
     }
-
-    /// Build canonical events query string with automatic class fetching if not provided
-    pub fn canonical_events_with_classes(&self) -> String {
-        self.canonical_events_with_classes_from_server("localhost", 5600)
-    }
-
-    /// Build canonical events query string with automatic class fetching from custom server
-    pub fn canonical_events_with_classes_from_server(&self, host: &str, port: u16) -> String {
-        match self {
-            QueryParams::Desktop(params) => {
-                let mut params_with_classes = params.clone();
-                if params_with_classes.base.classes.is_empty() {
-                    params_with_classes.base.classes =
-                        crate::classes::get_classes_from_server(host, port);
-                }
-                build_desktop_canonical_events(&params_with_classes)
-            }
-            QueryParams::Android(params) => {
-                let mut params_with_classes = params.clone();
-                if params_with_classes.base.classes.is_empty() {
-                    params_with_classes.base.classes =
-                        crate::classes::get_classes_from_server(host, port);
-                }
-                build_android_canonical_events(&params_with_classes)
-            }
-        }
-    }
 }
 
 /// Helper function to serialize classes in the format expected by the categorize function
@@ -201,7 +172,7 @@ fn serialize_classes(classes: &[ClassRule]) -> String {
     format!("[{}]", parts.join(", "))
 }
 
-fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
+pub fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
     let mut query = Vec::new();
 
     // Fetch window events
@@ -256,7 +227,7 @@ not_afk = period_union(not_afk, audible_events)"
     query.join(";\n")
 }
 
-fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
+pub fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
     let mut query = Vec::new();
 
     // Fetch app events
@@ -287,7 +258,7 @@ fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
     query.join(";\n")
 }
 
-fn build_browser_events(params: &DesktopQueryParams) -> String {
+pub fn build_browser_events(params: &DesktopQueryParams) -> String {
     let mut query = String::from("browser_events = [];");
 
     for browser_bucket in &params.base.bid_browsers {
@@ -311,38 +282,34 @@ browser_events = sort_by_timestamp(browser_events)",
     query
 }
 
-/// Build a full desktop query
+/// Build a full desktop query using default localhost:5600 configuration
 pub fn full_desktop_query(params: &DesktopQueryParams) -> String {
-    let mut query = QueryParams::Desktop(params.clone()).canonical_events_with_classes();
+    let mut query = QueryParams::Desktop(params.clone()).canonical_events();
 
     // Add basic event aggregations
-    query.push_str(&format!(
-        "
+    query.push_str(
+        &"
         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, {});
-        title_events = limit_events(title_events, {});
         duration = sum_durations(events);
-        ",
-        DEFAULT_LIMIT, DEFAULT_LIMIT
-    ));
+        "
+        .to_string(),
+    );
 
     // Add browser-specific query parts if browser buckets exist
     if !params.base.bid_browsers.is_empty() {
-        query.push_str(&format!(
-            "
+        query.push_str(
+            &"
             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, {});
             browser_domains = merge_events_by_keys(browser_events, [\"$domain\"]);
             browser_domains = sort_by_duration(browser_domains);
-            browser_domains = limit_events(browser_domains, {});
             browser_duration = sum_durations(browser_events);
-            ",
-            DEFAULT_LIMIT, DEFAULT_LIMIT
-        ));
+            "
+            .to_string(),
+        );
     } else {
         query.push_str(
             "
@@ -490,4 +457,85 @@ mod tests {
         assert!(query.contains("events = categorize"));
         assert!(query.contains("test"));
     }
+
+    #[test]
+    fn test_canonical_events_with_client_config() {
+        use crate::AwClient;
+
+        let params = DesktopQueryParams {
+            base: QueryParamsBase {
+                bid_browsers: vec![],
+                classes: vec![], // Empty - would fetch from server if available
+                filter_classes: vec![],
+                filter_afk: true,
+                include_audible: true,
+            },
+            bid_window: "test-window".to_string(),
+            bid_afk: "test-afk".to_string(),
+        };
+
+        // Test with custom port client
+        if let Ok(client) = AwClient::new("localhost", 8080, "test-client") {
+            let query_params = QueryParams::Desktop(params.clone());
+            let query = query_params.canonical_events_with_classes_from_client(&client);
+
+            // Should contain basic query structure
+            assert!(query.contains("events = flood"));
+            assert!(query.contains("test-window"));
+        }
+
+        // Test with blocking client
+        use crate::blocking::AwClient as BlockingClient;
+        if let Ok(blocking_client) = BlockingClient::new("localhost", 9090, "test-blocking-client")
+        {
+            let query_params = QueryParams::Desktop(params);
+            let query =
+                query_params.canonical_events_with_classes_from_blocking_client(&blocking_client);
+
+            // Should contain basic query structure
+            assert!(query.contains("events = flood"));
+            assert!(query.contains("test-window"));
+        }
+    }
+
+    #[test]
+    fn test_full_desktop_query_from_client() {
+        use crate::AwClient;
+
+        let params = DesktopQueryParams {
+            base: QueryParamsBase {
+                bid_browsers: vec!["aw-watcher-web-chrome".to_string()],
+                classes: vec![],
+                filter_classes: vec![],
+                filter_afk: true,
+                include_audible: true,
+            },
+            bid_window: "test-window".to_string(),
+            bid_afk: "test-afk".to_string(),
+        };
+
+        // Test the client-aware full desktop query
+        if let Ok(client) = AwClient::new("localhost", 8080, "test-client") {
+            let query = full_desktop_query_from_client(&params, &client);
+
+            // Should contain all expected parts
+            assert!(query.contains("events = flood"));
+            assert!(query.contains("title_events = sort_by_duration"));
+            assert!(query.contains("browser_events"));
+            assert!(query.contains("RETURN"));
+        }
+
+        // Test the blocking client version
+        use crate::blocking::AwClient as BlockingClient;
+        if let Ok(blocking_client) = BlockingClient::new("localhost", 9090, "test-blocking-client")
+        {
+            let query = full_desktop_query_from_blocking_client(&params, &blocking_client);
+
+            // Should contain all expected parts
+            assert!(query.contains("events = flood"));
+            assert!(query.contains("title_events = sort_by_duration"));
+            assert!(query.contains("browser_events"));
+            assert!(query.contains("RETURN"));
+        }
+    }
 }