akar_common/
extension_utils.rs

1//! Shared helpers for connector extensions (httpfs, azure, delta, iceberg,
2//! unity-catalog, sqlite, duckdb).
3//!
4//! Centralizes the scan-closure copy-paste that used to live in each connector:
5//! filling a single-string-column chunk, quoting SQL identifiers, and managing
6//! the lifecycle of temp files downloaded by `*_scan` table functions.
7
8use crate::data_chunk::DataChunk;
9use crate::types::PhysicalTypeID;
10use arrow::array::{ArrayRef, StringArray};
11use std::collections::VecDeque;
12use std::path::PathBuf;
13use std::sync::{Arc, Mutex};
14
15/// Fill a `DataChunk` with a single String column, replacing any existing schema.
16///
17/// The chunk `size` is set to `values.len()`. This mirrors the repeated
18/// clear -> push StringArray -> set names/types -> size sequence that every
19/// connector scan closure used to hand-roll.
20pub fn fill_chunk_with_strings(chunk: &mut DataChunk, field_name: &str, values: &[String]) {
21    let array: ArrayRef = Arc::new(StringArray::from_iter_values(values.iter().map(String::as_str)));
22    chunk.fields.clear();
23    chunk.field_types.clear();
24    chunk.field_names.clear();
25    chunk.fields.push(array);
26    chunk.field_types.push(PhysicalTypeID::String);
27    chunk.field_names.push(field_name.to_string());
28    chunk.size = values.len();
29}
30
31/// Quote a single SQL identifier using double quotes (escaping embedded quotes).
32pub fn quote_sql_identifier(name: &str) -> String {
33    format!("\"{}\"", name.replace('"', "\"\""))
34}
35
36/// Quote a possibly-qualified SQL table name (`catalog.schema.table`) part-by-part.
37///
38/// Each dot-separated component is quoted independently so a user-supplied table
39/// name cannot break out of the query (prevents SQL injection in DuckDB-delegated
40/// scans).
41pub fn quote_sql_table_name(name: &str) -> String {
42    name.split('.').map(quote_sql_identifier).collect::<Vec<_>>().join(".")
43}
44
45/// Maximum number of retained connector temp files before the oldest are evicted.
46const MAX_RETAINED_TEMP_FILES: usize = 64;
47
48/// Registry of temp files downloaded by connector scans and kept past the scan
49/// closure so downstream reads can open them by path.
50static RETAINED_TEMP_FILES: Mutex<VecDeque<PathBuf>> = Mutex::new(VecDeque::new());
51
52/// Track a kept temp file so it is eventually removed.
53///
54/// Connector scans previously called `NamedTempFile::keep()` and leaked the file
55/// permanently (unbounded accumulation). The registry bounds the number of
56/// retained files: when the cap is exceeded the oldest file is deleted from disk.
57/// Returns the path as a string for the scan result row.
58pub fn retain_temp_file(path: PathBuf) -> String {
59    let mut queue = RETAINED_TEMP_FILES.lock().unwrap_or_else(|p| p.into_inner());
60    queue.push_back(path.clone());
61    while queue.len() > MAX_RETAINED_TEMP_FILES {
62        if let Some(oldest) = queue.pop_front() {
63            let _ = std::fs::remove_file(&oldest);
64        }
65    }
66    path.to_string_lossy().to_string()
67}
68
69/// Remove all currently retained temp files (used by tests and clean shutdown).
70pub fn clear_retained_temp_files() {
71    let mut queue = RETAINED_TEMP_FILES.lock().unwrap_or_else(|p| p.into_inner());
72    for path in queue.drain(..) {
73        let _ = std::fs::remove_file(&path);
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use crate::vector::DataChunk;
81
82    #[test]
83    fn test_fill_chunk_with_strings() {
84        let mut chunk = DataChunk::new(Vec::new(), Vec::new());
85        fill_chunk_with_strings(&mut chunk, "path", &["a.parquet".into(), "b.parquet".into()]);
86        assert_eq!(chunk.size, 2);
87        assert_eq!(chunk.num_fields(), 1);
88        assert_eq!(chunk.field_names, vec!["path".to_string()]);
89        assert_eq!(chunk.field_types, vec![PhysicalTypeID::String]);
90    }
91
92    #[test]
93    fn test_fill_chunk_with_strings_replaces_schema() {
94        let mut chunk = DataChunk::new(Vec::new(), Vec::new());
95        fill_chunk_with_strings(&mut chunk, "first", &["a".into()]);
96        fill_chunk_with_strings(&mut chunk, "second", &["x".into(), "y".into(), "z".into()]);
97        assert_eq!(chunk.size, 3);
98        assert_eq!(chunk.num_fields(), 1);
99        assert_eq!(chunk.field_names, vec!["second".to_string()]);
100    }
101
102    #[test]
103    fn test_quote_sql_identifier() {
104        assert_eq!(quote_sql_identifier("tbl"), "\"tbl\"");
105        assert_eq!(quote_sql_identifier("we\"ird"), "\"we\"\"ird\"");
106    }
107
108    #[test]
109    fn test_quote_sql_table_name() {
110        assert_eq!(
111            quote_sql_table_name("main.default.people"),
112            "\"main\".\"default\".\"people\""
113        );
114        assert_eq!(
115            quote_sql_table_name("x.\"y\"; DROP TABLE t"),
116            "\"x\".\"\"\"y\"\"; DROP TABLE t\""
117        );
118    }
119
120    #[test]
121    fn test_retain_temp_file_evicts_oldest() {
122        clear_retained_temp_files();
123        let dir = std::env::temp_dir();
124        for i in 0..(MAX_RETAINED_TEMP_FILES + 5) {
125            let p = dir.join(format!("akar_retain_test_{i}.tmp"));
126            std::fs::write(&p, b"x").unwrap();
127            retain_temp_file(p);
128        }
129        // Oldest 5 files must have been evicted (deleted).
130        for i in 0..5 {
131            let p = dir.join(format!("akar_retain_test_{i}.tmp"));
132            assert!(!p.exists(), "file {i} should have been evicted");
133        }
134        // Newest MAX files still present.
135        for i in 5..(MAX_RETAINED_TEMP_FILES + 5) {
136            let p = dir.join(format!("akar_retain_test_{i}.tmp"));
137            assert!(p.exists(), "file {i} should still be retained");
138        }
139        clear_retained_temp_files();
140    }
141}
akar_common/extension_utils.rs

akar_common/
extension_utils.rs
