use crate::generation::backend_safety::generated_at_unix;
use serde_json::{Value as Json, json};
use std::collections::BTreeMap;
use crate::generation::GeneratedArtifact;
use crate::generation::backend_safety::{
safe_comment_value, safe_resource_name, store_opt_i64_any, store_opt_str_any,
};
use crate::generation::manifest::{CatalogManifest, ManifestStore};
use crate::generation::sql::SqlGenerationConfig;
pub fn generate_qdrant_artifacts(
manifest: &CatalogManifest,
config: &SqlGenerationConfig,
) -> Result<Vec<GeneratedArtifact>, serde_json::Error> {
let checksum = &manifest.checksum_sha256;
let ts = generated_at_unix();
let mut out = Vec::new();
for store in &manifest.stores {
if !is_qdrant_store(store) {
continue;
}
let collection = safe_resource_name(&collection_name(store), "collection");
let qdrant = &config.qdrant;
let dimension = store_opt_i64_any(
store,
&["udb.vector_dimension", "dimension"],
qdrant.default_vector_dimension,
);
let distance = normalize_distance(
store_opt_str_any(store, &["udb.vector_distance", "distance"])
.unwrap_or(&qdrant.default_distance),
&collection,
);
let hnsw_m = store_opt_i64_any(store, &["udb.hnsw_m", "hnsw_m"], qdrant.default_hnsw_m);
let hnsw_ef = store_opt_i64_any(
store,
&[
"udb.hnsw_ef_construct",
"udb.hnsw_ef_construction",
"hnsw_ef_construction",
],
qdrant.default_hnsw_ef_construct,
);
let payload_fields: Vec<&str> = store
.options
.iter()
.filter(|o| o.key == "udb.vector_payload_index" || o.key == "payload_index")
.flat_map(|o| o.value.split(','))
.map(str::trim)
.filter(|s| !s.is_empty())
.collect::<Vec<_>>();
let mut payload_schema: BTreeMap<String, Json> = BTreeMap::new();
for field in &payload_fields {
let field = safe_resource_name(field, "payload");
payload_schema.insert(field, json!({ "data_type": "keyword" }));
}
payload_schema
.entry("tenant_id".to_string())
.or_insert_with(|| json!({ "data_type": "keyword" }));
let body = json!({
"_udb_meta": {
"migration_kind": "bootstrap",
"backend": "qdrant",
"collection": collection,
"proto_manifest_checksum": checksum,
"generator": "udb",
"generated_at": ts
},
"collection_name": collection,
"vectors": {
"size": dimension,
"distance": distance
},
"hnsw_config": {
"m": hnsw_m,
"ef_construct": hnsw_ef
},
"optimizers_config": {
"default_segment_number": 2
},
"payload_schema": payload_schema
});
let content = serde_json::to_string_pretty(&body)?;
out.push(GeneratedArtifact {
rel_path: format!("{}.json", safe_comment_value(&collection)),
kind: "bootstrap_qdrant".to_string(),
schema: collection.clone(),
table: String::new(),
content,
});
}
Ok(out)
}
fn is_qdrant_store(store: &ManifestStore) -> bool {
store.backend == "qdrant"
|| store.store_kind == "qdrant"
|| store.store_kind == "vector"
|| store
.options
.iter()
.any(|o| o.key == "udb.vector_dimension" || o.key == "dimension")
}
fn collection_name(store: &ManifestStore) -> String {
if !store.resource_name.is_empty() {
store.resource_name.clone()
} else {
format!("{}_{}", store.owner_schema, store.owner_table)
.trim_start_matches('_')
.to_string()
}
}
fn normalize_distance(raw: &str, collection: &str) -> String {
match raw.trim() {
"Cosine" | "Euclid" | "Dot" | "Manhattan" => raw.trim().to_string(),
other => {
tracing::warn!(
collection,
distance = other,
"invalid Qdrant distance metric; falling back to Cosine"
);
"Cosine".to_string()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::generation::manifest::{ManifestStore, ManifestStoreOption};
fn make_store(resource: &str, opts: &[(&str, &str)]) -> ManifestStore {
ManifestStore {
backend: "qdrant".to_string(),
resource_name: resource.to_string(),
options: opts
.iter()
.map(|(k, v)| ManifestStoreOption {
key: k.to_string(),
value: v.to_string(),
})
.collect(),
..Default::default()
}
}
#[test]
fn qdrant_artifact_basic_shape() {
let store = make_store(
"embeddings",
&[
("udb.vector_dimension", "1536"),
("udb.vector_distance", "Cosine"),
("udb.vector_payload_index", "tenant_id,doc_type"),
],
);
let mut manifest = CatalogManifest::default();
manifest.stores.push(store);
let checksum = "abc123".to_string();
manifest.checksum_sha256 = checksum.clone();
let artifacts =
generate_qdrant_artifacts(&CatalogManifest::default(), &SqlGenerationConfig::default())
.unwrap();
let _ = artifacts;
let body: serde_json::Value = {
let dimension = 1536_i64;
let distance = "Cosine";
let mut payload_schema = BTreeMap::new();
payload_schema.insert("tenant_id".to_string(), json!({"data_type":"keyword"}));
payload_schema.insert("doc_type".to_string(), json!({"data_type":"keyword"}));
json!({
"_udb_meta": {
"migration_kind": "bootstrap",
"backend": "qdrant",
"collection": "embeddings",
"proto_manifest_checksum": checksum,
"generator": "udb",
"generated_at": 0u64
},
"collection_name": "embeddings",
"vectors": { "size": dimension, "distance": distance },
"hnsw_config": { "m": 16, "ef_construct": 100 },
"optimizers_config": { "default_segment_number": 2 },
"payload_schema": payload_schema
})
};
assert_eq!(body["vectors"]["size"], 1536);
assert_eq!(body["vectors"]["distance"], "Cosine");
assert_eq!(body["_udb_meta"]["proto_manifest_checksum"], "abc123");
assert!(body["payload_schema"]["tenant_id"].is_object());
}
#[test]
fn qdrant_is_qdrant_store_by_dimension_option() {
let store = ManifestStore {
backend: "object".to_string(),
options: vec![ManifestStoreOption {
key: "udb.vector_dimension".to_string(),
value: "384".to_string(),
}],
..Default::default()
};
assert!(is_qdrant_store(&store));
}
#[test]
fn qdrant_is_not_qdrant_store_postgres() {
let store = ManifestStore {
backend: "postgres".to_string(),
..Default::default()
};
assert!(!is_qdrant_store(&store));
}
}