20260718 新增查询服务 VRP 检索接口
This commit is contained in:
128
src/query_db.rs
128
src/query_db.rs
@@ -9,10 +9,11 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::query::artifact_manifest::build_artifact_manifest;
|
||||
use crate::query::report_stream::{self, ObjectLookup, ObjectScope, ReportSummary};
|
||||
use crate::query::vrp::{self, VrpLookup, VrpRecord};
|
||||
|
||||
use crate::blob_store::ExternalRepoBytesDb;
|
||||
|
||||
pub const QUERY_DB_SCHEMA_VERSION: u32 = 1;
|
||||
pub const QUERY_DB_SCHEMA_VERSION: u32 = 2;
|
||||
|
||||
pub const CF_META: &str = "meta";
|
||||
pub const CF_RUNS: &str = "runs";
|
||||
@@ -26,6 +27,7 @@ pub const CF_VALIDATION_EXPLAIN_CACHE: &str = "validation_explain_cache";
|
||||
pub const CF_EXPORT_JOBS: &str = "export_jobs";
|
||||
pub const CF_STATS: &str = "stats";
|
||||
pub const CF_REASON_INDEX: &str = "reason_index";
|
||||
pub const CF_VRPS: &str = "vrps";
|
||||
|
||||
pub const QUERY_DB_COLUMN_FAMILIES: &[&str] = &[
|
||||
CF_META,
|
||||
@@ -40,6 +42,7 @@ pub const QUERY_DB_COLUMN_FAMILIES: &[&str] = &[
|
||||
CF_EXPORT_JOBS,
|
||||
CF_STATS,
|
||||
CF_REASON_INDEX,
|
||||
CF_VRPS,
|
||||
];
|
||||
|
||||
const KEY_SCHEMA_VERSION: &[u8] = b"schema_version";
|
||||
@@ -92,6 +95,7 @@ pub struct QueryIndexSummary {
|
||||
pub object_instances_indexed: u64,
|
||||
pub object_projections_indexed: u64,
|
||||
pub stats_indexed: u64,
|
||||
pub vrps_indexed: u64,
|
||||
pub latest_ready_run: Option<String>,
|
||||
pub errors: Vec<String>,
|
||||
}
|
||||
@@ -369,6 +373,52 @@ impl QueryDb {
|
||||
self.list_json_by_prefix(CF_REPOS, &format!("repo/{run_id}/"), limit, cursor)
|
||||
}
|
||||
|
||||
pub fn lookup_vrps(
|
||||
&self,
|
||||
run_id: &str,
|
||||
lookup: &VrpLookup,
|
||||
asn: Option<u32>,
|
||||
limit: usize,
|
||||
cursor: Option<&str>,
|
||||
) -> QueryDbResult<QueryPage<VrpRecord>> {
|
||||
let cf = self.cf(CF_VRPS)?;
|
||||
let mut records = BTreeMap::<String, VrpRecord>::new();
|
||||
for candidate in lookup.candidate_prefixes() {
|
||||
let key_prefix = vrp::vrp_key_prefix(run_id, &candidate);
|
||||
let mode = IteratorMode::From(key_prefix.as_bytes(), rocksdb::Direction::Forward);
|
||||
for item in self.db.iterator_cf(cf, mode) {
|
||||
let (key, value) = item?;
|
||||
if !key.starts_with(key_prefix.as_bytes()) {
|
||||
break;
|
||||
}
|
||||
let key_string = String::from_utf8_lossy(&key).to_string();
|
||||
let record: VrpRecord = serde_json::from_slice(&value)?;
|
||||
if asn.is_some_and(|expected| expected != record.asn)
|
||||
|| !lookup
|
||||
.matches(&record)
|
||||
.map_err(QueryDbError::InvalidArtifact)?
|
||||
{
|
||||
continue;
|
||||
}
|
||||
records.insert(key_string, record);
|
||||
}
|
||||
}
|
||||
|
||||
let offset = vrp::parse_cursor(cursor).map_err(QueryDbError::InvalidArtifact)?;
|
||||
let bounded_limit = limit.clamp(1, 1000);
|
||||
let total = records.len();
|
||||
let data = records
|
||||
.into_values()
|
||||
.skip(offset)
|
||||
.take(bounded_limit)
|
||||
.collect::<Vec<_>>();
|
||||
Ok(QueryPage {
|
||||
data,
|
||||
next_cursor: vrp::next_cursor(offset, bounded_limit, total),
|
||||
limit: bounded_limit,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_repo(&self, run_id: &str, repo_id: &str) -> QueryDbResult<Option<RepositoryRecord>> {
|
||||
self.get_json_cf(CF_REPOS, repo_key(run_id, repo_id).as_bytes())
|
||||
}
|
||||
@@ -899,12 +949,20 @@ impl QueryDb {
|
||||
(CF_EXPORT_JOBS, format!("export/{}/", run.run_id)),
|
||||
(CF_STATS, format!("stats/{}/", run.run_id)),
|
||||
(CF_REASON_INDEX, format!("reason/{}/", run.run_id)),
|
||||
(CF_VRPS, format!("vrp/{}/", run.run_id)),
|
||||
] {
|
||||
self.delete_prefix_range(&mut batch, cf_name, &prefix)?;
|
||||
}
|
||||
self.write_batch(batch)
|
||||
}
|
||||
|
||||
fn clear_vrp_index(&self, run_id: &str) -> QueryDbResult<()> {
|
||||
let mut batch = WriteBatch::default();
|
||||
let prefix = format!("vrp/{run_id}/");
|
||||
self.delete_prefix_range(&mut batch, CF_VRPS, &prefix)?;
|
||||
self.write_batch(batch)
|
||||
}
|
||||
|
||||
fn cf(&self, name: &'static str) -> QueryDbResult<&rocksdb::ColumnFamily> {
|
||||
self.db
|
||||
.cf_handle(name)
|
||||
@@ -1080,6 +1138,7 @@ pub fn index_artifacts_with_open_db(
|
||||
summary.object_instances_indexed += run_summary.object_instances_indexed;
|
||||
summary.object_projections_indexed += run_summary.object_projections_indexed;
|
||||
summary.stats_indexed += run_summary.stats_indexed;
|
||||
summary.vrps_indexed += run_summary.vrps_indexed;
|
||||
summary.latest_ready_run = run_summary.latest_ready_run;
|
||||
}
|
||||
Err(err) => summary.errors.push(format!("{}: {err}", run_dir.display())),
|
||||
@@ -1152,6 +1211,7 @@ struct SingleRunIndexSummary {
|
||||
object_instances_indexed: u64,
|
||||
object_projections_indexed: u64,
|
||||
stats_indexed: u64,
|
||||
vrps_indexed: u64,
|
||||
latest_ready_run: Option<String>,
|
||||
}
|
||||
|
||||
@@ -1193,6 +1253,7 @@ fn index_run_dir(
|
||||
run_record.index_status = "building".to_string();
|
||||
run_record.index_error = None;
|
||||
db.put_json_cf(CF_RUNS, run_key(&run_id).as_bytes(), &run_record)?;
|
||||
db.clear_vrp_index(&run_id)?;
|
||||
|
||||
let indexed =
|
||||
match write_summary_index_records(db, &run_id, &report_summary, &artifact_manifest) {
|
||||
@@ -1236,6 +1297,7 @@ fn index_run_dir(
|
||||
object_instances_indexed: indexed.object_instances_indexed,
|
||||
object_projections_indexed: indexed.object_projections_indexed,
|
||||
stats_indexed: indexed.stats_indexed,
|
||||
vrps_indexed: indexed.vrps_indexed,
|
||||
latest_ready_run: if should_update_latest {
|
||||
Some(run_id)
|
||||
} else {
|
||||
@@ -1275,6 +1337,7 @@ struct IndexWriteSummary {
|
||||
object_instances_indexed: u64,
|
||||
object_projections_indexed: u64,
|
||||
stats_indexed: u64,
|
||||
vrps_indexed: u64,
|
||||
}
|
||||
|
||||
fn write_summary_index_records(
|
||||
@@ -1287,6 +1350,7 @@ fn write_summary_index_records(
|
||||
let mut batch = WriteBatch::default();
|
||||
let mut pending = 0usize;
|
||||
let mut summary = IndexWriteSummary::default();
|
||||
let mut seen_vrp_keys = BTreeSet::new();
|
||||
|
||||
macro_rules! flush_if_needed {
|
||||
() => {
|
||||
@@ -1324,6 +1388,18 @@ fn write_summary_index_records(
|
||||
flush_if_needed!();
|
||||
}
|
||||
|
||||
for vrp_record in &report_summary.vrps {
|
||||
let key = vrp_record
|
||||
.index_key()
|
||||
.map_err(QueryDbError::InvalidArtifact)?;
|
||||
if seen_vrp_keys.insert(key.clone()) {
|
||||
put_json_batch(&mut batch, db, CF_VRPS, key.as_bytes(), vrp_record)?;
|
||||
pending += 1;
|
||||
summary.vrps_indexed += 1;
|
||||
flush_if_needed!();
|
||||
}
|
||||
}
|
||||
|
||||
let artifacts_value = serde_json::to_value(artifact_manifest)?;
|
||||
let stats = report_stream::stats_records_from_summary(run_id, report_summary, artifacts_value);
|
||||
for record in &stats {
|
||||
@@ -1584,6 +1660,56 @@ mod tests {
|
||||
assert_eq!(db.count_cf(CF_OBJECT_INSTANCES).unwrap(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vrp_index_supports_ip_prefix_asn_and_pagination_queries() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let run_dir = temp.path().join("runs/run_0001");
|
||||
fs::create_dir_all(&run_dir).expect("run dir");
|
||||
write_sample_run(&run_dir, "run_0001", 1);
|
||||
let mut report = read_json_file(&run_dir.join("report.json")).expect("read report");
|
||||
let duplicate = report["vrps"][0].clone();
|
||||
report["vrps"].as_array_mut().unwrap().push(duplicate);
|
||||
fs::write(
|
||||
run_dir.join("report.json"),
|
||||
serde_json::to_vec(&report).unwrap(),
|
||||
)
|
||||
.expect("write duplicate VRP report");
|
||||
|
||||
let query_db_path = temp.path().join("query-db");
|
||||
let summary = index_artifacts(&ArtifactIndexerConfig {
|
||||
query_db_path: query_db_path.clone(),
|
||||
run_root: Some(temp.path().to_path_buf()),
|
||||
run_dir: None,
|
||||
repo_bytes_db_path: None,
|
||||
projection_entry_limit: 50,
|
||||
min_run_seq: None,
|
||||
retain_indexed_runs: None,
|
||||
})
|
||||
.expect("index");
|
||||
assert_eq!(summary.vrps_indexed, 1);
|
||||
|
||||
let db = QueryDb::open(&query_db_path).expect("query db");
|
||||
assert_eq!(db.count_cf(CF_VRPS).unwrap(), 1);
|
||||
let ip = VrpLookup::parse_ip("192.0.2.1").expect("ip lookup");
|
||||
let page = db
|
||||
.lookup_vrps("run_0001", &ip, None, 1, None)
|
||||
.expect("ip query");
|
||||
assert_eq!(page.data.len(), 1);
|
||||
assert_eq!(page.data[0].asn, 64496);
|
||||
assert_eq!(page.data[0].prefix, "192.0.2.0/24");
|
||||
assert!(page.next_cursor.is_none());
|
||||
|
||||
let prefix = VrpLookup::parse_prefix("192.0.2.0/24").expect("prefix lookup");
|
||||
let filtered = db
|
||||
.lookup_vrps("run_0001", &prefix, Some(64497), 100, None)
|
||||
.expect("asn query");
|
||||
assert!(filtered.data.is_empty());
|
||||
let latest = db
|
||||
.lookup_vrps("run_0001", &prefix, Some(64496), 100, None)
|
||||
.expect("filtered prefix query");
|
||||
assert_eq!(latest.data.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_run_does_not_replace_previous_latest() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
|
||||
Reference in New Issue
Block a user