Skip to content

Commit 3197f0e

Browse files
0xbrayo0xbrayo
committed
savepoint
1 parent b705f04 commit 3197f0e

File tree

3 files changed

+219
-92
lines changed

3 files changed

+219
-92
lines changed

‎aw-client-rust/src/classes.rs

Lines changed: 0 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -146,61 +146,3 @@ pub fn default_classes() -> Vec<(CategoryId, CategorySpec)> {
146146
),
147147
]
148148
}
149-
150-
/// Get classes from server-side settings using default localhost:5600.
151-
/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
152-
pub fn get_classes() -> Vec<(CategoryId, CategorySpec)> {
153-
get_classes_from_server("localhost", 5600)
154-
}
155-
156-
/// Get classes from server-side settings with custom host and port.
157-
/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
158-
pub fn get_classes_from_server(host: &str, port: u16) -> Vec<(CategoryId, CategorySpec)> {
159-
let mut rng = rand::rng();
160-
let random_int = rng.random_range(0..10001);
161-
let client_id = format!("get-setting-{}", random_int);
162-
163-
// Create a client with a random ID, similar to the Python implementation
164-
let awc = match ActivityWatchClient::new(host, port, &client_id) {
165-
Ok(client) => client,
166-
Err(_) => {
167-
warn!(
168-
"Failed to create ActivityWatch client for {}:{}, using default classes",
169-
host, port
170-
);
171-
return default_classes();
172-
}
173-
};
174-
175-
awc.get_setting("classes")
176-
.map(|setting_value| {
177-
// Try to deserialize the setting into Vec<ClassSetting>
178-
if setting_value.is_null() {
179-
return default_classes();
180-
}
181-
182-
let class_settings: Vec<ClassSetting> = match serde_json::from_value(setting_value) {
183-
Ok(classes) => classes,
184-
Err(e) => {
185-
warn!(
186-
"Failed to deserialize classes setting: {}, using default classes",
187-
e
188-
);
189-
return default_classes();
190-
}
191-
};
192-
193-
// Convert ClassSetting to (CategoryId, CategorySpec) format
194-
class_settings
195-
.into_iter()
196-
.map(|class| (class.name, class.rule))
197-
.collect()
198-
})
199-
.unwrap_or_else(|_| {
200-
warn!(
201-
"Failed to get classes from server {}:{}, using default classes as fallback",
202-
host, port
203-
);
204-
default_classes()
205-
})
206-
}

‎aw-client-rust/src/lib.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,6 @@ impl AwClient {
123123
.map(|(start, stop)| format!("{}/{}", start, stop))
124124
.collect();
125125

126-
let query_lines: Vec<&str> = query.split('\n').collect();
127-
128126
// Result is a sequence, one element per timeperiod
129127
self.client
130128
.post(url)

‎aw-client-rust/src/queries.rs

Lines changed: 219 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
//! ```
3737
3838
use crate::classes::{CategoryId, CategorySpec};
39+
use crate::AwClient;
3940
use serde::{Deserialize, Serialize};
4041

4142
/// Browser application names mapped by browser type
@@ -78,7 +79,6 @@ pub static BROWSER_APPNAMES: phf::Map<&'static str, &'static [&'static str]> = p
7879
"vivaldi" => &["Vivaldi-stable", "Vivaldi-snapshot", "vivaldi.exe"],
7980
};
8081

81-
pub const DEFAULT_LIMIT: u32 = 100;
8282

8383
/// Type alias for categorization classes
8484
pub type ClassRule = (CategoryId, CategorySpec);
@@ -135,33 +135,6 @@ impl QueryParams {
135135
QueryParams::Android(params) => build_android_canonical_events(params),
136136
}
137137
}
138-
139-
/// Build canonical events query string with automatic class fetching if not provided
140-
pub fn canonical_events_with_classes(&self) -> String {
141-
self.canonical_events_with_classes_from_server("localhost", 5600)
142-
}
143-
144-
/// Build canonical events query string with automatic class fetching from custom server
145-
pub fn canonical_events_with_classes_from_server(&self, host: &str, port: u16) -> String {
146-
match self {
147-
QueryParams::Desktop(params) => {
148-
let mut params_with_classes = params.clone();
149-
if params_with_classes.base.classes.is_empty() {
150-
params_with_classes.base.classes =
151-
crate::classes::get_classes_from_server(host, port);
152-
}
153-
build_desktop_canonical_events(&params_with_classes)
154-
}
155-
QueryParams::Android(params) => {
156-
let mut params_with_classes = params.clone();
157-
if params_with_classes.base.classes.is_empty() {
158-
params_with_classes.base.classes =
159-
crate::classes::get_classes_from_server(host, port);
160-
}
161-
build_android_canonical_events(&params_with_classes)
162-
}
163-
}
164-
}
165138
}
166139

