diff --git a/src/indexes.rs b/src/indexes.rs index 2b8ebe0..30e68f2 100644 --- a/src/indexes.rs +++ b/src/indexes.rs @@ -62,6 +62,69 @@ fn connection_lookup(api: &Api) -> HashMap { connection_label_to_id_map(&refs) } +/// Pick the connection id to address a per-table index call with during a +/// connection-wide scan. +/// +/// Prefers the caller-supplied `--connection-id`: it always resolves, including +/// for a database-scoped connection whose `information_schema` `label` +/// (`__db_*`) is absent from `connections list` (that listing hides +/// database-scoped connections, so `name_to_id` can't map it — #161). The scan's +/// tables are already filtered to that connection, so the supplied id is correct +/// for every row. With no `--connection-id` (the list-everything case), maps the +/// label back to an id, falling back to the label itself. +fn scan_connection_id<'a>( + supplied: Option<&'a str>, + label: &'a str, + name_to_id: &'a HashMap, +) -> &'a str { + supplied + .or_else(|| name_to_id.get(label).map(String::as_str)) + .unwrap_or(label) +} + +/// Gather index rows across a connection's (or the workspace's) tables — the +/// `indexes list` path when no full `connection.schema.table` triple is given. +/// +/// Enumerates tables via `information_schema`, then fetches each table's indexes +/// by the connection id chosen in [`scan_connection_id`] (the supplied +/// `--connection-id` wins, fixing the database-scoped case in #161). Skipped +/// connections / missing tables surface as no rows for that table, not an error. +fn collect_connection_wide( + api: &Api, + connection_id: Option<&str>, + schema: Option<&str>, + table: Option<&str>, +) -> Vec { + let tables = collect_tables(api, connection_id, schema, table); + // See `scan_connection_id`: a supplied `--connection-id` is used directly, + // so the name→id map is only needed (and only worth a round-trip) for the + // unscoped, list-everything case (#161). + let conn_ids = if connection_id.is_some() { + HashMap::new() + } else { + connection_lookup(api) + }; + let per_table: Vec<(String, Vec)> = tables + .par_iter() + .map(|t| { + let cid = scan_connection_id(connection_id, &t.connection, &conn_ids); + let full = format!("{}.{}.{}", t.connection, t.schema, t.table); + let indexes = list_one_table_scan(api, cid, &t.schema, &t.table); + (full, indexes) + }) + .collect(); + let mut rows: Vec = Vec::new(); + for (full, indexes) in per_table { + for i in indexes { + rows.push(IndexRow { + inner: i, + table: Some(full.clone()), + }); + } + } + rows +} + /// How to continue after merging one `/information_schema` page. fn information_schema_followup( has_more: bool, @@ -264,33 +327,10 @@ pub fn list( .collect(); (rows, false) } - _ => { - let tables = collect_tables(&api, connection_id, schema, table); - let conn_ids = connection_lookup(&api); - let api = api.clone(); - let per_table: Vec<(String, Vec)> = tables - .par_iter() - .map(|t| { - let cid = conn_ids - .get(&t.connection) - .map(String::as_str) - .unwrap_or(t.connection.as_str()); - let full = format!("{}.{}.{}", t.connection, t.schema, t.table); - let indexes = list_one_table_scan(&api, cid, &t.schema, &t.table); - (full, indexes) - }) - .collect(); - let mut rows: Vec = Vec::new(); - for (full, indexes) in per_table { - for i in indexes { - rows.push(IndexRow { - inner: i, - table: Some(full.clone()), - }); - } - } - (rows, true) - } + _ => ( + collect_connection_wide(&api, connection_id, schema, table), + true, + ), }; match format { @@ -602,6 +642,81 @@ mod tests { assert!(!m.contains_key("conn-id")); } + #[test] + fn scan_connection_id_prefers_supplied_id_over_label_map() { + // #161: a managed database's catalog surfaces under an internal + // `__db_*` label that `connections list` hides, so the name→id map is + // empty for it. The supplied --connection-id must win regardless. + let empty = HashMap::new(); + assert_eq!( + scan_connection_id(Some("conn-real"), "__db_jz50abc", &empty), + "conn-real" + ); + // Even when the label *is* in the map, the supplied id takes precedence. + let mut m = HashMap::new(); + m.insert("__db_jz50abc".to_string(), "conn-mapped".to_string()); + assert_eq!( + scan_connection_id(Some("conn-real"), "__db_jz50abc", &m), + "conn-real" + ); + } + + #[test] + fn scan_connection_id_maps_label_when_no_supplied_id() { + let mut m = HashMap::new(); + m.insert("Warehouse".to_string(), "conn-id".to_string()); + assert_eq!(scan_connection_id(None, "Warehouse", &m), "conn-id"); + } + + #[test] + fn scan_connection_id_falls_back_to_label_when_unmapped() { + let empty = HashMap::new(); + assert_eq!(scan_connection_id(None, "Warehouse", &empty), "Warehouse"); + } + + #[test] + fn collect_connection_wide_uses_supplied_id_for_db_scoped_label() { + // #161 regression: information_schema reports a managed database's + // catalog under an internal `__db_*` label, but the per-table index + // call must use the supplied --connection-id. The indexes endpoint is + // mocked ONLY for the real id (`conn-real`); had the scan used the + // `__db_*` label (the old behavior), it would miss this mock. No + // `connections list` mock is needed — a supplied id skips that lookup. + let mut server = mockito::Server::new(); + let info = server + .mock("GET", "/v1/information_schema") + .match_query(mockito::Matcher::Any) + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + r#"{"count":1,"limit":100,"tables":[ + {"connection":"__db_abc","schema":"public","table":"vec_mid","synced":true} + ],"has_more":false,"next_cursor":null}"#, + ) + .create(); + let idx = server + .mock( + "GET", + "/v1/connections/conn-real/tables/public/vec_mid/indexes", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + r#"{"indexes":[{"index_name":"vec_mid_idx","index_type":"vector", + "columns":["c"],"metric":"cosine","status":"ready", + "created_at":"2020-01-01T00:00:00Z","updated_at":"2020-01-01T00:00:00Z"}]}"#, + ) + .create(); + + let api = Api::test_new(&server.url(), "k", Some("ws")); + let rows = collect_connection_wide(&api, Some("conn-real"), None, None); + info.assert(); + idx.assert(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].inner.index_name, "vec_mid_idx"); + assert_eq!(rows[0].table.as_deref(), Some("__db_abc.public.vec_mid")); + } + #[test] fn collect_tables_single_page() { let mut server = mockito::Server::new();