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 pub async fn cancel_job(&self, job_id: &str) -> Result<()> {
208 let url = format!("{}/api/jobs/{job_id}", self.base_url);
209 let response = self
210 .client
211 .delete(&url)
212 .send()
213 .await
214 .context(format!("Failed to cancel job {job_id}"))?;
215
216 ensure_success(response).await?;
217
218 Ok(())
219 }
220
221 async fn fetch_query_result(
222 &self,
223 url: &str,
224 fetch_context: &str,
225 ) -> Result<crate::models::QueryResult> {
226 let result_response: crate::models::QueryResultResponse =
227 self.get_json(url, fetch_context).await?;
228
229 Ok(result_response.query_result)
230 }
231
232 pub async fn get_query_result(
233 &self,
234 query_id: u64,
235 result_id: u64,
236 ) -> Result<crate::models::QueryResult> {
237 let url = format!(
238 "{}/api/queries/{query_id}/results/{result_id}.json",
239 self.base_url
240 );
241
242 self.fetch_query_result(&url, &format!("result {result_id} for query {query_id}"))
243 .await
244 }
245
246 pub async fn refresh_adhoc_query(
247 &self,
248 sql: &str,
249 data_source_id: u64,
250 parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
251 ) -> Result<crate::models::Job> {
252 let url = format!("{}/api/query_results", self.base_url);
253
254 let request = crate::models::AdhocRefreshRequest {
255 sql: sql.to_string(),
256 data_source_id,
257 max_age: 0,
258 parameters,
259 };
260
261 let job_response: crate::models::JobResponse = self
262 .post_json(&url, &request, "ad-hoc query refresh")
263 .await?;
264
265 Ok(job_response.job)
266 }
267
268 pub async fn get_adhoc_query_result(
269 &self,
270 result_id: u64,
271 ) -> Result<crate::models::QueryResult> {
272 let url = format!("{}/api/query_results/{result_id}.json", self.base_url);
273
274 self.fetch_query_result(&url, &format!("ad-hoc result {result_id}"))
275 .await
276 }
277
278 async fn poll_job_to_completion(
279 &self,
280 job: crate::models::Job,
281 timeout_secs: u64,
282 poll_interval_ms: u64,
283 ) -> Result<u64> {
284 let job_id = job.id.clone();
285
286 let outcome = tokio::select! {
287 res = self.poll_loop(job, timeout_secs, poll_interval_ms) => res,
288 () = wait_for_ctrl_c() => Err(anyhow::anyhow!("Interrupted; cancelling query")),
289 };
290
291 if outcome.is_err() {
295 eprintln!("\nCancelling job {job_id}...");
296 if let Err(e) = self.cancel_job(&job_id).await {
297 eprintln!("Warning: failed to cancel job {job_id}: {e}");
298 }
299 }
300
301 outcome
302 }
303
304 async fn poll_loop(
305 &self,
306 job: crate::models::Job,
307 timeout_secs: u64,
308 poll_interval_ms: u64,
309 ) -> Result<u64> {
310 use crate::models::JobStatus;
311 use tokio::time::{Duration, sleep};
312
313 let start = std::time::Instant::now();
314 let timeout = Duration::from_secs(timeout_secs);
315 let poll_interval = Duration::from_millis(poll_interval_ms);
316
317 let mut current_job = job;
318 loop {
319 if start.elapsed() > timeout {
320 anyhow::bail!("Query execution timed out after {timeout_secs} seconds");
321 }
322
323 match JobStatus::from_u8(current_job.status)? {
324 JobStatus::Success => {
325 return current_job
326 .query_result_id
327 .context("Job succeeded but no result_id returned");
328 }
329 JobStatus::Failure => {
330 let error = current_job
331 .error
332 .unwrap_or_else(|| "Unknown error".to_string());
333 anyhow::bail!("Query execution failed: {error}");
334 }
335 JobStatus::Cancelled => {
336 anyhow::bail!("Query execution was cancelled");
337 }
338 JobStatus::Pending | JobStatus::Started => {
339 eprint!(".");
340 sleep(poll_interval).await;
341 current_job = self.poll_job(¤t_job.id).await?;
342 }
343 }
344 }
345 }
346
347 pub async fn execute_query_with_polling(
348 &self,
349 query_id: u64,
350 parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
351 timeout_secs: u64,
352 poll_interval_ms: u64,
353 ) -> Result<crate::models::QueryResult> {
354 eprintln!("Executing query {query_id}...");
355 let job = self.refresh_query(query_id, parameters).await?;
356 let result_id = self
357 .poll_job_to_completion(job, timeout_secs, poll_interval_ms)
358 .await?;
359
360 eprintln!("Query completed, fetching results...");
361 self.get_query_result(query_id, result_id).await
362 }
363
364 pub async fn execute_adhoc_with_polling(
365 &self,
366 sql: &str,
367 data_source_id: u64,
368 parameters: Option<std::collections::HashMap<String, serde_json::Value>>,
369 timeout_secs: u64,
370 poll_interval_ms: u64,
371 ) -> Result<crate::models::QueryResult> {
372 eprintln!("Executing SQL on data source {data_source_id}...");
373 let job = self
374 .refresh_adhoc_query(sql, data_source_id, parameters)
375 .await?;
376 let result_id = self
377 .poll_job_to_completion(job, timeout_secs, poll_interval_ms)
378 .await?;
379
380 eprintln!("Query completed, fetching results...");
381 self.get_adhoc_query_result(result_id).await
382 }
383
384 pub async fn archive_query(&self, query_id: u64) -> Result<Query> {
385 let url = format!("{}/api/queries/{query_id}", self.base_url);
386 let payload = serde_json::json!({"is_archived": true});
387 self.post_json(&url, &payload, &format!("query {query_id} archive"))
388 .await
389 }
390
391 pub async fn unarchive_query(&self, query_id: u64) -> Result<Query> {
392 let url = format!("{}/api/queries/{query_id}", self.base_url);
393 let payload = serde_json::json!({"is_archived": false});
394 self.post_json(&url, &payload, &format!("query {query_id} unarchive"))
395 .await
396 }
397
398 pub async fn list_query_snippets(&self) -> Result<Vec<QuerySnippet>> {
399 let url = format!("{}/api/query_snippets", self.base_url);
400 self.get_json(&url, "query snippets").await
401 }
402
403 pub async fn get_query_snippet(&self, snippet_id: u64) -> Result<QuerySnippet> {
404 let url = format!("{}/api/query_snippets/{snippet_id}", self.base_url);
405 self.get_json(&url, &format!("query snippet {snippet_id}"))
406 .await
407 }
408
409 pub async fn create_query_snippet(&self, create: &CreateQuerySnippet) -> Result<QuerySnippet> {
410 let url = format!("{}/api/query_snippets", self.base_url);
411 self.post_json(&url, create, "new query snippet").await
412 }
413
414 pub async fn update_query_snippet(&self, snippet: &QuerySnippet) -> Result<QuerySnippet> {
415 let url = format!("{}/api/query_snippets/{}", self.base_url, snippet.id);
416 self.post_json(
417 &url,
418 snippet,
419 &format!("query snippet {} update", snippet.id),
420 )
421 .await
422 }
423
424 pub async fn delete_query_snippet(&self, snippet_id: u64) -> Result<()> {
425 let url = format!("{}/api/query_snippets/{snippet_id}", self.base_url);
426 let response = self
427 .client
428 .delete(&url)
429 .send()
430 .await
431 .context(format!("Failed to delete query snippet {snippet_id}"))?;
432
433 ensure_success(response).await?;
434
435 Ok(())
436 }
437
438 pub async fn create_dashboard(&self, dashboard: &CreateDashboard) -> Result<Dashboard> {
439 let url = format!("{}/api/dashboards", self.base_url);
440 self.post_json(&url, dashboard, "new dashboard").await
441 }
442
443 pub async fn list_favorite_dashboards(
444 &self,
445 page: u32,
446 page_size: u32,
447 ) -> Result<DashboardsResponse> {
448 let url = format!(
449 "{}/api/dashboards/favorites?page={page}&page_size={page_size}",
450 self.base_url
451 );
452 self.get_json(&url, "favorite dashboards").await
453 }
454
455 pub async fn get_dashboard(&self, slug_or_id: &str) -> Result<Dashboard> {
456 let url = format!("{}/api/dashboards/{slug_or_id}", self.base_url);
457 self.get_json(&url, &format!("dashboard {slug_or_id}"))
458 .await
459 }
460
461 pub async fn update_dashboard(&self, dashboard: &Dashboard) -> Result<Dashboard> {
462 let url = format!("{}/api/dashboards/{}", self.base_url, dashboard.id);
463 self.post_json(
464 &url,
465 dashboard,
466 &format!("dashboard {} update", dashboard.id),
467 )
468 .await
469 }
470
471 pub async fn archive_dashboard(&self, dashboard_id: u64) -> Result<()> {
472 let url = format!("{}/api/dashboards/{dashboard_id}", self.base_url);
473 let payload = serde_json::json!({"is_archived": true});
474 let response = self
475 .client
476 .post(&url)
477 .json(&payload)
478 .send()
479 .await
480 .context(format!("Failed to archive dashboard {dashboard_id}"))?;
481
482 ensure_success(response).await?;
483
484 Ok(())
485 }
486
487 pub async fn unarchive_dashboard(&self, dashboard_id: u64) -> Result<Dashboard> {
488 let url = format!("{}/api/dashboards/{dashboard_id}", self.base_url);
489 let payload = serde_json::json!({"is_archived": false});
490 self.post_json(
491 &url,
492 &payload,
493 &format!("dashboard {dashboard_id} unarchive"),
494 )
495 .await
496 }
497
498 pub async fn create_widget(&self, widget: &CreateWidget) -> Result<crate::models::Widget> {
499 let url = format!("{}/api/widgets", self.base_url);
500 self.post_json(&url, widget, "new widget").await
501 }
502
503 pub async fn update_widget(
504 &self,
505 widget_id: u64,
506 widget: &CreateWidget,
507 ) -> Result<crate::models::Widget> {
508 let url = format!("{}/api/widgets/{widget_id}", self.base_url);
509 self.post_json(&url, widget, &format!("widget {widget_id} update"))
510 .await
511 }
512
513 pub async fn delete_widget(&self, widget_id: u64) -> Result<()> {
514 let url = format!("{}/api/widgets/{widget_id}", self.base_url);
515 let response = self
516 .client
517 .delete(&url)
518 .send()
519 .await
520 .context(format!("Failed to delete widget {widget_id}"))?;
521
522 ensure_success(response).await?;
523
524 Ok(())
525 }
526
527 pub async fn favorite_dashboard(&self, slug: &str) -> Result<()> {
528 let url = format!("{}/api/dashboards/{slug}/favorite", self.base_url);
529 let response = self
530 .client
531 .post(&url)
532 .json(&serde_json::json!({}))
533 .send()
534 .await
535 .context(format!("Failed to favorite dashboard {slug}"))?;
536
537 ensure_success(response).await?;
538
539 Ok(())
540 }
541
542 pub async fn fetch_favorite_dashboards(&self) -> Result<Vec<DashboardSummary>> {
543 let mut all_dashboards = Vec::new();
544 let mut page = 1;
545 let page_size = 100;
546
547 loop {
548 let response = self.list_favorite_dashboards(page, page_size).await?;
549
550 if response.results.is_empty() {
551 break;
552 }
553
554 all_dashboards.extend(response.results);
555 eprintln!(
556 "Fetched {} / {} dashboards...",
557 all_dashboards.len(),
558 response.count
559 );
560
561 #[allow(clippy::cast_possible_truncation)]
562 if all_dashboards.len() >= response.count as usize {
563 break;
564 }
565
566 page += 1;
567 }
568
569 Ok(all_dashboards)
570 }
571
572 async fn get_with_retry(
573 &self,
574 url: &str,
575 params: &[(&str, String)],
576 ) -> Result<reqwest::Response> {
577 use tokio::time::{Duration, sleep};
578
579 const MAX_ATTEMPTS: u32 = 4;
580 let base_delays = [250u64, 500, 1000, 2000];
581
582 let mut last_error = anyhow::anyhow!("No attempts made");
583 for attempt in 0..MAX_ATTEMPTS {
584 let response = self
585 .client
586 .get(url)
587 .query(params)
588 .send()
589 .await
590 .with_context(|| format!("Failed to GET {url}"))?;
591
592 let status = response.status();
593 let should_retry = status.as_u16() == 429 || status.is_server_error();
594
595 match ensure_success(response).await {
596 Ok(response) => return Ok(response),
597 Err(err) if !should_retry || attempt + 1 == MAX_ATTEMPTS => return Err(err),
598 Err(err) => last_error = err,
599 }
600
601 let delay_ms = base_delays[attempt as usize];
602 sleep(Duration::from_millis(delay_ms)).await;
603 }
604
605 Err(last_error)
606 }
607
608 async fn list_queries(&self, q: &str, page: u32, page_size: u32) -> Result<QueriesResponse> {
609 let url = format!("{}/api/queries", self.base_url);
610 let params = [
611 ("page", page.to_string()),
612 ("page_size", page_size.to_string()),
613 ("q", q.to_string()),
614 ];
615 self.get_with_retry(&url, ¶ms)
616 .await?
617 .json()
618 .await
619 .context("Failed to parse queries response")
620 }
621
622 async fn list_dashboards(
623 &self,
624 q: &str,
625 page: u32,
626 page_size: u32,
627 ) -> Result<DashboardsResponse> {
628 let url = format!("{}/api/dashboards", self.base_url);
629 let params = [
630 ("page", page.to_string()),
631 ("page_size", page_size.to_string()),
632 ("q", q.to_string()),
633 ];
634 self.get_with_retry(&url, ¶ms)
635 .await?
636 .json()
637 .await
638 .context("Failed to parse dashboards response")
639 }
640
641 pub async fn search_queries(&self, q: &str, limit: usize) -> Result<Vec<Query>> {
642 const PAGE_SIZE: usize = 250;
643
644 let mut results: Vec<Query> = Vec::new();
645 let mut page = 1u32;
646
647 loop {
648 let remaining = limit - results.len();
649 #[allow(clippy::cast_possible_truncation)]
650 let page_size = remaining.min(PAGE_SIZE) as u32;
651 let response = self.list_queries(q, page, page_size).await?;
652
653 results.extend(response.results);
654
655 #[allow(clippy::cast_possible_truncation)]
656 if results.len() >= limit || results.len() >= response.count as usize {
657 break;
658 }
659
660 page += 1;
661 }
662
663 results.truncate(limit);
664 Ok(results)
665 }
666
667 pub async fn search_dashboards(&self, q: &str, limit: usize) -> Result<Vec<DashboardSummary>> {
668 const PAGE_SIZE: usize = 250;
669
670 let mut results: Vec<DashboardSummary> = Vec::new();
671 let mut page = 1u32;
672
673 loop {
674 let remaining = limit - results.len();
675 #[allow(clippy::cast_possible_truncation)]
676 let page_size = remaining.min(PAGE_SIZE) as u32;
677 let response = self.list_dashboards(q, page, page_size).await?;
678
679 results.extend(response.results);
680
681 #[allow(clippy::cast_possible_truncation)]
682 if results.len() >= limit || results.len() >= response.count as usize {
683 break;
684 }
685
686 page += 1;
687 }
688
689 results.truncate(limit);
690 Ok(results)
691 }
692}
693
694async fn ensure_success(response: reqwest::Response) -> Result<reqwest::Response> {
695 let status = response.status();
696 if !status.is_success() {
697 let body = response.text().await.unwrap_or_default();
698 anyhow::bail!("API error {status}: {body}");
699 }
700 Ok(response)
701}
702
703async fn wait_for_ctrl_c() {
704 let _ = tokio::signal::ctrl_c().await;
705}