1use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11use super::common::MetabaseId;
12
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct DatasetQuery {
16 pub database: MetabaseId,
18
19 #[serde(rename = "type")]
21 pub query_type: String,
22
23 pub query: Value,
25
26 #[serde(skip_serializing_if = "Option::is_none")]
28 pub parameters: Option<Vec<QueryParameter>>,
29
30 #[serde(skip_serializing_if = "Option::is_none")]
32 pub constraints: Option<QueryConstraints>,
33}
34
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct NativeQuery {
38 pub query: String,
40
41 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
43 #[serde(rename = "template-tags")]
44 pub template_tags: HashMap<String, TemplateTag>,
45
46 #[serde(skip_serializing_if = "Option::is_none")]
48 pub collection: Option<String>,
49}
50
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
53pub struct QueryParameter {
54 pub id: String,
56
57 #[serde(rename = "type")]
59 pub parameter_type: String,
60
61 pub value: Value,
63
64 #[serde(skip_serializing_if = "Option::is_none")]
66 pub target: Option<Value>,
67}
68
69#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
71pub struct TemplateTag {
72 pub id: String,
74
75 pub name: String,
77
78 #[serde(rename = "display-name")]
80 pub display_name: String,
81
82 #[serde(rename = "type")]
84 pub tag_type: String,
85
86 #[serde(default)]
88 pub required: bool,
89
90 #[serde(skip_serializing_if = "Option::is_none")]
92 pub default: Option<Value>,
93}
94
95#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97pub struct QueryConstraints {
98 #[serde(skip_serializing_if = "Option::is_none")]
100 pub max_results: Option<i32>,
101
102 #[serde(skip_serializing_if = "Option::is_none")]
104 pub max_execution_time: Option<i32>,
105}
106
107#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
109pub struct QueryResult {
110 pub data: QueryData,
112
113 pub database_id: MetabaseId,
115
116 pub started_at: DateTime<Utc>,
118
119 #[serde(skip_serializing_if = "Option::is_none")]
121 pub finished_at: Option<DateTime<Utc>>,
122
123 pub json_query: Value,
125
126 pub status: QueryStatus,
128
129 #[serde(skip_serializing_if = "Option::is_none")]
131 pub row_count: Option<i32>,
132
133 #[serde(skip_serializing_if = "Option::is_none")]
135 pub running_time: Option<i32>,
136}
137
138#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
140#[serde(rename_all = "lowercase")]
141pub enum QueryStatus {
142 Running,
144 Completed,
146 Failed,
148 Cancelled,
150}
151
152#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
154pub struct QueryData {
155 pub cols: Vec<Column>,
157
158 pub rows: Vec<Vec<Value>>,
160
161 #[serde(skip_serializing_if = "Option::is_none")]
163 pub native_form: Option<Value>,
164
165 #[serde(default, skip_serializing_if = "Vec::is_empty")]
167 pub insights: Vec<Insight>,
168}
169
170#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
172pub struct Column {
173 pub name: String,
175
176 pub display_name: String,
178
179 pub base_type: String,
181
182 #[serde(skip_serializing_if = "Option::is_none")]
184 pub effective_type: Option<String>,
185
186 #[serde(skip_serializing_if = "Option::is_none")]
188 pub semantic_type: Option<String>,
189
190 #[serde(skip_serializing_if = "Option::is_none")]
192 pub field_ref: Option<Value>,
193}
194
195#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
197pub struct Insight {
198 #[serde(rename = "type")]
200 pub insight_type: String,
201
202 pub value: Value,
204}
205
206#[derive(Debug, Clone, Serialize)]
208pub struct ExecuteQueryRequest {
209 pub dataset_query: DatasetQuery,
211
212 #[serde(skip_serializing_if = "Option::is_none")]
214 pub visualization_settings: Option<Value>,
215
216 #[serde(skip_serializing_if = "Option::is_none")]
218 pub display: Option<String>,
219}
220
221#[derive(Debug, Clone, Serialize)]
223pub struct ExecuteNativeQueryRequest {
224 pub database: MetabaseId,
226
227 pub native: NativeQuery,
229
230 #[serde(skip_serializing_if = "Option::is_none")]
232 pub parameters: Option<Vec<QueryParameter>>,
233}
234
235impl DatasetQuery {
236 pub fn builder(database: MetabaseId) -> DatasetQueryBuilder {
238 DatasetQueryBuilder::new(database)
239 }
240}
241
242impl NativeQuery {
243 pub fn new(sql: impl Into<String>) -> Self {
245 Self {
246 query: sql.into(),
247 template_tags: HashMap::new(),
248 collection: None,
249 }
250 }
251
252 pub fn builder(sql: impl Into<String>) -> NativeQueryBuilder {
254 NativeQueryBuilder::new(sql)
255 }
256
257 pub fn with_param(mut self, name: &str, value: Value) -> Self {
259 let tag = TemplateTag {
260 id: uuid::Uuid::new_v4().to_string(),
261 name: name.to_string(),
262 display_name: name.to_string(),
263 tag_type: match &value {
264 Value::String(_) => "text",
265 Value::Number(_) => "number",
266 Value::Bool(_) => "text",
267 _ => "text",
268 }
269 .to_string(),
270 required: false,
271 default: Some(value),
272 };
273 self.template_tags.insert(name.to_string(), tag);
274 self
275 }
276}
277
278pub struct DatasetQueryBuilder {
280 database: MetabaseId,
281 query_type: String,
282 query: Value,
283 parameters: Option<Vec<QueryParameter>>,
284 constraints: Option<QueryConstraints>,
285}
286
287impl DatasetQueryBuilder {
288 pub fn new(database: MetabaseId) -> Self {
290 Self {
291 database,
292 query_type: "query".to_string(),
293 query: Value::Null,
294 parameters: None,
295 constraints: None,
296 }
297 }
298
299 pub fn query_type(mut self, query_type: impl Into<String>) -> Self {
301 self.query_type = query_type.into();
302 self
303 }
304
305 pub fn query(mut self, query: Value) -> Self {
307 self.query = query;
308 self
309 }
310
311 pub fn parameters(mut self, params: Vec<QueryParameter>) -> Self {
313 self.parameters = Some(params);
314 self
315 }
316
317 pub fn constraints(mut self, constraints: QueryConstraints) -> Self {
319 self.constraints = Some(constraints);
320 self
321 }
322
323 pub fn build(self) -> DatasetQuery {
325 DatasetQuery {
326 database: self.database,
327 query_type: self.query_type,
328 query: self.query,
329 parameters: self.parameters,
330 constraints: self.constraints,
331 }
332 }
333}
334
335pub struct NativeQueryBuilder {
337 query: String,
338 template_tags: HashMap<String, TemplateTag>,
339 collection: Option<String>,
340}
341
342impl NativeQueryBuilder {
343 pub fn new(sql: impl Into<String>) -> Self {
345 Self {
346 query: sql.into(),
347 template_tags: HashMap::new(),
348 collection: None,
349 }
350 }
351
352 pub fn add_param(mut self, name: &str, param_type: &str, value: Value) -> Self {
354 let tag = TemplateTag {
355 id: uuid::Uuid::new_v4().to_string(),
356 name: name.to_string(),
357 display_name: name.to_string(),
358 tag_type: param_type.to_string(),
359 required: false,
360 default: Some(value),
361 };
362 self.template_tags.insert(name.to_string(), tag);
363 self
364 }
365
366 pub fn add_text_param(self, name: &str, value: &str) -> Self {
368 self.add_param(name, "text", Value::String(value.to_string()))
369 }
370
371 pub fn add_number_param(self, name: &str, value: f64) -> Self {
373 self.add_param(name, "number", serde_json::json!(value))
374 }
375
376 pub fn add_date_param(self, name: &str, value: &str) -> Self {
378 self.add_param(name, "date", Value::String(value.to_string()))
379 }
380
381 pub fn with_params(mut self, params: HashMap<String, Value>) -> Self {
383 for (name, value) in params {
384 let param_type = match &value {
385 Value::String(_) => "text",
386 Value::Number(_) => "number",
387 Value::Bool(_) => "text",
388 _ => "text",
389 };
390
391 let tag = TemplateTag {
392 id: uuid::Uuid::new_v4().to_string(),
393 name: name.clone(),
394 display_name: name.clone(),
395 tag_type: param_type.to_string(),
396 required: false,
397 default: Some(value),
398 };
399 self.template_tags.insert(name, tag);
400 }
401 self
402 }
403
404 pub fn collection(mut self, collection: impl Into<String>) -> Self {
406 self.collection = Some(collection.into());
407 self
408 }
409
410 pub fn build(self) -> NativeQuery {
412 NativeQuery {
413 query: self.query,
414 template_tags: self.template_tags,
415 collection: self.collection,
416 }
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use super::*;
423 use serde_json::json;
424
425 #[test]
426 fn test_dataset_query_builder() {
427 let query = DatasetQuery::builder(MetabaseId(1))
428 .query_type("native")
429 .query(json!({"query": "SELECT * FROM users"}))
430 .build();
431
432 assert_eq!(query.database, MetabaseId(1));
433 assert_eq!(query.query_type, "native");
434 assert_eq!(query.query, json!({"query": "SELECT * FROM users"}));
435 }
436
437 #[test]
438 fn test_native_query() {
439 let mut template_tags = HashMap::new();
440 template_tags.insert(
441 "date".to_string(),
442 TemplateTag {
443 id: "test-id".to_string(),
444 name: "date".to_string(),
445 display_name: "Date".to_string(),
446 tag_type: "date".to_string(),
447 required: true,
448 default: None,
449 },
450 );
451
452 let native = NativeQuery {
453 query: "SELECT * FROM orders WHERE created_at > {{date}}".to_string(),
454 template_tags,
455 collection: None,
456 };
457
458 assert_eq!(native.template_tags.len(), 1);
459 assert!(native.template_tags.contains_key("date"));
460 assert!(native.template_tags["date"].required);
461 }
462
463 #[test]
464 fn test_native_query_builder() {
465 let query = NativeQuery::builder("SELECT * FROM orders WHERE status = {{status}}")
466 .add_text_param("status", "completed")
467 .build();
468
469 assert_eq!(
470 query.query,
471 "SELECT * FROM orders WHERE status = {{status}}"
472 );
473 assert!(query.template_tags.contains_key("status"));
474 assert_eq!(query.template_tags["status"].tag_type, "text");
475 assert_eq!(
476 query.template_tags["status"].default,
477 Some(json!("completed"))
478 );
479 }
480
481 #[test]
482 fn test_query_result() {
483 let result = QueryResult {
484 data: QueryData {
485 cols: vec![Column {
486 name: "id".to_string(),
487 display_name: "ID".to_string(),
488 base_type: "type/Integer".to_string(),
489 effective_type: None,
490 semantic_type: None,
491 field_ref: None,
492 }],
493 rows: vec![vec![json!(1)], vec![json!(2)]],
494 native_form: None,
495 insights: vec![],
496 },
497 database_id: MetabaseId(1),
498 started_at: Utc::now(),
499 finished_at: Some(Utc::now()),
500 json_query: json!({}),
501 status: QueryStatus::Completed,
502 row_count: Some(2),
503 running_time: Some(150),
504 };
505
506 assert_eq!(result.status, QueryStatus::Completed);
507 assert_eq!(result.row_count, Some(2));
508 assert_eq!(result.data.rows.len(), 2);
509 }
510
511 #[test]
512 fn test_query_constraints() {
513 let constraints = QueryConstraints {
514 max_results: Some(1000),
515 max_execution_time: Some(60),
516 };
517
518 assert_eq!(constraints.max_results, Some(1000));
519 assert_eq!(constraints.max_execution_time, Some(60));
520 }
521}