use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use arrow_schema::{DataType, Field, Schema};
use serde_json::json;
use uuid::Uuid;
use wiremock::matchers::{body_partial_json, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::{VectorIndexMetric, VectorIndexOptions};
use crate::catalog_types::ProjectCommitId;
use crate::{Client, RemoteScandal, Table};
async fn source(dtype: DataType, stable_fields: bool) -> (MockServer, Client, Table, Uuid) {
let server = MockServer::start().await;
let client = Client::builder()
.api_url(server.uri())
.bearer_token("caller")
.scandal(Arc::new(
RemoteScandal::new(server.uri().parse().unwrap()).unwrap(),
))
.build()
.unwrap();
let table = client.table_by_id(Uuid::new_v4(), Uuid::now_v7());
let field_id = Uuid::new_v4();
let identity = table.identity();
Mock::given(method("GET"))
.and(path(format!(
"/api/tables/{}/versions/{}",
identity.resource_id, identity.version_id
)))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(json!({"rootUri":"spfsu://source/", "seqNo":"1"})),
)
.mount(&server)
.await;
let expires = SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + Duration::from_mins(15);
Mock::given(method("POST"))
.and(path(format!("/api/tables/{}/storage-binding", identity.resource_id)))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rootUri":"spfsu://source/", "mode":"read_only", "expiresAt":expires.as_secs().to_string(),
"access":{"type":"signed", "signed":{"endpoint":"", "token":"source-grant"}},
}))).mount(&server).await;
let schema = Schema::new(vec![Field::new("embedding", dtype, false)]);
Mock::given(method("POST"))
.and(path("/scans/plan"))
.and(body_partial_json(json!({"scan":{"table":{"resourceId":identity.resource_id,"versionId":identity.version_id}, "columns":["embedding"], "limit":"0"}})))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"info": {"schema":schema, "filterClauses":[], "sourceFields":stable_fields.then(|| BTreeMap::from([("embedding", field_id)]))},
"splits":[],
}))).expect(1).mount(&server).await;
(server, client, table, field_id)
}
fn vector_type() -> DataType {
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 4)
}
#[tokio::test]
async fn creation_resolves_the_source_field_and_stages_a_vector_resource() {
let (server, client, source, field_id) = source(vector_type(), true).await;
let project = client.project("demo").unwrap();
let resource_id = Uuid::new_v4();
let version_id = Uuid::now_v7();
Mock::given(method("POST"))
.and(path("/api/vector-indexes"))
.and(body_partial_json(json!({
"projectId":"demo", "definition": {
"sourceResourceId":source.identity().resource_id,"sourceTableId":source.identity().resource_id,
"sourceFieldId":field_id,"dimensions":4,"metric":"dot","engineId":"vortex.ivf-rabitq",
"engineOptions":{"partitions":"32"},
},
})))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"vectorIndexId":resource_id,"versionId":version_id,"rootUri":"spfsu://index/"})))
.expect(1).mount(&server).await;
let index = project
.create_vector_index(
&source,
"embedding",
VectorIndexOptions {
metric: VectorIndexMetric::Dot,
engine_options: BTreeMap::from([("partitions".to_owned(), "32".to_owned())]),
},
)
.await
.unwrap();
assert_eq!(
(index.resource_id(), index.version_id()),
(resource_id, version_id)
);
Mock::given(method("GET"))
.and(path("/api/catalog/projects/demo/branches/by-name"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"branchId":Uuid::now_v7(),"name":"main","headCommitId":Uuid::now_v7(),
"lastBranchEventId":Uuid::now_v7(),"createdAtMicros":"0",
})))
.mount(&server)
.await;
let mut commit = project.branch("main").await.unwrap().commit();
commit.create_vector_index("frame-search", &index).unwrap();
let commit_id = Uuid::now_v7();
Mock::given(method("POST"))
.and(path(format!("/api/catalog/projects/demo/commits/{commit_id}/namespace/resolve")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"commitId": commit_id,
"resource": {"descriptor": {"resourceId": Uuid::new_v4(), "kind": "vector-index", "service": "vector-index", "namespaceManaged": false}, "versionId": Uuid::now_v7()},
})))
.expect(1).mount(&server).await;
let existing = project
.snapshot_at(ProjectCommitId::new(commit_id).unwrap())
.vector_index("existing")
.await
.unwrap();
assert!(commit.create_vector_index("existing", &existing).is_err());
let saved = serde_json::to_value(commit.freeze().unwrap().save()).unwrap();
assert_eq!(
saved["publication"]["resources"],
json!([{
"type":"create", "create":{"descriptor":{"resourceId":resource_id,"service":"vector-index","kind":"vector-index","namespaceManaged":false}, "versionId":version_id},
}])
);
assert_eq!(
saved["publication"]["namespace"][0]["bind"]["name"],
"frame-search"
);
server.verify().await;
}
#[tokio::test]
async fn unsupported_source_metadata_fails_before_resource_creation() {
for (dtype, stable_fields, message) in [
(vector_type(), false, "stable source field identities"),
(DataType::Int64, true, "fixed-size list"),
] {
let (server, client, source, _) = source(dtype, stable_fields).await;
Mock::given(path("/api/vector-indexes"))
.respond_with(ResponseTemplate::new(200))
.expect(0)
.mount(&server)
.await;
let error = client
.project("demo")
.unwrap()
.create_vector_index(&source, "embedding", VectorIndexOptions::default())
.await
.err()
.unwrap();
assert!(error.to_string().contains(message), "{error}");
server.verify().await;
}
}