athena_cli/commands/history/
list.rs

1use super::fields::{get_field_value, HistoryField};
2use crate::cli::HistoryArgs;
3use crate::context::Context;
4use anyhow::Result;
5use prettytable::{Cell, Row};
6use std::collections::HashMap;
7
8pub async fn list(ctx: &Context, args: &HistoryArgs) -> Result<()> {
9    let client = ctx.create_athena_client();
10    let workgroup = ctx.workgroup();
11
12    // Use limit from CLI args if provided, otherwise from config
13    let limit = args.limit.unwrap_or_else(|| ctx.history_size());
14
15    let result = client
16        .list_query_executions()
17        .work_group(&workgroup)
18        .max_results(limit)
19        .send()
20        .await?;
21
22    // Get query IDs
23    let query_ids = result.query_execution_ids();
24    if query_ids.is_empty() {
25        println!("No queries found in workgroup: {}", workgroup);
26        return Ok(());
27    }
28
29    println!(
30        "Found {} queries in workgroup: {}",
31        query_ids.len(),
32        workgroup
33    );
34
35    // Get details for all queries in a single batch request
36    let details = client
37        .batch_get_query_execution()
38        .set_query_execution_ids(Some(query_ids.to_vec()))
39        .send()
40        .await?;
41
42    // Create a map of query ID to execution for quick lookup
43    let executions_map: HashMap<String, &aws_sdk_athena::types::QueryExecution> = details
44        .query_executions()
45        .iter()
46        .filter_map(|exec| exec.query_execution_id().map(|id| (id.to_string(), exec)))
47        .collect();
48
49    // Only fetch row counts if the RowCount field is being displayed
50    let fields = super::fields::get_history_fields();
51    let mut row_counts: HashMap<String, String> = HashMap::new();
52
53    if fields.contains(&HistoryField::RowCount) {
54        // Get only SUCCEEDED query IDs to minimize API calls
55        let succeeded_query_ids: Vec<String> = query_ids
56            .iter()
57            .filter(|&id| {
58                if let Some(execution) = executions_map.get(id) {
59                    if let Some(status) = execution.status().and_then(|s| s.state()) {
60                        return status.as_str() == "SUCCEEDED";
61                    }
62                }
63                false
64            })
65            .map(|id| id.to_string())
66            .collect();
67
68        // Fetch row counts for successful queries in batches to reduce API calls
69        for chunk in succeeded_query_ids.chunks(10) {
70            for query_id in chunk {
71                match client
72                    .get_query_runtime_statistics()
73                    .query_execution_id(query_id)
74                    .send()
75                    .await
76                {
77                    Ok(stats) => {
78                        if let Some(rows) = stats.query_runtime_statistics().and_then(|s| s.rows())
79                        {
80                            if let Some(output_rows) = rows.output_rows() {
81                                row_counts.insert(query_id.clone(), output_rows.to_string());
82                            }
83                        }
84                    }
85                    Err(e) => {
86                        // Log the error but continue processing
87                        eprintln!("Failed to get row count for query {}: {}", query_id, e);
88                    }
89                }
90            }
91        }
92    }
93
94    // Process query IDs in the original order
95    // Create a new table
96    let mut table = prettytable::Table::new();
97
98    // Add header row with styling
99    let header_cells = fields
100        .iter()
101        .map(|field| prettytable::Cell::new(&field.to_string()).style_spec("Fb"))
102        .collect();
103    table.add_row(prettytable::Row::new(header_cells));
104
105    // Add data rows in the original order from query_ids
106    for query_id in query_ids {
107        if let Some(execution) = executions_map.get(query_id) {
108            // Filter by status if specified
109            if let Some(status_filter) = &args.status {
110                if let Some(status) = execution.status().and_then(|s| s.state()) {
111                    if status.as_str() != status_filter.to_uppercase() {
112                        continue;
113                    }
114                }
115            }
116
117            // Collect field values
118            let row_values: Vec<String> = fields
119                .iter()
120                .map(|&field| {
121                    if field == HistoryField::RowCount {
122                        // Use the row count from our map if available
123                        if let Some(count) =
124                            row_counts.get(execution.query_execution_id().unwrap_or_default())
125                        {
126                            count.clone()
127                        } else {
128                            "-".to_string()
129                        }
130                    } else {
131                        get_field_value(execution, field)
132                    }
133                })
134                .collect();
135
136            // Create cells for the row
137            let cells: Vec<Cell> = row_values.iter().map(|val| Cell::new(val)).collect();
138            table.add_row(Row::new(cells));
139        }
140    }
141
142    table.printstd();
143    Ok(())
144}
athena_cli/commands/history/list.rs

athena_cli/commands/history/
list.rs
