Skip to main content

stmo_cli/
api.rs

1#![allow(clippy::missing_errors_doc)]
2
3use crate::models::{
4    CreateDashboard, CreateQuery, CreateQuerySnippet, CreateWidget, Dashboard, DashboardSummary,
5    DashboardsResponse, DataSource, DataSourceSchema, QueriesResponse, Query, QuerySnippet,
6};
7use anyhow::{Context, Result};
8use reqwest::{Client, header};
9use serde::Serialize;
10use serde::de::DeserializeOwned;
11
12pub struct RedashClient {
13    client: Client,
14    base_url: String,
15}
16
17impl RedashClient {
18    pub fn new(base_url: String, api_key: &str) -> Result<Self> {
19        let mut headers = header::HeaderMap::new();
20        headers.insert(
21            "Authorization",
22            header::HeaderValue::from_str(&format!("Key {api_key}"))
23                .context("Invalid API key format")?,
24        );
25
26        let client = Client::builder()
27            .default_headers(headers)
28            .build()
29            .context("Failed to build HTTP client")?;
30
31        Ok(Self { client, base_url })
32    }
33
34    #[must_use]
35    pub fn base_url(&self) -> &str {
36        &self.base_url
37    }
38
39    async fn get_json<T: DeserializeOwned>(&self, url: &str, ctx: &str) -> Result<T> {
40        let response = self
41            .client
42            .get(url)
43            .send()
44            .await
45            .with_context(|| format!("Failed to request {ctx}"))?;
46
47        let response = ensure_success(response).await?;
48
49        response
50            .json()
51            .await
52            .with_context(|| format!("Failed to parse {ctx} response"))
53    }
54
55    async fn post_json<T: DeserializeOwned, B: Serialize + ?Sized>(
56        &self,
57        url: &str,
58        body: &B,
59        ctx: &str,
60    ) -> Result<T> {
61        let response = self
62            .client
63            .post(url)
64            .json(body)
65            .send()
66            .await
67            .with_context(|| format!("Failed to request {ctx}"))?;
68
69        let response = ensure_success(response).await?;
70
71        response
72            .json()
73            .await
74            .with_context(|| format!("Failed to parse {ctx} response"))
75    }
76
77    pub async fn list_my_queries(&self, page: u32, page_size: u32) -> Result<QueriesResponse> {
78        let url = format!(
79            "{}/api/queries/my?page={page}&page_size={page_size}",
80            self.base_url
81        );
82        self.get_json(&url, "my queries").await
83    }
84
85    pub async fn get_query(&self, query_id: u64) -> Result<Query> {
86        let url = format!("{}/api/queries/{query_id}", self.base_url);
87        self.get_json(&url, &format!("query {query_id}")).await
88    }
89
90    pub async fn list_data_sources(&self) -> Result<Vec<DataSource>> {
91        let url = format!("{}/api/data_sources", self.base_url);
92        self.get_json(&url, "data sources").await
93    }
94
95    pub async fn get_data_source(&self, data_source_id: u64) -> Result<DataSource> {
96        let url = format!("{}/api/data_sources/{data_source_id}", self.base_url);
97        self.get_json(&url, &format!("data source {data_source_id}"))
98            .await
99    }
100
101    pub async fn get_data_source_schema(
102        &self,
103        data_source_id: u64,
104        refresh: bool,
105    ) -> Result<DataSourceSchema> {
106        let url = if refresh {
107            format!(
108                "{}/api/data_sources/{data_source_id}/schema?refresh=true",
109                self.base_url
110            )
111        } else {
112            format!("{}/api/data_sources/{data_source_id}/schema", self.base_url)
113        };
114
115        self.get_json(&url, &format!("schema for data source {data_source_id}"))
116            .await
117    }
118
119    pub async fn create_query(&self, create_query: &CreateQuery) -> Result<Query> {
120        let url = format!("{}/api/queries", self.base_url);
121        self.post_json(&url, create_query, "new query").await
122    }
123
124    pub async fn create_or_update_query(&self, query: &Query) -> Result<Query> {
125        let url = format!("{}/api/queries/{}", self.base_url, query.id);
126        self.post_json(&url, query, &format!("query {} update", query.id))
127            .await
128    }
129
130    pub async fn create_visualization(
131        &self,
132        query_id: u64,
133        viz: &crate::models::CreateVisualization,
134    ) -> Result<crate::models::Visualization> {
135        let url = format!("{}/api/visualizations", self.base_url);
136        self.post_json(&url, viz, &format!("visualization for query {query_id}"))
137            .await
138    }
139
140    pub async fn update_visualization(
141        &self,
142        viz: &crate::models::Visualization,
143    ) -> Result<crate::models::Visualization> {
144        let url = format!("{}/api/visualizations/{}", self.base_url, viz.id);
145        self.post_json(&url, viz, &format!("visualization {} update", viz.id))
146            .await
147    }
148
149    pub async fn fetch_all_queries(&self) -> Result<Vec<Query>> {
150        let mut all_queries = Vec::new();
151        let mut page = 1;
152        let page_size = 100;
153
154        loop {
155            let response = self.list_my_queries(page, page_size).await?;
156
157            if response.results.is_empty() {
158                break;
159            }
160
161            all_queries.extend(response.results);
162            eprintln!(
163                "Fetched {} / {} queries...",
164                all_queries.len(),
165                response.count
166            );
167
168            #[allow(clippy::cast_possible_truncation)]
169            if all_queries.len() >= response.count as usize {
170                break;
171            }
172
173            page += 1;
174        }
175
176        Ok(all_queries)
177    }
178
179    pub async fn refresh_query(
180        &self,
181        query_id: u64,
182        parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
183    ) -> Result<crate::models::Job> {
184        let url = format!("{}/api/queries/{query_id}/results", self.base_url);
185
186        let request = crate::models::RefreshRequest {
187            max_age: 0,
188            parameters,
189        };
190
191        let job_response: crate::models::JobResponse = self
192            .post_json(&url, &request, &format!("query {query_id} refresh"))
193            .await?;
194
195        Ok(job_response.job)
196    }
197
198    pub async fn poll_job(&self, job_id: &str) -> Result<crate::models::Job> {
199        let url = format!("{}/api/jobs/{job_id}", self.base_url);
200
201        let job_response: crate::models::JobResponse =
202            self.get_json(&url, &format!("job {job_id}")).await?;
203
204        Ok(job_response.job)
205    }
206
207    async fn fetch_query_result(
208        &self,
209        url: &str,
210        fetch_context: &str,
211    ) -> Result<crate::models::QueryResult> {
212        let result_response: crate::models::QueryResultResponse =
213            self.get_json(url, fetch_context).await?;
214
215        Ok(result_response.query_result)
216    }
217
218    pub async fn get_query_result(
219        &self,
220        query_id: u64,
221        result_id: u64,
222    ) -> Result<crate::models::QueryResult> {
223        let url = format!(
224            "{}/api/queries/{query_id}/results/{result_id}.json",
225            self.base_url
226        );
227
228        self.fetch_query_result(&url, &format!("result {result_id} for query {query_id}"))
229            .await
230    }
231
232    pub async fn refresh_adhoc_query(
233        &self,
234        sql: &str,
235        data_source_id: u64,
236        parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
237    ) -> Result<crate::models::Job> {
238        let url = format!("{}/api/query_results", self.base_url);
239
240        let request = crate::models::AdhocRefreshRequest {
241            sql: sql.to_string(),
242            data_source_id,
243            max_age: 0,
244            parameters,
245        };
246
247        let job_response: crate::models::JobResponse = self
248            .post_json(&url, &request, "ad-hoc query refresh")
249            .await?;
250
251        Ok(job_response.job)
252    }
253
254    pub async fn get_adhoc_query_result(
255        &self,
256        result_id: u64,
257    ) -> Result<crate::models::QueryResult> {
258        let url = format!("{}/api/query_results/{result_id}.json", self.base_url);
259
260        self.fetch_query_result(&url, &format!("ad-hoc result {result_id}"))
261            .await
262    }
263
264    async fn poll_job_to_completion(
265        &self,
266        job: crate::models::Job,
267        timeout_secs: u64,
268        poll_interval_ms: u64,
269    ) -> Result<u64> {
270        use crate::models::JobStatus;
271        use tokio::time::{Duration, sleep};
272
273        let start = std::time::Instant::now();
274        let timeout = Duration::from_secs(timeout_secs);
275        let poll_interval = Duration::from_millis(poll_interval_ms);
276
277        let mut current_job = job;
278        loop {
279            if start.elapsed() > timeout {
280                anyhow::bail!("Query execution timed out after {timeout_secs} seconds");
281            }
282
283            match JobStatus::from_u8(current_job.status)? {
284                JobStatus::Success => {
285                    return current_job
286                        .query_result_id
287                        .context("Job succeeded but no result_id returned");
288                }
289                JobStatus::Failure => {
290                    let error = current_job
291                        .error
292                        .unwrap_or_else(|| "Unknown error".to_string());
293                    anyhow::bail!("Query execution failed: {error}");
294                }
295                JobStatus::Cancelled => {
296                    anyhow::bail!("Query execution was cancelled");
297                }
298                JobStatus::Pending | JobStatus::Started => {
299                    eprint!(".");
300                    sleep(poll_interval).await;
301                    current_job = self.poll_job(&current_job.id).await?;
302                }
303            }
304        }
305    }
306
307    pub async fn execute_query_with_polling(
308        &self,
309        query_id: u64,
310        parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
311        timeout_secs: u64,
312        poll_interval_ms: u64,
313    ) -> Result<crate::models::QueryResult> {
314        eprintln!("Executing query {query_id}...");
315        let job = self.refresh_query(query_id, parameters).await?;
316        let result_id = self
317            .poll_job_to_completion(job, timeout_secs, poll_interval_ms)
318            .await?;
319
320        eprintln!("Query completed, fetching results...");
321        self.get_query_result(query_id, result_id).await
322    }
323
324    pub async fn execute_adhoc_with_polling(
325        &self,
326        sql: &str,
327        data_source_id: u64,
328        parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
329        timeout_secs: u64,
330        poll_interval_ms: u64,
331    ) -> Result<crate::models::QueryResult> {
332        eprintln!("Executing SQL on data source {data_source_id}...");
333        let job = self
334            .refresh_adhoc_query(sql, data_source_id, parameters)
335            .await?;
336        let result_id = self
337            .poll_job_to_completion(job, timeout_secs, poll_interval_ms)
338            .await?;
339
340        eprintln!("Query completed, fetching results...");
341        self.get_adhoc_query_result(result_id).await
342    }
343
344    pub async fn archive_query(&self, query_id: u64) -> Result<Query> {
345        let url = format!("{}/api/queries/{query_id}", self.base_url);
346        let payload = serde_json::json!({"is_archived": true});
347        self.post_json(&url, &payload, &format!("query {query_id} archive"))
348            .await
349    }
350
351    pub async fn unarchive_query(&self, query_id: u64) -> Result<Query> {
352        let url = format!("{}/api/queries/{query_id}", self.base_url);
353        let payload = serde_json::json!({"is_archived": false});
354        self.post_json(&url, &payload, &format!("query {query_id} unarchive"))
355            .await
356    }
357
358    pub async fn list_query_snippets(&self) -> Result<Vec<QuerySnippet>> {
359        let url = format!("{}/api/query_snippets", self.base_url);
360        self.get_json(&url, "query snippets").await
361    }
362
363    pub async fn get_query_snippet(&self, snippet_id: u64) -> Result<QuerySnippet> {
364        let url = format!("{}/api/query_snippets/{snippet_id}", self.base_url);
365        self.get_json(&url, &format!("query snippet {snippet_id}"))
366            .await
367    }
368
369    pub async fn create_query_snippet(&self, create: &CreateQuerySnippet) -> Result<QuerySnippet> {
370        let url = format!("{}/api/query_snippets", self.base_url);
371        self.post_json(&url, create, "new query snippet").await
372    }
373
374    pub async fn update_query_snippet(&self, snippet: &QuerySnippet) -> Result<QuerySnippet> {
375        let url = format!("{}/api/query_snippets/{}", self.base_url, snippet.id);
376        self.post_json(
377            &url,
378            snippet,
379            &format!("query snippet {} update", snippet.id),
380        )
381        .await
382    }
383
384    pub async fn delete_query_snippet(&self, snippet_id: u64) -> Result<()> {
385        let url = format!("{}/api/query_snippets/{snippet_id}", self.base_url);
386        let response = self
387            .client
388            .delete(&url)
389            .send()
390            .await
391            .context(format!("Failed to delete query snippet {snippet_id}"))?;
392
393        ensure_success(response).await?;
394
395        Ok(())
396    }
397
398    pub async fn create_dashboard(&self, dashboard: &CreateDashboard) -> Result<Dashboard> {
399        let url = format!("{}/api/dashboards", self.base_url);
400        self.post_json(&url, dashboard, "new dashboard").await
401    }
402
403    pub async fn list_favorite_dashboards(
404        &self,
405        page: u32,
406        page_size: u32,
407    ) -> Result<DashboardsResponse> {
408        let url = format!(
409            "{}/api/dashboards/favorites?page={page}&page_size={page_size}",
410            self.base_url
411        );
412        self.get_json(&url, "favorite dashboards").await
413    }
414
415    pub async fn get_dashboard(&self, slug_or_id: &str) -> Result<Dashboard> {
416        let url = format!("{}/api/dashboards/{slug_or_id}", self.base_url);
417        self.get_json(&url, &format!("dashboard {slug_or_id}"))
418            .await
419    }
420
421    pub async fn update_dashboard(&self, dashboard: &Dashboard) -> Result<Dashboard> {
422        let url = format!("{}/api/dashboards/{}", self.base_url, dashboard.id);
423        self.post_json(
424            &url,
425            dashboard,
426            &format!("dashboard {} update", dashboard.id),
427        )
428        .await
429    }
430
431    pub async fn archive_dashboard(&self, dashboard_id: u64) -> Result<()> {
432        let url = format!("{}/api/dashboards/{dashboard_id}", self.base_url);
433        let payload = serde_json::json!({"is_archived": true});
434        let response = self
435            .client
436            .post(&url)
437            .json(&payload)
438            .send()
439            .await
440            .context(format!("Failed to archive dashboard {dashboard_id}"))?;
441
442        ensure_success(response).await?;
443
444        Ok(())
445    }
446
447    pub async fn unarchive_dashboard(&self, dashboard_id: u64) -> Result<Dashboard> {
448        let url = format!("{}/api/dashboards/{dashboard_id}", self.base_url);
449        let payload = serde_json::json!({"is_archived": false});
450        self.post_json(
451            &url,
452            &payload,
453            &format!("dashboard {dashboard_id} unarchive"),
454        )
455        .await
456    }
457
458    pub async fn create_widget(&self, widget: &CreateWidget) -> Result<crate::models::Widget> {
459        let url = format!("{}/api/widgets", self.base_url);
460        self.post_json(&url, widget, "new widget").await
461    }
462
463    pub async fn update_widget(
464        &self,
465        widget_id: u64,
466        widget: &CreateWidget,
467    ) -> Result<crate::models::Widget> {
468        let url = format!("{}/api/widgets/{widget_id}", self.base_url);
469        self.post_json(&url, widget, &format!("widget {widget_id} update"))
470            .await
471    }
472
473    pub async fn delete_widget(&self, widget_id: u64) -> Result<()> {
474        let url = format!("{}/api/widgets/{widget_id}", self.base_url);
475        let response = self
476            .client
477            .delete(&url)
478            .send()
479            .await
480            .context(format!("Failed to delete widget {widget_id}"))?;
481
482        ensure_success(response).await?;
483
484        Ok(())
485    }
486
487    pub async fn favorite_dashboard(&self, slug: &str) -> Result<()> {
488        let url = format!("{}/api/dashboards/{slug}/favorite", self.base_url);
489        let response = self
490            .client
491            .post(&url)
492            .json(&serde_json::json!({}))
493            .send()
494            .await
495            .context(format!("Failed to favorite dashboard {slug}"))?;
496
497        ensure_success(response).await?;
498
499        Ok(())
500    }
501
502    pub async fn fetch_favorite_dashboards(&self) -> Result<Vec<DashboardSummary>> {
503        let mut all_dashboards = Vec::new();
504        let mut page = 1;
505        let page_size = 100;
506
507        loop {
508            let response = self.list_favorite_dashboards(page, page_size).await?;
509
510            if response.results.is_empty() {
511                break;
512            }
513
514            all_dashboards.extend(response.results);
515            eprintln!(
516                "Fetched {} / {} dashboards...",
517                all_dashboards.len(),
518                response.count
519            );
520
521            #[allow(clippy::cast_possible_truncation)]
522            if all_dashboards.len() >= response.count as usize {
523                break;
524            }
525
526            page += 1;
527        }
528
529        Ok(all_dashboards)
530    }
531
532    async fn get_with_retry(
533        &self,
534        url: &str,
535        params: &[(&str, String)],
536    ) -> Result<reqwest::Response> {
537        use tokio::time::{Duration, sleep};
538
539        const MAX_ATTEMPTS: u32 = 4;
540        let base_delays = [250u64, 500, 1000, 2000];
541
542        let mut last_error = anyhow::anyhow!("No attempts made");
543        for attempt in 0..MAX_ATTEMPTS {
544            let response = self
545                .client
546                .get(url)
547                .query(params)
548                .send()
549                .await
550                .with_context(|| format!("Failed to GET {url}"))?;
551
552            let status = response.status();
553            let should_retry = status.as_u16() == 429 || status.is_server_error();
554
555            match ensure_success(response).await {
556                Ok(response) => return Ok(response),
557                Err(err) if !should_retry || attempt + 1 == MAX_ATTEMPTS => return Err(err),
558                Err(err) => last_error = err,
559            }
560
561            let delay_ms = base_delays[attempt as usize];
562            sleep(Duration::from_millis(delay_ms)).await;
563        }
564
565        Err(last_error)
566    }
567
568    async fn list_queries(&self, q: &str, page: u32, page_size: u32) -> Result<QueriesResponse> {
569        let url = format!("{}/api/queries", self.base_url);
570        let params = [
571            ("page", page.to_string()),
572            ("page_size", page_size.to_string()),
573            ("q", q.to_string()),
574        ];
575        self.get_with_retry(&url, &params)
576            .await?
577            .json()
578            .await
579            .context("Failed to parse queries response")
580    }
581
582    async fn list_dashboards(
583        &self,
584        q: &str,
585        page: u32,
586        page_size: u32,
587    ) -> Result<DashboardsResponse> {
588        let url = format!("{}/api/dashboards", self.base_url);
589        let params = [
590            ("page", page.to_string()),
591            ("page_size", page_size.to_string()),
592            ("q", q.to_string()),
593        ];
594        self.get_with_retry(&url, &params)
595            .await?
596            .json()
597            .await
598            .context("Failed to parse dashboards response")
599    }
600
601    pub async fn search_queries(&self, q: &str, limit: usize) -> Result<Vec<Query>> {
602        const PAGE_SIZE: usize = 250;
603
604        let mut results: Vec<Query> = Vec::new();
605        let mut page = 1u32;
606
607        loop {
608            let remaining = limit - results.len();
609            #[allow(clippy::cast_possible_truncation)]
610            let page_size = remaining.min(PAGE_SIZE) as u32;
611            let response = self.list_queries(q, page, page_size).await?;
612
613            results.extend(response.results);
614
615            #[allow(clippy::cast_possible_truncation)]
616            if results.len() >= limit || results.len() >= response.count as usize {
617                break;
618            }
619
620            page += 1;
621        }
622
623        results.truncate(limit);
624        Ok(results)
625    }
626
627    pub async fn search_dashboards(&self, q: &str, limit: usize) -> Result<Vec<DashboardSummary>> {
628        const PAGE_SIZE: usize = 250;
629
630        let mut results: Vec<DashboardSummary> = Vec::new();
631        let mut page = 1u32;
632
633        loop {
634            let remaining = limit - results.len();
635            #[allow(clippy::cast_possible_truncation)]
636            let page_size = remaining.min(PAGE_SIZE) as u32;
637            let response = self.list_dashboards(q, page, page_size).await?;
638
639            results.extend(response.results);
640
641            #[allow(clippy::cast_possible_truncation)]
642            if results.len() >= limit || results.len() >= response.count as usize {
643                break;
644            }
645
646            page += 1;
647        }
648
649        results.truncate(limit);
650        Ok(results)
651    }
652}
653
654async fn ensure_success(response: reqwest::Response) -> Result<reqwest::Response> {
655    let status = response.status();
656    if !status.is_success() {
657        let body = response.text().await.unwrap_or_default();
658        anyhow::bail!("API error {status}: {body}");
659    }
660    Ok(response)
661}