167140
/// Helper function to serialize classes in the format expected by the categorize function
@@ -201,7 +174,7 @@ fn serialize_classes(classes: &[ClassRule]) -> String {
201174
format!("[{}]", parts.join(", "))
202175
}
203176

204-
fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
177+
pub fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
205178
let mut query = Vec::new();
206179

207180
// Fetch window events
@@ -256,7 +229,7 @@ not_afk = period_union(not_afk, audible_events)"
256229
query.join(";\n")
257230
}
258231

259-
fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
232+
pub fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
260233
let mut query = Vec::new();
261234

262235
// Fetch app events
@@ -287,7 +260,7 @@ fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
287260
query.join(";\n")
288261
}
289262

290-
fn build_browser_events(params: &DesktopQueryParams) -> String {
263+
pub fn build_browser_events(params: &DesktopQueryParams) -> String {
291264
let mut query = String::from("browser_events = [];");
292265

293266
for browser_bucket in &params.base.bid_browsers {
@@ -311,10 +284,143 @@ browser_events = sort_by_timestamp(browser_events)",
311284
query
312285
}
313286

314-
/// Build a full desktop query
287+
/// Build a full desktop query using default localhost:5600 configuration
315288
pub fn full_desktop_query(params: &DesktopQueryParams) -> String {
316289
let mut query = QueryParams::Desktop(params.clone()).canonical_events_with_classes();
317290

291+
// Add basic event aggregations
292+
query.push_str(&format!(
293+
"
294+
title_events = sort_by_duration(merge_events_by_keys(events, [\"app\", \"title\"]));
295+
app_events = sort_by_duration(merge_events_by_keys(title_events, [\"app\"]));
296+
cat_events = sort_by_duration(merge_events_by_keys(events, [\"$category\"]));
297+
duration = sum_durations(events);
298+
",
299+
));
300+
301+
// Add browser-specific query parts if browser buckets exist
302+
if !params.base.bid_browsers.is_empty() {
303+
query.push_str(&format!(
304+
"
305+
browser_events = split_url_events(browser_events);
306+
browser_urls = merge_events_by_keys(browser_events, [\"url\"]);
307+
browser_urls = sort_by_duration(browser_urls);
308+
browser_domains = merge_events_by_keys(browser_events, [\"$domain\"]);
309+
browser_domains = sort_by_duration(browser_domains);
310+
browser_duration = sum_durations(browser_events);
311+
",
312+
));
313+
} else {
314+
query.push_str(
315+
"
316+
browser_events = [];
317+
browser_urls = [];
318+
browser_domains = [];
319+
browser_duration = 0;
320+
",
321+
);
322+
}
323+
324+
// Add return statement
325+
query.push_str(
326+
"
327+
RETURN = {
328+
\"events\": events,
329+
\"window\": {
330+
\"app_events\": app_events,
331+
\"title_events\": title_events,
332+
\"cat_events\": cat_events,
333+
\"active_events\": not_afk,
334+
\"duration\": duration
335+
},
336+
\"browser\": {
337+
\"domains\": browser_domains,
338+
\"urls\": browser_urls,
339+
\"duration\": browser_duration
340+
}
341+
};
342+
",
343+
);
344+
345+
query
346+
}
347+
348+
/// Build a full desktop query using client configuration
349+
pub fn full_desktop_query_from_client(
350+
params: &DesktopQueryParams,
351+
client: &crate::AwClient,
352+
) -> String {
353+
let mut query =
354+
QueryParams::Desktop(params.clone()).canonical_events_with_classes_from_client(client);
355+
356+
// Add basic event aggregations
357+
query.push_str(&format!(
358+
"
359+
title_events = sort_by_duration(merge_events_by_keys(events, [\"app\", \"title\"]));
360+
app_events = sort_by_duration(merge_events_by_keys(title_events, [\"app\"]));
361+
cat_events = sort_by_duration(merge_events_by_keys(events, [\"$category\"]));
362+
app_events = limit_events(app_events, {});
363+
title_events = limit_events(title_events, {});
364+
duration = sum_durations(events);
365+
",
366+
DEFAULT_LIMIT, DEFAULT_LIMIT
367+
));
368+
369+
// Add browser-specific query parts if browser buckets exist
370+
if !params.base.bid_browsers.is_empty() {
371+
query.push_str(&format!(
372+
"
373+
browser_events = split_url_events(browser_events);
374+
browser_urls = merge_events_by_keys(browser_events, [\"url\"]);
375+
browser_urls = sort_by_duration(browser_urls);
376+
browser_domains = merge_events_by_keys(browser_events, [\"$domain\"]);
377+
browser_domains = sort_by_duration(browser_domains);
378+
browser_duration = sum_durations(browser_events);
379+
"
380+
));
381+
} else {
382+
query.push_str(
383+
"
384+
browser_events = [];
385+
browser_urls = [];
386+
browser_domains = [];
387+
browser_duration = 0;
388+
",
389+
);
390+
}
391+
392+
// Add return statement
393+
query.push_str(
394+
"
395+
RETURN = {
396+
\"events\": events,
397+
\"window\": {
398+
\"app_events\": app_events,
399+
\"title_events\": title_events,
400+
\"cat_events\": cat_events,
401+
\"active_events\": not_afk,
402+
\"duration\": duration
403+
},
404+
\"browser\": {
405+
\"domains\": browser_domains,
406+
\"urls\": browser_urls,
407+
\"duration\": browser_duration
408+
}
409+
};
410+
",
411+
);
412+
413+
query
414+
}
415+
416+
/// Build a full desktop query using blocking client configuration
417+
pub fn full_desktop_query_from_blocking_client(
418+
params: &DesktopQueryParams,
419+
client: &crate::blocking::AwClient,
420+
) -> String {
421+
let mut query = QueryParams::Desktop(params.clone())
422+
.canonical_events_with_classes_from_blocking_client(client);
423+
318424
// Add basic event aggregations
319425
query.push_str(&format!(
320426
"
@@ -490,4 +596,85 @@ mod tests {
490596
assert!(query.contains("events = categorize"));
491597
assert!(query.contains("test"));
492598
}
599+
600+
#[test]
601+
fn test_canonical_events_with_client_config() {
602+
use crate::AwClient;
603+
604+
let params = DesktopQueryParams {
605+
base: QueryParamsBase {
606+
bid_browsers: vec![],
607+
classes: vec![], // Empty - would fetch from server if available
608+
filter_classes: vec![],
609+
filter_afk: true,
610+
include_audible: true,
611+
},
612+
bid_window: "test-window".to_string(),
613+
bid_afk: "test-afk".to_string(),
614+
};
615+
616+
// Test with custom port client
617+
if let Ok(client) = AwClient::new("localhost", 8080, "test-client") {
618+
let query_params = QueryParams::Desktop(params.clone());
619+
let query = query_params.canonical_events_with_classes_from_client(&client);
620+
621+
// Should contain basic query structure
622+
assert!(query.contains("events = flood"));
623+
assert!(query.contains("test-window"));
624+
}
625+
626+
// Test with blocking client
627+
use crate::blocking::AwClient as BlockingClient;
628+
if let Ok(blocking_client) = BlockingClient::new("localhost", 9090, "test-blocking-client")
629+
{
630+
let query_params = QueryParams::Desktop(params);
631+
let query =
632+
query_params.canonical_events_with_classes_from_blocking_client(&blocking_client);
633+
634+
// Should contain basic query structure
635+
assert!(query.contains("events = flood"));
636+
assert!(query.contains("test-window"));
637+
}
638+
}
639+
640+
#[test]
641+
fn test_full_desktop_query_from_client() {
642+
use crate::AwClient;
643+
644+
let params = DesktopQueryParams {
645+
base: QueryParamsBase {
646+
bid_browsers: vec!["aw-watcher-web-chrome".to_string()],
647+
classes: vec![],
648+
filter_classes: vec![],
649+
filter_afk: true,
650+
include_audible: true,
651+
},
652+
bid_window: "test-window".to_string(),
653+
bid_afk: "test-afk".to_string(),
654+
};
655+
656+
// Test the client-aware full desktop query
657+
if let Ok(client) = AwClient::new("localhost", 8080, "test-client") {
658+
let query = full_desktop_query_from_client(&params, &client);
659+
660+
// Should contain all expected parts
661+
assert!(query.contains("events = flood"));
662+
assert!(query.contains("title_events = sort_by_duration"));
663+
assert!(query.contains("browser_events"));
664+
assert!(query.contains("RETURN"));
665+
}
666+
667+
// Test the blocking client version
668+
use crate::blocking::AwClient as BlockingClient;
669+
if let Ok(blocking_client) = BlockingClient::new("localhost", 9090, "test-blocking-client")
670+
{
671+
let query = full_desktop_query_from_blocking_client(&params, &blocking_client);
672+
673+
// Should contain all expected parts
674+
assert!(query.contains("events = flood"));
675+
assert!(query.contains("title_events = sort_by_duration"));
676+
assert!(query.contains("browser_events"));
677+
assert!(query.contains("RETURN"));
678+
}
679+
}
493680
}

0 commit comments

Comments
 (0)