Skip to main content

systemprompt_cli/commands/infrastructure/db/
helpers.rs

1//! Shared db-command helpers: known tables, byte formatting.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6const KNOWN_TABLES: &[&str] = &[
7    "logs",
8    "ai_requests",
9    "mcp_tool_executions",
10    "agent_tasks",
11    "users",
12    "tenants",
13    "sessions",
14    "agent_execution_steps",
15    "agent_artifacts",
16    "credentials",
17    "mcp_servers",
18    "workflow_states",
19    "blog_posts",
20    "categories",
21    "sources",
22];
23
24pub fn format_bytes(bytes: i64) -> String {
25    const KB: i64 = 1024;
26    const MB: i64 = KB * 1024;
27    const GB: i64 = MB * 1024;
28
29    if bytes >= GB {
30        format!("{:.2} GB", bytes as f64 / GB as f64)
31    } else if bytes >= MB {
32        format!("{:.2} MB", bytes as f64 / MB as f64)
33    } else if bytes >= KB {
34        format!("{:.2} KB", bytes as f64 / KB as f64)
35    } else {
36        format!("{} bytes", bytes)
37    }
38}
39
40pub fn extract_relation_name(msg: &str) -> String {
41    if let Some(start) = msg.find('"')
42        && let Some(end) = msg[start + 1..].find('"')
43    {
44        return msg[start + 1..start + 1 + end].to_string();
45    }
46    "unknown".to_owned()
47}
48
49pub fn suggest_table_name(input: &str) -> Option<String> {
50    let input_lower = input.to_lowercase();
51    let input_parts: Vec<&str> = input_lower.split('_').collect();
52
53    KNOWN_TABLES
54        .iter()
55        .filter(|&&table| {
56            let table_lower = table.to_lowercase();
57            let table_parts: Vec<&str> = table_lower.split('_').collect();
58
59            table_lower.contains(&input_lower)
60                || input_lower.contains(&table_lower)
61                || levenshtein_distance(&input_lower, &table_lower) <= 4
62                || shares_prefix_parts(&input_parts, &table_parts, 2)
63        })
64        .min_by_key(|&&table| levenshtein_distance(&input_lower, &table.to_lowercase()))
65        .map(|&s| s.to_owned())
66}
67
68fn shares_prefix_parts(a: &[&str], b: &[&str], min_shared: usize) -> bool {
69    let shared = a.iter().zip(b.iter()).take_while(|(x, y)| x == y).count();
70    shared >= min_shared
71}
72
73fn levenshtein_distance(s1: &str, s2: &str) -> usize {
74    let len1 = s1.chars().count();
75    let len2 = s2.chars().count();
76
77    if len1 == 0 {
78        return len2;
79    }
80    if len2 == 0 {
81        return len1;
82    }
83
84    let mut prev_row: Vec<usize> = (0..=len2).collect();
85    let mut curr_row: Vec<usize> = vec![0; len2 + 1];
86
87    for (i, c1) in s1.chars().enumerate() {
88        curr_row[0] = i + 1;
89
90        for (j, c2) in s2.chars().enumerate() {
91            let cost = usize::from(c1 != c2);
92            curr_row[j + 1] = (prev_row[j + 1] + 1)
93                .min(curr_row[j] + 1)
94                .min(prev_row[j] + cost);
95        }
96
97        std::mem::swap(&mut prev_row, &mut curr_row);
98    }
99
100    prev_row[len2]
101}