query_forge/
engine.rs

1use std::{collections::HashSet, sync::OnceLock};
2
3use anyhow::{Context, Result, anyhow, bail};
4use regex::Regex;
5use rusqlite::{
6    Connection, params_from_iter,
7    types::ValueRef,
8};
9
10use crate::input;
11use crate::value::to_sql_value;
12use crate::{
13    ExtractionOptions, InputNormalizationOptions, QueryParam, QueryResult, QueryValue,
14    TypeInferenceOptions, WorkbookInput,
15};
16
17const SQLITE_MAX_INSERT_VARIABLES: usize = 999;
18const SQLITE_MAX_INSERT_ROWS: usize = 250;
19
20pub fn run_query(
21    workbook_path: &std::path::Path,
22    sheet_name: Option<&str>,
23    query: &str,
24    has_headers: bool,
25) -> Result<QueryResult> {
26    run_query_with_params(workbook_path, sheet_name, query, &[], has_headers)
27}
28
29pub fn run_query_with_params(
30    workbook_path: &std::path::Path,
31    sheet_name: Option<&str>,
32    query: &str,
33    params: &[QueryParam],
34    has_headers: bool,
35) -> Result<QueryResult> {
36    run_query_with_params_multi(&[workbook_path], sheet_name, query, params, has_headers)
37}
38
39pub fn run_query_with_params_multi(
40    workbook_paths: &[&std::path::Path],
41    sheet_name: Option<&str>,
42    query: &str,
43    params: &[QueryParam],
44    has_headers: bool,
45) -> Result<QueryResult> {
46    let workbook_inputs = workbook_paths
47        .iter()
48        .map(|path| WorkbookInput {
49            path,
50            sheet_name,
51            table_name: None,
52        })
53        .collect::<Vec<_>>();
54
55    run_query_with_params_multi_inputs(&workbook_inputs, query, params, has_headers)
56}
57
58pub fn run_query_with_params_multi_inputs(
59    workbook_inputs: &[WorkbookInput<'_>],
60    query: &str,
61    params: &[QueryParam],
62    has_headers: bool,
63) -> Result<QueryResult> {
64    run_query_with_params_multi_inputs_and_options(
65        workbook_inputs,
66        query,
67        params,
68        &TypeInferenceOptions::default(),
69        has_headers,
70    )
71}
72
73pub fn run_query_with_params_multi_inputs_and_options(
74    workbook_inputs: &[WorkbookInput<'_>],
75    query: &str,
76    params: &[QueryParam],
77    inference_options: &TypeInferenceOptions,
78    has_headers: bool,
79) -> Result<QueryResult> {
80    run_query_with_params_multi_inputs_and_options_and_normalization(
81        workbook_inputs,
82        query,
83        params,
84        inference_options,
85        &InputNormalizationOptions::default(),
86        &ExtractionOptions::default(),
87        has_headers,
88    )
89}
90
91pub fn run_query_with_params_multi_inputs_and_options_and_normalization(
92    workbook_inputs: &[WorkbookInput<'_>],
93    query: &str,
94    params: &[QueryParam],
95    inference_options: &TypeInferenceOptions,
96    normalization_options: &InputNormalizationOptions,
97    extraction_options: &ExtractionOptions,
98    has_headers: bool,
99) -> Result<QueryResult> {
100    if workbook_inputs.is_empty() {
101        bail!("at least one workbook input is required");
102    }
103
104    let connection =
105        Connection::open_in_memory().context("failed to create in-memory SQLite database")?;
106    let mut registered_sheet_views = HashSet::new();
107
108    for (index, workbook_input) in workbook_inputs.iter().enumerate() {
109        let mut sheet = input::load_input(
110            workbook_input.path,
111            workbook_input.sheet_name,
112            inference_options,
113            extraction_options,
114            has_headers,
115        )?;
116        input::apply_input_normalization(&mut sheet, normalization_options);
117        let table_name = workbook_input
118            .table_name
119            .map(str::to_owned)
120            .unwrap_or_else(|| {
121                if index == 0 {
122                    "table".to_owned()
123                } else {
124                    format!("table{}", index + 1)
125                }
126            });
127
128        register_sheet(
129            &connection,
130            &sheet,
131            &table_name,
132            &mut registered_sheet_views,
133        )?;
134    }
135
136    execute_query(&connection, query, params)
137}
138
139fn register_sheet(
140    connection: &Connection,
141    sheet: &input::SheetData,
142    table_name: &str,
143    registered_views: &mut HashSet<String>,
144) -> Result<()> {
145    let columns = sheet
146        .columns
147        .iter()
148        .map(|column| quote_identifier(column))
149        .collect::<Vec<_>>()
150        .join(", ");
151
152    connection
153        .execute(
154            &format!("CREATE TABLE {} ({columns})", quote_identifier(table_name)),
155            [],
156        )
157        .context("failed to create sheet table")?;
158
159    let table1_key = normalize_view_key("table1");
160    if table_name == "table" && !registered_views.contains(&table1_key) {
161        connection
162            .execute(
163                &format!(
164                    "CREATE VIEW {} AS SELECT * FROM {}",
165                    quote_identifier("table1"),
166                    quote_identifier("table")
167                ),
168                [],
169            )
170            .context("failed to register alias view table1 for first input")?;
171        registered_views.insert(table1_key);
172    }
173
174    let sanitized_sheet_name = sanitize_table_name(&sheet.original_name);
175    let sanitized_sheet_key = normalize_view_key(&sanitized_sheet_name);
176    if sanitized_sheet_name != table_name && !registered_views.contains(&sanitized_sheet_key) {
177        connection
178            .execute(
179                &format!(
180                    "CREATE VIEW {} AS SELECT * FROM {}",
181                    quote_identifier(&sanitized_sheet_name),
182                    quote_identifier(table_name)
183                ),
184                [],
185            )
186            .with_context(|| {
187                format!("failed to register view for sheet {}", sheet.original_name)
188            })?;
189        registered_views.insert(sanitized_sheet_key);
190    }
191
192    if sheet.rows.is_empty() {
193        return Ok(());
194    }
195
196    insert_sheet_rows(connection, table_name, sheet)
197}
198
199fn insert_sheet_rows(connection: &Connection, table_name: &str, sheet: &input::SheetData) -> Result<()> {
200    connection
201        .execute_batch("BEGIN IMMEDIATE TRANSACTION")
202        .context("failed to begin sheet insert transaction")?;
203
204    let insert_result = (|| {
205        let batch_size = calculate_insert_batch_size(sheet.columns.len());
206        for rows_chunk in sheet.rows.chunks(batch_size) {
207            let insert_sql =
208                build_batched_insert_sql(table_name, sheet.columns.len(), rows_chunk.len());
209            let mut values = Vec::with_capacity(rows_chunk.len() * sheet.columns.len());
210            for row in rows_chunk {
211                for index in 0..sheet.columns.len() {
212                    let value = row.get(index).unwrap_or(&QueryValue::Null);
213                    values.push(to_sql_value(value));
214                }
215            }
216
217            connection
218                .execute(&insert_sql, params_from_iter(values))
219                .context("failed to insert sheet rows")?;
220        }
221
222        Ok(())
223    })();
224
225    match insert_result {
226        Ok(()) => connection
227            .execute_batch("COMMIT")
228            .context("failed to commit sheet insert transaction"),
229        Err(error) => {
230            let _ = connection.execute_batch("ROLLBACK");
231            Err(error)
232        }
233    }
234}
235
236fn calculate_insert_batch_size(column_count: usize) -> usize {
237    let clamped_column_count = column_count.max(1);
238    (SQLITE_MAX_INSERT_VARIABLES / clamped_column_count)
239        .max(1)
240        .min(SQLITE_MAX_INSERT_ROWS)
241}
242
243fn build_batched_insert_sql(table_name: &str, column_count: usize, row_count: usize) -> String {
244    let row_placeholders = format!("({})", vec!["?"; column_count].join(", "));
245    let values = vec![row_placeholders; row_count].join(", ");
246    format!("INSERT INTO {} VALUES {values}", quote_identifier(table_name))
247}
248
249fn execute_query(
250    connection: &Connection,
251    query: &str,
252    params: &[QueryParam],
253) -> Result<QueryResult> {
254    let normalized_query = normalize_query_for_reserved_identifiers(query);
255    let mut statement = connection
256        .prepare(&normalized_query)
257        .map_err(|error| enrich_query_prepare_error(connection, query, error))?;
258
259    if statement.column_count() == 0 {
260        bail!("query must return rows");
261    }
262
263    let columns = statement
264        .column_names()
265        .into_iter()
266        .map(str::to_owned)
267        .collect::<Vec<_>>();
268
269    bind_query_params(&mut statement, params)?;
270
271    let column_count = statement.column_count();
272    let mut rows = statement.raw_query();
273    let mut result_rows = Vec::new();
274
275    while let Some(row) = rows.next().context("failed to fetch query row")? {
276        let mut values = Vec::with_capacity(column_count);
277
278        for index in 0..column_count {
279            values.push(match row.get_ref(index)? {
280                ValueRef::Null => QueryValue::Null,
281                ValueRef::Integer(value) => QueryValue::Integer(value),
282                ValueRef::Real(value) => QueryValue::Real(value),
283                ValueRef::Text(value) => {
284                    QueryValue::Text(String::from_utf8_lossy(value).into_owned())
285                }
286                ValueRef::Blob(value) => QueryValue::Text(format_blob(value)),
287            });
288        }
289
290        result_rows.push(values);
291    }
292
293    Ok(QueryResult {
294        columns,
295        rows: result_rows,
296    })
297}
298
299fn enrich_query_prepare_error(
300    connection: &Connection,
301    query: &str,
302    error: rusqlite::Error,
303) -> anyhow::Error {
304    let raw_error = error.to_string();
305
306    if let Some(table) = raw_error.strip_prefix("no such table: ").map(str::trim) {
307        let table = simplify_sqlite_missing_identifier(table);
308        let available_tables = list_loaded_table_names(connection);
309        let available_suffix = if available_tables.is_empty() {
310            String::new()
311        } else {
312            format!(" Available tables/views: {}.", available_tables.join(", "))
313        };
314
315        return anyhow!(
316            "query references unknown table '{table}'.{available_suffix} Check your table names or run `qf tables --input ...` to inspect available tables.\nQuery: {query}"
317        );
318    }
319
320    if let Some(column) = raw_error.strip_prefix("no such column: ").map(str::trim) {
321        let column = simplify_sqlite_missing_identifier(column);
322        return anyhow!(
323            "query references unknown column '{column}'. Check your column names or run `qf schema --input ...` to inspect available columns.\nQuery: {query}"
324        );
325    }
326
327    anyhow!(error).context(format!("failed to prepare query: {query}"))
328}
329
330fn simplify_sqlite_missing_identifier(identifier: &str) -> &str {
331    identifier
332        .split_once(" in ")
333        .map(|(name, _)| name)
334        .unwrap_or(identifier)
335        .trim()
336}
337
338fn list_loaded_table_names(connection: &Connection) -> Vec<String> {
339    let mut statement = match connection.prepare(
340        "SELECT name FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name",
341    ) {
342        Ok(statement) => statement,
343        Err(_) => return Vec::new(),
344    };
345
346    let names = match statement.query_map([], |row| row.get::<_, String>(0)) {
347        Ok(names) => names,
348        Err(_) => return Vec::new(),
349    };
350
351    names.filter_map(|name| name.ok()).collect()
352}
353
354fn normalize_query_for_reserved_identifiers(query: &str) -> String {
355    static RESERVED_TABLE_RE: OnceLock<Regex> = OnceLock::new();
356
357    let re = RESERVED_TABLE_RE.get_or_init(|| {
358        Regex::new(r"(?i)\b(from|join|update|into)\s+table\b")
359            .expect("reserved table regex should compile")
360    });
361
362    re.replace_all(query, |captures: &regex::Captures<'_>| {
363        format!("{} \"table\"", &captures[1])
364    })
365    .into_owned()
366}
367
368fn bind_query_params(statement: &mut rusqlite::Statement<'_>, params: &[QueryParam]) -> Result<()> {
369    for param in params {
370        let mut bound = false;
371
372        for prefix in [":", "@", "$"] {
373            let parameter_name = format!("{prefix}{}", param.name);
374            if let Some(index) = statement
375                .parameter_index(&parameter_name)
376                .with_context(|| format!("failed to inspect parameter {parameter_name}"))?
377            {
378                statement
379                    .raw_bind_parameter(index, to_sql_value(&param.value))
380                    .with_context(|| format!("failed to bind parameter {parameter_name}"))?;
381                bound = true;
382                break;
383            }
384        }
385
386        if !bound {
387            bail!("query does not contain parameter :{}", param.name);
388        }
389    }
390
391    Ok(())
392}
393
394fn format_blob(value: &[u8]) -> String {
395    use std::fmt::Write as _;
396
397    let mut formatted = String::from("0x");
398    for byte in value {
399        let _ = write!(&mut formatted, "{byte:02x}");
400    }
401
402    formatted
403}
404
405fn quote_identifier(identifier: &str) -> String {
406    format!("\"{}\"", identifier.replace('"', "\"\""))
407}
408
409fn sanitize_table_name(name: &str) -> String {
410    let sanitized = name
411        .chars()
412        .map(|character| {
413            if character.is_alphanumeric() || character == '_' {
414                character
415            } else {
416                '_'
417            }
418        })
419        .collect::<String>()
420        .trim_matches('_')
421        .to_string();
422
423    if sanitized.is_empty() {
424        "table_view".to_owned()
425    } else {
426        sanitized
427    }
428}
429
430fn normalize_view_key(name: &str) -> String {
431    name.to_ascii_lowercase()
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn calculates_safe_insert_batch_size() {
440        assert_eq!(calculate_insert_batch_size(1), SQLITE_MAX_INSERT_ROWS);
441        assert_eq!(calculate_insert_batch_size(3), SQLITE_MAX_INSERT_ROWS);
442        assert_eq!(calculate_insert_batch_size(400), 2);
443        assert_eq!(calculate_insert_batch_size(2_000), 1);
444    }
445
446    #[test]
447    fn builds_batched_insert_sql_for_multiple_rows() {
448        assert_eq!(
449            build_batched_insert_sql("table", 2, 3),
450            "INSERT INTO \"table\" VALUES (?, ?), (?, ?), (?, ?)"
451        );
452    }
453}
query_forge/engine.rs

query_forge/
engine.rs
