dataprof_db/connectors/
common.rs

1//! Common utilities and shared logic for database connectors
2//!
3//! This module provides reusable functions to reduce code duplication across
4//! PostgreSQL, MySQL, and SQLite connectors.
5
6use crate::DataProfilerError;
7use crate::security::{validate_base_query, validate_sql_identifier};
8
9/// Generate "not connected to database" error
10#[allow(dead_code)]
11pub fn not_connected_error() -> DataProfilerError {
12    DataProfilerError::database_connection("Not connected to database")
13}
14
15/// Generate feature-not-enabled error for a specific database
16#[allow(dead_code)]
17pub fn feature_not_enabled_error(db_name: &str, feature: &str) -> DataProfilerError {
18    DataProfilerError::database_feature_disabled(db_name, feature)
19}
20
21/// Render one column of one row as the string the profiler ingests.
22///
23/// The profiler consumes columns as `Vec<String>` and re-infers types from the
24/// textual form, so every SQL type has to be turned into a faithful string.
25/// sqlx offers no backend-agnostic "decode as whatever this is" call, so we try
26/// concrete types in order and take the first that decodes.
27///
28/// Order is load-bearing:
29///
30/// * `Option<String>` first. sqlx skips the type-compatibility check when the
31///   value is NULL, so a NULL of *any* SQL type decodes here as `Ok(None)`.
32///   That makes this arm the NULL detector as well as the text arm; everything
33///   below it is known to be a non-null value.
34/// * Integers before `bool`. SQLite stores booleans as INTEGER and will happily
35///   decode `42` as `true`, so trying `bool` early would render integers as
36///   "true"/"false".
37/// * Widest integer first, so an INT8/BIGINT is not truncated by a narrower arm.
38///
39/// A value whose type matches none of these arms (NUMERIC/DECIMAL, dates and
40/// times without the corresponding sqlx feature, BLOB) yields `None` and is
41/// recorded as a null, which is the pre-existing behaviour for those types.
42#[macro_export]
43macro_rules! db_column_to_string {
44    ($row:expr, $index:expr) => {{
45        let row = $row;
46        let index = $index;
47
48        if let Ok(v) = row.try_get::<Option<String>, _>(index) {
49            v
50        } else if let Ok(v) = row.try_get::<Option<i64>, _>(index) {
51            v.map(|x| x.to_string())
52        } else if let Ok(v) = row.try_get::<Option<i32>, _>(index) {
53            v.map(|x| x.to_string())
54        } else if let Ok(v) = row.try_get::<Option<i16>, _>(index) {
55            v.map(|x| x.to_string())
56        } else if let Ok(v) = row.try_get::<Option<f64>, _>(index) {
57            // `{:?}` keeps the decimal point on integral floats ("100.0", not
58            // "100"), so a REAL column of whole numbers is still inferred as a
59            // float downstream. It also stays compact at the extremes ("1e300").
60            v.map(|x| format!("{:?}", x))
61        } else if let Ok(v) = row.try_get::<Option<f32>, _>(index) {
62            v.map(|x| format!("{:?}", x))
63        } else if let Ok(v) = row.try_get::<Option<bool>, _>(index) {
64            v.map(|x| x.to_string())
65        } else {
66            None
67        }
68    }};
69}
70
71/// Macro to generate the streaming batch loop for profiling queries.
72#[macro_export]
73macro_rules! streaming_profile_loop {
74    ($pool:expr, $query:expr, $batch_size:expr, $total_rows:expr, $db_name:literal) => {{
75        use sqlx::{Column, Row};
76        use $crate::connectors::common::build_batch_query;
77        use $crate::streaming::{StreamingProgress, merge_column_batches};
78
79        let mut progress = StreamingProgress::new(Some($total_rows as u64));
80        let mut all_batches: Vec<$crate::QueryColumns> = Vec::new();
81        let mut offset = 0usize;
82
83        loop {
84            let batch_query = build_batch_query($query, $batch_size, offset)?;
85            let rows = sqlx::query(&batch_query)
86                .fetch_all($pool)
87                .await
88                .map_err(|e| $crate::DataProfilerError::DatabaseQueryError {
89                    message: format!("Batch query execution failed: {}", e),
90                })?;
91
92            if rows.is_empty() {
93                break;
94            }
95
96            let columns = rows[0].columns();
97            let column_names: Vec<String> = columns
98                .iter()
99                .map(|column| column.name().to_string())
100                .collect();
101            dataprof_core::validate_unique_column_names(
102                &column_names,
103                concat!($db_name, " query result"),
104            )?;
105            // Built from the driver's column list, so the batch carries the
106            // query's column order and values are filed by position.
107            let mut batch_result =
108                $crate::QueryColumns::with_names(column_names.clone(), rows.len());
109
110            for row in &rows {
111                for i in 0..columns.len() {
112                    let value: Option<String> = $crate::db_column_to_string!(row, i);
113                    // decode-audit: no-data — None is SQL NULL (or a type
114                    // db_column_to_string documents as unsupported); "" is
115                    // the profiler's textual null.
116                    batch_result.push_value(i, value.unwrap_or_default());
117                }
118            }
119
120            let batch_size_actual = rows.len();
121            all_batches.push(batch_result);
122            progress.update(batch_size_actual as u64);
123
124            if let Some(percentage) = progress.percentage() {
125                log::info!(
126                    "{} streaming progress: {:.1}% ({}/{} rows)",
127                    $db_name,
128                    percentage,
129                    progress.processed_rows,
130                    $total_rows
131                );
132            }
133
134            offset += $batch_size;
135            if batch_size_actual < $batch_size {
136                break;
137            }
138        }
139
140        Ok(merge_column_batches(all_batches))
141    }};
142}
143
144/// Macro to process rows into column-oriented results, in query column order.
145#[macro_export]
146macro_rules! process_rows_to_columns {
147    ($rows:expr) => {{
148        use sqlx::{Column, Row};
149
150        if $rows.is_empty() {
151            Ok($crate::QueryColumns::new())
152        } else {
153            let columns = $rows[0].columns();
154            let column_names: Vec<String> = columns
155                .iter()
156                .map(|column| column.name().to_string())
157                .collect();
158            match dataprof_core::validate_unique_column_names(
159                &column_names,
160                "database query result",
161            ) {
162                Err(error) => Err(error),
163                Ok(()) => {
164                    // Built from the driver's column list, so the result carries
165                    // the query's column order and values are filed by position.
166                    let mut result = $crate::QueryColumns::with_names(column_names, $rows.len());
167
168                    for row in &$rows {
169                        for i in 0..columns.len() {
170                            let value: Option<String> = $crate::db_column_to_string!(row, i);
171                            // decode-audit: no-data — None is SQL NULL (or a type
172                            // db_column_to_string documents as unsupported); "" is
173                            // the profiler's textual null.
174                            result.push_value(i, value.unwrap_or_default());
175                        }
176                    }
177
178                    Ok(result)
179                }
180            }
181        }
182    }};
183}
184
185/// Build a count query for a given table or query
186#[allow(dead_code)]
187pub fn build_count_query(query: &str) -> Result<String, DataProfilerError> {
188    if query.trim().to_uppercase().starts_with("SELECT") {
189        let validated_query = validate_base_query(query)?;
190        Ok(format!(
191            "SELECT COUNT(*) FROM ({}) as count_subquery",
192            validated_query
193        ))
194    } else {
195        validate_sql_identifier(query)?;
196        Ok(format!("SELECT COUNT(*) FROM {}", query))
197    }
198}
199
200/// Build a batch query with LIMIT and OFFSET
201#[allow(dead_code)]
202pub fn build_batch_query(
203    query: &str,
204    batch_size: usize,
205    offset: usize,
206) -> Result<String, DataProfilerError> {
207    let validated_query = if query.trim().to_uppercase().starts_with("SELECT") {
208        validate_base_query(query)?
209    } else {
210        validate_sql_identifier(query)?;
211        format!("SELECT * FROM {}", query)
212    };
213    Ok(format!(
214        "{} LIMIT {} OFFSET {}",
215        validated_query, batch_size, offset
216    ))
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn test_build_count_query_table() {
225        let result = build_count_query("users").unwrap();
226        assert_eq!(result, "SELECT COUNT(*) FROM users");
227    }
228
229    #[test]
230    fn test_build_count_query_select() {
231        let result = build_count_query("SELECT * FROM users WHERE active = true").unwrap();
232        assert!(result.contains("SELECT COUNT(*) FROM"));
233        assert!(result.contains("count_subquery"));
234    }
235
236    #[test]
237    fn test_build_batch_query() {
238        let result = build_batch_query("users", 100, 0).unwrap();
239        assert_eq!(result, "SELECT * FROM users LIMIT 100 OFFSET 0");
240    }
241}
dataprof_db/connectors/common.rs

dataprof_db/connectors/
common.rs
