zino_core/model/
context.rs

1use crate::Uuid;
2use std::time::Instant;
3
4/// Data associated with a query.
5#[derive(Debug, Clone)]
6pub struct QueryContext {
7    /// Model name.
8    model_name: &'static str,
9    /// Start time.
10    start_time: Instant,
11    /// Query ID.
12    query_id: Uuid,
13    /// A query.
14    query: String,
15    /// Arguments.
16    arguments: Vec<String>,
17    /// Last insert ID.
18    last_insert_id: Option<i64>,
19    /// Number of rows affected.
20    rows_affected: Option<u64>,
21    /// Indicates the query execution is successful or not.
22    success: bool,
23    /// Indicates the query execution is cancelled or not.
24    cancelled: bool,
25}
26
27impl QueryContext {
28    /// Creates a new instance.
29    #[inline]
30    pub fn new(model_name: &'static str) -> Self {
31        Self {
32            model_name,
33            start_time: Instant::now(),
34            query_id: Uuid::now_v7(),
35            query: String::new(),
36            arguments: Vec::new(),
37            last_insert_id: None,
38            rows_affected: None,
39            success: false,
40            cancelled: false,
41        }
42    }
43
44    /// Sets the query.
45    #[inline]
46    pub fn set_query(&mut self, query: impl Into<String>) {
47        self.query = query.into();
48    }
49
50    /// Adds an argument to the list of query arguments.
51    #[inline]
52    pub fn add_argument(&mut self, value: impl ToString) {
53        self.arguments.push(value.to_string());
54    }
55
56    /// Appends the query arguments.
57    #[inline]
58    pub fn append_arguments(&mut self, arguments: &mut Vec<String>) {
59        self.arguments.append(arguments);
60    }
61
62    /// Sets the last insert ID.
63    #[inline]
64    pub fn set_last_insert_id(&mut self, last_insert_id: i64) {
65        self.last_insert_id = Some(last_insert_id);
66    }
67
68    /// Sets the query result.
69    #[inline]
70    pub fn set_query_result(&mut self, rows_affected: impl Into<Option<u64>>, success: bool) {
71        self.rows_affected = rows_affected.into();
72        self.success = success;
73        self.cancelled = false;
74    }
75
76    /// Cancells the query execution.
77    #[inline]
78    pub fn cancel(&mut self) {
79        self.cancelled = true;
80    }
81
82    /// Returns the model name.
83    #[inline]
84    pub fn model_name(&self) -> &'static str {
85        self.model_name
86    }
87
88    /// Returns the start time.
89    #[inline]
90    pub fn start_time(&self) -> Instant {
91        self.start_time
92    }
93
94    /// Returns the query ID.
95    #[inline]
96    pub fn query_id(&self) -> Uuid {
97        self.query_id
98    }
99
100    /// Returns the query.
101    #[inline]
102    pub fn query(&self) -> &str {
103        &self.query
104    }
105
106    /// Returns the query arguments.
107    #[inline]
108    pub fn arguments(&self) -> &[String] {
109        &self.arguments
110    }
111
112    /// Returns the last insert ID.
113    #[inline]
114    pub fn last_insert_id(&self) -> Option<i64> {
115        self.last_insert_id
116    }
117
118    /// Returns the number of rows affected.
119    #[inline]
120    pub fn rows_affected(&self) -> Option<u64> {
121        self.rows_affected
122    }
123
124    /// Returns `true` if the query execution is cancelled.
125    #[inline]
126    pub fn is_cancelled(&self) -> bool {
127        self.cancelled
128    }
129
130    /// Returns `true` if the query execution is successful.
131    #[inline]
132    pub fn is_success(&self) -> bool {
133        self.success
134    }
135
136    /// Formats the query arguments as a `String` if they exist.
137    #[inline]
138    pub fn format_arguments(&self) -> Option<String> {
139        let arguments = self.arguments();
140        (!arguments.is_empty()).then(|| arguments.join(", "))
141    }
142
143    /// Records an error message for the query.
144    pub fn record_error(&self, message: impl AsRef<str>) {
145        fn inner(ctx: &QueryContext, message: &str) {
146            let model_name = ctx.model_name();
147            let query_id = ctx.query_id().to_string();
148            let query = ctx.query();
149            let arguments = ctx.format_arguments();
150            if ctx.is_cancelled() {
151                tracing::warn!(
152                    cancelled = true,
153                    model_name,
154                    query_id,
155                    query,
156                    arguments,
157                    message,
158                );
159            } else {
160                tracing::error!(model_name, query_id, query, arguments, message);
161            }
162        }
163        inner(self, message.as_ref())
164    }
165
166    /// Emits the metrics for the query.
167    #[cfg(feature = "metrics")]
168    #[inline]
169    pub fn emit_metrics(&self, action: impl Into<crate::SharedString>) {
170        fn inner(ctx: &QueryContext, action: crate::SharedString) {
171            metrics::histogram!(
172                "zino_model_query_duration_seconds",
173                "model_name" => ctx.model_name(),
174                "action" => action,
175            )
176            .record(ctx.start_time().elapsed().as_secs_f64());
177        }
178        inner(self, action.into())
179    }
180}