20260718 新增查询服务 VRP 检索接口

This commit is contained in:
2026-07-20 10:04:43 +08:00
parent e80fba44c3
commit 6d9eedd680
5 changed files with 532 additions and 3 deletions

View File

@@ -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");