fix: enrichment table query (#2603)

Author: hengfeiyang

src/common/infra/db/sqlite.rs

@@ -245,6 +245,9 @@ impl super::Db for SqliteDb {
             return Err(e.into());
         }
 
+        // release lock
+        drop(client);
+
         // event watch
         if need_watch {
             if let Err(e) = CHANNEL

src/common/infra/file_list/sqlite.rs

@@ -96,9 +96,9 @@ INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records
             return Ok(());
         }
         let chunks = files.chunks(100);
-        let client = CLIENT_RW.clone();
-        let client = client.lock().await;
         for files in chunks {
+            let client = CLIENT_RW.clone();
+            let client = client.lock().await;
             let mut tx = client.begin().await?;
             let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
                 "INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size)",
@@ -142,6 +142,9 @@ INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records
                     log::error!("[SQLITE] rollback file_list batch add error: {}", e);
                     return Err(e.into());
                 }
+                // release lock
+                drop(client);
+                // add file one by one
                 for item in files {
                     if let Err(e) = self.add(&item.key, &item.meta).await {
                         log::error!("[SQLITE] single insert file_list add error: {}", e);
@@ -161,10 +164,10 @@ INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records
             return Ok(());
         }
         let chunks = files.chunks(100);
-        let client = CLIENT_RW.clone();
-        let client = client.lock().await;
         for files in chunks {
             // get ids of the files
+            let client = CLIENT_RW.clone();
+            let client = client.lock().await;
             let pool = client.clone();
             let mut ids = Vec::with_capacity(files.len());
             for file in files {
@@ -211,9 +214,9 @@ INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records
             return Ok(());
         }
         let chunks = files.chunks(100);
-        let client = CLIENT_RW.clone();
-        let client = client.lock().await;
         for files in chunks {
+            let client = CLIENT_RW.clone();
+            let client = client.lock().await;
             let mut tx = client.begin().await?;
             let mut query_builder: QueryBuilder<Sqlite> = QueryBuilder::new(
                 "INSERT INTO file_list_deleted (org, stream, date, file, created_at)",
@@ -246,10 +249,10 @@ INSERT INTO file_list (org, stream, date, file, deleted, min_ts, max_ts, records
             return Ok(());
         }
         let chunks = files.chunks(100);
-        let client = CLIENT_RW.clone();
-        let client = client.lock().await;
         for files in chunks {
             // get ids of the files
+            let client = CLIENT_RW.clone();
+            let client = client.lock().await;
             let pool = client.clone();
             let mut ids = Vec::with_capacity(files.len());
             for file in files {

src/common/meta/search.rs

@@ -176,6 +176,7 @@ pub struct Response {
     #[serde(skip_serializing)]
     pub file_count: usize,
     pub scan_size: usize,
+    pub scan_records: usize,
     #[serde(default)]
     #[serde(skip_serializing_if = "String::is_empty")]
     pub response_type: String,
@@ -204,6 +205,7 @@ impl Response {
             size,
             file_count: 0,
             scan_size: 0,
+            scan_records: 0,
             columns: Vec::new(),
             hits: Vec::new(),
             aggs: HashMap::new(),
@@ -252,6 +254,10 @@ impl Response {
         self.scan_size = val;
     }
 
+    pub fn set_scan_records(&mut self, val: usize) {
+        self.scan_records = val;
+    }
+
     pub fn set_session_id(&mut self, session_id: String) {
         self.session_id = session_id;
     }

src/ingester/src/wal.rs

@@ -119,7 +119,13 @@ pub(crate) async fn replay_wal_files() -> Result<()> {
             .unwrap_or_default();
         let key = WriterKey::new(org_id, stream_type);
         let mut memtable = memtable::MemTable::new();
-        let mut reader = wal::Reader::from_path(wal_file).context(WalSnafu)?;
+        let mut reader = match wal::Reader::from_path(wal_file) {
+            Ok(v) => v,
+            Err(e) => {
+                log::error!("Unable to open the wal file err: {}, skip", e);
+                continue;
+            }
+        };
         let mut total = 0;
         let mut i = 0;
         loop {

src/job/compact.rs

@@ -90,14 +90,8 @@ async fn run_sync_to_db() -> Result<(), anyhow::Error> {
     interval.tick().await; // trigger the first run
     loop {
         interval.tick().await;
-        let ret = service::db::compact::files::sync_cache_to_db().await;
-        if ret.is_err() {
-            log::error!(
-                "[COMPACTOR] run offset sync cache to db error: {}",
-                ret.err().unwrap()
-            );
-        } else {
-            log::info!("[COMPACTOR] run offset sync cache to db done");
+        if let Err(e) = service::db::compact::files::sync_cache_to_db().await {
+            log::error!("[COMPACTOR] run offset sync cache to db error: {}", e);
         }
     }
 }

src/service/compact/file_list_deleted.rs

@@ -156,10 +156,10 @@ pub async fn load_prefix_from_s3(
     let prefix = format!("file_list_deleted/{prefix}/");
     let files = storage::list(&prefix).await?;
     let files_num = files.len();
-    log::info!("Load file_list_deleted gets {} files", files_num);
     if files.is_empty() {
         return Ok(HashMap::default());
     }
+    log::info!("Load file_list_deleted gets {} files", files_num);
 
     let mut tasks = Vec::with_capacity(CONFIG.limit.query_thread_num + 1);
     let chunk_size = std::cmp::max(1, files_num / CONFIG.limit.query_thread_num);

src/service/db/enrichment_table.rs

@@ -40,7 +40,7 @@ pub async fn get(org_id: &str, name: &str) -> Result<Vec<vrl::value::Value>, any
     let query = meta::search::Query {
         sql: format!("SELECT * FROM \"{name}\" limit {rec_num}"),
         start_time: BASE_TIME.timestamp_micros(),
-        end_time: Utc::now().timestamp(),
+        end_time: Utc::now().timestamp_micros(),
         sql_mode: "full".to_owned(),
         ..Default::default()
     };

src/service/search/grpc/wal.rs

@@ -93,6 +93,7 @@ pub async fn search_parquet(
                 let parquet_meta = read_metadata_from_bytes(&file_data)
                     .await
                     .unwrap_or_default();
+                scan_stats.records += parquet_meta.records;
                 scan_stats.original_size += parquet_meta.original_size;
                 if let Some((min_ts, max_ts)) = sql.meta.time_range {
                     if parquet_meta.min_ts <= max_ts && parquet_meta.max_ts >= min_ts {
@@ -334,7 +335,10 @@ pub async fn search_memtable(
     let mut batch_groups: HashMap<Arc<Schema>, Vec<RecordBatch>> = HashMap::with_capacity(2);
     for (schema, batch) in batches {
         let entry = batch_groups.entry(schema).or_default();
-        scan_stats.original_size += batch.iter().map(|r| r.data_json_size).sum::<usize>() as i64;
+        for r in batch.iter() {
+            scan_stats.records += r.data.num_rows() as i64;
+            scan_stats.original_size += r.data_json_size as i64;
+        }
         entry.extend(batch.into_iter().map(|r| r.data.clone()));
     }
 

src/service/search/mod.rs

@@ -587,6 +587,7 @@ async fn search_in_cluster(mut req: cluster_rpc::SearchRequest) -> Result<search
     result.set_cluster_took(start.elapsed().as_millis() as usize, took_wait);
     result.set_file_count(scan_stats.files as usize);
     result.set_scan_size(scan_stats.original_size as usize);
+    result.set_scan_records(scan_stats.records as usize);
 
     if query_type == "table" {
         result.response_type = "table".to_string();