Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 50 additions & 0 deletions src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,18 @@ impl ApiClient {
}
}

/// Test-only client (no config load). Used with a local mock HTTP server.
#[cfg(test)]
pub(crate) fn test_new(api_url: &str, api_key: &str, workspace_id: Option<&str>) -> Self {
Self {
client: reqwest::blocking::Client::new(),
api_key: api_key.to_string(),
api_url: api_url.to_string(),
workspace_id: workspace_id.map(String::from),
sandbox_id: None,
}
}

fn debug_headers(&self) -> Vec<(&str, String)> {
let masked = if self.api_key.len() > 4 {
format!("Bearer ...{}", &self.api_key[self.api_key.len()-4..])
Expand Down Expand Up @@ -336,6 +348,44 @@ fn format_fail_message(
mod tests {
use super::*;
use auth::AuthStatus;
use serde::Deserialize;

#[derive(Deserialize)]
struct Probe {
n: i32,
}

#[test]
fn get_none_if_not_found_returns_none_on_404() {
let mut server = mockito::Server::new();
let mock = server
.mock("GET", "/missing")
.match_header("Authorization", "Bearer test-key")
.with_status(404)
.create();

let api = ApiClient::test_new(&server.url(), "test-key", None);
let got: Option<Probe> = api.get_none_if_not_found("/missing");
assert!(got.is_none());
mock.assert();
}

#[test]
fn get_none_if_not_found_returns_some_on_200() {
let mut server = mockito::Server::new();
let mock = server
.mock("GET", "/ok")
.match_header("Authorization", "Bearer test-key")
.match_header("X-Workspace-Id", "ws-1")
.with_status(200)
.with_body(r#"{"n":7}"#)
.create();

let api = ApiClient::test_new(&server.url(), "test-key", Some("ws-1"));
let got: Option<Probe> = api.get_none_if_not_found("/ok");
assert_eq!(got.unwrap().n, 7);
mock.assert();
}

#[test]
fn format_fail_message_401_with_invalid_key_shows_reauth_hint() {
Expand Down
200 changes: 186 additions & 14 deletions src/indexes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::api::ApiClient;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::ops::ControlFlow;

#[derive(Deserialize, Serialize)]
struct Index {
Expand Down Expand Up @@ -52,16 +53,42 @@ struct ConnectionsBody {
connections: Vec<ConnectionRef>,
}

fn connection_lookup(api: &ApiClient) -> HashMap<String, String> {
let body: ConnectionsBody = api.get("/connections");
fn connection_label_to_id_map(connections: &[ConnectionRef]) -> HashMap<String, String> {
let mut m = HashMap::new();
for c in body.connections {
m.insert(c.id.clone(), c.id.clone());
m.insert(c.name.clone(), c.id);
for c in connections {
m.insert(c.name.clone(), c.id.clone());
}
m
}

fn connection_lookup(api: &ApiClient) -> HashMap<String, String> {
let body: ConnectionsBody = api.get("/connections");
connection_label_to_id_map(&body.connections)
}

/// How to continue after merging one `/information_schema` page.
fn information_schema_followup(
has_more: bool,
next_cursor: Option<String>,
) -> ControlFlow<(), String> {
if !has_more {
return ControlFlow::Break(());
}
let Some(c) = next_cursor else {
return ControlFlow::Break(());
};
ControlFlow::Continue(c)
}

fn sort_info_tables(tables: &mut [InfoTable]) {
tables.sort_by(|a, b| {
a.connection
.cmp(&b.connection)
.then_with(|| a.schema.cmp(&b.schema))
.then_with(|| a.table.cmp(&b.table))
});
}

fn collect_tables(
api: &ApiClient,
connection_id: Option<&str>,
Expand All @@ -86,17 +113,12 @@ fn collect_tables(
}
let body: InfoListResponse = api.get_with_params("/information_schema", &params);
out.extend(body.tables);
if !body.has_more {
break;
match information_schema_followup(body.has_more, body.next_cursor) {
ControlFlow::Break(()) => break,
ControlFlow::Continue(c) => cursor = Some(c),
}
cursor = body.next_cursor;
}
out.sort_by(|a, b| {
a.connection
.cmp(&b.connection)
.then_with(|| a.schema.cmp(&b.schema))
.then_with(|| a.table.cmp(&b.table))
});
sort_info_tables(&mut out);
out
}

Expand Down Expand Up @@ -266,3 +288,153 @@ pub fn create(
println!("{}", "Index created.".green());
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn information_schema_followup_terminates_when_not_has_more() {
assert!(matches!(
information_schema_followup(false, Some("c".into())),
ControlFlow::Break(())
));
}

#[test]
fn information_schema_followup_breaks_when_more_but_no_cursor() {
assert!(matches!(
information_schema_followup(true, None),
ControlFlow::Break(())
));
}

#[test]
fn information_schema_followup_continues_with_cursor() {
assert!(matches!(
information_schema_followup(true, Some("next".into())),
ControlFlow::Continue(ref s) if s == "next"
));
}

#[test]
fn sort_info_tables_orders_by_connection_schema_table() {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: The test only exercises connection-level ordering (entries differ at connection). Schema and table tiebreaking are never reached, so the test name oversells the coverage. Consider adding a second pair that shares the same connection but differs on schema or table to actually verify the composite key. (not blocking)

let mut tables = vec![
InfoTable {
connection: "b".into(),
schema: "s".into(),
table: "t2".into(),
},
InfoTable {
connection: "a".into(),
schema: "s".into(),
table: "t1".into(),
},
];
sort_info_tables(&mut tables);
assert_eq!(tables[0].table, "t1");
assert_eq!(tables[1].table, "t2");
}

#[test]
fn connection_label_to_id_map_maps_names_only() {
let connections = vec![
ConnectionRef {
id: "conn-id".into(),
name: "Warehouse".into(),
},
ConnectionRef {
id: "other".into(),
name: "Lake".into(),
},
];
let m = connection_label_to_id_map(&connections);
assert_eq!(m.get("Warehouse").map(String::as_str), Some("conn-id"));
assert_eq!(m.get("Lake").map(String::as_str), Some("other"));
assert!(!m.contains_key("conn-id"));
}

#[test]
fn collect_tables_single_page() {
let mut server = mockito::Server::new();
let mock = server
.mock("GET", "/information_schema")
.match_header("Authorization", "Bearer k")
.match_header("X-Workspace-Id", "ws1")
.with_status(200)
.with_body(
r#"{"tables":[
{"connection":"c1","schema":"public","table":"z"},
{"connection":"c1","schema":"public","table":"a"}
],"has_more":false,"next_cursor":null}"#,
)
.create();

let api = ApiClient::test_new(&server.url(), "k", Some("ws1"));
let tables = collect_tables(&api, None, None, None);
mock.assert();
assert_eq!(tables.len(), 2);
assert_eq!(tables[0].table, "a");
assert_eq!(tables[1].table, "z");
}

#[test]
fn list_one_table_scan_returns_empty_on_404() {
let mut server = mockito::Server::new();
let mock = server
.mock(
"GET",
mockito::Matcher::Regex(r"^/connections/.+/tables/.+/.+/indexes$".into()),
)
.match_header("Authorization", "Bearer k")
.with_status(404)
.create();

let api = ApiClient::test_new(&server.url(), "k", Some("ws"));
let rows = list_one_table_scan(&api, "cid", "sch", "tbl");
mock.assert();
assert!(rows.is_empty());
}

#[test]
fn list_one_table_returns_indexes() {
let mut server = mockito::Server::new();
let mock = server
.mock("GET", "/connections/cid/tables/sch/tbl/indexes")
.match_header("Authorization", "Bearer k")
.with_status(200)
.with_body(
r#"{"indexes":[{
"index_name":"ix1",
"index_type":"btree",
"columns":["c1"],
"metric":null,
"status":"ready",
"created_at":"2020-01-01T00:00:00Z",
"updated_at":"2020-01-01T00:00:00Z"
}]}"#,
)
.create();

let api = ApiClient::test_new(&server.url(), "k", None);
let rows = list_one_table(&api, "cid", "sch", "tbl");
mock.assert();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].index_name, "ix1");
}

#[test]
fn list_one_table_scan_returns_indexes_on_200() {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

super nit: Unlike the other mockito tests in this file, this one doesn't call .match_header("Authorization", "Bearer k"), so the header is not verified. Consistent header matching would make the test stronger. (not blocking)

let mut server = mockito::Server::new();
let mock = server
.mock("GET", "/connections/x/tables/s/t/indexes")
.with_status(200)
.with_body(r#"{"indexes":[]}"#)
.create();

let api = ApiClient::test_new(&server.url(), "k", None);
let rows = list_one_table_scan(&api, "x", "s", "t");
mock.assert();
assert!(rows.is_empty());
}
}
Loading