metabase_api_rs/core/models/
query.rs

1//! Query models for executing queries against Metabase
2//!
3//! This module provides data structures for dataset queries,
4//! native SQL queries, and their results.
5
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11use super::common::MetabaseId;
12
13/// Represents a dataset query to be executed
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct DatasetQuery {
16    /// Database ID to query against
17    pub database: MetabaseId,
18
19    /// Query type (e.g., "query", "native")
20    #[serde(rename = "type")]
21    pub query_type: String,
22
23    /// The actual query (MBQL or native)
24    pub query: Value,
25
26    /// Optional parameters for the query
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub parameters: Option<Vec<QueryParameter>>,
29
30    /// Optional constraints (e.g., max rows)
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub constraints: Option<QueryConstraints>,
33}
34
35/// Represents a native SQL query
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct NativeQuery {
38    /// The SQL query string
39    pub query: String,
40
41    /// Template tags for parameterized queries
42    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
43    #[serde(rename = "template-tags")]
44    pub template_tags: HashMap<String, TemplateTag>,
45
46    /// Collection of tables used in the query
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub collection: Option<String>,
49}
50
51/// Parameter for a query
52#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
53pub struct QueryParameter {
54    /// Parameter ID or name
55    pub id: String,
56
57    /// Parameter type
58    #[serde(rename = "type")]
59    pub parameter_type: String,
60
61    /// Parameter value
62    pub value: Value,
63
64    /// Target field or variable
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub target: Option<Value>,
67}
68
69/// Template tag for native queries
70#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
71pub struct TemplateTag {
72    /// Unique identifier for the tag
73    pub id: String,
74
75    /// Tag name
76    pub name: String,
77
78    /// Display name
79    #[serde(rename = "display-name")]
80    pub display_name: String,
81
82    /// Tag type (e.g., "text", "number", "date")
83    #[serde(rename = "type")]
84    pub tag_type: String,
85
86    /// Whether the tag is required
87    #[serde(default)]
88    pub required: bool,
89
90    /// Default value
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub default: Option<Value>,
93}
94
95/// Query constraints
96#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97pub struct QueryConstraints {
98    /// Maximum number of rows to return
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub max_results: Option<i32>,
101
102    /// Maximum execution time in seconds
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub max_execution_time: Option<i32>,
105}
106
107/// Result of a query execution
108#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
109pub struct QueryResult {
110    /// The actual data returned
111    pub data: QueryData,
112
113    /// Database that was queried
114    pub database_id: MetabaseId,
115
116    /// When the query started
117    pub started_at: DateTime<Utc>,
118
119    /// When the query finished
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub finished_at: Option<DateTime<Utc>>,
122
123    /// The query that was executed
124    pub json_query: Value,
125
126    /// Query status
127    pub status: QueryStatus,
128
129    /// Row count
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub row_count: Option<i32>,
132
133    /// Running time in milliseconds
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub running_time: Option<i32>,
136}
137
138/// Query execution status
139#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
140#[serde(rename_all = "lowercase")]
141pub enum QueryStatus {
142    /// Query is still running
143    Running,
144    /// Query completed successfully
145    Completed,
146    /// Query failed with an error
147    Failed,
148    /// Query was cancelled
149    Cancelled,
150}
151
152/// The actual data from a query result
153#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
154pub struct QueryData {
155    /// Column information
156    pub cols: Vec<Column>,
157
158    /// Row data
159    pub rows: Vec<Vec<Value>>,
160
161    /// Native form of the results
162    #[serde(skip_serializing_if = "Option::is_none")]
163    pub native_form: Option<Value>,
164
165    /// Insights from the query
166    #[serde(default, skip_serializing_if = "Vec::is_empty")]
167    pub insights: Vec<Insight>,
168}
169
170/// Column metadata
171#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
172pub struct Column {
173    /// Column name
174    pub name: String,
175
176    /// Display name
177    pub display_name: String,
178
179    /// Base type (e.g., "type/Text", "type/Integer")
180    pub base_type: String,
181
182    /// Effective type
183    #[serde(skip_serializing_if = "Option::is_none")]
184    pub effective_type: Option<String>,
185
186    /// Semantic type
187    #[serde(skip_serializing_if = "Option::is_none")]
188    pub semantic_type: Option<String>,
189
190    /// Field reference
191    #[serde(skip_serializing_if = "Option::is_none")]
192    pub field_ref: Option<Value>,
193}
194
195/// Query insight
196#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
197pub struct Insight {
198    /// Insight type
199    #[serde(rename = "type")]
200    pub insight_type: String,
201
202    /// Insight value
203    pub value: Value,
204}
205
206/// Request to execute a dataset query
207#[derive(Debug, Clone, Serialize)]
208pub struct ExecuteQueryRequest {
209    /// The dataset query to execute
210    pub dataset_query: DatasetQuery,
211
212    /// Visualization settings
213    #[serde(skip_serializing_if = "Option::is_none")]
214    pub visualization_settings: Option<Value>,
215
216    /// Display type
217    #[serde(skip_serializing_if = "Option::is_none")]
218    pub display: Option<String>,
219}
220
221/// Request to execute a native SQL query
222#[derive(Debug, Clone, Serialize)]
223pub struct ExecuteNativeQueryRequest {
224    /// Database to execute against
225    pub database: MetabaseId,
226
227    /// The native query
228    pub native: NativeQuery,
229
230    /// Parameters for the query
231    #[serde(skip_serializing_if = "Option::is_none")]
232    pub parameters: Option<Vec<QueryParameter>>,
233}
234
235impl DatasetQuery {
236    /// Creates a new dataset query builder
237    pub fn builder(database: MetabaseId) -> DatasetQueryBuilder {
238        DatasetQueryBuilder::new(database)
239    }
240}
241
242impl NativeQuery {
243    /// Creates a new NativeQuery with the given SQL
244    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    /// Creates a new NativeQuery builder
253    pub fn builder(sql: impl Into<String>) -> NativeQueryBuilder {
254        NativeQueryBuilder::new(sql)
255    }
256
257    /// Add a parameter to the query
258    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
278/// Builder for creating DatasetQuery instances
279pub 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    /// Creates a new query builder
289    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    /// Sets the query type
300    pub fn query_type(mut self, query_type: impl Into<String>) -> Self {
301        self.query_type = query_type.into();
302        self
303    }
304
305    /// Sets the query
306    pub fn query(mut self, query: Value) -> Self {
307        self.query = query;
308        self
309    }
310
311    /// Sets the parameters
312    pub fn parameters(mut self, params: Vec<QueryParameter>) -> Self {
313        self.parameters = Some(params);
314        self
315    }
316
317    /// Sets the constraints
318    pub fn constraints(mut self, constraints: QueryConstraints) -> Self {
319        self.constraints = Some(constraints);
320        self
321    }
322
323    /// Builds the DatasetQuery
324    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
335/// Builder for creating NativeQuery instances
336pub struct NativeQueryBuilder {
337    query: String,
338    template_tags: HashMap<String, TemplateTag>,
339    collection: Option<String>,
340}
341
342impl NativeQueryBuilder {
343    /// Creates a new NativeQuery builder
344    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    /// Adds a generic parameter
353    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    /// Adds a text parameter
367    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    /// Adds a number parameter
372    pub fn add_number_param(self, name: &str, value: f64) -> Self {
373        self.add_param(name, "number", serde_json::json!(value))
374    }
375
376    /// Adds a date parameter
377    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    /// Adds parameters from a HashMap
382    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    /// Sets the collection
405    pub fn collection(mut self, collection: impl Into<String>) -> Self {
406        self.collection = Some(collection.into());
407        self
408    }
409
410    /// Builds the NativeQuery
411    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}