use serde_json::{Value as Json, json};
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::ast::ProtoSchema;
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(
schemas: &[ProtoSchema],
_config: &SqlGenerationConfig,
) -> Result<Vec<GeneratedArtifact>, serde_json::Error> {
let manifest = CatalogManifest::from_schemas(schemas)?;
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 dimension = store_opt_i64_any(store, &["udb.vector_dimension", "dimension"], 1536);
let distance = match store_opt_str_any(store, &["udb.vector_distance", "distance"])
.unwrap_or("Cosine")
{
"Cosine" | "Euclid" | "Dot" | "Manhattan" => {
store_opt_str_any(store, &["udb.vector_distance", "distance"]).unwrap_or("Cosine")
}
_ => "Cosine",
};
let hnsw_m = store_opt_i64_any(store, &["udb.hnsw_m", "hnsw_m"], 16);
let hnsw_ef = store_opt_i64_any(
store,
&[
"udb.hnsw_ef_construct",
"udb.hnsw_ef_construction",
"hnsw_ef_construction",
],
100,
);
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 generated_at_unix() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or_default()
}
#[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(&[], &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));
}
}