wp_knowledge/mem/
query_util.rs

1use orion_error::ErrorOwe;
2use rusqlite::Params;
3use std::collections::hash_map::DefaultHasher;
4use std::future::Future;
5use std::hash::{Hash, Hasher};
6use std::num::NonZeroUsize;
7use std::sync::{Arc, RwLock};
8use wp_error::KnowledgeResult;
9use wp_log::debug_kdb;
10use wp_model_core::model::{self, DataField};
11
12use lazy_static::lazy_static;
13use lru::LruCache;
14
15use crate::loader::ProviderKind;
16use crate::mem::RowData;
17use crate::runtime::{DatasourceId, Generation, MetadataCacheScope, runtime};
18use crate::telemetry::{
19    CacheLayer, CacheOutcome, CacheTelemetryEvent, telemetry, telemetry_enabled,
20};
21
22#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23pub struct MetadataCacheKey {
24    pub datasource_id: DatasourceId,
25    pub generation: Generation,
26    pub query_hash: u64,
27}
28
29lazy_static! {
30    /// Global column metadata cache keyed by datasource/generation/query hash.
31    pub static ref COLNAME_CACHE: RwLock<LruCache<MetadataCacheKey, Arc<Vec<String>>>> =
32        RwLock::new(LruCache::new(
33            NonZeroUsize::new(512).expect("non-zero metadata cache size")
34        ));
35}
36
37pub fn column_metadata_cache_snapshot() -> (usize, usize) {
38    COLNAME_CACHE
39        .read()
40        .map(|cache| (cache.len(), cache.cap().get()))
41        .unwrap_or((0, 0))
42}
43
44fn stable_hash(value: &str) -> u64 {
45    let mut hasher = DefaultHasher::new();
46    value.hash(&mut hasher);
47    hasher.finish()
48}
49
50#[cfg(test)]
51pub(crate) fn metadata_cache_key_for_current_scope(sql: &str) -> MetadataCacheKey {
52    let scope = runtime().current_metadata_scope();
53    metadata_cache_key_for_scope(&scope, sql)
54}
55
56pub(crate) fn metadata_cache_key_for_scope(
57    scope: &MetadataCacheScope,
58    sql: &str,
59) -> MetadataCacheKey {
60    MetadataCacheKey {
61        datasource_id: scope.datasource_id.clone(),
62        generation: scope.generation,
63        query_hash: stable_hash(sql),
64    }
65}
66
67pub(crate) fn metadata_cache_get_or_try_init<F>(sql: &str, load: F) -> KnowledgeResult<Vec<String>>
68where
69    F: FnOnce() -> KnowledgeResult<Option<Vec<String>>>,
70{
71    let scope = runtime().current_metadata_scope();
72    let provider_kind = runtime().current_provider_kind();
73    metadata_cache_get_or_try_init_for_scope(&scope, provider_kind, sql, load)
74}
75
76pub(crate) fn metadata_cache_get_or_try_init_for_scope<F>(
77    scope: &MetadataCacheScope,
78    provider_kind: Option<ProviderKind>,
79    sql: &str,
80    load: F,
81) -> KnowledgeResult<Vec<String>>
82where
83    F: FnOnce() -> KnowledgeResult<Option<Vec<String>>>,
84{
85    let cache_key = metadata_cache_key_for_scope(scope, sql);
86    if let Some(names) = COLNAME_CACHE
87        .read()
88        .ok()
89        .and_then(|m| m.peek(&cache_key).cloned())
90    {
91        runtime().record_metadata_cache_hit();
92        if telemetry_enabled() {
93            telemetry().on_cache(&CacheTelemetryEvent {
94                layer: CacheLayer::Metadata,
95                outcome: CacheOutcome::Hit,
96                provider_kind: provider_kind.clone(),
97            });
98        }
99        debug_kdb!(
100            "[kdb] metadata cache hit datasource_id={} generation={}",
101            cache_key.datasource_id.0,
102            cache_key.generation.0
103        );
104        return Ok((*names).clone());
105    }
106
107    runtime().record_metadata_cache_miss();
108    if telemetry_enabled() {
109        telemetry().on_cache(&CacheTelemetryEvent {
110            layer: CacheLayer::Metadata,
111            outcome: CacheOutcome::Miss,
112            provider_kind,
113        });
114    }
115    debug_kdb!(
116        "[kdb] metadata cache miss datasource_id={} generation={}",
117        cache_key.datasource_id.0,
118        cache_key.generation.0
119    );
120
121    let Some(names) = load()? else {
122        return Ok(Vec::new());
123    };
124    if let Ok(mut m) = COLNAME_CACHE.write() {
125        m.put(cache_key, Arc::new(names.clone()));
126    }
127    Ok(names)
128}
129
130pub(crate) async fn metadata_cache_get_or_try_init_async_for_scope<F, Fut>(
131    scope: &MetadataCacheScope,
132    provider_kind: Option<ProviderKind>,
133    sql: &str,
134    load: F,
135) -> KnowledgeResult<Vec<String>>
136where
137    F: FnOnce() -> Fut,
138    Fut: Future<Output = KnowledgeResult<Option<Vec<String>>>>,
139{
140    metadata_cache_get_or_try_init_async_for_scope_typed(scope, provider_kind, sql, load).await
141}
142
143pub(crate) async fn metadata_cache_get_or_try_init_async_for_scope_typed<F, Fut, E>(
144    scope: &MetadataCacheScope,
145    provider_kind: Option<ProviderKind>,
146    sql: &str,
147    load: F,
148) -> Result<Vec<String>, E>
149where
150    F: FnOnce() -> Fut,
151    Fut: Future<Output = Result<Option<Vec<String>>, E>>,
152{
153    let cache_key = metadata_cache_key_for_scope(scope, sql);
154    if let Some(names) = COLNAME_CACHE
155        .read()
156        .ok()
157        .and_then(|m| m.peek(&cache_key).cloned())
158    {
159        runtime().record_metadata_cache_hit();
160        if telemetry_enabled() {
161            telemetry().on_cache(&CacheTelemetryEvent {
162                layer: CacheLayer::Metadata,
163                outcome: CacheOutcome::Hit,
164                provider_kind: provider_kind.clone(),
165            });
166        }
167        return Ok((*names).clone());
168    }
169
170    runtime().record_metadata_cache_miss();
171    if telemetry_enabled() {
172        telemetry().on_cache(&CacheTelemetryEvent {
173            layer: CacheLayer::Metadata,
174            outcome: CacheOutcome::Miss,
175            provider_kind,
176        });
177    }
178
179    let Some(names) = load().await? else {
180        return Ok(Vec::new());
181    };
182    if let Ok(mut m) = COLNAME_CACHE.write() {
183        m.put(cache_key, Arc::new(names.clone()));
184    }
185    Ok(names)
186}
187
188/// 将一行数据映射为 RowData
189fn map_row(row: &rusqlite::Row<'_>, col_names: &[String]) -> KnowledgeResult<RowData> {
190    let mut result = Vec::with_capacity(col_names.len());
191    for (i, col_name) in col_names.iter().enumerate() {
192        let value = row.get_ref(i).owe_rule()?;
193        let field = match value {
194            rusqlite::types::ValueRef::Null => {
195                DataField::new(model::DataType::default(), col_name, model::Value::Null)
196            }
197            rusqlite::types::ValueRef::Integer(v) => DataField::from_digit(col_name, v),
198            rusqlite::types::ValueRef::Real(v) => DataField::from_float(col_name, v),
199            rusqlite::types::ValueRef::Text(v) => {
200                DataField::from_chars(col_name, String::from_utf8(v.to_vec()).owe_rule()?)
201            }
202            rusqlite::types::ValueRef::Blob(v) => {
203                DataField::from_chars(col_name, String::from_utf8_lossy(v).to_string())
204            }
205        };
206        result.push(field);
207    }
208    Ok(result)
209}
210
211/// 从 statement 获取列名（普通版，带 debug 日志）
212fn extract_col_names(stmt: &rusqlite::Statement<'_>) -> Vec<String> {
213    let col_cnt = stmt.column_count();
214    debug_kdb!("[memdb] col_cnt={}", col_cnt);
215    let mut col_names = Vec::with_capacity(col_cnt);
216    for i in 0..col_cnt {
217        let name = stmt.column_name(i).unwrap_or("").to_string();
218        debug_kdb!("[memdb] col[{}] name='{}'", i, name);
219        col_names.push(name);
220    }
221    col_names
222}
223
224/// 从 statement 获取列名（cached 版，使用全局缓存）
225fn extract_col_names_cached(
226    stmt: &rusqlite::Statement<'_>,
227    sql: &str,
228) -> KnowledgeResult<Vec<String>> {
229    metadata_cache_get_or_try_init(sql, || {
230        let col_cnt = stmt.column_count();
231        let mut names = Vec::with_capacity(col_cnt);
232        for i in 0..col_cnt {
233            names.push(stmt.column_name(i).owe_rule()?.to_string());
234        }
235        Ok(Some(names))
236    })
237}
238
239fn extract_col_names_cached_with_scope(
240    stmt: &rusqlite::Statement<'_>,
241    scope: &MetadataCacheScope,
242    provider_kind: Option<ProviderKind>,
243    sql: &str,
244) -> KnowledgeResult<Vec<String>> {
245    metadata_cache_get_or_try_init_for_scope(scope, provider_kind, sql, || {
246        let col_cnt = stmt.column_count();
247        let mut names = Vec::with_capacity(col_cnt);
248        for i in 0..col_cnt {
249            names.push(stmt.column_name(i).owe_rule()?.to_string());
250        }
251        Ok(Some(names))
252    })
253}
254
255pub fn query<P: Params>(
256    conn: &rusqlite::Connection,
257    sql: &str,
258    params: P,
259) -> KnowledgeResult<Vec<RowData>> {
260    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
261    let col_names = extract_col_names(&stmt);
262    let mut rows = stmt.query(params).owe_rule()?;
263    let mut all_result = Vec::new();
264    while let Some(row) = rows.next().owe_rule()? {
265        all_result.push(map_row(row, &col_names)?);
266    }
267    Ok(all_result)
268}
269
270/// Query first row and map columns into RowData with column names preserved.
271pub fn query_first_row<P: Params>(
272    conn: &rusqlite::Connection,
273    sql: &str,
274    params: P,
275) -> KnowledgeResult<RowData> {
276    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
277    let col_names = extract_col_names(&stmt);
278    let mut rows = stmt.query(params).owe_rule()?;
279    if let Some(row) = rows.next().owe_rule()? {
280        map_row(row, &col_names)
281    } else {
282        debug_kdb!("[memdb] no row for sql");
283        Ok(Vec::new())
284    }
285}
286
287pub fn query_cached<P: Params>(
288    conn: &rusqlite::Connection,
289    sql: &str,
290    params: P,
291) -> KnowledgeResult<Vec<RowData>> {
292    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
293    // Column names cache (per SQL)
294    let col_names = extract_col_names_cached(&stmt, sql)?;
295    let mut rows = stmt.query(params).owe_rule()?;
296    let mut all_result = Vec::new();
297    while let Some(row) = rows.next().owe_rule()? {
298        all_result.push(map_row(row, &col_names)?);
299    }
300    Ok(all_result)
301}
302
303pub fn query_cached_with_scope<P: Params>(
304    conn: &rusqlite::Connection,
305    scope: &MetadataCacheScope,
306    provider_kind: Option<ProviderKind>,
307    sql: &str,
308    params: P,
309) -> KnowledgeResult<Vec<RowData>> {
310    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
311    let col_names = extract_col_names_cached_with_scope(&stmt, scope, provider_kind, sql)?;
312    let mut rows = stmt.query(params).owe_rule()?;
313    let mut all_result = Vec::new();
314    while let Some(row) = rows.next().owe_rule()? {
315        all_result.push(map_row(row, &col_names)?);
316    }
317    Ok(all_result)
318}
319
320/// Same as `query_first_row` but with a shared column-names cache to reduce metadata lookups.
321pub fn query_first_row_cached<P: Params>(
322    conn: &rusqlite::Connection,
323    sql: &str,
324    params: P,
325) -> KnowledgeResult<RowData> {
326    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
327    let col_names = extract_col_names_cached(&stmt, sql)?;
328    let mut rows = stmt.query(params).owe_rule()?;
329    if let Some(row) = rows.next().owe_rule()? {
330        map_row(row, &col_names)
331    } else {
332        Ok(Vec::new())
333    }
334}
335
336pub fn query_first_row_cached_with_scope<P: Params>(
337    conn: &rusqlite::Connection,
338    scope: &MetadataCacheScope,
339    provider_kind: Option<ProviderKind>,
340    sql: &str,
341    params: P,
342) -> KnowledgeResult<RowData> {
343    let mut stmt = conn.prepare_cached(sql).owe_rule()?;
344    let col_names = extract_col_names_cached_with_scope(&stmt, scope, provider_kind, sql)?;
345    let mut rows = stmt.query(params).owe_rule()?;
346    if let Some(row) = rows.next().owe_rule()? {
347        map_row(row, &col_names)
348    } else {
349        Ok(Vec::new())
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356    use rusqlite::Connection;
357
358    fn setup_test_db() -> Connection {
359        let conn = Connection::open_in_memory().unwrap();
360        conn.execute(
361            "CREATE TABLE test (id INTEGER, name TEXT, score REAL, data BLOB, empty)",
362            [],
363        )
364        .unwrap();
365        conn
366    }
367
368    #[test]
369    fn test_query_returns_all_rows() {
370        let conn = setup_test_db();
371        let rows = query(&conn, "SELECT * FROM test", []).unwrap();
372        assert!(rows.is_empty());
373        conn.execute("INSERT INTO test (id, name) VALUES (1, 'alice')", [])
374            .unwrap();
375        conn.execute("INSERT INTO test (id, name) VALUES (2, 'bob')", [])
376            .unwrap();
377        conn.execute("INSERT INTO test (id, name) VALUES (3, 'charlie')", [])
378            .unwrap();
379
380        let rows = query(&conn, "SELECT id, name FROM test ORDER BY id", []).unwrap();
381        assert_eq!(rows.len(), 3);
382    }
383
384    #[test]
385    fn test_query_first_row_returns_single_row() {
386        let conn = setup_test_db();
387        let row = query_first_row(&conn, "SELECT * FROM test", []).unwrap();
388        assert!(row.is_empty());
389        conn.execute("INSERT INTO test (id, name) VALUES (1, 'first')", [])
390            .unwrap();
391        conn.execute("INSERT INTO test (id, name) VALUES (2, 'second')", [])
392            .unwrap();
393
394        let row = query_first_row(&conn, "SELECT id, name FROM test ORDER BY id", []).unwrap();
395        assert_eq!(row.len(), 2);
396        assert_eq!(row[0].to_string(), "digit(1)");
397        assert_eq!(row[1].to_string(), "chars(first)");
398    }
399
400    #[test]
401    fn test_map_row_handles_all_types() {
402        let conn = setup_test_db();
403        conn.execute(
404            "INSERT INTO test (id, name, score, data, empty) VALUES (42, 'hello', 3.14, X'414243', NULL)",
405            [],
406        )
407        .unwrap();
408
409        let row =
410            query_first_row(&conn, "SELECT id, name, score, data, empty FROM test", []).unwrap();
411        assert_eq!(row.len(), 5);
412    }
413
414    #[test]
415    fn test_extract_col_names_preserves_aliases() {
416        let conn = setup_test_db();
417        conn.execute("INSERT INTO test (id, name) VALUES (1, 'x')", [])
418            .unwrap();
419
420        let row = query_first_row(
421            &conn,
422            "SELECT id AS user_id, name AS user_name FROM test",
423            [],
424        )
425        .unwrap();
426        assert_eq!(row[0].get_name(), "user_id");
427        assert_eq!(row[1].get_name(), "user_name");
428    }
429
430    #[test]
431    fn test_query_cached_uses_cache() {
432        let _guard = crate::runtime::runtime_test_guard()
433            .lock()
434            .expect("runtime test guard");
435        let conn = setup_test_db();
436        conn.execute("INSERT INTO test (id) VALUES (1)", [])
437            .unwrap();
438
439        let sql = "SELECT id FROM test WHERE id = 1";
440        // 第一次查询，填充缓存
441        let _ = query_cached(&conn, sql, []).unwrap();
442        // 第二次查询，应命中缓存
443        let rows = query_cached(&conn, sql, []).unwrap();
444        assert_eq!(rows.len(), 1);
445
446        // 验证缓存已填充
447        let cache = COLNAME_CACHE.read().unwrap();
448        assert!(cache.contains(&metadata_cache_key_for_current_scope(sql)));
449    }
450
451    #[test]
452    fn test_query_with_params() {
453        let conn = setup_test_db();
454        conn.execute("INSERT INTO test (id, name) VALUES (1, 'alice')", [])
455            .unwrap();
456        conn.execute("INSERT INTO test (id, name) VALUES (2, 'bob')", [])
457            .unwrap();
458
459        let rows = query(&conn, "SELECT name FROM test WHERE id = ?1", [2]).unwrap();
460        assert_eq!(rows.len(), 1);
461        assert_eq!(rows[0][0].to_string(), "chars(bob)");
462    }
463
464    #[test]
465    fn test_metadata_cache_key_for_scope_is_explicit() {
466        let sql = "SELECT id FROM test";
467        let scope_a = MetadataCacheScope {
468            datasource_id: DatasourceId("postgres:aaaa".to_string()),
469            generation: Generation(1),
470        };
471        let scope_b = MetadataCacheScope {
472            datasource_id: DatasourceId("postgres:bbbb".to_string()),
473            generation: Generation(2),
474        };
475        let key_a = metadata_cache_key_for_scope(&scope_a, sql);
476        let key_b = metadata_cache_key_for_scope(&scope_b, sql);
477        assert_ne!(key_a, key_b);
478    }
479}
wp_knowledge/mem/query_util.rs

wp_knowledge/mem/
query_util.rs
