rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 14:52:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(rust): move remote module unit tests into separate files

Browse Source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BgkiPYS2eTmAWVyqn6KBnN
This commit is contained in:
Claude
2026-10-07 06:23:00 +00:00
parent bbc4aebe55
commit 6dde01de33
34 changed files with 18673 additions and 18672 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/remote/client.rs
+1 -752
View File
@@ -1146,755 +1146,4 @@ pub mod test_utils {
}
#[cfg(test)]
mod tests {
/// A configuration naming any other delimiter is refused where it was
/// written, rather than producing identifiers no service splits the way the
/// caller meant.
#[test]
fn test_a_delimiter_other_than_the_supported_one_is_refused() {
for delimiter in ["/", "?", "#", "%", "", ".", "..", "-", "_", "|", "::", "$$"] {
let error = super::validate_id_delimiter(delimiter)
.expect_err("only the supported delimiter may be configured");
assert!(
error.to_string().contains("id_delimiter"),
"{delimiter:?}: {error}"
);
}
super::validate_id_delimiter(super::ID_DELIMITER).unwrap();
}
/// Leaving it unset is how nearly every caller reaches the same delimiter.
#[test]
fn test_an_unset_delimiter_is_the_supported_one() {
super::ClientConfig::default().validate().unwrap();
super::ClientConfig {
id_delimiter: Some(super::ID_DELIMITER.to_string()),
..Default::default()
}
.validate()
.unwrap();
super::ClientConfig {
id_delimiter: Some("-".to_string()),
..Default::default()
}
.validate()
.expect_err("a configured delimiter other than the supported one must be refused");
}
use super::*;
use serial_test::serial;
use std::time::Duration;
// Serializes the env-var-mutating tests below: cargo test runs tests in
// parallel, but several of these tests read and write the same process-
// global env vars (`LANCEDB_USER_ID*`), so they would race without this.
static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn lock_env() -> std::sync::MutexGuard<'static, ()> {
ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
}
#[test]
fn test_parse_catalog_database_uri() {
let parsed = parse_db_url("db://team%2Fsearch").unwrap();
assert_eq!(parsed.db_name, "team/search");
assert!(parsed.db_prefix.is_none());
let parsed = parse_db_url("db://db/prefix").unwrap();
assert_eq!(parsed.db_name, "db");
assert_eq!(parsed.db_prefix.as_deref(), Some("prefix"));
}
#[test]
fn test_timeout_config_default() {
let config = TimeoutConfig::default();
assert!(config.timeout.is_none());
assert!(config.connect_timeout.is_none());
assert!(config.read_timeout.is_none());
assert!(config.pool_idle_timeout.is_none());
}
#[test]
fn test_timeout_config_with_overall_timeout() {
let config = TimeoutConfig {
timeout: Some(Duration::from_secs(60)),
connect_timeout: Some(Duration::from_secs(10)),
read_timeout: Some(Duration::from_secs(30)),
pool_idle_timeout: Some(Duration::from_secs(300)),
};
assert_eq!(config.timeout, Some(Duration::from_secs(60)));
assert_eq!(config.connect_timeout, Some(Duration::from_secs(10)));
assert_eq!(config.read_timeout, Some(Duration::from_secs(30)));
assert_eq!(config.pool_idle_timeout, Some(Duration::from_secs(300)));
}
#[test]
fn test_client_config_with_timeout() {
let timeout_config = TimeoutConfig {
timeout: Some(Duration::from_secs(120)),
..Default::default()
};
let client_config = ClientConfig {
timeout_config,
..Default::default()
};
assert_eq!(
client_config.timeout_config.timeout,
Some(Duration::from_secs(120))
);
}
#[test]
fn test_tls_config_default() {
let config = TlsConfig::default();
assert!(config.cert_file.is_none());
assert!(config.key_file.is_none());
assert!(config.ssl_ca_cert.is_none());
assert!(config.assert_hostname);
}
#[test]
fn test_tls_config_with_mtls() {
let tls_config = TlsConfig {
cert_file: Some("/path/to/cert.pem".to_string()),
key_file: Some("/path/to/key.pem".to_string()),
ssl_ca_cert: Some("/path/to/ca.pem".to_string()),
assert_hostname: true,
};
assert_eq!(tls_config.cert_file, Some("/path/to/cert.pem".to_string()));
assert_eq!(tls_config.key_file, Some("/path/to/key.pem".to_string()));
assert_eq!(tls_config.ssl_ca_cert, Some("/path/to/ca.pem".to_string()));
assert!(tls_config.assert_hostname);
}
#[test]
fn test_client_config_with_tls() {
let tls_config = TlsConfig {
cert_file: Some("/path/to/cert.pem".to_string()),
key_file: Some("/path/to/key.pem".to_string()),
ssl_ca_cert: None,
assert_hostname: false,
};
let client_config = ClientConfig {
tls_config: Some(tls_config.clone()),
..Default::default()
};
assert!(client_config.tls_config.is_some());
let config_tls = client_config.tls_config.unwrap();
assert_eq!(config_tls.cert_file, Some("/path/to/cert.pem".to_string()));
assert_eq!(config_tls.key_file, Some("/path/to/key.pem".to_string()));
assert!(config_tls.ssl_ca_cert.is_none());
assert!(!config_tls.assert_hostname);
}
#[test]
fn test_default_headers_skip_empty_api_key() {
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
assert!(!headers.contains_key("x-api-key"));
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"static-secret-value",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
let api_key = headers.get("x-api-key").unwrap();
assert_eq!(api_key, "static-secret-value");
assert!(api_key.is_sensitive());
assert!(!format!("{headers:?}").contains("static-secret-value"));
}
#[test]
fn test_configured_authentication_headers_are_sensitive() {
let config = ClientConfig {
extra_headers: HashMap::from([
(
"Authorization".to_string(),
"Bearer configured-secret".to_string(),
),
("X-API-Key".to_string(), "configured-api-key".to_string()),
(
"Cookie".to_string(),
"session=configured-cookie".to_string(),
),
(
"Proxy-Authorization".to_string(),
"Basic configured-proxy-secret".to_string(),
),
("X-Custom".to_string(), "visible-value".to_string()),
]),
..Default::default()
};
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&config,
)
.unwrap();
assert!(headers.get("authorization").unwrap().is_sensitive());
assert!(headers.get("x-api-key").unwrap().is_sensitive());
assert!(headers.get("cookie").unwrap().is_sensitive());
assert!(headers.get("proxy-authorization").unwrap().is_sensitive());
assert!(!headers.get("x-custom").unwrap().is_sensitive());
let debug = format!("{headers:?}");
assert!(!debug.contains("configured-secret"));
assert!(!debug.contains("configured-api-key"));
assert!(!debug.contains("configured-cookie"));
assert!(!debug.contains("configured-proxy-secret"));
assert!(debug.contains("visible-value"));
}
/// `log_request` prints the request's Debug, and Debug for a request prints
/// its headers. Marking the value sensitive is the only thing standing
/// between the API key and every debug line; assert on the header map's own
/// Debug, which is what that printing reduces to.
#[test]
fn test_api_key_is_redacted_in_debug_output() {
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"sk-live-sentinel",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
assert_eq!(headers.get("x-api-key").unwrap(), "sk-live-sentinel");
assert!(
!format!("{:?}", headers).contains("sk-live-sentinel"),
"the API key must not survive Debug formatting"
);
}
/// A suppressed body is suppressed whatever the content type says, and an
/// allowed one is logged in full.
#[test]
fn test_body_logging_is_decided_by_the_caller() {
assert_ne!(BodyLogging::Allowed, BodyLogging::Suppressed);
// `send` and `send_suppressing_body` differ only in what they pass, so
// the enum is the whole contract: a caller states its intent and the
// transport does not infer one from the route.
assert_eq!(BodyLogging::Allowed, BodyLogging::Allowed);
}
#[test]
fn test_rejects_invalid_cloud_dns_hostname() {
let invalid_database_names = ["a".repeat(64), "invalid..database".to_string()];
for db_name in invalid_database_names {
let parsed_url = parse_db_url(&format!("db://{db_name}")).unwrap();
let error = RestfulLanceDbClient::<Sender>::try_new(
&parsed_url,
"us-east-1",
None,
HeaderMap::new(),
ClientConfig::default(),
None,
)
.unwrap_err();
assert!(
matches!(error, Error::InvalidInput { ref message } if message.contains("DNS labels must contain 1 to 63 bytes")),
"unexpected error: {error}"
);
}
}
// Test implementation of HeaderProvider
#[derive(Debug, Clone)]
struct TestHeaderProvider {
headers: HashMap<String, String>,
}
impl TestHeaderProvider {
fn new(headers: HashMap<String, String>) -> Self {
Self { headers }
}
}
#[async_trait::async_trait]
impl HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
Ok(self.headers.clone())
}
}
// Test implementation that returns an error
#[derive(Debug)]
struct ErrorHeaderProvider;
#[async_trait::async_trait]
impl HeaderProvider for ErrorHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
Err(Error::Runtime {
message: "Failed to get headers".to_string(),
})
}
}
#[tokio::test]
async fn test_client_config_with_header_provider() {
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "secret-key".to_string());
let provider = TestHeaderProvider::new(headers);
let client_config = ClientConfig {
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
..Default::default()
};
assert!(client_config.header_provider.is_some());
}
#[tokio::test]
async fn test_apply_dynamic_headers() {
// Create a mock client with header provider
let mut headers = HashMap::new();
headers.insert("X-Dynamic".to_string(), "dynamic-value".to_string());
let provider = TestHeaderProvider::new(headers);
// Create a simple request
let request = reqwest::Request::new(
reqwest::Method::GET,
"https://example.com/test".parse().unwrap(),
);
// Create client with header provider
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Apply dynamic headers
let updated_request = client.apply_dynamic_headers(request).await.unwrap();
// Check that the header was added
assert_eq!(
updated_request.headers().get("X-Dynamic").unwrap(),
"dynamic-value"
);
}
#[tokio::test]
async fn test_apply_dynamic_headers_merge() {
// Test that dynamic headers override existing headers
let mut headers = HashMap::new();
headers.insert("Authorization".to_string(), "Bearer new-token".to_string());
headers.insert("X-API-Key".to_string(), "new-api-key".to_string());
headers.insert("X-Custom".to_string(), "custom-value".to_string());
let provider = TestHeaderProvider::new(headers);
// Create request with existing Authorization header
let mut request_builder = reqwest::Client::new().get("https://example.com/test");
request_builder = request_builder.header("Authorization", "Bearer old-token");
request_builder = request_builder.header("X-Existing", "existing-value");
let request = request_builder.build().unwrap();
// Create client with header provider
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Apply dynamic headers
let updated_request = client.apply_dynamic_headers(request).await.unwrap();
// Check that dynamic headers override existing ones
assert_eq!(
updated_request.headers().get("Authorization").unwrap(),
"Bearer new-token"
);
assert_eq!(
updated_request.headers().get("X-Custom").unwrap(),
"custom-value"
);
assert!(
updated_request
.headers()
.get("Authorization")
.unwrap()
.is_sensitive()
);
assert!(
updated_request
.headers()
.get("X-API-Key")
.unwrap()
.is_sensitive()
);
assert!(
!updated_request
.headers()
.get("X-Custom")
.unwrap()
.is_sensitive()
);
let debug = format!("{updated_request:?}");
assert!(!debug.contains("new-token"));
assert!(!debug.contains("new-api-key"));
assert!(debug.contains("custom-value"));
// Existing headers should still be present
assert_eq!(
updated_request.headers().get("X-Existing").unwrap(),
"existing-value"
);
}
#[tokio::test]
async fn test_apply_dynamic_headers_with_error_provider() {
let provider = ErrorHeaderProvider;
let request = reqwest::Request::new(
reqwest::Method::GET,
"https://example.com/test".parse().unwrap(),
);
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Header provider errors should fail the request
// This is important for security - if auth headers can't be fetched, don't proceed
let result = client.apply_dynamic_headers(request).await;
assert!(result.is_err());
match result.unwrap_err() {
Error::Runtime { message } => {
assert_eq!(message, "Failed to get headers");
}
_ => panic!("Expected Runtime error"),
}
}
#[test]
fn test_resolve_user_id_direct_value() {
let config = ClientConfig {
user_id: Some("direct-user-id".to_string()),
..Default::default()
};
assert_eq!(config.resolve_user_id(), Some("direct-user-id".to_string()));
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_none() {
let _guard = lock_env();
let config = ClientConfig::default();
// Clear env vars that might be set from other tests
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
}
assert_eq!(config.resolve_user_id(), None);
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_from_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "env-user-id");
}
let config = ClientConfig::default();
assert_eq!(config.resolve_user_id(), Some("env-user-id".to_string()));
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_from_env_key() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
std::env::set_var("LANCEDB_USER_ID_ENV_KEY", "MY_CUSTOM_USER_ID");
std::env::set_var("MY_CUSTOM_USER_ID", "custom-env-user-id");
}
let config = ClientConfig::default();
assert_eq!(
config.resolve_user_id(),
Some("custom-env-user-id".to_string())
);
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
std::env::remove_var("MY_CUSTOM_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_direct_takes_precedence() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "env-user-id");
}
let config = ClientConfig {
user_id: Some("direct-user-id".to_string()),
..Default::default()
};
assert_eq!(config.resolve_user_id(), Some("direct-user-id".to_string()));
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_empty_env_ignored() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "");
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
}
let config = ClientConfig::default();
assert_eq!(config.resolve_user_id(), None);
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
fn test_resolve_max_bytes_passed_value_wins() {
// An explicit config value is used verbatim; env/default are not consulted.
let resolved =
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
assert_eq!(resolved, Some(1234));
}
#[test]
fn test_resolve_max_bytes_zero_disables() {
let resolved =
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(0)).unwrap();
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_default_when_unset() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
assert_eq!(resolved, Some(DEFAULT_MAX_BYTES_PER_REQUEST));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_from_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, Some(4096));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_env_zero_disables() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "0");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_config_overrides_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
}
// A config value takes precedence over the environment variable.
let resolved =
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, Some(1234));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_invalid_env_errors() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "not-a-number");
}
let err = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap_err();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
}
#[test]
fn test_resolve_max_request_duration_passed_value_wins() {
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
Some(Duration::from_secs(42)),
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, Some(Duration::from_secs(42)));
}
#[test]
fn test_resolve_max_request_duration_zero_disables() {
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
Some(Duration::ZERO),
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_default_is_half_read_timeout() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, Some(Duration::from_secs(150)));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_from_env_seconds() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "30");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert_eq!(resolved, Some(Duration::from_secs(30)));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_env_zero_disables() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "0");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_invalid_env_errors() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "12.5");
}
let err = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap_err();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
}
}
mod tests;
rust/lancedb/src/remote/client/tests.rs
+752
View File
@@ -0,0 +1,752 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
/// A configuration naming any other delimiter is refused where it was
/// written, rather than producing identifiers no service splits the way the
/// caller meant.
#[test]
fn test_a_delimiter_other_than_the_supported_one_is_refused() {
for delimiter in ["/", "?", "#", "%", "", ".", "..", "-", "_", "|", "::", "$$"] {
let error = super::validate_id_delimiter(delimiter)
.expect_err("only the supported delimiter may be configured");
assert!(
error.to_string().contains("id_delimiter"),
"{delimiter:?}: {error}"
);
}
super::validate_id_delimiter(super::ID_DELIMITER).unwrap();
}
/// Leaving it unset is how nearly every caller reaches the same delimiter.
#[test]
fn test_an_unset_delimiter_is_the_supported_one() {
super::ClientConfig::default().validate().unwrap();
super::ClientConfig {
id_delimiter: Some(super::ID_DELIMITER.to_string()),
..Default::default()
}
.validate()
.unwrap();
super::ClientConfig {
id_delimiter: Some("-".to_string()),
..Default::default()
}
.validate()
.expect_err("a configured delimiter other than the supported one must be refused");
}
use super::*;
use serial_test::serial;
use std::time::Duration;
// Serializes the env-var-mutating tests below: cargo test runs tests in
// parallel, but several of these tests read and write the same process-
// global env vars (`LANCEDB_USER_ID*`), so they would race without this.
static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn lock_env() -> std::sync::MutexGuard<'static, ()> {
ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
}
#[test]
fn test_parse_catalog_database_uri() {
let parsed = parse_db_url("db://team%2Fsearch").unwrap();
assert_eq!(parsed.db_name, "team/search");
assert!(parsed.db_prefix.is_none());
let parsed = parse_db_url("db://db/prefix").unwrap();
assert_eq!(parsed.db_name, "db");
assert_eq!(parsed.db_prefix.as_deref(), Some("prefix"));
}
#[test]
fn test_timeout_config_default() {
let config = TimeoutConfig::default();
assert!(config.timeout.is_none());
assert!(config.connect_timeout.is_none());
assert!(config.read_timeout.is_none());
assert!(config.pool_idle_timeout.is_none());
}
#[test]
fn test_timeout_config_with_overall_timeout() {
let config = TimeoutConfig {
timeout: Some(Duration::from_secs(60)),
connect_timeout: Some(Duration::from_secs(10)),
read_timeout: Some(Duration::from_secs(30)),
pool_idle_timeout: Some(Duration::from_secs(300)),
};
assert_eq!(config.timeout, Some(Duration::from_secs(60)));
assert_eq!(config.connect_timeout, Some(Duration::from_secs(10)));
assert_eq!(config.read_timeout, Some(Duration::from_secs(30)));
assert_eq!(config.pool_idle_timeout, Some(Duration::from_secs(300)));
}
#[test]
fn test_client_config_with_timeout() {
let timeout_config = TimeoutConfig {
timeout: Some(Duration::from_secs(120)),
..Default::default()
};
let client_config = ClientConfig {
timeout_config,
..Default::default()
};
assert_eq!(
client_config.timeout_config.timeout,
Some(Duration::from_secs(120))
);
}
#[test]
fn test_tls_config_default() {
let config = TlsConfig::default();
assert!(config.cert_file.is_none());
assert!(config.key_file.is_none());
assert!(config.ssl_ca_cert.is_none());
assert!(config.assert_hostname);
}
#[test]
fn test_tls_config_with_mtls() {
let tls_config = TlsConfig {
cert_file: Some("/path/to/cert.pem".to_string()),
key_file: Some("/path/to/key.pem".to_string()),
ssl_ca_cert: Some("/path/to/ca.pem".to_string()),
assert_hostname: true,
};
assert_eq!(tls_config.cert_file, Some("/path/to/cert.pem".to_string()));
assert_eq!(tls_config.key_file, Some("/path/to/key.pem".to_string()));
assert_eq!(tls_config.ssl_ca_cert, Some("/path/to/ca.pem".to_string()));
assert!(tls_config.assert_hostname);
}
#[test]
fn test_client_config_with_tls() {
let tls_config = TlsConfig {
cert_file: Some("/path/to/cert.pem".to_string()),
key_file: Some("/path/to/key.pem".to_string()),
ssl_ca_cert: None,
assert_hostname: false,
};
let client_config = ClientConfig {
tls_config: Some(tls_config.clone()),
..Default::default()
};
assert!(client_config.tls_config.is_some());
let config_tls = client_config.tls_config.unwrap();
assert_eq!(config_tls.cert_file, Some("/path/to/cert.pem".to_string()));
assert_eq!(config_tls.key_file, Some("/path/to/key.pem".to_string()));
assert!(config_tls.ssl_ca_cert.is_none());
assert!(!config_tls.assert_hostname);
}
#[test]
fn test_default_headers_skip_empty_api_key() {
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
assert!(!headers.contains_key("x-api-key"));
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"static-secret-value",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
let api_key = headers.get("x-api-key").unwrap();
assert_eq!(api_key, "static-secret-value");
assert!(api_key.is_sensitive());
assert!(!format!("{headers:?}").contains("static-secret-value"));
}
#[test]
fn test_configured_authentication_headers_are_sensitive() {
let config = ClientConfig {
extra_headers: HashMap::from([
(
"Authorization".to_string(),
"Bearer configured-secret".to_string(),
),
("X-API-Key".to_string(), "configured-api-key".to_string()),
(
"Cookie".to_string(),
"session=configured-cookie".to_string(),
),
(
"Proxy-Authorization".to_string(),
"Basic configured-proxy-secret".to_string(),
),
("X-Custom".to_string(), "visible-value".to_string()),
]),
..Default::default()
};
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&config,
)
.unwrap();
assert!(headers.get("authorization").unwrap().is_sensitive());
assert!(headers.get("x-api-key").unwrap().is_sensitive());
assert!(headers.get("cookie").unwrap().is_sensitive());
assert!(headers.get("proxy-authorization").unwrap().is_sensitive());
assert!(!headers.get("x-custom").unwrap().is_sensitive());
let debug = format!("{headers:?}");
assert!(!debug.contains("configured-secret"));
assert!(!debug.contains("configured-api-key"));
assert!(!debug.contains("configured-cookie"));
assert!(!debug.contains("configured-proxy-secret"));
assert!(debug.contains("visible-value"));
}
/// `log_request` prints the request's Debug, and Debug for a request prints
/// its headers. Marking the value sensitive is the only thing standing
/// between the API key and every debug line; assert on the header map's own
/// Debug, which is what that printing reduces to.
#[test]
fn test_api_key_is_redacted_in_debug_output() {
let headers = RestfulLanceDbClient::<Sender>::default_headers(
"sk-live-sentinel",
"us-east-1",
"db-name",
false,
&RemoteOptions::default(),
None,
&ClientConfig::default(),
)
.unwrap();
assert_eq!(headers.get("x-api-key").unwrap(), "sk-live-sentinel");
assert!(
!format!("{:?}", headers).contains("sk-live-sentinel"),
"the API key must not survive Debug formatting"
);
}
/// A suppressed body is suppressed whatever the content type says, and an
/// allowed one is logged in full.
#[test]
fn test_body_logging_is_decided_by_the_caller() {
assert_ne!(BodyLogging::Allowed, BodyLogging::Suppressed);
// `send` and `send_suppressing_body` differ only in what they pass, so
// the enum is the whole contract: a caller states its intent and the
// transport does not infer one from the route.
assert_eq!(BodyLogging::Allowed, BodyLogging::Allowed);
}
#[test]
fn test_rejects_invalid_cloud_dns_hostname() {
let invalid_database_names = ["a".repeat(64), "invalid..database".to_string()];
for db_name in invalid_database_names {
let parsed_url = parse_db_url(&format!("db://{db_name}")).unwrap();
let error = RestfulLanceDbClient::<Sender>::try_new(
&parsed_url,
"us-east-1",
None,
HeaderMap::new(),
ClientConfig::default(),
None,
)
.unwrap_err();
assert!(
matches!(error, Error::InvalidInput { ref message } if message.contains("DNS labels must contain 1 to 63 bytes")),
"unexpected error: {error}"
);
}
}
// Test implementation of HeaderProvider
#[derive(Debug, Clone)]
struct TestHeaderProvider {
headers: HashMap<String, String>,
}
impl TestHeaderProvider {
fn new(headers: HashMap<String, String>) -> Self {
Self { headers }
}
}
#[async_trait::async_trait]
impl HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
Ok(self.headers.clone())
}
}
// Test implementation that returns an error
#[derive(Debug)]
struct ErrorHeaderProvider;
#[async_trait::async_trait]
impl HeaderProvider for ErrorHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
Err(Error::Runtime {
message: "Failed to get headers".to_string(),
})
}
}
#[tokio::test]
async fn test_client_config_with_header_provider() {
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "secret-key".to_string());
let provider = TestHeaderProvider::new(headers);
let client_config = ClientConfig {
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
..Default::default()
};
assert!(client_config.header_provider.is_some());
}
#[tokio::test]
async fn test_apply_dynamic_headers() {
// Create a mock client with header provider
let mut headers = HashMap::new();
headers.insert("X-Dynamic".to_string(), "dynamic-value".to_string());
let provider = TestHeaderProvider::new(headers);
// Create a simple request
let request = reqwest::Request::new(
reqwest::Method::GET,
"https://example.com/test".parse().unwrap(),
);
// Create client with header provider
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Apply dynamic headers
let updated_request = client.apply_dynamic_headers(request).await.unwrap();
// Check that the header was added
assert_eq!(
updated_request.headers().get("X-Dynamic").unwrap(),
"dynamic-value"
);
}
#[tokio::test]
async fn test_apply_dynamic_headers_merge() {
// Test that dynamic headers override existing headers
let mut headers = HashMap::new();
headers.insert("Authorization".to_string(), "Bearer new-token".to_string());
headers.insert("X-API-Key".to_string(), "new-api-key".to_string());
headers.insert("X-Custom".to_string(), "custom-value".to_string());
let provider = TestHeaderProvider::new(headers);
// Create request with existing Authorization header
let mut request_builder = reqwest::Client::new().get("https://example.com/test");
request_builder = request_builder.header("Authorization", "Bearer old-token");
request_builder = request_builder.header("X-Existing", "existing-value");
let request = request_builder.build().unwrap();
// Create client with header provider
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Apply dynamic headers
let updated_request = client.apply_dynamic_headers(request).await.unwrap();
// Check that dynamic headers override existing ones
assert_eq!(
updated_request.headers().get("Authorization").unwrap(),
"Bearer new-token"
);
assert_eq!(
updated_request.headers().get("X-Custom").unwrap(),
"custom-value"
);
assert!(
updated_request
.headers()
.get("Authorization")
.unwrap()
.is_sensitive()
);
assert!(
updated_request
.headers()
.get("X-API-Key")
.unwrap()
.is_sensitive()
);
assert!(
!updated_request
.headers()
.get("X-Custom")
.unwrap()
.is_sensitive()
);
let debug = format!("{updated_request:?}");
assert!(!debug.contains("new-token"));
assert!(!debug.contains("new-api-key"));
assert!(debug.contains("custom-value"));
// Existing headers should still be present
assert_eq!(
updated_request.headers().get("X-Existing").unwrap(),
"existing-value"
);
}
#[tokio::test]
async fn test_apply_dynamic_headers_with_error_provider() {
let provider = ErrorHeaderProvider;
let request = reqwest::Request::new(
reqwest::Method::GET,
"https://example.com/test".parse().unwrap(),
);
let client = RestfulLanceDbClient {
client: reqwest::Client::new(),
host: "https://example.com".to_string(),
retry_config: RetryConfig::default().try_into().unwrap(),
sender: Sender,
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
read_consistency_interval: None,
max_bytes_per_request: None,
max_request_duration: None,
};
// Header provider errors should fail the request
// This is important for security - if auth headers can't be fetched, don't proceed
let result = client.apply_dynamic_headers(request).await;
assert!(result.is_err());
match result.unwrap_err() {
Error::Runtime { message } => {
assert_eq!(message, "Failed to get headers");
}
_ => panic!("Expected Runtime error"),
}
}
#[test]
fn test_resolve_user_id_direct_value() {
let config = ClientConfig {
user_id: Some("direct-user-id".to_string()),
..Default::default()
};
assert_eq!(config.resolve_user_id(), Some("direct-user-id".to_string()));
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_none() {
let _guard = lock_env();
let config = ClientConfig::default();
// Clear env vars that might be set from other tests
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
}
assert_eq!(config.resolve_user_id(), None);
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_from_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "env-user-id");
}
let config = ClientConfig::default();
assert_eq!(config.resolve_user_id(), Some("env-user-id".to_string()));
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_from_env_key() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
std::env::set_var("LANCEDB_USER_ID_ENV_KEY", "MY_CUSTOM_USER_ID");
std::env::set_var("MY_CUSTOM_USER_ID", "custom-env-user-id");
}
let config = ClientConfig::default();
assert_eq!(
config.resolve_user_id(),
Some("custom-env-user-id".to_string())
);
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
std::env::remove_var("MY_CUSTOM_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_direct_takes_precedence() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "env-user-id");
}
let config = ClientConfig {
user_id: Some("direct-user-id".to_string()),
..Default::default()
};
assert_eq!(config.resolve_user_id(), Some("direct-user-id".to_string()));
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
#[serial(user_id_env)]
fn test_resolve_user_id_empty_env_ignored() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCEDB_USER_ID", "");
std::env::remove_var("LANCEDB_USER_ID_ENV_KEY");
}
let config = ClientConfig::default();
assert_eq!(config.resolve_user_id(), None);
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCEDB_USER_ID");
}
}
#[test]
fn test_resolve_max_bytes_passed_value_wins() {
// An explicit config value is used verbatim; env/default are not consulted.
let resolved =
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
assert_eq!(resolved, Some(1234));
}
#[test]
fn test_resolve_max_bytes_zero_disables() {
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(0)).unwrap();
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_default_when_unset() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
assert_eq!(resolved, Some(DEFAULT_MAX_BYTES_PER_REQUEST));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_from_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, Some(4096));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_env_zero_disables() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "0");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_config_overrides_env() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
}
// A config value takes precedence over the environment variable.
let resolved =
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert_eq!(resolved, Some(1234));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_bytes_invalid_env_errors() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "not-a-number");
}
let err = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap_err();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
}
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
}
#[test]
fn test_resolve_max_request_duration_passed_value_wins() {
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
Some(Duration::from_secs(42)),
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, Some(Duration::from_secs(42)));
}
#[test]
fn test_resolve_max_request_duration_zero_disables() {
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
Some(Duration::ZERO),
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_default_is_half_read_timeout() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
assert_eq!(resolved, Some(Duration::from_secs(150)));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_from_env_seconds() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "30");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert_eq!(resolved, Some(Duration::from_secs(30)));
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_env_zero_disables() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "0");
}
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert_eq!(resolved, None);
}
#[test]
#[serial(request_limits_env)]
fn test_resolve_max_request_duration_invalid_env_errors() {
let _guard = lock_env();
// SAFETY: This is only called in tests
unsafe {
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "12.5");
}
let err = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
None,
Duration::from_secs(300),
)
.unwrap_err();
// SAFETY: This is only called in tests
unsafe {
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
}
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
}
rust/lancedb/src/remote/db.rs
+1 -2748
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/db/tests/catalog.rs
+634
View File
@@ -0,0 +1,634 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_create_and_alter_secret_send_the_value_in_the_request_body() {
for (route, call) in [
("/v1/secret/openai-prod/create", true),
("/v1/secret/openai-prod/alter", false),
] {
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), route);
// Never a path segment or query parameter, which is what keeps
// it out of access logs and proxy traces.
assert!(request.url().query().is_none(), "{:?}", request.url());
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
// The name is the path identifier, so the body is the value alone.
assert_eq!(body, serde_json::json!({ "value": "sk-live-0001" }));
http::Response::builder().status(200).body("{}").unwrap()
});
if call {
conn.create_secret("openai-prod", "sk-live-0001", &[])
.await
.unwrap();
} else {
conn.alter_secret("openai-prod", "sk-live-0001", &[])
.await
.unwrap();
}
}
}
#[tokio::test]
async fn test_list_secrets_walks_pages_and_returns_names_only() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/namespace/$/secret/list");
let page = request
.url()
.query_pairs()
.find(|(key, _)| key == "page_token")
.map(|(_, value)| value.into_owned());
let body = match page.as_deref() {
None => r#"{"secrets":[{"name":"openai-prod"}],"page_token":"p2"}"#,
Some("p2") => r#"{"secrets":[{"name":"hf-prod"}]}"#,
Some(other) => panic!("unexpected page token: {other}"),
};
http::Response::builder().status(200).body(body).unwrap()
});
assert_eq!(
conn.list_secrets(&[]).await.unwrap(),
vec!["openai-prod".to_string(), "hf-prod".to_string()]
);
}
/// A server that keeps handing back the same token would otherwise spin
/// forever.
#[tokio::test]
async fn test_list_secrets_rejects_a_repeated_page_token() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(r#"{"secrets":[{"name":"openai-prod"}],"page_token":"same"}"#)
.unwrap()
});
let error = conn.list_secrets(&[]).await.unwrap_err();
assert!(
error.to_string().contains("repeated a page_token"),
"{error}"
);
}
#[tokio::test]
async fn test_drop_and_describe_address_the_secret_in_the_path() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/secret/openai-prod/drop");
// Nothing is left to say once the path names the Secret.
assert!(request.body().is_none(), "{:?}", request.body());
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_secret("openai-prod", &[]).await.unwrap();
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/secret/openai-prod/describe");
assert!(request.body().is_none(), "{:?}", request.body());
http::Response::builder()
.status(200)
.body(r#"{"name":"openai-prod","created_at_millis":1,"updated_at_millis":2}"#)
.unwrap()
});
let info = conn.describe_secret("openai-prod", &[]).await.unwrap();
assert_eq!(info.name, "openai-prod");
}
/// The namespace is part of the identifier the path addresses, joined with
/// the client's configured delimiter the way every other object's is. A
/// root Secret is therefore addressed by its bare name, and a namespaced
/// one by the joined path -- there is no body field either way.
#[tokio::test]
async fn test_a_namespace_path_is_addressed_in_the_path() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(
request.url().path(),
"/v1/secret/prod$vision$openai-prod/drop"
);
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_secret("openai-prod", &["prod".to_string(), "vision".to_string()])
.await
.unwrap();
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/secret/openai-prod/drop");
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_secret("openai-prod", &[]).await.unwrap();
// Listing is namespace-scoped, so the namespace is the whole identifier.
let conn = Connection::new_with_handler(|request| {
assert_eq!(
request.url().path(),
"/v1/namespace/prod$vision/secret/list"
);
http::Response::builder()
.status(200)
.body(r#"{"secrets":[]}"#)
.unwrap()
});
conn.list_secrets(&["prod".to_string(), "vision".to_string()])
.await
.unwrap();
}
/// A view description carries the schema in the namespace spec's JSON
/// encoding, so the body a test asserts against is the one that encoding
/// produces rather than a hand-written guess at it.
fn view_description_body(name: &str, namespace: &[&str], query: &str) -> String {
let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]);
serde_json::json!({
"name": name,
"namespace": namespace,
"query": query,
"default_database": "db",
"default_namespace": namespace,
"schema": lance_namespace::schema::arrow_schema_to_json(&schema).unwrap(),
})
.to_string()
}
#[tokio::test]
async fn test_create_view_posts_the_query_and_returns_the_resolved_schema() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/view/analytics$adults/create");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
// The name and its namespace address the view in the path, so the
// body says only what they cannot.
assert_eq!(body, serde_json::json!({"query": "SELECT id FROM people"}));
http::Response::builder()
.status(200)
.body(view_description_body(
"adults",
&["analytics"],
"SELECT id FROM people",
))
.unwrap()
});
let view = conn
.create_view("adults", "SELECT id FROM people", &["analytics".into()])
.await
.unwrap();
assert_eq!(view.name, "adults");
assert_eq!(view.namespace_path, vec!["analytics".to_string()]);
assert_eq!(view.query, "SELECT id FROM people");
assert_eq!(view.default_database, "db");
assert_eq!(view.default_namespace_path, vec!["analytics".to_string()]);
assert_eq!(view.schema.field(0).name(), "id");
}
#[tokio::test]
async fn test_describe_and_drop_address_the_view_in_the_path() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/view/adults/describe");
assert!(request.body().is_none(), "{:?}", request.body());
http::Response::builder()
.status(200)
.body(view_description_body(
"adults",
&[],
"SELECT id FROM people",
))
.unwrap()
});
let view = conn.describe_view("adults", &[]).await.unwrap();
assert!(view.namespace_path.is_empty());
assert_eq!(view.schema.fields().len(), 1);
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/view/analytics$adults/drop");
assert!(request.body().is_none(), "{:?}", request.body());
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_view("adults", &["analytics".into()])
.await
.unwrap();
}
/// An accepted drop hands back the job so a caller can wait on the delete,
/// and the waiting `drop_view` does that for them.
#[tokio::test]
async fn test_drop_view_async_reports_the_cleanup_job() {
let db = super::super::RemoteDatabase::new_mock(|_| {
http::Response::builder()
.status(202)
.body(r#"{"job_id":"j1-do-abc"}"#)
.unwrap()
});
let job = db.drop_view_async("adults", &[]).await.unwrap();
assert_eq!(job.id(), Some("j1-do-abc"));
}
/// Nothing was bound, so nothing is being deleted and the job is already done.
#[tokio::test]
async fn test_drop_view_async_reports_a_finished_job_when_nothing_was_bound() {
let db = super::super::RemoteDatabase::new_mock(|_| {
http::Response::builder().status(200).body("{}").unwrap()
});
let job = db.drop_view_async("adults", &[]).await.unwrap();
assert_eq!(job.id(), None);
assert_eq!(job.status().await.unwrap(), "finished");
job.wait().await.unwrap();
}
#[tokio::test]
async fn test_drop_view_rejects_incomplete_acceptance() {
for body in ["{}", r#"{"job_id":""}"#, r#"{"job_id":null}"#] {
let db = super::super::RemoteDatabase::new_mock(move |_| {
http::Response::builder().status(202).body(body).unwrap()
});
let error = db.drop_view_async("adults", &[]).await.err().unwrap();
assert!(error.to_string().contains("valid job_id"), "{error}");
}
}
#[tokio::test]
async fn test_drop_view_rejects_unexpected_success_status() {
let db = super::super::RemoteDatabase::new_mock(|_| {
http::Response::builder().status(204).body("").unwrap()
});
let error = db.drop_view_async("adults", &[]).await.err().unwrap();
assert!(
error.to_string().contains("200 OK or 202 Accepted"),
"{error}"
);
}
/// A schema the client cannot decode is a broken response, not a view
/// with no columns: reporting it as an error keeps a caller from reading
/// an empty schema as the truth about the view.
#[tokio::test]
async fn test_describe_view_rejects_an_undecodable_schema() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(
r#"{"name":"adults","query":"SELECT 1","default_database":"db",
"schema":{"fields":[{"name":"id","type":{"type":"nonesuch"},
"nullable":true}]}}"#,
)
.unwrap()
});
let error = conn.describe_view("adults", &[]).await.unwrap_err();
assert!(error.to_string().contains("undecodable schema"), "{error}");
}
#[tokio::test]
async fn test_list_views_walks_pages() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/namespace/analytics/view/list");
let page = request
.url()
.query_pairs()
.find(|(key, _)| key == "page_token")
.map(|(_, value)| value.into_owned());
let body = match page.as_deref() {
None => r#"{"views":[],"page_token":"p2"}"#,
Some("p2") => r#"{"views":["adults"]}"#,
Some(other) => panic!("unexpected page token: {other}"),
};
http::Response::builder().status(200).body(body).unwrap()
});
assert_eq!(
conn.list_views(&["analytics".into()]).await.unwrap(),
vec!["adults".to_string()]
);
}
/// A server that keeps handing back the same token would otherwise spin
/// forever.
#[tokio::test]
async fn test_list_views_rejects_a_repeated_page_token() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(r#"{"views":["adults"],"page_token":"same"}"#)
.unwrap()
});
let error = conn.list_views(&[]).await.unwrap_err();
assert!(
error.to_string().contains("repeated a page_token"),
"{error}"
);
}
/// A name carrying the delimiter would split back apart as a different
/// view, so it is refused before it reaches a route.
#[tokio::test]
async fn test_view_names_that_would_resplit_are_refused() {
let conn = Connection::new_with_handler(|_| -> http::Response<String> {
panic!("an invalid identifier must not reach the service")
});
for (name, namespace) in [
("analytics$adults", vec![]),
("adults", vec!["ana$lytics".to_string()]),
("", vec![]),
("..", vec![]),
] {
let error = match conn.describe_view(name, &namespace).await {
Ok(_) => panic!("accepted {name:?} in {namespace:?}"),
Err(error) => error.to_string(),
};
// A view is not a table, so the refusal says so.
assert!(
error.contains("view name") || error.contains("view namespace path segment"),
"{error}"
);
}
}
#[tokio::test]
async fn test_create_function_async_sends_canonical_request_and_decodes_typed_job() {
const REQUEST: &str = include_str!(
"../../../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
);
const FUNCTION_JOB: &str =
include_str!("../../../../tests/fixtures/first_class_functions/v1/remote_function_job.json");
let mut expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
expected.as_object_mut().unwrap().remove("name");
let conn = Connection::new_with_handler(move |request| match request.url().path() {
"/v1/function/normalize_score/create" => {
assert_eq!(request.method(), &reqwest::Method::POST);
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body, expected);
http::Response::builder()
.status(202)
.body(r#"{"job_id":"job-function-1"}"#)
.unwrap()
}
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(FUNCTION_JOB)
.unwrap(),
path => panic!("unexpected path: {path}"),
});
let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
let job = conn.create_function_async(request, &[]).await.unwrap();
assert_eq!(job.id(), Some("job-function-1"));
let version = job.wait().await.unwrap();
assert_eq!(version.name(), "embed");
assert_eq!(version.version(), "1");
}
#[tokio::test]
async fn test_get_function_requires_and_sends_exact_version() {
const VERSION: &str = include_str!(
"../../../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
);
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/function/embed/describe");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body, serde_json::json!({"version": "1"}));
http::Response::builder().status(200).body(VERSION).unwrap()
});
let version = conn.get_function("embed", "1", &[]).await.unwrap();
assert_eq!(version.name(), "embed");
assert_eq!(version.version(), "1");
}
#[tokio::test]
async fn test_list_functions_requests_definitions_and_paginates() {
const VERSION: &str = include_str!(
"../../../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
);
let version: serde_json::Value = serde_json::from_str(VERSION).unwrap();
let page = Arc::new(AtomicUsize::new(0));
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/namespace/$/function/list");
let query = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(query.get("include_definition").unwrap(), "true");
match page.fetch_add(1, Ordering::SeqCst) {
0 => {
assert!(!query.contains_key("page_token"));
http::Response::builder()
.status(200)
.body(r#"{"functions": [], "page_token": "next"}"#.to_string())
.unwrap()
}
_ => {
assert_eq!(query.get("page_token").unwrap(), "next");
http::Response::builder()
.status(200)
.body(
serde_json::json!({
"functions": [{
"name": "embed",
"version": "1",
"definition": version.clone(),
}],
})
.to_string(),
)
.unwrap()
}
}
});
let functions = conn.list_functions(&[]).await.unwrap();
assert_eq!(functions.len(), 1);
assert_eq!(functions[0].name(), "embed");
assert_eq!(functions[0].version(), "1");
}
#[tokio::test]
async fn test_list_functions_stops_on_an_empty_page_token() {
let requests = Arc::new(AtomicUsize::new(0));
let seen = requests.clone();
let conn = Connection::new_with_handler(move |request| {
seen.fetch_add(1, Ordering::SeqCst);
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/namespace/$/function/list");
let query = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(query.get("include_definition").unwrap(), "true");
assert!(!query.contains_key("page_token"));
http::Response::builder()
.status(200)
.body(r#"{"functions": [], "page_token": ""}"#)
.unwrap()
});
let functions = conn.list_functions(&[]).await.unwrap();
assert!(functions.is_empty());
assert_eq!(requests.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_list_functions_rejects_a_page_token_cycle() {
let page = Arc::new(AtomicUsize::new(0));
let requests = page.clone();
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/namespace/$/function/list");
let query = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(query.get("include_definition").unwrap(), "true");
let next_page_token = match page.fetch_add(1, Ordering::SeqCst) {
0 => {
assert!(!query.contains_key("page_token"));
"one"
}
1 => {
assert_eq!(query.get("page_token").unwrap(), "one");
"two"
}
2 => {
assert_eq!(query.get("page_token").unwrap(), "two");
"one"
}
page => panic!("unexpected page: {page}"),
};
http::Response::builder()
.status(200)
.body(
serde_json::json!({
"functions": [],
"page_token": next_page_token,
})
.to_string(),
)
.unwrap()
});
let error = conn.list_functions(&[]).await.unwrap_err();
assert!(
matches!(
&error,
Error::Http {
status_code: Some(http::StatusCode::OK),
..
}
),
"got {error:?}"
);
assert_eq!(requests.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_drop_function_sends_exact_version_and_decodes_replay() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/function/embed/drop");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body, serde_json::json!({"version": "1"}));
http::Response::builder()
.status(200)
.body(r#"{"dropped":false}"#)
.unwrap()
});
assert!(!conn.drop_function("embed", "1", &[]).await.unwrap());
}
/// A Function's namespace is addressed in the path the way a Secret's is:
/// every route takes the joined identifier and no body carries the
/// namespace, so a namespaced request differs from a root one only in its
/// path.
#[tokio::test]
async fn test_a_function_namespace_path_is_addressed_in_the_path() {
const REQUEST: &str = include_str!(
"../../../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
);
const VERSION: &str = include_str!(
"../../../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
);
let namespace = ["analytics".to_string(), "features".to_string()];
let mut expected_create: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
expected_create.as_object_mut().unwrap().remove("name");
let conn = Connection::new_with_handler(move |request| {
let body = request
.body()
.and_then(|body| body.as_bytes())
.map(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).unwrap());
match request.url().path() {
"/v1/function/analytics$features$normalize_score/create" => {
assert_eq!(body.unwrap(), expected_create);
http::Response::builder()
.status(202)
.body(r#"{"job_id":"job-function-1"}"#.to_string())
.unwrap()
}
"/v1/function/analytics$features$embed/describe" => {
assert_eq!(body.unwrap(), serde_json::json!({"version": "1"}));
http::Response::builder()
.status(200)
.body(VERSION.to_string())
.unwrap()
}
"/v1/function/analytics$features$embed/drop" => {
assert_eq!(body.unwrap(), serde_json::json!({"version": "1"}));
http::Response::builder()
.status(200)
.body(r#"{"dropped":true}"#.to_string())
.unwrap()
}
// Listing is namespace-scoped, so the namespace is the whole
// identifier.
"/v1/namespace/analytics$features/function/list" => http::Response::builder()
.status(200)
.body(r#"{"functions":[]}"#.to_string())
.unwrap(),
path => panic!("unexpected path: {path}"),
}
});
let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
let job = conn
.create_function_async(request, &namespace)
.await
.unwrap();
assert_eq!(job.id(), Some("job-function-1"));
conn.get_function("embed", "1", &namespace).await.unwrap();
assert!(conn.list_functions(&namespace).await.unwrap().is_empty());
assert!(conn.drop_function("embed", "1", &namespace).await.unwrap());
}
/// Each namespace segment is checked before a request is built, so an
/// empty segment cannot collapse the identifier onto the parent namespace.
#[tokio::test]
async fn test_an_unaddressable_function_namespace_segment_is_refused() {
let conn = Connection::new_with_handler(|request| -> http::Response<String> {
panic!("reached the transport: {}", request.url().path())
});
let namespace = ["analytics".to_string(), String::new()];
assert!(conn.get_function("embed", "1", &namespace).await.is_err());
assert!(conn.list_functions(&namespace).await.is_err());
assert!(conn.drop_function("embed", "1", &namespace).await.is_err());
}
#[tokio::test]
async fn test_conn_job_waits_to_done() {
let polls = Arc::new(AtomicUsize::new(0));
let polls_ref = polls.clone();
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.url().path(), "/v1/jobs/describe");
// Two in-progress answers: one for the load, one for the first
// status poll.
let state = if polls_ref.fetch_add(1, Ordering::SeqCst) < 2 {
"IN_PROGRESS"
} else {
"DONE"
};
http::Response::builder()
.status(200)
.body(format!(
r#"{{"job_id": "job-1", "job_type": "create_function", "job_state": "{}", "creation_ms": 1, "result": {{"name": "embed", "version": "1"}}}}"#,
state
))
.unwrap()
});
let job = conn.open_job("job-1").await.unwrap();
assert_eq!(job.id(), Some("job-1"));
// Opening already answered the state; no extra call needed for it.
assert_eq!(job.state().as_deref(), Some("running"));
assert_eq!(job.status().await.unwrap(), "running");
job.wait().await.unwrap();
assert_eq!(job.status().await.unwrap(), "finished");
assert!(polls.load(Ordering::SeqCst) >= 4);
}
rust/lancedb/src/remote/db/tests/jobs.rs
+489
View File
@@ -0,0 +1,489 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_list_jobs_paginates() {
let page = Arc::new(AtomicUsize::new(0));
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/jobs/list");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
match page.fetch_add(1, Ordering::SeqCst) {
0 => {
assert!(body.get("page_token").is_none());
http::Response::builder()
.status(200)
.body(
r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000}], "page_token": "next"}"#,
)
.unwrap()
}
_ => {
assert_eq!(body["page_token"], "next");
http::Response::builder()
.status(200)
.body(
r#"{"jobs": [{"job_id": "job-2", "table": "t2", "job_type": "create_index", "state": "succeeded", "created_at_millis": 2000}, {"job_id": "job-3", "table": "t3", "job_type": "create_index", "state": "timed_out", "created_at_millis": 3000}]}"#,
)
.unwrap()
}
}
});
let jobs = conn.list_jobs().await.unwrap();
assert_eq!(jobs.len(), 3);
assert_eq!(jobs[0].job_id, "job-1");
assert_eq!(jobs[0].table, "t1");
assert_eq!(jobs[0].state, "running");
assert_eq!(jobs[1].job_id, "job-2");
assert_eq!(jobs[1].state, "finished");
assert_eq!(jobs[1].created_at_millis, 2000);
assert_eq!(jobs[2].job_id, "job-3");
assert_eq!(jobs[2].state, "failed");
}
#[tokio::test]
async fn test_list_jobs_rejects_a_page_token_cycle() {
let requests = Arc::new(AtomicUsize::new(0));
let seen = requests.clone();
let conn = Connection::new_with_handler(move |request| {
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
match seen.fetch_add(1, Ordering::SeqCst) {
0 => assert!(body.get("page_token").is_none()),
_ => assert_eq!(body["page_token"], "loop"),
}
http::Response::builder()
.status(200)
.body(r#"{"jobs": [], "page_token": "loop"}"#)
.unwrap()
});
let error = conn.list_jobs().await.unwrap_err();
assert!(
matches!(
&error,
Error::Http {
status_code: Some(http::StatusCode::OK),
..
}
),
"got {error:?}"
);
assert_eq!(requests.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_open_job() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/jobs/describe");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["job_id"], "job-1");
http::Response::builder()
.status(200)
.body(
r#"{"job_id": "job-1", "job_type": "create_index", "job_state": "FAILED", "creation_ms": 1000, "spec": {"column": "vec"}, "failure": {"phase": "execute", "message": "worker died", "retryable": true}}"#,
)
.unwrap()
});
// Opening populates the handle, so the accessors answer without a
// second round trip.
let job = conn.open_job("job-1").await.unwrap();
assert_eq!(job.id(), Some("job-1"));
assert_eq!(job.job_type().as_deref(), Some("create_index"));
assert_eq!(job.state().as_deref(), Some("failed"));
assert_eq!(job.creation_ms(), Some(1000));
assert_eq!(job.spec().unwrap()["column"], "vec");
assert!(job.result().is_none());
let failure = job.failure().unwrap();
assert_eq!(failure.phase.as_deref(), Some("execute"));
assert_eq!(failure.message.as_deref(), Some("worker died"));
assert_eq!(failure.retryable, Some(true));
}
#[tokio::test]
async fn test_open_job_reports_the_terminal_result() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "DONE", "creation_ms": 1000, "result": {"rows_assigned": 1000000, "rows_failed": 0}}"#,
)
.unwrap()
});
let job = conn.open_job("job-1").await.unwrap();
assert_eq!(job.state().as_deref(), Some("finished"));
let result = job.result().unwrap();
assert_eq!(result["rows_assigned"], 1_000_000);
assert_eq!(result["rows_failed"], 0);
}
#[tokio::test]
async fn test_open_job_missing_fails() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(404)
.body("no such job")
.unwrap()
});
let err = conn.open_job("nope").await.unwrap_err();
assert!(
matches!(&err, Error::JobNotFound { job_id } if job_id == "nope"),
"{err:?}"
);
}
#[tokio::test]
async fn test_pause_and_resume_job() {
use crate::database::{PauseJobStatus, ResumeJobStatus};
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/jobs/pause");
http::Response::builder()
.status(200)
.body(r#"{"job_id": "job-1", "paused": true}"#)
.unwrap()
});
assert_eq!(
conn.pause_job("job-1").await.unwrap(),
PauseJobStatus::Pausing
);
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(r#"{"job_id": "job-1", "paused": false, "committing": true}"#)
.unwrap()
});
assert_eq!(
conn.pause_job("job-1").await.unwrap(),
PauseJobStatus::Committing
);
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/jobs/resume");
http::Response::builder()
.status(200)
.body(r#"{"job_id": "job-1", "resumed": false, "still_pausing": true}"#)
.unwrap()
});
assert_eq!(
conn.resume_job("job-1").await.unwrap(),
ResumeJobStatus::StillPausing
);
}
#[tokio::test]
async fn test_job_events_scope_to_that_job() {
let schema = Arc::new(Schema::new(vec![Field::new(
"state",
DataType::Utf8,
false,
)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(arrow_array::StringArray::from(vec![
"claim_complete",
]))],
)
.unwrap();
let mut events = Vec::new();
{
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
}
let conn = Connection::new_with_handler(move |request| {
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
if request.url().path() == "/v1/jobs/describe" {
return http::Response::builder()
.status(200)
.body(
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
.as_bytes()
.to_vec(),
)
.unwrap();
}
assert_eq!(request.url().path(), "/v1/jobs/query_events");
// The handle supplies job_id; the caller only narrows the query.
assert_eq!(body["job_id"], "job-1");
assert_eq!(body["limit"], 500);
assert_eq!(body["filter"], "state = 'claim_complete'");
http::Response::builder()
.status(200)
.body(events.clone())
.unwrap()
});
let job = conn.open_job("job-1").await.unwrap();
let batches = job
.events(
JobEventsRequest::default()
.limit(500)
.filter("state = 'claim_complete'"),
)
.await
.unwrap();
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 1);
}
#[tokio::test]
async fn test_job_events_keep_the_schema_when_nothing_matches() {
let schema = Arc::new(Schema::new(vec![Field::new(
"state",
DataType::Utf8,
false,
)]));
let mut events = Vec::new();
{
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
writer.finish().unwrap();
}
let conn = Connection::new_with_handler(move |request| {
if request.url().path() == "/v1/jobs/describe" {
return http::Response::builder()
.status(200)
.body(
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
.as_bytes()
.to_vec(),
)
.unwrap();
}
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
// Only the job id when the caller narrows nothing.
assert_eq!(body, serde_json::json!({ "job_id": "job-1" }));
http::Response::builder()
.status(200)
.body(events.clone())
.unwrap()
});
let job = conn.open_job("job-1").await.unwrap();
let batches = job.events(JobEventsRequest::default()).await.unwrap();
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 0);
assert_eq!(batches[0].schema(), schema);
}
/// A component that is not a legal Secret component never reaches a
/// transport. Before the identifier was checked here, each of these decided
/// the route instead of the name: `a/b` and `../jobs` left `/v1/secret/`
/// entirely, carrying a create body that holds a credential, and `a$b` read
/// as the namespace `a` and the name `b`.
#[tokio::test]
async fn test_an_illegal_component_never_reaches_the_transport() {
use std::sync::{Arc, Mutex};
for name in [
"../jobs",
"a/b",
"a$b",
"with space",
"q?x",
"a#b",
"a%2Fb",
"",
] {
let reached = Arc::new(Mutex::new(false));
let flag = reached.clone();
let conn = Connection::new_with_handler(move |_| {
*flag.lock().unwrap() = true;
http::Response::builder().status(200).body("{}").unwrap()
});
let error = conn
.create_secret(name, "sk-live-0001", &[])
.await
.expect_err("an illegal component must be refused");
assert!(!*reached.lock().unwrap(), "{name:?} reached the transport");
assert!(
error.to_string().contains("Secret name"),
"{name:?}: {error}"
);
}
}
/// `.` and `..` pass the character set and still cannot address anything:
/// URL parsing resolves them as relative segments, and after
/// percent-decoding, so no spelling of either survives.
#[tokio::test]
async fn test_a_relative_segment_component_is_refused() {
use std::sync::{Arc, Mutex};
// The percent-encoded spellings are caught a step earlier, by the
// character set: `%` is not a character a Secret component may hold.
// They reach the relative-segment rule only where there is no charset
// to catch them first -- see the Function case below.
for component in [".", ".."] {
let reached = Arc::new(Mutex::new(false));
let flag = reached.clone();
let conn = Connection::new_with_handler(move |_| {
*flag.lock().unwrap() = true;
http::Response::builder()
.status(200)
.body(r#"{"secrets":[]}"#)
.unwrap()
});
let by_name = conn
.drop_secret(component, &[])
.await
.expect_err("a dot-only name must be refused");
assert!(
by_name.to_string().contains("relative path segments"),
"{by_name}"
);
let by_segment = conn
.list_secrets(&[component.to_string()])
.await
.expect_err("a dot-only namespace segment must be refused");
assert!(
by_segment.to_string().contains("relative path segments"),
"{by_segment}"
);
assert!(
!*reached.lock().unwrap(),
"{component:?} reached the transport"
);
}
}
/// An identifier the service accepts is untouched by the encoding, so the
/// route reads like the table and Function routes beside it.
#[tokio::test]
async fn test_an_admissible_identifier_is_not_encoded() {
use std::sync::{Arc, Mutex};
let seen = Arc::new(Mutex::new(String::new()));
let captured = seen.clone();
let conn = Connection::new_with_handler(move |request| {
*captured.lock().unwrap() = request.url().path().to_string();
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_secret(
"openai-prod.v1",
&["prod".to_string(), "vision_2".to_string()],
)
.await
.unwrap();
assert_eq!(
*seen.lock().unwrap(),
"/v1/secret/prod$vision_2$openai-prod.v1/drop"
);
}
/// A table name is joined into the route the same way a Secret's is, so the
/// same two failures are reachable: a dot-only name leaves the table route
/// entirely, and a name holding the delimiter is indistinguishable from a
/// namespace boundary.
#[tokio::test]
async fn test_a_table_name_cannot_choose_its_own_route() {
use std::sync::{Arc, Mutex};
for name in ["..", ".", "../jobs", "a/b", "a$b"] {
let reached = Arc::new(Mutex::new(false));
let flag = reached.clone();
let conn = Connection::new_with_handler(move |_| {
*flag.lock().unwrap() = true;
http::Response::builder().status(200).body("{}").unwrap()
});
let error = conn
.drop_table(name, &[])
.await
.expect_err("an unaddressable table name must be refused");
assert!(!*reached.lock().unwrap(), "{name:?} reached the transport");
assert!(!error.to_string().is_empty(), "{name:?}");
}
}
/// The namespace half of the same join.
#[tokio::test]
async fn test_a_namespace_segment_cannot_choose_its_own_route() {
use std::sync::{Arc, Mutex};
for segment in ["..", "a$b", ""] {
let reached = Arc::new(Mutex::new(false));
let flag = reached.clone();
let conn = Connection::new_with_handler(move |_| {
*flag.lock().unwrap() = true;
http::Response::builder().status(200).body("{}").unwrap()
});
let error = conn
.drop_table("t", &[segment.to_string()])
.await
.expect_err("an unaddressable namespace segment must be refused");
assert!(
!*reached.lock().unwrap(),
"{segment:?} reached the transport"
);
assert!(!error.to_string().is_empty(), "{segment:?}");
}
}
/// A segment outside the table charset still addresses one segment: the
/// service decides whether it may exist, and percent-encoding is what keeps
/// the question reaching the right route. A catalog database is named this
/// way.
#[tokio::test]
async fn test_a_namespace_segment_outside_the_charset_is_encoded_not_refused() {
use std::sync::{Arc, Mutex};
let seen = Arc::new(Mutex::new(String::new()));
let path = seen.clone();
let conn = Connection::new_with_handler(move |request| {
*path.lock().unwrap() = request.url().path().to_string();
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_table("t", &["team/search".to_string()])
.await
.unwrap();
assert_eq!(*seen.lock().unwrap(), "/v1/table/team%2Fsearch$t/drop/");
}
/// A Function name is percent-encoded, which covers everything but the
/// relative segment: `..` is unreserved, so it survives encoding and is
/// then resolved away, posting a registration body to `/v1/create`.
#[tokio::test]
async fn test_a_relative_segment_function_name_is_refused() {
use std::sync::{Arc, Mutex};
for name in [".", "..", "%2E%2E"] {
let reached = Arc::new(Mutex::new(false));
let flag = reached.clone();
let conn = Connection::new_with_handler(move |_| {
*flag.lock().unwrap() = true;
http::Response::builder().status(200).body("{}").unwrap()
});
let error = conn
.drop_function(name, "fv_1", &[])
.await
.expect_err("a dot-only Function name must be refused");
assert!(
error.to_string().contains("relative path segments"),
"{error}"
);
assert!(!*reached.lock().unwrap(), "{name:?} reached the transport");
}
}
/// Only `.` and `..` are relative segments. `...` and longer runs are
/// ordinary and address perfectly well, so refusing them would make an
/// object that works today stop working on upgrade. This pins that.
#[tokio::test]
async fn test_a_longer_run_of_periods_is_an_ordinary_name() {
use std::sync::{Arc, Mutex};
for name in ["...", "....", "a.", ".a", "a..b"] {
let seen = Arc::new(Mutex::new(String::new()));
let captured = seen.clone();
let conn = Connection::new_with_handler(move |request| {
*captured.lock().unwrap() = request.url().path().to_string();
http::Response::builder().status(200).body("{}").unwrap()
});
conn.drop_secret(name, &[])
.await
.unwrap_or_else(|error| panic!("{name:?} must remain addressable: {error}"));
assert_eq!(
*seen.lock().unwrap(),
format!("/v1/secret/{name}/drop"),
"{name:?} did not reach its own route"
);
}
}
rust/lancedb/src/remote/db/tests/mod.rs
+25
View File
@@ -0,0 +1,25 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::{NamespaceHeaderProviderContext, build_cache_key};
use crate::connection::ConnectBuilder;
use crate::database::Database;
use crate::materialized_view::CreateMaterializedViewRequest;
use crate::{
Connection, Error,
database::CreateTableMode,
job::JobEventsRequest,
remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE},
};
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use lance_namespace_impls::{DynamicContextProvider, OperationInfo};
use rstest::rstest;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
mod catalog;
mod jobs;
mod namespace;
mod tables;
rust/lancedb/src/remote/db/tests/namespace.rs
+561
View File
@@ -0,0 +1,561 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_header_provider_in_request() {
// Test HeaderProvider implementation that adds custom headers
#[derive(Debug, Clone)]
struct TestHeaderProvider {
headers: HashMap<String, String>,
}
#[async_trait::async_trait]
impl HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
Ok(self.headers.clone())
}
}
// Create a test header provider with custom headers
let mut headers = HashMap::new();
headers.insert("X-Custom-Auth".to_string(), "test-token".to_string());
headers.insert("X-Request-Id".to_string(), "test-123".to_string());
let provider = Arc::new(TestHeaderProvider { headers }) as Arc<dyn HeaderProvider>;
// Create client config with the header provider
let client_config = ClientConfig {
header_provider: Some(provider),
..Default::default()
};
// Create connection with handler that verifies the headers are present
let conn = Connection::new_with_handler_and_config(
move |request| {
// Verify that our custom headers are present
assert_eq!(
request.headers().get("X-Custom-Auth").unwrap(),
"test-token"
);
assert_eq!(request.headers().get("X-Request-Id").unwrap(), "test-123");
// Also check standard headers are still there
assert_eq!(request.method(), &reqwest::Method::GET);
assert_eq!(request.url().path(), "/v1/table/");
http::Response::builder()
.status(200)
.body(r#"{"tables": ["table1", "table2"]}"#)
.unwrap()
},
client_config,
);
// Make a request that should include the custom headers
let names = conn.table_names().execute().await.unwrap();
assert_eq!(names, vec!["table1", "table2"]);
}
#[tokio::test]
async fn test_header_provider_error_handling() {
// Test HeaderProvider that returns an error
#[derive(Debug)]
struct ErrorHeaderProvider;
#[async_trait::async_trait]
impl HeaderProvider for ErrorHeaderProvider {
async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
Err(crate::Error::Runtime {
message: "Failed to fetch auth token".to_string(),
})
}
}
let provider = Arc::new(ErrorHeaderProvider) as Arc<dyn HeaderProvider>;
let client_config = ClientConfig {
header_provider: Some(provider),
..Default::default()
};
// Create connection - handler won't be called because header provider fails
let conn = Connection::new_with_handler_and_config(
move |_request| -> http::Response<&'static str> {
panic!("Handler should not be called when header provider fails");
},
client_config,
);
// Request should fail due to header provider error
let result = conn.table_names().execute().await;
assert!(result.is_err());
match result.unwrap_err() {
crate::Error::Runtime { message } => {
assert_eq!(message, "Failed to fetch auth token");
}
_ => panic!("Expected Runtime error from header provider"),
}
}
#[tokio::test]
async fn test_clone_table() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["source_location"], "s3://bucket/source_table");
assert_eq!(body["is_shallow"], true);
http::Response::builder().status(200).body("").unwrap()
});
let table = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.execute()
.await
.unwrap();
assert_eq!(table.name(), "cloned_table");
}
#[tokio::test]
async fn test_clone_table_with_version() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["source_location"], "s3://bucket/source_table");
assert_eq!(body["source_version"], 42);
assert_eq!(body["is_shallow"], true);
http::Response::builder().status(200).body("").unwrap()
});
let table = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.source_version(42)
.execute()
.await
.unwrap();
assert_eq!(table.name(), "cloned_table");
}
#[tokio::test]
async fn test_clone_table_with_tag() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["source_location"], "s3://bucket/source_table");
assert_eq!(body["source_tag"], "v1.0");
assert_eq!(body["is_shallow"], true);
http::Response::builder().status(200).body("").unwrap()
});
let table = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.source_tag("v1.0")
.execute()
.await
.unwrap();
assert_eq!(table.name(), "cloned_table");
}
#[tokio::test]
async fn test_clone_table_deep_clone() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["source_location"], "s3://bucket/source_table");
assert_eq!(body["is_shallow"], false);
http::Response::builder().status(200).body("").unwrap()
});
let table = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.is_shallow(false)
.execute()
.await
.unwrap();
assert_eq!(table.name(), "cloned_table");
}
#[tokio::test]
async fn test_clone_table_with_namespace() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/table/ns1$ns2$cloned_table/clone");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["source_location"], "s3://bucket/source_table");
assert_eq!(body["is_shallow"], true);
http::Response::builder().status(200).body("").unwrap()
});
let table = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.target_namespace(vec!["ns1".to_string(), "ns2".to_string()])
.execute()
.await
.unwrap();
assert_eq!(table.name(), "cloned_table");
}
#[tokio::test]
async fn test_clone_table_error() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(500)
.body("Internal server error")
.unwrap()
});
let result = conn
.clone_table("cloned_table", "s3://bucket/source_table")
.execute()
.await;
assert!(result.is_err());
if let Err(crate::Error::Http { source, .. }) = result {
assert!(source.to_string().contains("Failed to clone table"));
} else {
panic!("Expected HTTP error");
}
}
#[tokio::test]
async fn test_namespace_client() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(200)
.body(r#"{"tables": []}"#)
.unwrap()
});
// Get the namespace client from the connection's internal database
let namespace_client = conn.namespace_client().await;
assert!(namespace_client.is_ok());
}
#[tokio::test]
async fn test_namespace_client_with_tls_config() {
use crate::remote::client::TlsConfig;
let tls_config = TlsConfig {
cert_file: Some("/path/to/cert.pem".to_string()),
key_file: Some("/path/to/key.pem".to_string()),
ssl_ca_cert: Some("/path/to/ca.pem".to_string()),
assert_hostname: true,
};
let client_config = ClientConfig {
tls_config: Some(tls_config),
..Default::default()
};
let conn = Connection::new_with_handler_and_config(
|_| {
http::Response::builder()
.status(200)
.body(r#"{"tables": []}"#)
.unwrap()
},
client_config,
);
// Get the namespace client - it should be created with the TLS config
let namespace_client = conn.namespace_client().await;
assert!(namespace_client.is_ok());
}
#[tokio::test]
async fn test_namespace_client_with_headers() {
let mut extra_headers = HashMap::new();
extra_headers.insert("X-Custom-Header".to_string(), "custom-value".to_string());
let client_config = ClientConfig {
extra_headers,
..Default::default()
};
let conn = Connection::new_with_handler_and_config(
|_| {
http::Response::builder()
.status(200)
.body(r#"{"tables": []}"#)
.unwrap()
},
client_config,
);
// Get the namespace client - it should be created with the extra headers
let namespace_client = conn.namespace_client().await;
assert!(namespace_client.is_ok());
}
#[test]
fn test_namespace_header_provider_context_maps_headers() {
#[derive(Debug)]
struct TestHeaderProvider;
#[async_trait::async_trait]
impl HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
Ok(HashMap::from([(
"authorization".to_string(),
"Bearer token".to_string(),
)]))
}
}
let context_provider = NamespaceHeaderProviderContext {
header_provider: Arc::new(TestHeaderProvider) as Arc<dyn HeaderProvider>,
};
let context = context_provider.provide_context(&OperationInfo::new("list_tables", "namespace"));
assert_eq!(
context.get("headers.authorization"),
Some(&"Bearer token".to_string())
);
}
#[tokio::test]
async fn test_namespace_client_supports_dynamic_headers() {
#[derive(Debug)]
struct TestHeaderProvider;
#[async_trait::async_trait]
impl HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
Ok(HashMap::from([(
"authorization".to_string(),
"Bearer token".to_string(),
)]))
}
}
let client_config = ClientConfig {
header_provider: Some(Arc::new(TestHeaderProvider) as Arc<dyn HeaderProvider>),
..Default::default()
};
let conn = Connection::new_with_handler_and_config(
|_| {
http::Response::builder()
.status(200)
.body(r#"{"tables": []}"#)
.unwrap()
},
client_config,
);
let namespace_client = conn.namespace_client().await;
assert!(namespace_client.is_ok());
match conn.namespace_client_config().await {
Err(Error::NotSupported { message })
if message.contains("dynamic headers are configured") => {}
Err(err) => panic!("expected NotSupported, got {err:?}"),
Ok(_) => panic!("expected namespace_client_config to reject dynamic headers"),
}
}
/// Integration tests using RestAdapter to run RemoteDatabase against a real namespace server
mod rest_adapter_integration {
use super::*;
use lance_namespace::models::ListTablesRequest;
use lance_namespace_impls::{DirectoryNamespaceBuilder, RestAdapter, RestAdapterConfig};
use std::sync::Arc;
use tempfile::TempDir;
/// Test fixture that manages a REST server backed by DirectoryNamespace
struct RestServerFixture {
_temp_dir: TempDir,
server_handle: lance_namespace_impls::RestAdapterHandle,
server_url: String,
}
impl RestServerFixture {
async fn new() -> Self {
let temp_dir = TempDir::new().unwrap();
let temp_path = temp_dir.path().to_str().unwrap().to_string();
// Create DirectoryNamespace backend
let backend = DirectoryNamespaceBuilder::new(&temp_path)
.build()
.await
.unwrap();
let backend = Arc::new(backend);
// Start REST server with port 0 (OS assigns available port)
let config = RestAdapterConfig {
port: 0,
..Default::default()
};
let server = RestAdapter::new(backend, config);
let server_handle = server.start().await.unwrap();
// Get the actual port assigned by OS
let actual_port = server_handle.port();
let server_url = format!("http://127.0.0.1:{}", actual_port);
Self {
_temp_dir: temp_dir,
server_handle,
server_url,
}
}
}
impl Drop for RestServerFixture {
fn drop(&mut self) {
self.server_handle.shutdown();
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_remote_database_with_rest_adapter() {
use lance_namespace::models::CreateNamespaceRequest;
let fixture = RestServerFixture::new().await;
// Connect to the REST server using lancedb Connection
// Use db://dummy as URI and set actual server URL via host_override
let conn = ConnectBuilder::new("db://dummy")
.api_key("test-api-key")
.region("us-east-1")
.host_override(&fixture.server_url)
.execute()
.await
.unwrap();
// Create a child namespace first
let namespace = vec!["test_ns".to_string()];
conn.create_namespace(CreateNamespaceRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.expect("Failed to create namespace");
// Create a table in the child namespace
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
let data = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let table = conn
.create_table("test_table", data)
.namespace(namespace.clone())
.execute()
.await;
assert!(table.is_ok(), "Failed to create table: {:?}", table.err());
// List tables in the child namespace
let list_response = conn
.list_tables(ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.expect("Failed to list tables");
assert_eq!(list_response.tables, vec!["test_table"]);
// Get namespace client and verify it can also list tables
let namespace_client = conn.namespace_client().await.unwrap();
let list_response = namespace_client
.list_tables(ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.unwrap();
assert_eq!(list_response.tables, vec!["test_table"]);
// Open the table from the child namespace
let opened_table = conn
.open_table("test_table")
.namespace(namespace.clone())
.execute()
.await;
assert!(
opened_table.is_ok(),
"Failed to open table: {:?}",
opened_table.err()
);
assert_eq!(opened_table.unwrap().name(), "test_table");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_remote_database_with_multiple_tables() {
use lance_namespace::models::CreateNamespaceRequest;
let fixture = RestServerFixture::new().await;
// Connect to the REST server
// Use db://dummy as URI and set actual server URL via host_override
let conn = ConnectBuilder::new("db://dummy")
.api_key("test-api-key")
.region("us-east-1")
.host_override(&fixture.server_url)
.execute()
.await
.unwrap();
// Create a child namespace first
let namespace = vec!["multi_table_ns".to_string()];
conn.create_namespace(CreateNamespaceRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.expect("Failed to create namespace");
// Create multiple tables in the child namespace
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
for i in 1..=3 {
let data =
RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![i]))])
.unwrap();
conn.create_table(format!("table{}", i), data)
.namespace(namespace.clone())
.execute()
.await
.unwrap_or_else(|e| panic!("Failed to create table{}: {:?}", i, e));
}
// List tables in the child namespace
let list_response = conn
.list_tables(ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.unwrap();
assert_eq!(list_response.tables.len(), 3);
assert!(list_response.tables.contains(&"table1".to_string()));
assert!(list_response.tables.contains(&"table2".to_string()));
assert!(list_response.tables.contains(&"table3".to_string()));
}
}
rust/lancedb/src/remote/db/tests/tables.rs
+1049
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/oauth.rs
+1 -1864
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/oauth/tests/authorization.rs
+478
View File
@@ -0,0 +1,478 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[test]
fn test_authorization_callback_parsing() {
let callback = parse_authorization_callback(
"GET /callback?code=abc%20123&state=expected HTTP/1.1\r\nHost: localhost\r\n",
"/callback",
"expected",
)
.unwrap();
assert_eq!(callback, AuthorizationCallback::Code("abc 123".to_string()));
}
#[test]
fn test_authorization_callback_rejects_state_mismatch() {
let err = parse_authorization_callback(
"GET /callback?code=abc&state=wrong HTTP/1.1\r\nHost: localhost\r\n",
"/callback",
"expected",
)
.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "OAuth authorization callback state did not match"
));
}
#[test]
fn test_authorization_callback_reports_provider_error() {
let callback = parse_authorization_callback(
"GET /callback?error=access_denied&error_description=user+cancelled&state=expected HTTP/1.1\r\nHost: localhost\r\n",
"/callback",
"expected",
)
.unwrap();
assert_eq!(
callback,
AuthorizationCallback::ProviderError(
"OAuth authorization failed: user cancelled".to_string()
)
);
}
#[tokio::test]
async fn test_authorization_callback_ignores_unrelated_connection_and_partial_read() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let source = AuthorizationCodeSource::new(
"http://127.0.0.1:1".to_string(),
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new().redirect_uri(format!("http://127.0.0.1:{port}/callback")),
)
.unwrap();
let browser = tokio::spawn(async move {
let mut unrelated = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
unrelated
.write_all(b"GET /favicon.ico HTTP/1.1\r\nHost: localhost\r\n\r\n")
.await
.unwrap();
let mut response = Vec::new();
unrelated.read_to_end(&mut response).await.unwrap();
assert!(
String::from_utf8(response)
.unwrap()
.starts_with("HTTP/1.1 400")
);
let mut callback = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
callback
.write_all(b"GET /callback?code=auth")
.await
.unwrap();
tokio::task::yield_now().await;
callback
.write_all(b"-code&state=expected HTTP/1.1\r\nHost: localhost\r\n\r\n")
.await
.unwrap();
});
assert_eq!(
source
.wait_for_callback(&listener, "expected")
.await
.unwrap(),
"auth-code"
);
browser.await.unwrap();
}
#[tokio::test]
async fn test_authorization_callback_read_respects_overall_deadline() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let client = tokio::spawn(async move {
let _stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
});
let (mut stream, _) = listener.accept().await.unwrap();
let err = read_authorization_callback(
&mut stream,
"/callback",
"expected",
TokioInstant::now() + Duration::from_millis(20),
)
.await
.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "Timed out reading the OAuth authorization callback"
));
client.abort();
}
#[tokio::test]
async fn test_authorization_request_uses_pkce_by_default() {
let (issuer_url, server) = spawn_discovery_server(1).await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string(), "profile".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
let request = source.build_authorization_request().await.unwrap();
let params: HashMap<_, _> = request.url.query_pairs().into_owned().collect();
assert_eq!(
params.get("response_type").map(String::as_str),
Some("code")
);
assert_eq!(
params.get("client_id").map(String::as_str),
Some("client-id")
);
assert_eq!(
params.get("redirect_uri").map(String::as_str),
Some("http://127.0.0.1:8400/callback")
);
assert_eq!(
params.get("scope").map(String::as_str),
Some("openid profile")
);
assert!(params.get("state").is_some_and(|state| state.len() >= 16));
assert_eq!(
params.get("code_challenge_method").map(String::as_str),
Some("S256")
);
assert!(params.contains_key("code_challenge"));
assert!(request.code_verifier.is_some());
server.await.unwrap();
}
#[tokio::test]
async fn test_authorization_request_rejects_plaintext_provider_endpoint() {
let (issuer_url, server) = spawn_insecure_authorization_discovery_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
let err = source.build_authorization_request().await.unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message.contains("authorization_endpoint must use https")
));
server.await.unwrap();
}
#[tokio::test]
async fn test_authorization_request_can_disable_pkce() {
let (issuer_url, server) = spawn_discovery_server(1).await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new().use_pkce(false),
)
.unwrap();
let request = source.build_authorization_request().await.unwrap();
let params: HashMap<_, _> = request.url.query_pairs().into_owned().collect();
assert!(!params.contains_key("code_challenge"));
assert!(!params.contains_key("code_challenge_method"));
assert!(request.code_verifier.is_none());
server.await.unwrap();
}
#[tokio::test]
async fn test_authorization_code_exchange_public_client_sends_client_id_only() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
let response = source
.exchange_code(
"auth-code",
Some(PkceCodeVerifier::new("verifier".to_string())),
)
.await
.unwrap();
assert_eq!(response.access_token.secret(), "token");
let request = request.lock().unwrap().take().unwrap();
assert_eq!(request.header("authorization"), None);
assert!(request.body.contains("grant_type=authorization_code"));
assert!(request.body.contains("code=auth-code"));
assert!(request.body.contains("code_verifier=verifier"));
assert!(request.body.contains("client_id=client-id"));
assert!(!request.body.contains("client_secret"));
server.await.unwrap();
}
#[tokio::test]
async fn test_authorization_code_exchange_uses_basic_auth_by_default() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
source
.exchange_code(
"auth-code",
Some(PkceCodeVerifier::new("verifier".to_string())),
)
.await
.unwrap();
let request = request.lock().unwrap().take().unwrap();
assert_eq!(
request.header("authorization").as_deref(),
Some(basic_authorization("client-id", "secret").as_str())
);
assert!(request.body.contains("grant_type=authorization_code"));
assert!(request.body.contains("code=auth-code"));
assert!(request.body.contains("code_verifier=verifier"));
assert!(!request.body.contains("client_secret"));
server.await.unwrap();
}
#[tokio::test]
async fn test_authorization_code_exchange_supports_client_secret_post() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretPost,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
source
.exchange_code(
"auth-code",
Some(PkceCodeVerifier::new("verifier".to_string())),
)
.await
.unwrap();
let request = request.lock().unwrap().take().unwrap();
assert_eq!(request.header("authorization"), None);
assert!(request.body.contains("client_id=client-id"));
assert!(request.body.contains("client_secret=secret"));
server.await.unwrap();
}
#[tokio::test]
async fn test_refresh_uses_basic_auth_by_default() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
assert!(matches!(
source.oidc.refresh_token("refresh-token").await.unwrap(),
RefreshResult::Refreshed(_)
));
let request = request.lock().unwrap().take().unwrap();
assert_eq!(
request.header("authorization").as_deref(),
Some(basic_authorization("client-id", "secret").as_str())
);
assert!(request.body.contains("grant_type=refresh_token"));
assert!(request.body.contains("refresh_token=refresh-token"));
assert!(!request.body.contains("client_secret"));
server.await.unwrap();
}
#[tokio::test]
async fn test_refresh_supports_client_secret_post_and_rotation() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretPost,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
let response = match source.oidc.refresh_token("old-refresh").await.unwrap() {
RefreshResult::Refreshed(response) => response,
other => panic!("expected refresh, got {other:?}"),
};
assert_eq!(response.refresh_token().unwrap().secret(), "refresh");
let request = request.lock().unwrap().take().unwrap();
assert_eq!(request.header("authorization"), None);
assert!(request.body.contains("grant_type=refresh_token"));
assert!(request.body.contains("refresh_token=old-refresh"));
assert!(request.body.contains("client_id=client-id"));
assert!(request.body.contains("client_secret=secret"));
server.await.unwrap();
}
#[tokio::test]
async fn test_refresh_invalid_grant_requires_reauthentication() {
let (issuer_url, server) =
spawn_refresh_error_server("400 Bad Request", r#"{"error":"invalid_grant"}"#).await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
assert!(matches!(
source.oidc.refresh_token("revoked").await.unwrap(),
RefreshResult::Reauthenticate
));
server.await.unwrap();
}
#[tokio::test]
async fn test_refresh_transient_failure_remains_retryable() {
let (issuer_url, server) = spawn_refresh_error_server(
"503 Service Unavailable",
r#"{"error":"temporarily_unavailable"}"#,
)
.await;
let source = AuthorizationCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
AuthorizationCodeOptions::new(),
)
.unwrap();
let err = source.oidc.refresh_token("still-valid").await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message.contains("503 Service Unavailable") && message.contains("transient")
));
server.await.unwrap();
}
#[tokio::test]
async fn test_token_request_rejects_insecure_redirect() {
let (issuer_url, server) = spawn_redirecting_token_server().await;
let source = ClientCredentialsSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["scope".to_string()],
None,
None,
)
.unwrap();
let err = TokenSource::fetch_token(&source).await.unwrap_err();
let Error::Runtime { message } = &err else {
panic!("expected runtime error, got {err:?}");
};
assert!(message.contains("redirect"));
// The insecure redirect target must never be contacted.
assert!(!message.contains("idp.example.com"));
server.await.unwrap();
}
#[tokio::test]
async fn test_malformed_token_response_error_does_not_leak_body() {
let (issuer_url, server) = spawn_malformed_token_server().await;
let source = ClientCredentialsSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["scope".to_string()],
None,
None,
)
.unwrap();
let err = TokenSource::fetch_token(&source).await.unwrap_err();
let message = format!("{err:?}");
assert!(!message.contains("leak-marker"));
assert!(matches!(
err,
Error::Runtime { message }
if message.contains("could not be parsed") || message.contains("Content-Type")
));
server.await.unwrap();
}
rust/lancedb/src/remote/oauth/tests/config.rs
+255
View File
@@ -0,0 +1,255 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
fn test_config(flow: OAuthFlow, client_secret: Option<String>) -> OAuthConfig {
OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret,
scopes: vec!["scope".to_string()],
resource: None,
audience: None,
flow,
client_auth_method: None,
refresh_buffer_secs: None,
token_cache: None,
}
}
#[test]
fn test_oauth_config_debug_redacts_client_secret() {
let mut config = test_config(
OAuthFlow::ClientCredentials,
Some("super-secret".to_string()),
);
config.client_auth_method = Some(ClientAuthMethod::ClientSecretBasic);
let debug = format!("{config:?}");
assert!(!debug.contains("super-secret"));
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
assert!(debug.contains("client_auth_method"));
}
#[test]
fn test_oauth_header_provider_debug_redacts_client_secret() {
let config = test_config(
OAuthFlow::ClientCredentials,
Some("super-secret".to_string()),
);
let provider = OAuthHeaderProvider::new(config).unwrap();
let debug = format!("{provider:?}");
assert!(!debug.contains("super-secret"));
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
}
#[test]
fn test_managed_identity_resource_from_default_scope() {
assert_eq!(
AzureImdsSource::resource_from_scopes(&["api://test/.default".to_string()]).unwrap(),
"api://test"
);
}
#[test]
fn test_managed_identity_resource_without_default_suffix() {
assert_eq!(
AzureImdsSource::resource_from_scopes(&["api://test".to_string()]).unwrap(),
"api://test"
);
}
#[test]
fn test_managed_identity_rejects_multiple_scopes() {
let config = OAuthConfig {
issuer_url: "https://login.microsoftonline.com/tenant/v2.0".to_string(),
client_id: "app-id".to_string(),
client_secret: None,
scopes: vec![
"api://test-a/.default".to_string(),
"api://test-b/.default".to_string(),
],
resource: None,
audience: None,
flow: OAuthFlow::AzureManagedIdentity { client_id: None },
client_auth_method: None,
refresh_buffer_secs: None,
token_cache: None,
};
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[tokio::test]
async fn test_token_endpoint_requires_discovery_success() {
let (issuer_url, server) = spawn_discovery_error_server().await;
let source = ClientCredentialsSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["scope".to_string()],
None,
None,
)
.unwrap();
let err = source.oidc.get_token_endpoint().await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message.contains("OIDC discovery failed with status 503")
));
server.await.unwrap();
}
#[test]
fn test_client_credentials_requires_secret() {
let config = test_config(OAuthFlow::ClientCredentials, None);
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[test]
fn test_client_credentials_rejects_insecure_non_loopback_issuer() {
let mut config = test_config(OAuthFlow::ClientCredentials, Some("secret".to_string()));
config.issuer_url = "http://issuer.example.com".to_string();
let err = OAuthHeaderProvider::new(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message
== "OAuth issuer_url must use https, except for http on a loopback host"
));
}
#[test]
fn test_empty_scopes_rejected() {
let mut config = test_config(OAuthFlow::AzureManagedIdentity { client_id: None }, None);
config.scopes = vec![];
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[tokio::test]
async fn test_client_credentials_token_lifecycle() {
let (issuer_url, token_requests, server) = spawn_oauth_server().await;
let config = OAuthConfig {
issuer_url,
client_id: "client-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
resource: None,
audience: None,
flow: OAuthFlow::ClientCredentials,
client_auth_method: None,
refresh_buffer_secs: Some(0),
token_cache: None,
};
let provider = OAuthHeaderProvider::new(config).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-1");
assert_eq!(token_requests.load(Ordering::SeqCst), 1);
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-1");
assert_eq!(token_requests.load(Ordering::SeqCst), 1);
provider.token_state.write().await.expires_at = Some(Instant::now() - Duration::from_secs(1));
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-2");
assert_eq!(token_requests.load(Ordering::SeqCst), 2);
server.await.unwrap();
}
#[tokio::test]
async fn test_client_credentials_supports_client_secret_post() {
let (issuer_url, request, server) = spawn_captured_token_server().await;
let config = OAuthConfig {
issuer_url,
client_id: "client-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
resource: None,
audience: None,
flow: OAuthFlow::ClientCredentials,
client_auth_method: Some(ClientAuthMethod::ClientSecretPost),
refresh_buffer_secs: None,
token_cache: None,
};
let provider = OAuthHeaderProvider::new(config).unwrap();
provider.get_headers().await.unwrap();
let request = request.lock().unwrap().take().unwrap();
assert_eq!(request.header("authorization"), None);
assert!(request.body.contains("grant_type=client_credentials"));
assert!(request.body.contains("client_id=client-id"));
assert!(request.body.contains("client_secret=secret"));
assert!(request.body.contains("scope=scope"));
server.await.unwrap();
}
async fn spawn_oauth_server() -> (String, Arc<AtomicUsize>, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let token_requests = Arc::new(AtomicUsize::new(0));
let server_token_requests = Arc::clone(&token_requests);
let server = tokio::spawn(async move {
for _ in 0..3 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request.line.starts_with("POST /token ") {
assert_eq!(
request.header("authorization").as_deref(),
Some(basic_authorization("client-id", "secret").as_str())
);
assert!(request.body.contains("grant_type=client_credentials"));
assert!(request.body.contains("scope=scope"));
assert!(!request.body.contains("client_secret"));
assert!(!request.body.contains("client_id"));
let token_num = server_token_requests.fetch_add(1, Ordering::SeqCst) + 1;
let token = format!(
r#"{{"access_token":"token-{token_num}","expires_in":3600,"token_type":"Bearer"}}"#
);
write_json_response(&mut stream, "200 OK", &token).await;
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, token_requests, server)
}
async fn spawn_discovery_error_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
assert!(
request
.line
.starts_with("GET /.well-known/openid-configuration ")
);
write_json_response(&mut stream, "503 Service Unavailable", "{}").await;
});
(issuer_url, server)
}
rust/lancedb/src/remote/oauth/tests/device.rs
+300
View File
@@ -0,0 +1,300 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_device_authorization_polls_until_success() {
let (issuer_url, token_requests, server) = spawn_device_server().await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
ClientAuthMethod::ClientSecretBasic,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let device = source.request_device_authorization().await.unwrap();
let response = source.poll_for_token(&device).await.unwrap();
assert_eq!(response.access_token.secret(), "device-token");
assert_eq!(
response
.refresh_token()
.map(|token| token.secret().as_str()),
Some("device-refresh")
);
assert_eq!(token_requests.load(Ordering::SeqCst), 3);
server.await.unwrap();
}
#[tokio::test]
async fn test_device_authorization_rejects_plaintext_verification_uri() {
let (issuer_url, server) = spawn_insecure_device_verification_server().await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let Err(err) = source.request_device_authorization().await else {
panic!("expected insecure verification URI to be rejected");
};
assert!(matches!(
err,
Error::InvalidInput { message }
if message.contains("verification_uri must use https")
));
server.await.unwrap();
}
#[tokio::test]
async fn test_device_authorization_retries_transient_failures() {
let (issuer_url, token_requests, server) = spawn_device_transient_server().await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let device = test_device_authorization_response(60, 1);
let response = source.poll_for_token(&device).await.unwrap();
assert_eq!(response.access_token.secret(), "device-token");
assert_eq!(token_requests.load(Ordering::SeqCst), 4);
server.await.unwrap();
}
#[tokio::test]
async fn test_device_authorization_reports_access_denied() {
let (issuer_url, server) = spawn_device_error_server("access_denied").await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let device = test_device_authorization_response(60, 1);
let err = source.poll_for_token(&device).await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "Device authorization was denied by the user"
));
server.await.unwrap();
}
#[tokio::test]
async fn test_device_authorization_reports_provider_expiry() {
let (issuer_url, server) = spawn_device_error_server("expired_token").await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let device = test_device_authorization_response(60, 1);
let err = source.poll_for_token(&device).await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "Device authorization expired before authentication completed"
));
server.await.unwrap();
}
#[tokio::test]
async fn test_device_authorization_stops_at_local_deadline() {
let (issuer_url, server) = spawn_discovery_server(1).await;
let source = DeviceCodeSource::new(
issuer_url,
"client-id".to_string(),
None,
ClientAuthMethod::None,
vec!["openid".to_string()],
None,
None,
)
.unwrap();
let device = test_device_authorization_response(1, 1);
let err = source.poll_for_token(&device).await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "Device authorization expired before authentication completed"
));
server.await.unwrap();
}
#[derive(Debug)]
struct RefreshingTokenSource {
fetches: Arc<AtomicUsize>,
refreshes: Arc<AtomicUsize>,
}
#[async_trait]
impl TokenSource for RefreshingTokenSource {
async fn fetch_token(&self) -> Result<TokenResponse> {
self.fetches.fetch_add(1, Ordering::SeqCst);
Ok(token_response("initial", Some("refresh"), Some(3600)))
}
async fn refresh_token(&self, refresh_token: &str) -> Result<RefreshResult> {
assert_eq!(refresh_token, "refresh");
self.refreshes.fetch_add(1, Ordering::SeqCst);
Ok(RefreshResult::Refreshed(token_response(
"refreshed",
None,
Some(3600),
)))
}
}
#[tokio::test]
async fn test_header_provider_uses_and_retains_refresh_token() {
let fetches = Arc::new(AtomicUsize::new(0));
let refreshes = Arc::new(AtomicUsize::new(0));
let provider = OAuthHeaderProvider {
token_source: Box::new(RefreshingTokenSource {
fetches: Arc::clone(&fetches),
refreshes: Arc::clone(&refreshes),
}),
token_state: Arc::new(RwLock::new(TokenState::new())),
refresh_buffer: Duration::ZERO,
token_cache: None,
};
assert_eq!(provider.get_valid_token().await.unwrap(), "initial");
provider.token_state.write().await.expires_at = Some(Instant::now() - Duration::from_secs(1));
assert_eq!(provider.get_valid_token().await.unwrap(), "refreshed");
assert_eq!(fetches.load(Ordering::SeqCst), 1);
assert_eq!(refreshes.load(Ordering::SeqCst), 1);
assert_eq!(
provider.token_state.read().await.refresh_token.as_deref(),
Some("refresh")
);
}
#[derive(Debug)]
struct FailedRefreshTokenSource {
fetches: Arc<AtomicUsize>,
refreshes: Arc<AtomicUsize>,
}
#[async_trait]
impl TokenSource for FailedRefreshTokenSource {
async fn fetch_token(&self) -> Result<TokenResponse> {
self.fetches.fetch_add(1, Ordering::SeqCst);
Ok(token_response(
"reauthenticated",
Some("new-refresh"),
Some(3600),
))
}
async fn refresh_token(&self, refresh_token: &str) -> Result<RefreshResult> {
assert_eq!(refresh_token, "revoked-refresh");
self.refreshes.fetch_add(1, Ordering::SeqCst);
Ok(RefreshResult::Reauthenticate)
}
}
#[tokio::test]
async fn test_header_provider_reauthenticates_after_refresh_failure() {
let fetches = Arc::new(AtomicUsize::new(0));
let refreshes = Arc::new(AtomicUsize::new(0));
let provider = OAuthHeaderProvider {
token_source: Box::new(FailedRefreshTokenSource {
fetches: Arc::clone(&fetches),
refreshes: Arc::clone(&refreshes),
}),
token_state: Arc::new(RwLock::new(TokenState {
access_token: Some("expired".to_string()),
refresh_token: Some("revoked-refresh".to_string()),
expires_at: Some(Instant::now() - Duration::from_secs(1)),
})),
refresh_buffer: Duration::ZERO,
token_cache: None,
};
assert_eq!(provider.get_valid_token().await.unwrap(), "reauthenticated");
assert_eq!(fetches.load(Ordering::SeqCst), 1);
assert_eq!(refreshes.load(Ordering::SeqCst), 1);
assert_eq!(
provider.token_state.read().await.refresh_token.as_deref(),
Some("new-refresh")
);
}
#[derive(Debug)]
struct TransientRefreshFailureSource {
fetches: Arc<AtomicUsize>,
}
#[async_trait]
impl TokenSource for TransientRefreshFailureSource {
async fn fetch_token(&self) -> Result<TokenResponse> {
self.fetches.fetch_add(1, Ordering::SeqCst);
unreachable!("a transient refresh failure must not start an interactive flow")
}
async fn refresh_token(&self, refresh_token: &str) -> Result<RefreshResult> {
assert_eq!(refresh_token, "valid-refresh");
Err(Error::Runtime {
message: "token endpoint temporarily unavailable".to_string(),
})
}
}
#[tokio::test]
async fn test_header_provider_preserves_refresh_token_after_transient_failure() {
let fetches = Arc::new(AtomicUsize::new(0));
let provider = OAuthHeaderProvider {
token_source: Box::new(TransientRefreshFailureSource {
fetches: Arc::clone(&fetches),
}),
token_state: Arc::new(RwLock::new(TokenState {
access_token: Some("expired".to_string()),
refresh_token: Some("valid-refresh".to_string()),
expires_at: Some(Instant::now() - Duration::from_secs(1)),
})),
refresh_buffer: Duration::ZERO,
token_cache: None,
};
let err = provider.get_valid_token().await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message == "token endpoint temporarily unavailable"
));
assert_eq!(fetches.load(Ordering::SeqCst), 0);
assert_eq!(
provider.token_state.read().await.refresh_token.as_deref(),
Some("valid-refresh")
);
}
rust/lancedb/src/remote/oauth/tests/flows.rs
+361
View File
@@ -0,0 +1,361 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_target_parameters_across_oauth_flows() {
for (resource, audience) in [
(None, None),
(Some("https://api.example.com/a?x=1&y=two"), None),
(None, Some("audience + & / ü")),
(Some("urn:example:resource"), Some("audience + & / ü")),
] {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer = format!("http://{addr}");
let server = tokio::spawn(async move {
let mut grants = Vec::new();
// Three discovery requests and six form submissions.
for _ in 0..9 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
let response = if request.line.starts_with("GET ") {
serde_json::json!({
"token_endpoint": format!("http://{addr}/token"),
"authorization_endpoint": format!("http://{addr}/authorize"),
"device_authorization_endpoint": format!("http://{addr}/device"),
})
} else {
let params: Vec<_> =
url::form_urlencoded::parse(request.body.as_bytes()).collect();
for (key, expected) in [("resource", resource), ("audience", audience)] {
let values: Vec<_> = params
.iter()
.filter(|(name, _)| name == key)
.map(|(_, value)| value.as_ref())
.collect();
assert_eq!(values, expected.into_iter().collect::<Vec<_>>());
}
if request.line.starts_with("POST /device ") {
serde_json::json!({
"device_code": "device-code", "user_code": "ABCD",
"verification_uri": format!("http://{addr}/verify"),
"expires_in": 60, "interval": 1,
})
} else {
grants.push(
params
.iter()
.find(|(key, _)| key == "grant_type")
.unwrap()
.1
.to_string(),
);
serde_json::json!({"access_token": "access", "refresh_token": "refresh", "expires_in": 3600})
}
};
write_json_response(&mut stream, "200 OK", &response.to_string()).await;
}
assert_eq!(
grants,
[
"client_credentials",
"authorization_code",
"refresh_token",
"urn:ietf:params:oauth:grant-type:device_code",
"refresh_token"
]
);
});
let credentials = ClientCredentialsSource::new(
issuer.clone(),
"client".into(),
Some("secret".into()),
ClientAuthMethod::ClientSecretBasic,
vec!["scope".into()],
resource.map(str::to_owned),
audience.map(str::to_owned),
)
.unwrap();
credentials.fetch_token().await.unwrap();
let browser = AuthorizationCodeSource::new(
issuer.clone(),
"client".into(),
None,
ClientAuthMethod::None,
vec!["scope".into()],
resource.map(str::to_owned),
audience.map(str::to_owned),
AuthorizationCodeOptions::new(),
)
.unwrap();
let request = browser.build_authorization_request().await.unwrap();
for (key, expected) in [("resource", resource), ("audience", audience)] {
let values: Vec<_> = request
.url
.query_pairs()
.filter(|(name, _)| name == key)
.map(|(_, value)| value.into_owned())
.collect();
assert_eq!(
values,
expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
);
}
browser
.exchange_code("code", Some(PkceCodeVerifier::new("verifier".to_string())))
.await
.unwrap();
browser.refresh_token("refresh").await.unwrap();
let device = DeviceCodeSource::new(
issuer,
"client".into(),
None,
ClientAuthMethod::None,
vec!["scope".into()],
resource.map(str::to_owned),
audience.map(str::to_owned),
)
.unwrap();
let authorization = device.request_device_authorization().await.unwrap();
device.poll_for_token(&authorization).await.unwrap();
device.refresh_token("refresh").await.unwrap();
server.await.unwrap();
}
}
#[test]
fn test_managed_identity_rejects_target_parameters() {
for (resource, audience) in [(Some("urn:resource"), None), (None, Some("audience"))] {
let config = OAuthConfig {
issuer_url: "https://issuer.example.com".into(),
client_id: "client".into(),
client_secret: None,
scopes: vec!["api://app/.default".into()],
resource: resource.map(str::to_owned),
audience: audience.map(str::to_owned),
flow: OAuthFlow::AzureManagedIdentity { client_id: None },
client_auth_method: None,
refresh_buffer_secs: None,
token_cache: None,
};
let error = OAuthHeaderProvider::new(config).unwrap_err().to_string();
assert!(error.contains("resource and audience are not supported for AzureManagedIdentity"));
}
}
#[test]
fn test_token_state_expiry() {
let mut state = TokenState::new();
assert!(state.is_expired(Duration::from_secs(0)));
state.access_token = Some("tok".to_string());
state.expires_at = Some(Instant::now() + Duration::from_secs(600));
assert!(!state.is_expired(Duration::from_secs(300)));
assert!(state.is_expired(Duration::from_secs(601)));
state.expires_at = None;
assert!(state.is_expired(Duration::from_secs(0)));
}
#[test]
fn test_token_state_uses_default_expiry() {
let mut state = TokenState::new();
state.update(&token_response("tok", None, None));
assert!(!state.is_expired(Duration::from_secs(DEFAULT_TOKEN_TTL_SECS - 1)));
assert!(state.is_expired(Duration::from_secs(DEFAULT_TOKEN_TTL_SECS + 1)));
}
#[test]
fn test_token_state_retains_refresh_token_when_not_rotated() {
let mut state = TokenState::new();
state.update(&token_response("token-1", Some("refresh-1"), Some(60)));
state.update(&token_response("token-2", None, Some(60)));
assert_eq!(state.refresh_token.as_deref(), Some("refresh-1"));
}
#[test]
fn test_token_response_accepts_float_expires_in() {
let response: TokenResponse =
serde_json::from_str(r#"{"access_token":"tok","expires_in":3600.0}"#).unwrap();
assert_eq!(response.expires_in, Some(3600));
}
#[test]
fn test_token_response_rejects_negative_expires_in() {
let err = serde_json::from_str::<TokenResponse>(r#"{"access_token":"tok","expires_in":-1}"#)
.unwrap_err();
assert!(err.to_string().contains("invalid expires_in value: -1"));
}
#[test]
fn test_token_response_debug_redacts_access_token() {
let response = TokenResponse {
access_token: AccessToken::new("secret-token".to_string()),
refresh_token: Some(RefreshToken::new("secret-refresh-token".to_string())),
expires_in: Some(3600),
token_type: Some(BasicTokenType::Bearer),
};
let debug = format!("{response:?}");
assert!(!debug.contains("secret-token"));
assert!(!debug.contains("secret-refresh-token"));
assert!(debug.contains("access_token: \"<redacted>\""));
}
#[test]
fn test_client_auth_method_defaults() {
assert_eq!(
resolve_client_auth_method(None, Some("secret")).unwrap(),
ClientAuthMethod::ClientSecretBasic
);
assert_eq!(
resolve_client_auth_method(None, None).unwrap(),
ClientAuthMethod::None
);
}
#[test]
fn test_client_auth_method_explicit_values() {
for method in [
ClientAuthMethod::None,
ClientAuthMethod::ClientSecretBasic,
ClientAuthMethod::ClientSecretPost,
] {
let secret = (method != ClientAuthMethod::None).then_some("secret");
assert_eq!(
resolve_client_auth_method(Some(method), secret).unwrap(),
method
);
}
}
#[test]
fn test_client_auth_method_rejects_inconsistent_configuration() {
let err = resolve_client_auth_method(Some(ClientAuthMethod::None), Some("secret")).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "client_auth_method None cannot be combined with client_secret"
));
for method in [
ClientAuthMethod::ClientSecretBasic,
ClientAuthMethod::ClientSecretPost,
] {
let err = resolve_client_auth_method(Some(method), None).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == format!("client_auth_method {method:?} requires client_secret to be set")
));
}
}
#[test]
fn test_oauth_transport_requires_https_except_for_loopback() {
assert!(validate_oauth_url("https://idp.example.com/token", "endpoint").is_ok());
assert!(validate_oauth_url("http://localhost:8080/token", "endpoint").is_ok());
assert!(validate_oauth_url("http://127.0.0.1:8080/token", "endpoint").is_ok());
assert!(validate_oauth_url("http://[::1]:8080/token", "endpoint").is_ok());
let err = validate_oauth_url("http://idp.example.com/token", "endpoint").unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "OAuth endpoint must use https, except for http on a loopback host"
));
}
#[test]
fn test_interactive_prompts_use_default_visible_output() {
let authorization_url = Url::parse("https://idp.example.com/authorize?state=abc").unwrap();
let authorization = authorization_prompt(&authorization_url);
let device = device_prompt("https://idp.example.com/device", "ABCD-EFGH");
let mut output = Vec::new();
write_oauth_prompt(&mut output, &authorization);
write_oauth_prompt(&mut output, &device);
let output = String::from_utf8(output).unwrap();
assert!(output.contains(authorization_url.as_str()));
assert!(output.contains("https://idp.example.com/device"));
assert!(output.contains("ABCD-EFGH"));
}
#[test]
fn test_authorization_code_options_default_to_pkce() {
let options = AuthorizationCodeOptions::new();
assert!(options.use_pkce);
assert!(options.redirect_uri.is_none());
assert!(options.callback_port.is_none());
}
#[test]
fn test_authorization_redirect_defaults_to_ipv4_loopback() {
let redirect = ResolvedRedirect::new(&AuthorizationCodeOptions::new()).unwrap();
assert_eq!(
redirect.uri,
format!("http://127.0.0.1:{DEFAULT_CALLBACK_PORT}/callback")
);
assert!(redirect.bind_addr.ip().is_loopback());
assert_eq!(redirect.bind_addr.port(), DEFAULT_CALLBACK_PORT);
assert_eq!(redirect.callback_path, "/callback");
}
#[test]
fn test_authorization_redirect_rejects_non_loopback_host() {
let options =
AuthorizationCodeOptions::new().redirect_uri("https://client.example.com/oauth/callback");
let err = ResolvedRedirect::new(&options).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "OAuth redirect_uri must use http with a loopback host"
));
}
#[test]
fn test_authorization_redirect_rejects_mismatched_port() {
let options = AuthorizationCodeOptions::new()
.redirect_uri("http://127.0.0.1:8401/callback")
.callback_port(8400);
let err = ResolvedRedirect::new(&options).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message.contains("does not match redirect_uri port")
));
}
#[test]
fn test_authorization_redirect_requires_explicit_port() {
let options = AuthorizationCodeOptions::new().redirect_uri("http://127.0.0.1/callback");
let err = ResolvedRedirect::new(&options).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "OAuth redirect_uri must include a port"
));
}
#[test]
fn test_authorization_redirect_accepts_ipv6_loopback() {
let options = AuthorizationCodeOptions::new().redirect_uri("http://[::1]:8400/callback");
let redirect = ResolvedRedirect::new(&options).unwrap();
assert_eq!(
redirect.bind_addr,
"[::1]:8400".parse::<SocketAddr>().unwrap()
);
}
rust/lancedb/src/remote/oauth/tests/mod.rs
+485
View File
@@ -0,0 +1,485 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
use base64::Engine;
use oauth2::TokenResponse as _;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinHandle;
fn token_response(
access_token: &str,
refresh_token: Option<&str>,
expires_in: Option<u64>,
) -> TokenResponse {
TokenResponse {
access_token: AccessToken::new(access_token.to_string()),
refresh_token: refresh_token.map(|token| RefreshToken::new(token.to_string())),
expires_in,
token_type: None,
}
}
fn basic_authorization(client_id: &str, client_secret: &str) -> String {
format!(
"Basic {}",
base64::engine::general_purpose::STANDARD.encode(format!("{client_id}:{client_secret}"))
)
}
struct CapturedRequest {
line: String,
headers: String,
body: String,
}
impl CapturedRequest {
fn header(&self, name: &str) -> Option<String> {
self.headers.lines().find_map(|line| {
let (key, value) = line.split_once(':')?;
key.trim()
.eq_ignore_ascii_case(name)
.then(|| value.trim().to_string())
})
}
}
async fn spawn_discovery_server(expected_requests: usize) -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..expected_requests {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
assert!(
request
.line
.starts_with("GET /.well-known/openid-configuration ")
);
let discovery = format!(
r#"{{"token_endpoint":"http://{addr}/token","authorization_endpoint":"http://{addr}/authorize","device_authorization_endpoint":"http://{addr}/device"}}"#
);
write_json_response(&mut stream, "200 OK", &discovery).await;
}
});
(issuer_url, server)
}
async fn spawn_insecure_authorization_discovery_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
assert!(
request
.line
.starts_with("GET /.well-known/openid-configuration ")
);
write_json_response(
&mut stream,
"200 OK",
r#"{"token_endpoint":"https://idp.example.com/token","authorization_endpoint":"http://idp.example.com/authorize"}"#,
)
.await;
});
(issuer_url, server)
}
async fn spawn_insecure_device_verification_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(
r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"#
);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else {
assert!(request.line.starts_with("POST /device "));
write_json_response(
&mut stream,
"200 OK",
r#"{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://idp.example.com/device","expires_in":60}"#,
)
.await;
}
}
});
(issuer_url, server)
}
async fn spawn_captured_token_server() -> (
String,
Arc<std::sync::Mutex<Option<CapturedRequest>>>,
JoinHandle<()>,
) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let request = Arc::new(std::sync::Mutex::new(None));
let server_request = Arc::clone(&request);
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let captured = read_http_request(&mut stream).await;
if captured
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if captured.line.starts_with("POST /token ") {
*server_request.lock().unwrap() = Some(captured);
write_json_response(
&mut stream,
"200 OK",
r#"{"access_token":"token","refresh_token":"refresh","expires_in":3600,"token_type":"Bearer"}"#,
)
.await;
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, request, server)
}
async fn spawn_redirecting_token_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else {
assert!(request.line.starts_with("POST /token "));
let response = "HTTP/1.1 302 Found\r\nlocation: http://idp.example.com/steal\r\ncontent-length: 0\r\nconnection: close\r\n\r\n".to_string();
stream.write_all(response.as_bytes()).await.unwrap();
}
}
});
(issuer_url, server)
}
async fn spawn_malformed_token_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else {
assert!(request.line.starts_with("POST /token "));
write_json_response(
&mut stream,
"200 OK",
r#"{"access_token":{"nested":"leak-marker-12345"}}"#,
)
.await;
}
}
});
(issuer_url, server)
}
async fn spawn_refresh_error_server(
status: &'static str,
response_body: &'static str,
) -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request.line.starts_with("POST /token ") {
assert!(request.body.contains("grant_type=refresh_token"));
assert!(request.body.contains("refresh_token="));
write_json_response(&mut stream, status, response_body).await;
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, server)
}
async fn spawn_device_server() -> (String, Arc<AtomicUsize>, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let token_requests = Arc::new(AtomicUsize::new(0));
let server_token_requests = Arc::clone(&token_requests);
let server = tokio::spawn(async move {
for _ in 0..5 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(
r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"#
);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request.line.starts_with("POST /device ") {
// The resolved default for a confidential client is HTTP Basic.
assert_eq!(
request.header("authorization").as_deref(),
Some(basic_authorization("client-id", "secret").as_str())
);
assert!(request.body.contains("scope=openid"));
assert!(!request.body.contains("client_secret"));
let device = format!(
r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","verification_uri_complete":"http://{addr}/verify?user_code=ABCD-EFGH","expires_in":60,"interval":1}}"#
);
write_json_response(&mut stream, "200 OK", &device).await;
} else if request.line.starts_with("POST /token ") {
assert_eq!(
request.header("authorization").as_deref(),
Some(basic_authorization("client-id", "secret").as_str())
);
assert!(
request.body.contains(
"grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code"
)
);
assert!(request.body.contains("device_code=device-code"));
assert!(!request.body.contains("client_secret"));
let request = server_token_requests.fetch_add(1, Ordering::SeqCst);
match request {
0 => {
write_json_response(
&mut stream,
"400 Bad Request",
r#"{"error":"authorization_pending"}"#,
)
.await;
}
1 => {
write_json_response(
&mut stream,
"400 Bad Request",
r#"{"error":"slow_down"}"#,
)
.await;
}
_ => {
write_json_response(
&mut stream,
"200 OK",
r#"{"access_token":"device-token","refresh_token":"device-refresh","expires_in":3600,"token_type":"Bearer"}"#,
)
.await;
}
}
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, token_requests, server)
}
async fn spawn_device_transient_server() -> (String, Arc<AtomicUsize>, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let token_requests = Arc::new(AtomicUsize::new(0));
let server_token_requests = Arc::clone(&token_requests);
let server = tokio::spawn(async move {
for _ in 0..5 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
continue;
}
assert!(request.line.starts_with("POST /token "));
match server_token_requests.fetch_add(1, Ordering::SeqCst) {
0 => drop(stream),
1 => {
write_json_response(
&mut stream,
"503 Service Unavailable",
r#"{"error":"server_error"}"#,
)
.await;
}
2 => {
write_json_response(
&mut stream,
"400 Bad Request",
r#"{"error":"temporarily_unavailable"}"#,
)
.await;
}
_ => {
write_json_response(
&mut stream,
"200 OK",
r#"{"access_token":"device-token","expires_in":3600,"token_type":"Bearer"}"#,
)
.await;
}
}
}
});
(issuer_url, token_requests, server)
}
fn test_device_authorization_response(
expires_in: u64,
interval: u64,
) -> StandardDeviceAuthorizationResponse {
serde_json::from_str(&format!(
r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://127.0.0.1/verify","expires_in":{expires_in},"interval":{interval}}}"#
))
.unwrap()
}
async fn spawn_device_error_server(error: &'static str) -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
if request
.line
.starts_with("GET /.well-known/openid-configuration ")
{
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request.line.starts_with("POST /token ") {
write_json_response(
&mut stream,
"400 Bad Request",
&format!(r#"{{"error":"{error}"}}"#),
)
.await;
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, server)
}
async fn read_http_request(stream: &mut TcpStream) -> CapturedRequest {
let mut buffer = Vec::new();
let mut header_end = None;
while header_end.is_none() {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request headers");
buffer.extend_from_slice(&chunk[..read]);
header_end = find_subsequence(&buffer, b"\r\n\r\n").map(|pos| pos + 4);
}
let header_end = header_end.unwrap();
let headers = String::from_utf8_lossy(&buffer[..header_end]).to_string();
let line = headers.lines().next().unwrap_or_default().to_string();
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while buffer.len() < header_end + content_length {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request body");
buffer.extend_from_slice(&chunk[..read]);
}
let body =
String::from_utf8_lossy(&buffer[header_end..header_end + content_length]).to_string();
CapturedRequest {
line,
headers,
body,
}
}
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack
.windows(needle.len())
.position(|window| window == needle)
}
async fn write_json_response(stream: &mut TcpStream, status: &str, body: &str) {
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).await.unwrap();
}
mod authorization;
mod config;
mod device;
mod flows;
rust/lancedb/src/remote/table.rs
+1 -10428
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/blobs.rs
+1 -1238
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/blobs/tests.rs
+1234
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/insert.rs
+1 -727
View File
@@ -869,730 +869,4 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
}
#[cfg(test)]
mod tests {
use arrow_array::record_batch;
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
use datafusion::prelude::SessionContext;
use datafusion_catalog::MemTable;
use datafusion_common::{DataFusionError, Result as DataFusionResult};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use super::RemoteWriteExec;
use super::WriteOp;
use crate::Table;
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
use crate::remote::table::MergeInsertRequest;
use crate::table::datafusion::BaseTableAdapter;
fn schema_json() -> &'static str {
r#"{"fields": [{"name": "id", "type": {"type": "int32"}, "nullable": true}]}"#
}
#[tokio::test]
async fn test_remote_insert_exec_execute_empty() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
// Return schema for BaseTableAdapter::try_new
return http::Response::builder()
.status(200)
.body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json()))
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
request_count_clone.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
});
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// Create empty MemTable (no batches)
let source_table = MemTable::try_new(schema, vec![vec![]]).unwrap();
let ctx = SessionContext::new();
// Register the remote table as insert target
let provider = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
ctx.register_table("my_table", Arc::new(provider)).unwrap();
// Register empty source
ctx.register_table("empty_source", Arc::new(source_table))
.unwrap();
// Execute the INSERT
ctx.sql("INSERT INTO my_table SELECT * FROM empty_source")
.await
.unwrap()
.collect()
.await
.unwrap();
// Verify: should have made exactly one HTTP request even with empty input
assert_eq!(request_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_remote_insert_exec_multi_partition() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
// Return schema for BaseTableAdapter::try_new
return http::Response::builder()
.status(200)
.body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json()))
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
request_count_clone.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
});
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// Create MemTable with multiple partitions and multiple batches
let source_table = MemTable::try_new(
schema,
vec![
// Partition 0
vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
],
// Partition 1
vec![record_batch!(("id", Int32, [5, 6, 7])).unwrap()],
// Partition 2
vec![record_batch!(("id", Int32, [8])).unwrap()],
],
)
.unwrap();
let ctx = SessionContext::new();
// Register the remote table as insert target
let provider = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
ctx.register_table("my_table", Arc::new(provider)).unwrap();
// Register multi-partition source
ctx.register_table("multi_partition_source", Arc::new(source_table))
.unwrap();
// Get the physical plan and verify it includes a repartition to 1
let df = ctx
.sql("INSERT INTO my_table SELECT * FROM multi_partition_source")
.await
.unwrap();
let plan = df.clone().create_physical_plan().await.unwrap();
let plan_str = datafusion::physical_plan::displayable(plan.as_ref())
.indent(true)
.to_string();
// The plan should include a CoalescePartitionsExec to merge partitions
assert!(
plan_str.contains("CoalescePartitionsExec"),
"Expected CoalescePartitionsExec in plan:\n{}",
plan_str
);
// Execute the INSERT
df.collect().await.unwrap();
// Verify: should have made exactly one HTTP request despite multiple input partitions
assert_eq!(request_count.load(Ordering::SeqCst), 1);
}
/// Build a single-partition input plan from the given batches.
async fn input_plan_from_batches(
schema: Arc<ArrowSchema>,
batches: Vec<arrow_array::RecordBatch>,
) -> Arc<dyn ExecutionPlan> {
use datafusion_catalog::TableProvider;
let mem = MemTable::try_new(schema, vec![batches]).unwrap();
let ctx = SessionContext::new();
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
}
/// Build a single-partition input plan from the batches spread across the
/// given partitions.
async fn input_plan_from_partitions(
schema: Arc<ArrowSchema>,
partitions: Vec<Vec<arrow_array::RecordBatch>>,
) -> Arc<dyn ExecutionPlan> {
use datafusion_catalog::TableProvider;
let mem = MemTable::try_new(schema, partitions).unwrap();
let ctx = SessionContext::new();
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
}
fn counting_insert_client(
counter: Arc<AtomicUsize>,
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender>
{
crate::remote::client::test_utils::client_with_handler(move |request| {
let path = request.url().path();
assert_eq!(path, "/v1/table/my_table/insert/");
let query = request.url().query().unwrap_or("");
assert!(query.contains("upload_id=upload-1"), "query: {query}");
assert!(query.contains("upload_part_id="), "query: {query}");
counter.fetch_add(1, Ordering::SeqCst);
http::Response::builder()
.status(200)
.body(String::new())
.unwrap()
})
}
/// Insert handler that records the `upload_part_id` of every part request so
/// a test can assert the ids are distinct.
fn recording_insert_client(
part_ids: Arc<Mutex<Vec<String>>>,
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender>
{
crate::remote::client::test_utils::client_with_handler(move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/insert/");
let part_id = request
.url()
.query_pairs()
.find(|(k, _)| k == "upload_part_id")
.map(|(_, v)| v.into_owned())
.expect("upload_part_id query param");
part_ids.lock().unwrap().push(part_id);
http::Response::builder()
.status(200)
.body(String::new())
.unwrap()
})
}
/// Single-partition input plan that yields one good batch and then an error,
/// for exercising the mid-part input-error abort path in `send_one_part`.
#[derive(Debug)]
struct ErroringExec {
schema: Arc<ArrowSchema>,
properties: Arc<PlanProperties>,
}
impl ErroringExec {
fn new() -> Self {
let schema = record_batch!(("id", Int32, [1, 2])).unwrap().schema();
let properties = PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
datafusion_physical_plan::Partitioning::UnknownPartitioning(1),
datafusion_physical_plan::execution_plan::EmissionType::Incremental,
datafusion_physical_plan::execution_plan::Boundedness::Bounded,
);
Self {
schema,
properties: Arc::new(properties),
}
}
}
impl DisplayAs for ErroringExec {
fn fmt_as(
&self,
_t: DisplayFormatType,
f: &mut std::fmt::Formatter<'_>,
) -> std::fmt::Result {
write!(f, "ErroringExec")
}
}
impl ExecutionPlan for ErroringExec {
fn name(&self) -> &str {
"ErroringExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
Ok(self)
}
fn execute(
&self,
_partition: usize,
_context: Arc<TaskContext>,
) -> DataFusionResult<SendableRecordBatchStream> {
let batch = record_batch!(("id", Int32, [1, 2])).unwrap();
let stream = futures::stream::iter(vec![
Ok(batch),
Err(DataFusionError::Execution("boom".to_string())),
]);
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.schema.clone(),
stream,
)))
}
}
#[tokio::test]
async fn test_multipart_chunked_splits_into_parts() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A 1-byte budget forces every batch into its own part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_single_part_when_under_budget() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A large byte budget and no time limit keep the whole partition in a
// single part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_chunked_splits_by_duration() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A large byte budget but a tiny duration budget: writing and sending
// one batch already takes longer than the limit, so each batch is cut
// into its own part on the time check rather than the byte check.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
Some(std::time::Duration::from_nanos(1)),
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_empty_partition_stages_nothing() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// An empty partition should stage no parts; on the multipart path the
// write relies on another partition having data to commit.
let input = input_plan_from_batches(schema, vec![]).await;
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_multipart_chunked_uses_distinct_part_ids() {
use futures::StreamExt;
use std::collections::HashSet;
let part_ids = Arc::new(Mutex::new(Vec::new()));
let client = recording_insert_client(part_ids.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A 1-byte budget forces every batch into its own part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
let ids = part_ids.lock().unwrap().clone();
assert_eq!(ids.len(), 3, "expected one part id per part: {ids:?}");
assert!(
ids.iter().all(|id| !id.is_empty()),
"part ids must be non-empty: {ids:?}"
);
let unique: HashSet<&String> = ids.iter().collect();
assert_eq!(unique.len(), 3, "part ids must be distinct: {ids:?}");
}
#[tokio::test]
async fn test_multipart_chunks_each_partition_independently() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let partitions = vec![
// Partition 0: two batches, split into two parts by the 1-byte budget.
vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
],
// Partition 1: one batch, one part.
vec![record_batch!(("id", Int32, [5, 6])).unwrap()],
];
let input = input_plan_from_partitions(schema, partitions).await;
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
for partition in 0..2 {
let mut stream = exec
.execute(partition, Arc::new(TaskContext::default()))
.unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
}
// 2 parts from partition 0 + 1 part from partition 1.
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_input_error_surfaces_original() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
// A large byte budget keeps the good batch and the following error in
// the same part, exercising the mid-part abort path.
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
let mut err = None;
while let Some(item) = stream.next().await {
if let Err(e) = item {
err = Some(e);
break;
}
}
let err = err.expect("expected the input stream error to surface");
// The original DataFusion error must win over the HTTP error it induces.
assert!(
err.to_string().contains("boom"),
"expected original input error, got: {err}"
);
}
#[tokio::test]
async fn test_merge_insert_input_error_surfaces_original() {
// Regression test for #2339 on the single-request merge_insert path.
// When the input stream errors mid-body, Hyper masks it under HTTP2 as a
// generic "stream error sent by user" message. The error side-channel
// must recover and surface the original DataFusion error instead.
use futures::StreamExt;
let client = crate::remote::client::test_utils::client_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
http::Response::builder()
.status(200)
.body(
r#"{"version": 2, "num_updated_rows": 0, "num_inserted_rows": 0, "num_deleted_rows": 0}"#
.to_string(),
)
.unwrap()
});
let query = MergeInsertRequest {
on: vec!["id".to_string()],
when_matched_update_all: false,
when_matched_update_all_filt: None,
when_not_matched_insert_all: false,
when_not_matched_by_source_delete: false,
when_not_matched_by_source_delete_filt: None,
use_index: true,
use_lsm: None,
};
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
let exec = RemoteWriteExec::new(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
WriteOp::MergeInsert {
query,
timeout: None,
},
None,
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
let mut err = None;
while let Some(item) = stream.next().await {
if let Err(e) = item {
err = Some(e);
break;
}
}
let err = err.expect("expected the input stream error to surface");
assert!(
err.to_string().contains("boom"),
"expected original input error, got: {err}"
);
}
#[tokio::test]
async fn test_multipart_records_progress_within_a_part() {
use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker};
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
let observed = Arc::new(Mutex::new(Vec::<usize>::new()));
let observed_cb = observed.clone();
let callback: ProgressCallback = Arc::new(Mutex::new(move |p: &WriteProgress| {
observed_cb.lock().unwrap().push(p.output_bytes());
}));
let tracker = Arc::new(WriteProgressTracker::new(callback, None));
// A large byte budget keeps all three batches in one part; smooth
// progress therefore requires bytes to be reported per chunk rather than
// once when the part completes.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
Some(tracker),
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(
insert_count.load(Ordering::SeqCst),
1,
"batches should all land in a single part"
);
let observed = observed.lock().unwrap();
assert!(
observed.len() > 1,
"expected multiple incremental progress updates within the part: {observed:?}"
);
assert!(
observed.windows(2).all(|w| w[1] >= w[0]),
"progress bytes should be monotonic: {observed:?}"
);
assert!(
*observed.last().unwrap() > 0,
"final progress should report bytes: {observed:?}"
);
}
}
mod tests;
rust/lancedb/src/remote/table/insert/tests.rs
+722
View File
@@ -0,0 +1,722 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use arrow_array::record_batch;
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
use datafusion::prelude::SessionContext;
use datafusion_catalog::MemTable;
use datafusion_common::{DataFusionError, Result as DataFusionResult};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use super::RemoteWriteExec;
use super::WriteOp;
use crate::Table;
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
use crate::remote::table::MergeInsertRequest;
use crate::table::datafusion::BaseTableAdapter;
fn schema_json() -> &'static str {
r#"{"fields": [{"name": "id", "type": {"type": "int32"}, "nullable": true}]}"#
}
#[tokio::test]
async fn test_remote_insert_exec_execute_empty() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
// Return schema for BaseTableAdapter::try_new
return http::Response::builder()
.status(200)
.body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json()))
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
request_count_clone.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
});
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// Create empty MemTable (no batches)
let source_table = MemTable::try_new(schema, vec![vec![]]).unwrap();
let ctx = SessionContext::new();
// Register the remote table as insert target
let provider = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
ctx.register_table("my_table", Arc::new(provider)).unwrap();
// Register empty source
ctx.register_table("empty_source", Arc::new(source_table))
.unwrap();
// Execute the INSERT
ctx.sql("INSERT INTO my_table SELECT * FROM empty_source")
.await
.unwrap()
.collect()
.await
.unwrap();
// Verify: should have made exactly one HTTP request even with empty input
assert_eq!(request_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_remote_insert_exec_multi_partition() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
// Return schema for BaseTableAdapter::try_new
return http::Response::builder()
.status(200)
.body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json()))
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
request_count_clone.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
});
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// Create MemTable with multiple partitions and multiple batches
let source_table = MemTable::try_new(
schema,
vec![
// Partition 0
vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
],
// Partition 1
vec![record_batch!(("id", Int32, [5, 6, 7])).unwrap()],
// Partition 2
vec![record_batch!(("id", Int32, [8])).unwrap()],
],
)
.unwrap();
let ctx = SessionContext::new();
// Register the remote table as insert target
let provider = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
ctx.register_table("my_table", Arc::new(provider)).unwrap();
// Register multi-partition source
ctx.register_table("multi_partition_source", Arc::new(source_table))
.unwrap();
// Get the physical plan and verify it includes a repartition to 1
let df = ctx
.sql("INSERT INTO my_table SELECT * FROM multi_partition_source")
.await
.unwrap();
let plan = df.clone().create_physical_plan().await.unwrap();
let plan_str = datafusion::physical_plan::displayable(plan.as_ref())
.indent(true)
.to_string();
// The plan should include a CoalescePartitionsExec to merge partitions
assert!(
plan_str.contains("CoalescePartitionsExec"),
"Expected CoalescePartitionsExec in plan:\n{}",
plan_str
);
// Execute the INSERT
df.collect().await.unwrap();
// Verify: should have made exactly one HTTP request despite multiple input partitions
assert_eq!(request_count.load(Ordering::SeqCst), 1);
}
/// Build a single-partition input plan from the given batches.
async fn input_plan_from_batches(
schema: Arc<ArrowSchema>,
batches: Vec<arrow_array::RecordBatch>,
) -> Arc<dyn ExecutionPlan> {
use datafusion_catalog::TableProvider;
let mem = MemTable::try_new(schema, vec![batches]).unwrap();
let ctx = SessionContext::new();
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
}
/// Build a single-partition input plan from the batches spread across the
/// given partitions.
async fn input_plan_from_partitions(
schema: Arc<ArrowSchema>,
partitions: Vec<Vec<arrow_array::RecordBatch>>,
) -> Arc<dyn ExecutionPlan> {
use datafusion_catalog::TableProvider;
let mem = MemTable::try_new(schema, partitions).unwrap();
let ctx = SessionContext::new();
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
}
fn counting_insert_client(
counter: Arc<AtomicUsize>,
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender> {
crate::remote::client::test_utils::client_with_handler(move |request| {
let path = request.url().path();
assert_eq!(path, "/v1/table/my_table/insert/");
let query = request.url().query().unwrap_or("");
assert!(query.contains("upload_id=upload-1"), "query: {query}");
assert!(query.contains("upload_part_id="), "query: {query}");
counter.fetch_add(1, Ordering::SeqCst);
http::Response::builder()
.status(200)
.body(String::new())
.unwrap()
})
}
/// Insert handler that records the `upload_part_id` of every part request so
/// a test can assert the ids are distinct.
fn recording_insert_client(
part_ids: Arc<Mutex<Vec<String>>>,
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender> {
crate::remote::client::test_utils::client_with_handler(move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/insert/");
let part_id = request
.url()
.query_pairs()
.find(|(k, _)| k == "upload_part_id")
.map(|(_, v)| v.into_owned())
.expect("upload_part_id query param");
part_ids.lock().unwrap().push(part_id);
http::Response::builder()
.status(200)
.body(String::new())
.unwrap()
})
}
/// Single-partition input plan that yields one good batch and then an error,
/// for exercising the mid-part input-error abort path in `send_one_part`.
#[derive(Debug)]
struct ErroringExec {
schema: Arc<ArrowSchema>,
properties: Arc<PlanProperties>,
}
impl ErroringExec {
fn new() -> Self {
let schema = record_batch!(("id", Int32, [1, 2])).unwrap().schema();
let properties = PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
datafusion_physical_plan::Partitioning::UnknownPartitioning(1),
datafusion_physical_plan::execution_plan::EmissionType::Incremental,
datafusion_physical_plan::execution_plan::Boundedness::Bounded,
);
Self {
schema,
properties: Arc::new(properties),
}
}
}
impl DisplayAs for ErroringExec {
fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "ErroringExec")
}
}
impl ExecutionPlan for ErroringExec {
fn name(&self) -> &str {
"ErroringExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
Ok(self)
}
fn execute(
&self,
_partition: usize,
_context: Arc<TaskContext>,
) -> DataFusionResult<SendableRecordBatchStream> {
let batch = record_batch!(("id", Int32, [1, 2])).unwrap();
let stream = futures::stream::iter(vec![
Ok(batch),
Err(DataFusionError::Execution("boom".to_string())),
]);
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.schema.clone(),
stream,
)))
}
}
#[tokio::test]
async fn test_multipart_chunked_splits_into_parts() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A 1-byte budget forces every batch into its own part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_single_part_when_under_budget() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A large byte budget and no time limit keep the whole partition in a
// single part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_chunked_splits_by_duration() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A large byte budget but a tiny duration budget: writing and sending
// one batch already takes longer than the limit, so each batch is cut
// into its own part on the time check rather than the byte check.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
Some(std::time::Duration::from_nanos(1)),
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_empty_partition_stages_nothing() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
// An empty partition should stage no parts; on the multipart path the
// write relies on another partition having data to commit.
let input = input_plan_from_batches(schema, vec![]).await;
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(insert_count.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_multipart_chunked_uses_distinct_part_ids() {
use futures::StreamExt;
use std::collections::HashSet;
let part_ids = Arc::new(Mutex::new(Vec::new()));
let client = recording_insert_client(part_ids.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
// A 1-byte budget forces every batch into its own part.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
let ids = part_ids.lock().unwrap().clone();
assert_eq!(ids.len(), 3, "expected one part id per part: {ids:?}");
assert!(
ids.iter().all(|id| !id.is_empty()),
"part ids must be non-empty: {ids:?}"
);
let unique: HashSet<&String> = ids.iter().collect();
assert_eq!(unique.len(), 3, "part ids must be distinct: {ids:?}");
}
#[tokio::test]
async fn test_multipart_chunks_each_partition_independently() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let partitions = vec![
// Partition 0: two batches, split into two parts by the 1-byte budget.
vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
],
// Partition 1: one batch, one part.
vec![record_batch!(("id", Int32, [5, 6])).unwrap()],
];
let input = input_plan_from_partitions(schema, partitions).await;
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(1),
None,
);
for partition in 0..2 {
let mut stream = exec
.execute(partition, Arc::new(TaskContext::default()))
.unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
}
// 2 parts from partition 0 + 1 part from partition 1.
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_multipart_input_error_surfaces_original() {
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
// A large byte budget keeps the good batch and the following error in
// the same part, exercising the mid-part abort path.
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
None,
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
let mut err = None;
while let Some(item) = stream.next().await {
if let Err(e) = item {
err = Some(e);
break;
}
}
let err = err.expect("expected the input stream error to surface");
// The original DataFusion error must win over the HTTP error it induces.
assert!(
err.to_string().contains("boom"),
"expected original input error, got: {err}"
);
}
#[tokio::test]
async fn test_merge_insert_input_error_surfaces_original() {
// Regression test for #2339 on the single-request merge_insert path.
// When the input stream errors mid-body, Hyper masks it under HTTP2 as a
// generic "stream error sent by user" message. The error side-channel
// must recover and surface the original DataFusion error instead.
use futures::StreamExt;
let client = crate::remote::client::test_utils::client_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
http::Response::builder()
.status(200)
.body(
r#"{"version": 2, "num_updated_rows": 0, "num_inserted_rows": 0, "num_deleted_rows": 0}"#
.to_string(),
)
.unwrap()
});
let query = MergeInsertRequest {
on: vec!["id".to_string()],
when_matched_update_all: false,
when_matched_update_all_filt: None,
when_not_matched_insert_all: false,
when_not_matched_by_source_delete: false,
when_not_matched_by_source_delete_filt: None,
use_index: true,
use_lsm: None,
};
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
let exec = RemoteWriteExec::new(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
WriteOp::MergeInsert {
query,
timeout: None,
},
None,
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
let mut err = None;
while let Some(item) = stream.next().await {
if let Err(e) = item {
err = Some(e);
break;
}
}
let err = err.expect("expected the input stream error to surface");
assert!(
err.to_string().contains("boom"),
"expected original input error, got: {err}"
);
}
#[tokio::test]
async fn test_multipart_records_progress_within_a_part() {
use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker};
use futures::StreamExt;
let insert_count = Arc::new(AtomicUsize::new(0));
let client = counting_insert_client(insert_count.clone());
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"id",
DataType::Int32,
true,
)]));
let batches = vec![
record_batch!(("id", Int32, [1, 2])).unwrap(),
record_batch!(("id", Int32, [3, 4])).unwrap(),
record_batch!(("id", Int32, [5, 6])).unwrap(),
];
let input = input_plan_from_batches(schema, batches).await;
let observed = Arc::new(Mutex::new(Vec::<usize>::new()));
let observed_cb = observed.clone();
let callback: ProgressCallback = Arc::new(Mutex::new(move |p: &WriteProgress| {
observed_cb.lock().unwrap().push(p.output_bytes());
}));
let tracker = Arc::new(WriteProgressTracker::new(callback, None));
// A large byte budget keeps all three batches in one part; smooth
// progress therefore requires bytes to be reported per chunk rather than
// once when the part completes.
let exec = RemoteWriteExec::new_multipart(
"my_table".to_string(),
"my_table".to_string(),
client,
input,
false,
"upload-1".to_string(),
Some(tracker),
None,
Some(64 * 1024 * 1024),
None,
);
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
while stream.next().await.transpose().unwrap().is_some() {}
assert_eq!(
insert_count.load(Ordering::SeqCst),
1,
"batches should all land in a single part"
);
let observed = observed.lock().unwrap();
assert!(
observed.len() > 1,
"expected multiple incremental progress updates within the part: {observed:?}"
);
assert!(
observed.windows(2).all(|w| w[1] >= w[0]),
"progress bytes should be monotonic: {observed:?}"
);
assert!(
*observed.last().unwrap() > 0,
"final progress should report bytes: {observed:?}"
);
}
rust/lancedb/src/remote/table/tests/blobs.rs
+753
View File
@@ -0,0 +1,753 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_query_plain() {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let expected_body = serde_json::json!({
"filter": "`A` > 0",
"k": isize::MAX as usize,
"prefilter": true,
"vector": [], // Empty vector means no vector query.
"version": null,
});
assert_eq!(body, expected_body);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let query = AnyQuery::Query(QueryRequest {
filter: Some(QueryFilter::Sql(r#""A" > 0"#.into())),
..Default::default()
});
let data = table
.base_table()
.query(&query, Default::default())
.await
.unwrap()
.collect::<Vec<_>>()
.await;
assert_eq!(data.len(), 1);
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
fn blob_describe_response() -> http::Response<String> {
let schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
crate::blob("image", true),
Field::new("caption", DataType::Utf8, true),
crate::blob("thumbnail", true),
]);
let json_schema = JsonSchema::try_from(&schema).unwrap();
http::Response::builder()
.status(200)
.body(serde_json::json!({ "version": 1, "schema": json_schema }).to_string())
.unwrap()
}
#[rstest]
#[case(semver::Version::new(0, 1, 0))]
#[case(semver::Version::new(0, 5, 0))]
#[tokio::test]
async fn test_blob_columns_read_the_schema_on_any_server_version(#[case] version: semver::Version) {
let table = Table::new_with_handler_version("my_table", version, |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/describe/");
blob_describe_response()
});
let columns = table.blob_columns().await.unwrap();
assert_eq!(columns, vec!["image".to_string(), "thumbnail".to_string()]);
}
#[tokio::test]
async fn test_fetch_blobs_decodes_null_aligned_bytes() {
let mut builder = LargeBinaryBuilder::new();
builder.append_value(b"alpha");
builder.append_null();
builder.append_value(b"gamma");
let blobs = builder.finish();
let expected = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(blobs)],
)
.unwrap();
let expected_ref = expected.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["column"], "image");
assert_eq!(body["row_ids"], serde_json::json!([10, 20, 30]));
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(write_ipc_stream_uncompressed(&expected_ref))
.unwrap()
},
);
let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap();
assert_eq!(blobs.len(), 3);
assert_eq!(blobs.value(0), b"alpha");
assert!(blobs.is_null(1));
assert_eq!(blobs.value(2), b"gamma");
}
#[tokio::test]
async fn test_fetch_blobs_concatenates_multiple_batches() {
let schema = Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)]));
let mut first_builder = LargeBinaryBuilder::new();
first_builder.append_value(b"alpha");
first_builder.append_null();
let first_batch =
RecordBatch::try_new(schema.clone(), vec![Arc::new(first_builder.finish())]).unwrap();
let mut second_builder = LargeBinaryBuilder::new();
second_builder.append_value(b"gamma");
let second_batch =
RecordBatch::try_new(schema.clone(), vec![Arc::new(second_builder.finish())]).unwrap();
let mut body = Vec::new();
{
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap();
writer.write(&first_batch).unwrap();
writer.write(&second_batch).unwrap();
writer.finish().unwrap();
}
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/");
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(body.clone())
.unwrap()
},
);
let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap();
assert_eq!(blobs.len(), 3);
assert_eq!(blobs.value(0), b"alpha");
assert!(blobs.is_null(1));
assert_eq!(blobs.value(2), b"gamma");
}
#[tokio::test]
async fn test_fetch_blobs_splits_row_ids_at_one_version_and_preserves_order() {
let request_sizes = Arc::new(std::sync::Mutex::new(Vec::new()));
let seen = request_sizes.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
if request.url().path() == "/v1/table/my_table/describe/" {
return http::Response::builder()
.status(200)
.body(br#"{"version":7,"schema":{"fields":[]}}"#.to_vec())
.unwrap();
}
assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/");
let body = request_body_json(&request);
assert_eq!(body["version"], 7);
let ids = body["row_ids"].as_array().unwrap();
seen.lock().unwrap().push(ids.len());
if ids.len() > 1024 {
return http::Response::builder()
.status(400)
.body(b"fetch_blobs accepts at most 1024 row IDs".to_vec())
.unwrap();
}
let mut builder = LargeBinaryBuilder::new();
for id in ids {
let id = id.as_u64().unwrap();
if id == 1023 {
builder.append_null();
} else {
builder.append_value(id.to_string().as_bytes());
}
}
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(builder.finish())],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(write_ipc_stream_uncompressed(&batch))
.unwrap()
},
);
let ids: Vec<u64> = (0..1024).chain([42]).collect();
let blobs = table.fetch_blobs("image", &ids).await.unwrap();
assert_eq!(blobs.len(), 1025);
assert_eq!(blobs.value(0), b"0");
assert_eq!(blobs.value(1022), b"1022");
assert!(blobs.is_null(1023));
assert_eq!(blobs.value(1024), b"42");
assert_eq!(request_sizes.lock().unwrap().as_slice(), &[1024, 1]);
}
#[tokio::test]
async fn test_fetch_blobs_splits_byte_limited_requests_and_reads_large_blob_by_range() {
// Simulate a lower byte cap so this test exercises the same 400 response
// without allocating 64 MiB of blob data.
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let seen = requests.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return http::Response::builder()
.status(200)
.body(br#"{"version":42,"schema":{"fields":[]}}"#.to_vec())
.unwrap();
}
if path == "/v1/table/my_table/fetch_blobs/" {
let body = request_body_json(&request);
let ids = body["row_ids"].as_array().unwrap();
if body["version"].is_null() {
// Only the initial failed request may read live latest.
assert_eq!(ids.len(), 5);
} else {
assert_eq!(body["version"], 42);
}
seen.lock().unwrap().push(format!("POST {}", ids.len()));
let mut builder = LargeBinaryBuilder::new();
let mut total_bytes = 0;
for id in ids {
let value: Option<&[u8]> = match id.as_u64().unwrap() {
10 => Some(b"aaaa"),
20 => Some(b"bbb"),
30 => None,
40 => Some(b"0123456789"),
id => panic!("unexpected row id {id}"),
};
if let Some(value) = value {
total_bytes += value.len();
builder.append_value(value);
} else {
builder.append_null();
}
}
if total_bytes > 6 {
return http::Response::builder()
.status(400)
.body(br#"{"error":"Bad request: fetch_blobs accepts at most 67108864 total blob bytes"}"#.to_vec())
.unwrap();
}
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(builder.finish())],
)
.unwrap();
return http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(write_ipc_stream_uncompressed(&batch))
.unwrap();
}
assert_eq!(path, "/v1/table/my_table/blob/image/40/bytes");
assert!(request.url().query().unwrap().contains("version=42"));
let range = request
.headers()
.get(reqwest::header::RANGE)
.unwrap()
.to_str()
.unwrap();
seen.lock().unwrap().push(format!("GET {range}"));
match range {
"bytes=0-0" => http::Response::builder()
.status(206)
.header(reqwest::header::CONTENT_RANGE, "bytes 0-0/10")
.header(VERSION_HEADER, "42")
.body(b"0".to_vec())
.unwrap(),
"bytes=0-" => http::Response::builder()
.status(206)
.header(reqwest::header::CONTENT_RANGE, "bytes 0-9/10")
.body(b"0123456789".to_vec())
.unwrap(),
_ => panic!("unexpected range: {range}"),
}
},
);
let blobs = table
.fetch_blobs("image", &[10, 20, 30, 40, 20])
.await
.unwrap();
assert_eq!(blobs.len(), 5);
assert_eq!(blobs.value(0), b"aaaa");
assert_eq!(blobs.value(1), b"bbb");
assert!(blobs.is_null(2));
assert_eq!(blobs.value(3), b"0123456789");
assert_eq!(blobs.value(4), b"bbb");
let requests = requests.lock().unwrap();
assert!(requests.contains(&"GET bytes=0-0".to_string()));
assert!(requests.contains(&"GET bytes=0-".to_string()));
}
#[tokio::test]
async fn test_fetch_blobs_does_not_split_unrelated_bad_requests() {
let requests = Arc::new(AtomicUsize::new(0));
let seen = requests.clone();
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 5, 0), move |_| {
seen.fetch_add(1, Ordering::SeqCst);
http::Response::builder()
.status(400)
.body(b"unknown blob column".to_vec())
.unwrap()
});
assert_fetch_blobs_http_error(
table.fetch_blobs("missing", &[10, 20]).await.unwrap_err(),
"unknown blob column",
);
assert_eq!(requests.load(Ordering::SeqCst), 1);
}
fn table_with_fetch_blobs_response(body: Vec<u8>) -> Table {
table_with_fetch_blobs_content_type(Some(ARROW_STREAM_CONTENT_TYPE), body)
}
fn table_with_fetch_blobs_content_type(content_type: Option<&'static str>, body: Vec<u8>) -> Table {
Table::new_with_handler_version("my_table", semver::Version::new(0, 5, 0), move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/");
let mut response = http::Response::builder().status(200);
if let Some(content_type) = content_type {
response = response.header(CONTENT_TYPE, content_type);
}
response.body(body.clone()).unwrap()
})
}
fn one_row_blob_batch(column: &str) -> RecordBatch {
let mut builder = LargeBinaryBuilder::new();
builder.append_value(b"alpha");
RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
column,
DataType::LargeBinary,
true,
)])),
vec![Arc::new(builder.finish())],
)
.unwrap()
}
fn one_row_blob_ipc_stream(column: &str) -> Vec<u8> {
write_ipc_stream_uncompressed(&one_row_blob_batch(column))
}
#[tokio::test]
async fn test_remote_scan_order_is_not_deterministic() {
// A distributed scan answers in no fixed order, so callers that assign meaning
// to row position have to sort for themselves.
let table = Table::new_with_handler("my_table", |_| {
http::Response::builder()
.status(200)
.body(Vec::new())
.unwrap()
});
assert!(!table.base_table().scan_order_is_deterministic());
}
#[tokio::test]
async fn test_checkout_branch_pins_without_touching_the_original() {
let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
let recorder = seen.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(br#"{"version": 42, "schema": {"fields": []}}"#.to_vec())
.unwrap(),
"/v1/table/my_table/count_rows/" => {
let body = request_body_json(&request);
recorder.lock().unwrap().push(body["version"].clone());
http::Response::builder()
.status(200)
.body(b"0".to_vec())
.unwrap()
}
path => panic!("unexpected request path: {path}"),
},
);
let pinned = table.checkout_branch("main", Some(42)).await.unwrap();
pinned.count_rows(None).await.unwrap();
table.count_rows(None).await.unwrap();
let seen = seen.lock().unwrap();
assert_eq!(seen[0], 42, "the pinned handle must send its version");
assert!(
seen[1].is_null(),
"the original handle must still track latest, got {:?}",
seen[1]
);
}
#[tokio::test]
async fn test_fetch_blobs_sends_the_checked_out_version() {
let ipc = one_row_blob_ipc_stream("image");
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(br#"{"version": 42, "schema": {"fields": []}}"#.to_vec())
.unwrap(),
"/v1/table/my_table/fetch_blobs/" => {
let body = request_body_json(&request);
assert_eq!(
body["version"], 42,
"blob reads must use the same snapshot as the query"
);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(ipc.clone())
.unwrap()
}
path => panic!("unexpected request path: {path}"),
},
);
table.checkout(42).await.unwrap();
assert_eq!(table.fetch_blobs("image", &[10]).await.unwrap().len(), 1);
}
#[tokio::test]
async fn test_fetch_blobs_sends_the_checked_out_branch() {
use lance::dataset::refs::Ref;
let ipc = one_row_blob_ipc_stream("image");
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| match request.url().path() {
"/v1/table/my_table/branches/create/" => http::Response::builder()
.status(200)
.body(b"{}".to_vec())
.unwrap(),
"/v1/table/my_table/fetch_blobs/" => {
let body = request_body_json(&request);
assert_eq!(body["branch"], "exp", "blob reads must stay on the branch");
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(ipc.clone())
.unwrap()
}
path => panic!("unexpected request path: {path}"),
},
);
let branch = table
.create_branch("exp", Ref::Version(None, None))
.await
.unwrap();
assert_eq!(branch.fetch_blobs("image", &[10]).await.unwrap().len(), 1);
}
#[tokio::test]
async fn test_fetch_blobs_sends_a_nested_column_as_a_dotted_path() {
let ipc = one_row_blob_ipc_stream("info.blob");
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
let body = request_body_json(&request);
assert_eq!(body["column"], "info.blob");
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(ipc.clone())
.unwrap()
},
);
let blobs = table.fetch_blobs("info.blob", &[10]).await.unwrap();
assert_eq!(blobs.value(0), b"alpha");
}
fn write_empty_ipc_stream(schema: &Schema) -> Vec<u8> {
let mut body = Vec::new();
arrow_ipc::writer::StreamWriter::try_new(&mut body, schema)
.unwrap()
.finish()
.unwrap();
body
}
fn assert_fetch_blobs_http_error(error: Error, expected: &str) {
match error {
Error::Http {
source, request_id, ..
} => {
assert!(source.to_string().contains(expected));
assert!(!request_id.is_empty());
}
error => panic!("expected HTTP error, got {error}"),
}
}
fn assert_not_supported_error(error: Error, expected: &str) {
match error {
Error::NotSupported { message } => assert!(message.contains(expected)),
error => panic!("expected not-supported error, got {error}"),
}
}
#[tokio::test]
async fn test_fetch_blobs_rejects_missing_column() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"other",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(LargeBinaryArray::from(vec![Some(
b"value".as_slice(),
)]))],
)
.unwrap();
let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch));
let error = table.fetch_blobs("image", &[10]).await.unwrap_err();
assert_fetch_blobs_http_error(error, "missing the 'image' column");
}
#[tokio::test]
async fn test_fetch_blobs_rejects_wrong_column_type() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::Int32,
false,
)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch));
let error = table.fetch_blobs("image", &[10]).await.unwrap_err();
assert_fetch_blobs_http_error(
error,
"type Int32, expected Binary, LargeBinary, or BinaryView",
);
}
#[rstest]
#[case(DataType::Binary)]
#[case(DataType::LargeBinary)]
#[case(DataType::BinaryView)]
#[tokio::test]
async fn test_fetch_blobs_accepts_binary_large_binary_and_binary_view(#[case] data_type: DataType) {
let binary_values = BinaryArray::from(vec![
Some(b"alpha".as_slice()),
None,
Some(b"gamma".as_slice()),
]);
let typed_column = arrow::compute::cast(&binary_values, &data_type).unwrap();
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("image", data_type, true)])),
vec![typed_column],
)
.unwrap();
let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch));
let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap();
assert_eq!(blobs.len(), 3);
assert_eq!(blobs.value(0), b"alpha");
assert!(blobs.is_null(1));
assert_eq!(blobs.value(2), b"gamma");
}
#[tokio::test]
async fn test_fetch_blobs_rejects_row_count_mismatch() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(LargeBinaryArray::from(vec![Some(
b"value".as_slice(),
)]))],
)
.unwrap();
let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch));
let error = table.fetch_blobs("image", &[10, 20]).await.unwrap_err();
assert_fetch_blobs_http_error(error, "returned 1 rows for 2 row ids");
}
#[tokio::test]
async fn test_fetch_blobs_rejects_zero_batches_for_nonempty_row_ids() {
let schema = Schema::new(vec![Field::new("image", DataType::LargeBinary, true)]);
let table = table_with_fetch_blobs_response(write_empty_ipc_stream(&schema));
let error = table.fetch_blobs("image", &[10]).await.unwrap_err();
assert_fetch_blobs_http_error(error, "returned 0 rows for 1 row ids");
}
#[tokio::test]
async fn test_fetch_blobs_skips_the_request_for_empty_row_ids() {
let table = Table::new_with_handler("my_table", |_| -> http::Response<String> {
panic!("fetch_blobs must not call the server for an empty selection");
});
let blobs = table.fetch_blobs("image", &[]).await.unwrap();
assert!(blobs.is_empty());
}
#[tokio::test]
async fn test_blob_byte_apis_not_supported_on_old_server() {
let table = Table::new_with_handler("my_table", |_| -> http::Response<String> {
panic!("blob request must not reach a server without blob support");
});
assert_not_supported_error(
table.fetch_blobs("image", &[1]).await.unwrap_err(),
"fetch_blobs",
);
assert_not_supported_error(
table.fetch_blob_files("image", &[1]).await.unwrap_err(),
"requires LanceDB Cloud server 0.5.0 or newer",
);
}
#[rstest]
#[case(ARROW_STREAM_CONTENT_TYPE)]
#[case("application/vnd.apache.arrow.stream; charset=utf-8")]
#[case("APPLICATION/VND.APACHE.ARROW.STREAM")]
#[tokio::test]
async fn test_fetch_blobs_accepts_stream_content_type_variants(#[case] content_type: &'static str) {
let table = table_with_fetch_blobs_content_type(
Some(content_type),
write_ipc_stream_uncompressed(&one_row_blob_batch("image")),
);
let blobs = table.fetch_blobs("image", &[10]).await.unwrap();
assert_eq!(blobs.value(0), b"alpha");
}
// Shared decoder accepts file framing because /query still returns it.
// fetch_blobs wire contract is stream. Enforced on the server.
#[rstest]
#[case(ARROW_FILE_CONTENT_TYPE)]
#[case("application/vnd.apache.arrow.file; charset=utf-8")]
#[case("APPLICATION/VND.APACHE.ARROW.FILE")]
#[tokio::test]
async fn test_fetch_blobs_accepts_file_content_type_variants(#[case] content_type: &'static str) {
let table = table_with_fetch_blobs_content_type(
Some(content_type),
write_ipc_file(&one_row_blob_batch("image")),
);
let blobs = table.fetch_blobs("image", &[10]).await.unwrap();
assert_eq!(blobs.value(0), b"alpha");
}
#[rstest]
#[case(ARROW_STREAM_CONTENT_TYPE, write_ipc_file(&one_row_blob_batch("image")))]
#[case(
ARROW_FILE_CONTENT_TYPE,
write_ipc_stream_uncompressed(&one_row_blob_batch("image"))
)]
#[tokio::test]
async fn test_fetch_blobs_fails_when_the_body_contradicts_the_content_type(
#[case] content_type: &'static str,
#[case] body: Vec<u8>,
) {
let table = table_with_fetch_blobs_content_type(Some(content_type), body);
let error = table.fetch_blobs("image", &[10]).await.unwrap_err();
assert!(
matches!(error, Error::Arrow { .. }),
"expected an Arrow decode failure, got {error}"
);
}
#[tokio::test]
async fn test_fetch_blobs_rejects_a_response_that_is_not_arrow_ipc() {
let table =
table_with_fetch_blobs_content_type(Some("application/json"), br#"{"blobs": []}"#.to_vec());
let error = table.fetch_blobs("image", &[10]).await.unwrap_err();
assert_fetch_blobs_http_error(error, "got 'application/json'");
}
#[tokio::test]
async fn test_fetch_blobs_without_content_type_falls_back_to_file_framing() {
let table =
table_with_fetch_blobs_content_type(None, write_ipc_file(&one_row_blob_batch("image")));
let blobs = table.fetch_blobs("image", &[10]).await.unwrap();
assert_eq!(blobs.value(0), b"alpha");
}
rust/lancedb/src/remote/table/tests/branches.rs
+1223
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/tests/columns.rs
+1041
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/tests/freshness.rs
+717
View File
@@ -0,0 +1,717 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
// ---- Read freshness header tests ------------------------------------
#[test]
fn test_compute_min_timestamp_combines_baseline_and_interval() {
let now = SystemTime::now();
let baseline = now - Duration::from_secs(60);
// No interval, no baseline -> no header.
assert_eq!(
compute_min_timestamp(&FreshnessState::default(), None, now),
None
);
// Baseline only -> baseline.
let state = FreshnessState {
min_version: None,
checkout_baseline: Some(baseline),
min_read_version: None,
..FreshnessState::default()
};
assert_eq!(compute_min_timestamp(&state, None, now), Some(baseline));
// ZERO interval, no baseline -> now.
assert_eq!(
compute_min_timestamp(&FreshnessState::default(), Some(Duration::ZERO), now),
Some(now)
);
// Positive interval, no baseline -> now - interval.
assert_eq!(
compute_min_timestamp(
&FreshnessState::default(),
Some(Duration::from_secs(10)),
now
),
Some(now - Duration::from_secs(10))
);
// Both: pick the more-recent (i.e. tighter) constraint.
// baseline = now-60, now-interval = now-10. now-10 is newer.
let state = FreshnessState {
min_version: None,
checkout_baseline: Some(baseline),
min_read_version: None,
..FreshnessState::default()
};
assert_eq!(
compute_min_timestamp(&state, Some(Duration::from_secs(10)), now),
Some(now - Duration::from_secs(10))
);
// Both, baseline newer: pick baseline.
let recent_baseline = now - Duration::from_secs(5);
let state = FreshnessState {
min_version: None,
checkout_baseline: Some(recent_baseline),
min_read_version: None,
..FreshnessState::default()
};
assert_eq!(
compute_min_timestamp(&state, Some(Duration::from_secs(60)), now),
Some(recent_baseline)
);
}
/// Allowed slop when comparing a header timestamp against a locally
/// captured wall-clock bound. Tests run fast enough that 1s is plenty.
const FRESHNESS_TOLERANCE: Duration = Duration::from_secs(1);
fn capturing_handler<F>(
body_for: F,
) -> (
impl Fn(reqwest::Request) -> http::Response<String> + Clone + Send + Sync + 'static,
Arc<std::sync::Mutex<Option<http::HeaderMap>>>,
)
where
F: Fn(&str) -> String + Clone + Send + Sync + 'static,
{
let captured = Arc::new(std::sync::Mutex::new(None));
let captured_c = captured.clone();
let handler = move |request: reqwest::Request| {
*captured_c.lock().unwrap() = Some(request.headers().clone());
let path = request.url().path().to_string();
http::Response::builder()
.status(200)
.body(body_for(&path))
.unwrap()
};
(handler, captured)
}
fn parse_min_timestamp(headers: &http::HeaderMap) -> SystemTime {
let value = headers
.get("x-lancedb-min-timestamp")
.expect("expected x-lancedb-min-timestamp header")
.to_str()
.unwrap();
chrono::DateTime::parse_from_rfc3339(value)
.unwrap()
.with_timezone(&chrono::Utc)
.into()
}
#[tokio::test]
async fn test_freshness_default_sends_no_headers() {
let (handler, captured) = capturing_handler(|_| "42".to_string());
let table = Table::new_with_handler("my_table", handler);
let _ = table.count_rows(None).await.unwrap();
let headers = captured.lock().unwrap().clone().unwrap();
assert!(!headers.contains_key("x-lancedb-min-timestamp"));
assert!(!headers.contains_key("x-lancedb-min-version"));
}
#[tokio::test]
async fn test_freshness_zero_interval_sends_now() {
let (handler, captured) = capturing_handler(|_| "42".to_string());
let table =
Table::new_with_handler_and_interval("my_table", handler, Some(Duration::from_secs(0)));
let before = SystemTime::now();
table.count_rows(None).await.unwrap();
let after = SystemTime::now();
let headers = captured.lock().unwrap().clone().unwrap();
let sent = parse_min_timestamp(&headers);
assert!(
sent >= before - FRESHNESS_TOLERANCE && sent <= after + FRESHNESS_TOLERANCE,
"expected timestamp roughly equal to wall clock"
);
assert!(!headers.contains_key("x-lancedb-min-version"));
}
#[tokio::test]
async fn test_checkout_disables_read_consistency_interval() {
let (handler, captured) = capturing_handler(|path| match path {
"/v1/table/my_table/describe/" => r#"{"version":5,"schema":{"fields":[]}}"#.to_string(),
"/v1/table/my_table/count_rows/" => "42".to_string(),
_ => panic!("unexpected path: {}", path),
});
let table =
Table::new_with_handler_and_interval("my_table", handler, Some(Duration::from_secs(0)));
table.checkout(5).await.unwrap();
table.count_rows(None).await.unwrap();
let headers = captured.lock().unwrap().clone().unwrap();
assert!(!headers.contains_key("x-lancedb-min-timestamp"));
assert!(!headers.contains_key("x-lancedb-min-version"));
assert!(!headers.contains_key("x-lancedb-min-read-version"));
}
#[tokio::test]
async fn test_read_snapshot_keeps_selector_and_freshness_generation_bound() {
let table = RemoteTable::new_mock_with_consistency_interval(
"my_table".to_string(),
|_| {
http::Response::builder()
.status(200)
.body(r#"{"version":5,"schema":{"fields":[]}}"#.to_string())
.unwrap()
},
Some(Duration::ZERO),
);
let latest = table.snapshot_read_state().await;
table.checkout(5).await.unwrap();
let latest_request = latest
.freshness
.apply(
table
.client
.post("/v1/table/my_table/count_rows/")
.json(&serde_json::json!({ "version": latest.version })),
)
.build()
.unwrap();
assert!(request_body_json(&latest_request)["version"].is_null());
assert!(latest_request.headers().contains_key(MIN_TIMESTAMP_HEADER));
let pinned = table.snapshot_read_state().await;
let pinned_request = pinned
.freshness
.apply(
table
.client
.post("/v1/table/my_table/count_rows/")
.json(&serde_json::json!({ "version": pinned.version })),
)
.build()
.unwrap();
assert_eq!(request_body_json(&pinned_request)["version"], 5);
assert!(!pinned_request.headers().contains_key(MIN_TIMESTAMP_HEADER));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_cancelled_checkout_keeps_latest_freshness_enabled() {
let (described_tx, described_rx) = std::sync::mpsc::channel::<()>();
let table = Arc::new(RemoteTable::new_mock_with_consistency_interval(
"my_table".to_string(),
move |_| {
described_tx.send(()).unwrap();
http::Response::builder()
.status(200)
.body(r#"{"version":5,"schema":{"fields":[]}}"#.to_string())
.unwrap()
},
Some(Duration::from_secs(0)),
));
let version_guard = table.version.write().await;
let checkout = tokio::spawn({
let table = table.clone();
async move { table.checkout(5).await }
});
tokio::task::spawn_blocking(move || {
described_rx.recv_timeout(Duration::from_secs(10)).unwrap()
})
.await
.unwrap();
for _ in 0..100 {
if table.freshness.lock().unwrap().pinned {
break;
}
tokio::task::yield_now().await;
}
assert!(!checkout.is_finished());
checkout.abort();
assert!(checkout.await.unwrap_err().is_cancelled());
drop(version_guard);
assert_eq!(*table.version.read().await, None);
assert!(table.snapshot_freshness_headers().min_timestamp.is_some());
}
#[tokio::test]
async fn test_freshness_positive_interval_sends_now_minus_interval() {
let (handler, captured) = capturing_handler(|_| "42".to_string());
let interval = Duration::from_secs(30);
let table = Table::new_with_handler_and_interval("my_table", handler, Some(interval));
let before = SystemTime::now();
table.count_rows(None).await.unwrap();
let after = SystemTime::now();
let headers = captured.lock().unwrap().clone().unwrap();
let sent = parse_min_timestamp(&headers);
assert!(
sent >= before - interval - FRESHNESS_TOLERANCE
&& sent <= after - interval + FRESHNESS_TOLERANCE,
"expected timestamp roughly equal to now - interval"
);
}
#[tokio::test]
async fn test_freshness_checkout_latest_sets_baseline() {
let (handler, captured) = capturing_handler(|path| match path {
"/v1/table/my_table/count_rows/" => "42".to_string(),
_ => panic!("unexpected path: {}", path),
});
// No interval — only the baseline should drive the timestamp.
let table = Table::new_with_handler_and_interval("my_table", handler, None);
let before_checkout = SystemTime::now();
table.checkout_latest().await.unwrap();
let after_checkout = SystemTime::now();
table.count_rows(None).await.unwrap();
let headers = captured.lock().unwrap().clone().unwrap();
let sent = parse_min_timestamp(&headers);
assert!(
sent >= before_checkout - FRESHNESS_TOLERANCE
&& sent <= after_checkout + FRESHNESS_TOLERANCE,
"expected timestamp captured at checkout_latest() time"
);
assert!(!headers.contains_key("x-lancedb-min-version"));
}
#[tokio::test]
async fn test_freshness_min_version_tracked_after_write() {
let (handler, captured) = capturing_handler(|path| match path {
"/v1/table/my_table/update/" => r#"{"rows_updated":1,"version":7}"#.to_string(),
"/v1/table/my_table/count_rows/" => "42".to_string(),
_ => panic!("unexpected path: {}", path),
});
let table = Table::new_with_handler("my_table", handler);
let _ = table.update().column("a", "a + 1").execute().await.unwrap();
// Update headers also pass through captured; reset by reading after.
table.count_rows(None).await.unwrap();
let headers = captured.lock().unwrap().clone().unwrap();
assert_eq!(
headers
.get("x-lancedb-min-version")
.unwrap()
.to_str()
.unwrap(),
"7"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_inflight_write_result_cannot_cross_checkout_generation() {
let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
let release_rx = Arc::new(std::sync::Mutex::new(release_rx));
let (arrived_tx, arrived_rx) = std::sync::mpsc::channel::<()>();
let arrived_tx = Arc::new(std::sync::Mutex::new(arrived_tx));
let count_headers = Arc::new(std::sync::Mutex::new(None));
let captured = count_headers.clone();
let table = Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/update/" => {
arrived_tx.lock().unwrap().send(()).unwrap();
release_rx
.lock()
.unwrap()
.recv_timeout(Duration::from_secs(10))
.unwrap();
http::Response::builder()
.status(200)
.body(r#"{"rows_updated":1,"version":100}"#.to_string())
.unwrap()
}
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(r#"{"version":5,"schema":{"fields":[]}}"#.to_string())
.unwrap(),
"/v1/table/my_table/count_rows/" => {
*captured.lock().unwrap() = Some(request.headers().clone());
http::Response::builder()
.status(200)
.body("1".to_string())
.unwrap()
}
path => panic!("unexpected path: {path}"),
});
let update = tokio::spawn({
let table = table.clone();
async move { table.update().column("a", "a + 1").execute().await }
});
tokio::task::spawn_blocking(move || arrived_rx.recv_timeout(Duration::from_secs(10)).unwrap())
.await
.unwrap();
table.checkout(5).await.unwrap();
release_tx.send(()).unwrap();
update.await.unwrap().unwrap();
table.count_rows(None).await.unwrap();
let headers = count_headers.lock().unwrap();
let headers = headers.as_ref().unwrap();
assert!(!headers.contains_key("x-lancedb-min-version"));
assert!(!headers.contains_key("x-lancedb-min-read-version"));
}
/// A handler that records every request's headers and answers each read with
/// an `x-lancedb-version` response header taken from `versions` (by call
/// index, saturating at the last entry). An empty string means "no header".
fn read_version_handler(
versions: &'static [&'static str],
) -> (
impl Fn(reqwest::Request) -> http::Response<String> + Clone + Send + Sync + 'static,
Arc<std::sync::Mutex<Vec<http::HeaderMap>>>,
) {
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let requests_c = requests.clone();
let call = Arc::new(AtomicUsize::new(0));
let handler = move |request: reqwest::Request| {
requests_c.lock().unwrap().push(request.headers().clone());
let i = call.fetch_add(1, Ordering::SeqCst).min(versions.len() - 1);
let mut builder = http::Response::builder().status(200);
if !versions[i].is_empty() {
builder = builder.header("x-lancedb-version", versions[i]);
}
builder.body("42".to_string()).unwrap()
};
(handler, requests)
}
#[tokio::test]
async fn test_read_version_watermark_tracked_and_sent() {
let (handler, requests) = read_version_handler(&["100", "100"]);
let table = Table::new_with_handler("my_table", handler);
// First read has no watermark yet; the response advertises version 100,
// so the second read must floor the server at 100.
table.count_rows(None).await.unwrap();
table.count_rows(None).await.unwrap();
let reqs = requests.lock().unwrap();
assert!(!reqs[0].contains_key("x-lancedb-min-read-version"));
assert_eq!(
reqs[1]
.get("x-lancedb-min-read-version")
.unwrap()
.to_str()
.unwrap(),
"100"
);
}
#[tokio::test]
async fn test_schema_response_advances_read_watermark() {
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let captured = requests.clone();
let table = Table::new_with_handler("my_table", move |request| {
captured
.lock()
.unwrap()
.push((request.url().path().to_string(), request.headers().clone()));
match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.header("x-lancedb-version", "100")
.body(r#"{"version":100,"schema":{"fields":[]}}"#.to_string())
.unwrap(),
"/v1/table/my_table/count_rows/" => http::Response::builder()
.status(200)
.body("42".to_string())
.unwrap(),
path => panic!("unexpected path: {path}"),
}
});
table.schema().await.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 42);
let requests = requests.lock().unwrap();
let count_headers = &requests[1].1;
assert_eq!(
count_headers
.get("x-lancedb-min-read-version")
.and_then(|value| value.to_str().ok()),
Some("100")
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_inflight_schema_response_cannot_cross_checkout_generation() {
let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
let release_rx = Arc::new(std::sync::Mutex::new(release_rx));
let (arrived_tx, arrived_rx) = std::sync::mpsc::channel::<()>();
let arrived_tx = Arc::new(std::sync::Mutex::new(arrived_tx));
let count_headers = Arc::new(std::sync::Mutex::new(None));
let captured = count_headers.clone();
let table = Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => {
let body = request_body_json(&request);
if body["version"].is_null() {
arrived_tx.lock().unwrap().send(()).unwrap();
release_rx
.lock()
.unwrap()
.recv_timeout(Duration::from_secs(10))
.unwrap();
http::Response::builder()
.status(200)
.header("x-lancedb-version", "100")
.body(r#"{"version":100,"schema":{"fields":[]}}"#.to_string())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version":5,"schema":{"fields":[]}}"#.to_string())
.unwrap()
}
}
"/v1/table/my_table/count_rows/" => {
*captured.lock().unwrap() = Some(request.headers().clone());
http::Response::builder()
.status(200)
.body("1".to_string())
.unwrap()
}
path => panic!("unexpected path: {path}"),
});
let schema = tokio::spawn({
let table = table.clone();
async move { table.schema().await }
});
tokio::task::spawn_blocking(move || arrived_rx.recv_timeout(Duration::from_secs(10)).unwrap())
.await
.unwrap();
table.checkout(5).await.unwrap();
release_tx.send(()).unwrap();
schema.await.unwrap().unwrap();
table.count_rows(None).await.unwrap();
assert!(
!count_headers
.lock()
.unwrap()
.as_ref()
.unwrap()
.contains_key("x-lancedb-min-read-version")
);
}
#[tokio::test]
async fn test_streaming_write_uses_and_advances_read_watermark() {
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let describe_body = serde_json::to_string(&json!({
"version": 7,
"schema": JsonSchema::try_from(data.schema().as_ref()).unwrap(),
}))
.unwrap();
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let captured = requests.clone();
let table = Table::new_with_handler("my_table", move |request| {
captured
.lock()
.unwrap()
.push((request.url().path().to_string(), request.headers().clone()));
match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => http::Response::builder()
.status(200)
.header("x-lancedb-version", "8")
.body(r#"{"version":8}"#.to_string())
.unwrap(),
"/v1/table/my_table/count_rows/" => http::Response::builder()
.status(200)
.body("3".to_string())
.unwrap(),
path => panic!("unexpected path: {path}"),
}
});
assert_eq!(table.add(data).execute().await.unwrap().version, 8);
assert_eq!(table.count_rows(None).await.unwrap(), 3);
let requests = requests.lock().unwrap();
let insert_headers = &requests[1].1;
assert_eq!(
insert_headers
.get("x-lancedb-min-read-version")
.and_then(|value| value.to_str().ok()),
Some("7")
);
let count_headers = &requests[2].1;
assert_eq!(
count_headers
.get("x-lancedb-min-read-version")
.and_then(|value| value.to_str().ok()),
Some("8")
);
}
#[tokio::test]
async fn test_read_version_watermark_keeps_max() {
// Server reports 100 then a stale 50; the watermark must not regress.
let (handler, requests) = read_version_handler(&["100", "50", "50"]);
let table = Table::new_with_handler("my_table", handler);
table.count_rows(None).await.unwrap();
table.count_rows(None).await.unwrap();
table.count_rows(None).await.unwrap();
let reqs = requests.lock().unwrap();
assert_eq!(
reqs[2]
.get("x-lancedb-min-read-version")
.unwrap()
.to_str()
.unwrap(),
"100"
);
}
#[tokio::test]
async fn test_read_version_absent_header_no_watermark() {
// An old server that doesn't return the version header leaves the
// watermark unset, preserving backward compatibility.
let (handler, requests) = read_version_handler(&[""]);
let table = Table::new_with_handler("my_table", handler);
table.count_rows(None).await.unwrap();
table.count_rows(None).await.unwrap();
let reqs = requests.lock().unwrap();
assert!(!reqs[1].contains_key("x-lancedb-min-read-version"));
}
#[tokio::test]
async fn test_read_version_watermark_reset_on_checkout_latest() {
let (handler, requests) = read_version_handler(&["100", "100"]);
let table = Table::new_with_handler("my_table", handler);
table.count_rows(None).await.unwrap();
table.checkout_latest().await.unwrap();
table.count_rows(None).await.unwrap();
// The read after checkout_latest starts from a clean slate.
let reqs = requests.lock().unwrap();
assert!(
!reqs
.last()
.unwrap()
.contains_key("x-lancedb-min-read-version")
);
}
/// Like `capturing_handler`, but keeps a per-path snapshot of the headers
/// from every request so tests can assert on a specific endpoint.
#[allow(clippy::type_complexity)]
fn path_capturing_handler<F>(
body_for: F,
) -> (
impl Fn(reqwest::Request) -> http::Response<String> + Clone + Send + Sync + 'static,
Arc<std::sync::Mutex<HashMap<String, http::HeaderMap>>>,
)
where
F: Fn(&str) -> String + Clone + Send + Sync + 'static,
{
let captured: Arc<std::sync::Mutex<HashMap<String, http::HeaderMap>>> =
Arc::new(std::sync::Mutex::new(HashMap::new()));
let captured_c = captured.clone();
let handler = move |request: reqwest::Request| {
let path = request.url().path().to_string();
captured_c
.lock()
.unwrap()
.insert(path.clone(), request.headers().clone());
http::Response::builder()
.status(200)
.body(body_for(&path))
.unwrap()
};
(handler, captured)
}
#[tokio::test]
async fn test_freshness_checkout_validation_sends_no_min_version() {
// After a write bumps min_version, calling checkout(v) must not let
// that stale header ride along on the validating /describe/ request.
let (handler, captured) = path_capturing_handler(|path| match path {
"/v1/table/my_table/update/" => r#"{"rows_updated":1,"version":7}"#.to_string(),
"/v1/table/my_table/describe/" => r#"{"version":5,"schema":{"fields":[]}}"#.to_string(),
_ => panic!("unexpected path: {}", path),
});
let table = Table::new_with_handler("my_table", handler);
table.update().column("a", "a + 1").execute().await.unwrap();
table.checkout(5).await.unwrap();
let captured = captured.lock().unwrap();
let describe_headers = captured
.get("/v1/table/my_table/describe/")
.expect("describe should have been called by checkout(v)");
assert!(
!describe_headers.contains_key("x-lancedb-min-version"),
"checkout(v) describe must not carry stale min_version",
);
assert!(!describe_headers.contains_key("x-lancedb-min-timestamp"));
}
#[tokio::test]
async fn test_freshness_checkout_tag_resolve_sends_no_min_version() {
// Same invariant for checkout_tag: the tag-resolve request must not
// pick up a stale min_version from a prior write.
let (handler, captured) = path_capturing_handler(|path| match path {
"/v1/table/my_table/update/" => r#"{"rows_updated":1,"version":7}"#.to_string(),
"/v1/table/my_table/tags/version/" => r#"{"version":5}"#.to_string(),
_ => panic!("unexpected path: {}", path),
});
let table = Table::new_with_handler("my_table", handler);
table.update().column("a", "a + 1").execute().await.unwrap();
table.checkout_tag("v_initial").await.unwrap();
let captured = captured.lock().unwrap();
let resolve_headers = captured
.get("/v1/table/my_table/tags/version/")
.expect("tags/version should have been called by checkout_tag");
assert!(
!resolve_headers.contains_key("x-lancedb-min-version"),
"checkout_tag resolve must not carry stale min_version",
);
assert!(!resolve_headers.contains_key("x-lancedb-min-timestamp"));
}
#[tokio::test]
async fn test_freshness_checkout_clears_min_version() {
let (handler, captured) = capturing_handler(|path| match path {
"/v1/table/my_table/update/" => r#"{"rows_updated":1,"version":7}"#.to_string(),
// checkout(5) needs to describe version 5 first
"/v1/table/my_table/describe/" => r#"{"version":5,"schema":{"fields":[]}}"#.to_string(),
"/v1/table/my_table/count_rows/" => "42".to_string(),
_ => panic!("unexpected path: {}", path),
});
let table = Table::new_with_handler("my_table", handler);
table.update().column("a", "a + 1").execute().await.unwrap();
table.checkout(5).await.unwrap();
table.count_rows(None).await.unwrap();
let headers = captured.lock().unwrap().clone().unwrap();
assert!(!headers.contains_key("x-lancedb-min-version"));
assert!(!headers.contains_key("x-lancedb-min-timestamp"));
}
rust/lancedb/src/remote/table/tests/hybrid.rs
+250
View File
@@ -0,0 +1,250 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
// WAL-PK-FUSION: delete everything from here to the end of `mod tests`.
/// Verbatim from sophon's `LSM_WITH_ROW_ID_UNSUPPORTED`; the fallback keys
/// off this sentence, so a test that paraphrased it would prove nothing.
const WAL_ROW_ID_REFUSAL_BODY: &str = r#"{"code":13,"error":"Bad request: the MemWAL LSM scanner does not support with_row_id (the LSM scanner exposes _rowaddr, not a stable _rowid); set use_lsm(false) to read the base table only (results will exclude un-compacted MemWAL data)"}"#;
fn pk_schema() -> Schema {
use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
Schema::new(vec![
Field::new("id", DataType::Utf8, false).with_metadata(
[(
LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(),
"1".to_string(),
)]
.into_iter()
.collect(),
),
Field::new("text", DataType::Utf8, true),
])
}
/// A server that refuses `_rowid` the way a MemWAL table does, recording
/// every query body so a test can see which key each attempt asked for.
fn wal_refusing_table(bodies: Arc<std::sync::Mutex<Vec<serde_json::Value>>>) -> Table {
Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_response(&pk_schema()).into_bytes())
.unwrap(),
"/v1/table/my_table/query/" => {
let body = request_body_json(&request);
let asked_for_row_id = body["with_row_id"] == serde_json::Value::Bool(true);
let is_fts = body.get("full_text_query").is_some_and(|v| !v.is_null());
bodies.lock().unwrap().push(body);
if asked_for_row_id {
return http::Response::builder()
.status(400)
.body(WAL_ROW_ID_REFUSAL_BODY.as_bytes().to_vec())
.unwrap();
}
let batch = if is_fts {
leg("_score", vec!["a", "b"], false)
} else {
leg("_distance", vec!["a"], false)
};
http::Response::builder()
.status(200)
.body(write_ipc_file(&batch))
.unwrap()
}
path => panic!("unexpected request path: {path}"),
})
}
fn hybrid(table: &Table) -> crate::query::VectorQuery {
table
.query()
.full_text_search(FullTextSearchQuery::new("a".to_string()))
.nearest_to(&[0.0, 0.0])
.unwrap()
}
/// The optimistic path: ask for `_rowid`, and when a MemWAL table refuses,
/// fall back to the primary key rather than failing.
#[tokio::test]
async fn test_hybrid_falls_back_to_the_primary_key_when_row_id_is_refused() {
let bodies = Arc::new(std::sync::Mutex::new(Vec::new()));
let table = wal_refusing_table(bodies.clone());
let results = hybrid(&table)
.select(Select::columns(&["text"]))
.limit(10)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let bodies = bodies.lock().unwrap();
let (asked_row_id, asked_pk): (Vec<_>, Vec<_>) = bodies
.iter()
.partition(|b| b["with_row_id"] == serde_json::Value::Bool(true));
// The legs run concurrently under `try_join!`, which abandons the
// sibling as soon as one errors — so the refused attempt is one or two
// requests, never zero, and the retry always sends both.
assert!(
!asked_row_id.is_empty(),
"`_rowid` has to be tried before the fallback means anything"
);
assert_eq!(asked_pk.len(), 2, "both legs are retried on the key");
for body in asked_pk {
let columns = body["columns"].as_array().expect("a column projection");
assert!(
columns.iter().any(|c| c == "id"),
"the retry has to project the key: {body}"
);
}
let batch = &results[0];
assert!(batch.column_by_name(ROW_ID).is_none());
assert!(batch.column_by_name("id").is_none(), "injected key leaked");
assert_eq!(batch.num_rows(), 2, "`a` is in both legs and fuses");
}
/// A hybrid query that matches nothing must return nothing, not fail.
/// When both legs come back empty, `query_schemas` synthesizes a schema
/// that already carries `_rowid`, which the stamping has to tolerate.
#[tokio::test]
async fn test_hybrid_with_no_matches_returns_an_empty_result() {
let table = Table::new_with_handler("my_table", move |request| {
match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_response(&pk_schema()).into_bytes())
.unwrap(),
"/v1/table/my_table/query/" => {
let body = request_body_json(&request);
if body["with_row_id"] == serde_json::Value::Bool(true) {
return http::Response::builder()
.status(400)
.body(WAL_ROW_ID_REFUSAL_BODY.as_bytes().to_vec())
.unwrap();
}
// A leg that matched nothing: schema, no batches.
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
let mut body = Vec::new();
{
let mut writer =
arrow_ipc::writer::FileWriter::try_new(&mut body, &schema).unwrap();
writer.finish().unwrap();
}
http::Response::builder().status(200).body(body).unwrap()
}
path => panic!("unexpected request path: {path}"),
}
});
let results = hybrid(&table)
.limit(10)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(results.iter().map(|b| b.num_rows()).sum::<usize>(), 0);
}
/// The refusal is paid once: the second query goes straight to the key.
#[tokio::test]
async fn test_hybrid_remembers_the_refusal() {
let bodies = Arc::new(std::sync::Mutex::new(Vec::new()));
let table = wal_refusing_table(bodies.clone());
for _ in 0..2 {
hybrid(&table)
.limit(10)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
}
let refused = bodies
.lock()
.unwrap()
.iter()
.filter(|b| b["with_row_id"] == serde_json::Value::Bool(true))
.count();
assert!(
(1..=2).contains(&refused),
"only the first query should pay the refusal, saw {refused} attempts"
);
}
/// A table that never refuses must never pay for the fallback existing.
#[tokio::test]
async fn test_hybrid_on_a_base_table_sends_no_extra_requests() {
let paths = Arc::new(std::sync::Mutex::new(Vec::new()));
let seen = paths.clone();
let table = Table::new_with_handler("my_table", move |request| {
seen.lock().unwrap().push(request.url().path().to_string());
let batch = if request_body_json(&request)
.get("full_text_query")
.is_some_and(|v| !v.is_null())
{
leg("_score", vec!["a", "b"], true)
} else {
leg("_distance", vec!["a"], true)
};
http::Response::builder()
.status(200)
.body(write_ipc_file(&batch))
.unwrap()
});
hybrid(&table)
.limit(10)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let paths = paths.lock().unwrap();
assert!(
paths.iter().all(|p| p.ends_with("/query/")),
"a table that never refuses must not be probed, got {paths:?}"
);
assert_eq!(paths.len(), 2, "one request per leg, got {paths:?}");
}
/// A caller who asked for `_rowid` gets the server's refusal, not a
/// surrogate dressed up as one.
#[tokio::test]
async fn test_hybrid_does_not_fall_back_when_the_caller_asked_for_row_id() {
let bodies = Arc::new(std::sync::Mutex::new(Vec::new()));
let table = wal_refusing_table(bodies.clone());
let err = hybrid(&table)
.with_row_id()
.limit(10)
.execute()
.await
.err()
.expect("the refusal must reach the caller");
// Translated, not passed through: the caller gets the reason and the
// way out, which a bare 400 does not carry.
assert!(matches!(err, Error::NotSupported { .. }), "{err:?}");
assert!(err.to_string().contains("use_lsm(false)"), "{err}");
assert!(
bodies
.lock()
.unwrap()
.iter()
.all(|b| b["with_row_id"] == serde_json::Value::Bool(true)),
"no retry should have been attempted"
);
}
rust/lancedb/src/remote/table/tests/index.rs
+1280
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/tests/lsm.rs
+901
View File
@@ -0,0 +1,901 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_prewarm_index() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/prewarm/"
);
http::Response::builder().status(200).body("{}").unwrap()
});
table.prewarm_index("my_index").await.unwrap();
}
#[tokio::test]
async fn test_prewarm_index_not_found() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/prewarm/"
);
http::Response::builder().status(404).body("{}").unwrap()
});
let e = table.prewarm_index("my_index").await.unwrap_err();
assert!(matches!(e, Error::IndexNotFound { .. }));
}
#[tokio::test]
async fn test_prewarm_data() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/page_cache/prewarm/"
);
http::Response::builder().status(200).body("{}").unwrap()
});
table.prewarm_data(None).await.unwrap();
}
#[tokio::test]
async fn test_prewarm_data_with_columns() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/page_cache/prewarm/"
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["columns"], serde_json::json!(["col_a", "col_b"]));
http::Response::builder().status(200).body("{}").unwrap()
});
table
.prewarm_data(Some(vec!["col_a".into(), "col_b".into()]))
.await
.unwrap();
}
#[tokio::test]
async fn test_drop_index() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/drop/"
);
http::Response::builder().status(200).body("{}").unwrap()
});
table.drop_index("my_index").await.unwrap();
}
#[tokio::test]
async fn test_drop_index_not_exists() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/drop/"
);
http::Response::builder().status(404).body("{}").unwrap()
});
// Assert that the error is IndexNotFound
let e = table.drop_index("my_index").await.unwrap_err();
assert!(matches!(e, Error::IndexNotFound { .. }));
}
/// Index names are unvalidated, so reserved characters must be
/// percent-encoded or they restructure the request path.
#[tokio::test]
async fn test_per_index_paths_encode_reserved_characters() {
const NAME: &str = "my/index?a#b c";
const PREFIX: &str = "/v1/table/my_table/index/my%2Findex%3Fa%23b%20c";
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.url().path(), format!("{PREFIX}/stats/"));
let body = serde_json::json!({
"num_indexed_rows": 1,
"num_unindexed_rows": 0,
"index_type": "IVF_PQ",
"distance_type": "l2"
});
http::Response::builder()
.status(200)
.body(serde_json::to_string(&body).unwrap())
.unwrap()
});
assert!(table.index_stats(NAME).await.unwrap().is_some());
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.url().path(), format!("{PREFIX}/drop/"));
http::Response::builder().status(200).body("{}").unwrap()
});
table.drop_index(NAME).await.unwrap();
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.url().path(), format!("{PREFIX}/prewarm/"));
http::Response::builder().status(200).body("{}").unwrap()
});
table.prewarm_index(NAME).await.unwrap();
}
#[tokio::test]
async fn test_set_lsm_write_spec_unsharded() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/set_lsm_write_spec/"
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["sharding"], serde_json::json!({ "mode": "unsharded" }));
assert_eq!(body["maintained_indexes"], serde_json::json!(["id_idx"]));
assert_eq!(
body["writer_config_defaults"],
serde_json::json!({ "max_memtable_rows": "1000" })
);
http::Response::builder()
.status(200)
.body(r#"{"maintained_indexes":["id_idx"]}"#)
.unwrap()
});
let spec = crate::table::LsmWriteSpec::unsharded()
.with_maintained_indexes(vec!["id_idx".to_string()])
.with_writer_config_defaults([("max_memtable_rows", "1000")]);
table.set_lsm_write_spec(spec).await.unwrap();
}
#[tokio::test]
async fn test_set_lsm_write_spec_bucket() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/set_lsm_write_spec/"
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["sharding"],
serde_json::json!({ "mode": "bucket", "column": "id", "num_buckets": 16 })
);
// An unpinned maintained set sends null: resolve server-side.
assert_eq!(body["maintained_indexes"], serde_json::Value::Null);
http::Response::builder().status(200).body("{}").unwrap()
});
table
.set_lsm_write_spec(crate::table::LsmWriteSpec::bucket("id", 16))
.await
.unwrap();
}
/// `[]` (none) must stay distinguishable on the wire from null (all).
#[tokio::test]
async fn test_set_lsm_write_spec_no_maintained_indexes() {
let table = Table::new_with_handler("my_table", |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["maintained_indexes"], serde_json::json!([]));
http::Response::builder().status(200).body("{}").unwrap()
});
table
.set_lsm_write_spec(
crate::table::LsmWriteSpec::bucket("id", 16).with_maintained_indexes(Vec::new()),
)
.await
.unwrap();
}
#[tokio::test]
async fn test_set_lsm_write_spec_identity() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/set_lsm_write_spec/"
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["sharding"],
serde_json::json!({ "mode": "identity", "column": "tenant" })
);
http::Response::builder().status(200).body("{}").unwrap()
});
table
.set_lsm_write_spec(crate::table::LsmWriteSpec::identity("tenant"))
.await
.unwrap();
}
#[tokio::test]
async fn test_unset_lsm_write_spec() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/unset_lsm_write_spec/"
);
http::Response::builder().status(200).body("{}").unwrap()
});
table.unset_lsm_write_spec().await.unwrap();
}
#[tokio::test]
async fn test_get_lsm_write_spec() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/get_lsm_write_spec/"
);
// The server resolves the spec and re-encodes it into the same
// sophon-owned shape the set endpoint accepts (`Sharding` internally
// tagged on `mode`, wrapped in `lsm_write_spec`).
let response = serde_json::json!({
"lsm_write_spec": {
"sharding": { "mode": "bucket", "column": "id", "num_buckets": 4 },
"maintained_indexes": ["id_idx"],
"writer_config_defaults": { "durable_write": "false" },
}
});
http::Response::builder()
.status(200)
.body(response.to_string())
.unwrap()
});
let spec = table
.get_lsm_write_spec()
.await
.unwrap()
.expect("a spec should be reported");
match spec {
crate::table::LsmWriteSpec::Bucket {
column,
num_buckets,
maintained_indexes,
writer_config_defaults,
} => {
assert_eq!(column, "id");
assert_eq!(num_buckets, 4);
assert_eq!(maintained_indexes, Some(vec!["id_idx".to_string()]));
assert_eq!(
writer_config_defaults
.get("durable_write")
.map(String::as_str),
Some("false")
);
}
other => panic!("expected a bucket spec, got {:?}", other),
}
}
/// Every selection reads back as the server reported it: every index
/// (`null`), none (`[]`), and a named list.
#[rstest::rstest]
#[case::every_index(serde_json::Value::Null, None)]
#[case::no_index(serde_json::json!([]), Some(vec![]))]
#[case::named(serde_json::json!(["id_idx"]), Some(vec!["id_idx".to_string()]))]
#[tokio::test]
async fn test_get_lsm_write_spec_round_trips_the_selection(
#[case] reported: serde_json::Value,
#[case] expected: Option<Vec<String>>,
) {
let table = Table::new_with_handler("my_table", move |_| {
let response = serde_json::json!({
"lsm_write_spec": {
"sharding": { "mode": "unsharded" },
"maintained_indexes": reported,
"writer_config_defaults": {},
}
});
http::Response::builder()
.status(200)
.body(response.to_string())
.unwrap()
});
let spec = table
.get_lsm_write_spec()
.await
.unwrap()
.expect("a spec should be reported");
assert_eq!(
spec.maintained_indexes().map(<[String]>::to_vec),
expected,
"the selection the server reported must survive the read"
);
}
#[tokio::test]
async fn test_get_lsm_write_spec_absent() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(
request.url().path(),
"/v1/table/my_table/get_lsm_write_spec/"
);
// Null spec → the LSM write path is not enabled.
let response = serde_json::json!({ "lsm_write_spec": null });
http::Response::builder()
.status(200)
.body(response.to_string())
.unwrap()
});
assert!(table.get_lsm_write_spec().await.unwrap().is_none());
}
/// `flush_lsm` / `compact_lsm` answer 202 with no body at all.
fn accepted() -> http::Response<String> {
http::Response::builder()
.status(202)
.body(String::new())
.unwrap()
}
fn ok_json(body: String) -> http::Response<String> {
http::Response::builder().status(200).body(body).unwrap()
}
/// A flush landing in an empty L0 finishes on the opening stats read
/// alone. Asserting zero compacts is the point: "it returned Ok" is also
/// true of a loop that ran a pointless pass.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_short_circuits_on_empty_l0() {
let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = compacts.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("compact_lsm") {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
panic!("an already-converged table must issue no compact calls");
}
if path.contains("flush_lsm") {
return accepted();
}
assert_eq!(path, "/v1/table/my_table/get_lsm_stats/");
ok_json(stats_body(&[], false))
});
table.checkpoint_lsm().await.unwrap();
assert_eq!(compacts.load(std::sync::atomic::Ordering::SeqCst), 0);
}
/// The loop triggers compaction until every generation that existed at
/// the start is gone, one bounded prefix per pass.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_triggers_until_targets_are_drained() {
let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = compacts.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") {
return accepted();
}
if path.contains("compact_lsm") {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
return accepted();
}
// Each pass drains the oldest generation.
let drained = seen.load(std::sync::atomic::Ordering::SeqCst);
let left: Vec<u64> = [1u64, 2, 3].into_iter().skip(drained).collect();
ok_json(stats_body(&left, false))
});
table.checkpoint_lsm().await.unwrap();
assert_eq!(
compacts.load(std::sync::atomic::Ordering::SeqCst),
3,
"one trigger per generation prefix, then stop"
);
}
/// Generations created *during* the checkpoint are not waited on, which
/// is what lets the loop terminate on a table taking writes where "L0 is
/// empty" never becomes true.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_ignores_generations_created_while_it_runs() {
let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = compacts.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") {
return accepted();
}
if path.contains("compact_lsm") {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
return accepted();
}
// Target is 5. One pass drains it; a writer keeps adding above.
let n = seen.load(std::sync::atomic::Ordering::SeqCst);
let body = if n == 0 {
stats_body(&[5], false)
} else {
stats_body(&[6, 7], false)
};
ok_json(body)
});
table.checkpoint_lsm().await.unwrap();
assert_eq!(
compacts.load(std::sync::atomic::Ordering::SeqCst),
1,
"the loop must not chase generations written after it started"
);
}
/// Contention is a 429 and must be retried. The server keeps it off 503
/// precisely so the client can act on the status alone — reading it as
/// terminal stops the checkpoint early on a healthy node.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_retries_contention() {
let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = compacts.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") {
return accepted();
}
if path.contains("compact_lsm") {
// First two triggers: every bucket already latched.
if seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst) < 2 {
return http::Response::builder()
.status(429)
.body(r#"{"code":21,"error":"Too many concurrent writes"}"#.to_string())
.unwrap();
}
return accepted();
}
let accepted_triggers = seen
.load(std::sync::atomic::Ordering::SeqCst)
.saturating_sub(2);
let left: Vec<u64> = if accepted_triggers == 0 {
vec![1]
} else {
vec![]
};
ok_json(stats_body(&left, false))
});
table
.checkpoint_lsm()
.await
.expect("contention must not abort the checkpoint");
assert_eq!(
compacts.load(std::sync::atomic::Ordering::SeqCst),
3,
"assert the retry count, not just the outcome"
);
}
/// A transient fault on the poll must not abort the checkpoint. This route
/// meets the most contention — it runs every `POLL_INTERVAL` for the
/// checkpoint's whole life, with the transport retry layer disabled — yet
/// was the one call reached with a bare `?`.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_retries_a_contended_stats_poll() {
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = polls.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") || path.contains("compact_lsm") {
return accepted();
}
// The opening read lands; the next two polls are latched out.
let n = seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if (1..3).contains(&n) {
return http::Response::builder()
.status(429)
.body(r#"{"code":21,"error":"Too many concurrent writes"}"#.to_string())
.unwrap();
}
ok_json(stats_body(if n < 4 { &[1] } else { &[] }, false))
});
table
.checkpoint_lsm()
.await
.expect("a contended poll must be retried, not surfaced");
assert_eq!(
polls.load(std::sync::atomic::Ordering::SeqCst),
5,
"the two rejected polls must be re-issued, not skipped"
);
}
/// Contention and a lost claim draw on separate budgets: five straight
/// 429s on `flush`, more than `MAX_REISSUES`, must still converge. On one
/// shared counter this spent the re-issue cap and then reported a lost
/// claim nothing had ever reported.
#[tokio::test(start_paused = true)]
async fn test_contention_does_not_exhaust_the_reissue_budget() {
let flushes = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = flushes.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") {
if seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst) < 5 {
return http::Response::builder()
.status(429)
.body(r#"{"code":21,"error":"Too many concurrent writes"}"#.to_string())
.unwrap();
}
return accepted();
}
if path.contains("compact_lsm") {
return accepted();
}
ok_json(stats_body(&[], false))
});
table
.checkpoint_lsm()
.await
.expect("contention must not be reported as a lost claim");
assert_eq!(
flushes.load(std::sync::atomic::Ordering::SeqCst),
6,
"five retries against one seal, then it lands"
);
}
/// An exhausted retry budget surfaces the fault that consumed it, not a
/// message the loop invented: "429, nine times" points an operator at a
/// saturated pool, a generic runtime error points them nowhere.
#[tokio::test(start_paused = true)]
async fn test_exhausted_retries_surface_the_underlying_fault() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = calls.clone();
let table = Table::new_with_handler("my_table", move |_request| {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
http::Response::builder()
.status(429)
.body(r#"{"code":21,"error":"Too many concurrent writes"}"#.to_string())
.unwrap()
});
let err = table.checkpoint_lsm().await.unwrap_err();
assert!(
matches!(&err, Error::Http { status_code: Some(s), .. } if s.as_u16() == 429),
"the fault that spent the budget must be the one reported: {err:?}"
);
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
9,
"one call plus MAX_RETRIES — the re-issue budget is not spent on top"
);
}
/// A draining node is terminal, but the client does not know that from the
/// status: draining and a proxy blip are both 503, and telling them apart
/// takes parsing the body for a namespace code. So it spends the retry
/// budget and then reports what the server said — the drain gate never
/// releases, so the answer does not change, and the operator still reads
/// "WAL node draining" in the error.
#[tokio::test(start_paused = true)]
async fn test_draining_surfaces_after_the_retry_budget() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = calls.clone();
let table = Table::new_with_handler("my_table", move |_request| {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
http::Response::builder()
.status(503)
.body(r#"{"code":19,"error":"WAL node draining"}"#.to_string())
.unwrap()
});
let err = table.checkpoint_lsm().await.unwrap_err();
let message = err.to_string();
assert!(
matches!(&err, Error::Http { status_code: Some(s), .. } if s.as_u16() == 503),
"the 503 must surface as itself: {err:?}"
);
assert!(
message.contains("WAL node draining"),
"the server's own diagnosis must survive to the caller: {message}"
);
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
9,
"one call plus MAX_RETRIES, then it reports rather than spinning"
);
}
/// A long stall with nothing compacting must keep waiting, not fail. The
/// client cannot judge this: a checkpoint queued behind unrelated tables
/// on the pod-wide compactor pool reports exactly these numbers — flat
/// generations, an idle latch — as one whose merges are failing. The
/// deadline is the caller's.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_waits_out_a_long_stall_rather_than_failing() {
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = polls.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") || path.contains("compact_lsm") {
return accepted();
}
// Flat for far longer than any bound this loop ever had, with
// `compacting: false` throughout — then it drains.
let n = seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
ok_json(stats_body(if n < 40 { &[1, 2] } else { &[] }, false))
});
table
.checkpoint_lsm()
.await
.expect("a stall is the server being slow, not the client's call to make");
assert!(
polls.load(std::sync::atomic::Ordering::SeqCst) > 40,
"the loop must have kept polling well past the old ten-poll bound"
);
}
/// A pass already owns the latch on every outstanding bucket, so the loop
/// waits rather than piling on triggers it would only refuse. This is the
/// sole thing `compacting` is read for.
#[tokio::test(start_paused = true)]
async fn test_checkpoint_waits_while_a_pass_is_running() {
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen_polls = polls.clone();
let seen_compacts = compacts.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
if path.contains("flush_lsm") {
return accepted();
}
if path.contains("compact_lsm") {
seen_compacts.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
return accepted();
}
// Latched for many polls, then done.
let n = seen_polls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
ok_json(if n > 15 {
stats_body(&[], false)
} else {
stats_body(&[1], true)
})
});
table
.checkpoint_lsm()
.await
.expect("a running pass is progress, not a stall");
assert_eq!(
compacts.load(std::sync::atomic::Ordering::SeqCst),
0,
"never trigger against a bucket already compacting"
);
}
/// WAL off ⇒ `None`; WAL on ⇒ a fully populated `Some` with no field
/// defaulting to a zero it did not measure. `include_generation_rows`
/// rides in the body and is off unless asked for.
#[tokio::test]
async fn test_get_lsm_stats_round_trip() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/get_lsm_stats/");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["include_generation_rows"], true,
"the flag must reach the server, not be silently dropped"
);
let response = serde_json::json!({
"lsm_stats": {
"buckets": [{
"shard_id": "b0",
"status": "Active",
"writer_epoch": 3,
"manifest_version": 11,
"current_generation": 9,
"replay_after_wal_entry_position": 100,
"wal_entry_position_last_seen": 140,
"generations": [{ "generation": 8, "bytes": 4096, "rows": 30 }],
"compacting": false,
"memtables": [
{ "generation": 9, "rows": 12, "bytes": 900, "batches": 2,
"indexes": ["vec_idx"] }
],
}],
}
});
http::Response::builder()
.status(200)
.body(response.to_string())
.unwrap()
});
let stats = table
.get_lsm_stats(true)
.await
.unwrap()
.expect("a WAL-backed table reports Some");
let bucket = &stats.buckets[0];
assert_eq!(bucket.replay_after_wal_entry_position, 100);
assert_eq!(bucket.wal_entry_position_last_seen, 140);
assert!(!bucket.compacting);
assert_eq!(bucket.generations[0].generation, 8);
assert_eq!(bucket.generations[0].rows, Some(30));
// The line that answers "why is my fresh-tier vector search
// brute-force" — an absent index name is the whole explanation.
let memtables = bucket.memtables.as_ref().unwrap();
assert_eq!(memtables[0].indexes, vec!["vec_idx".to_string()]);
}
/// A 404 arrives as `TableNotFound`, not as a lost claim the loop
/// re-issues from flush until its cap. The two are distinguished by
/// status: 404 is "no such table", 421 is "this node holds no claim".
/// They shared 404 once, and the loop chased a name that never existed.
#[tokio::test(start_paused = true)]
async fn test_missing_table_is_not_read_as_a_lost_claim() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = calls.clone();
let table = Table::new_with_handler("my_table", move |_request| {
seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
http::Response::builder()
.status(404)
.body(r#"{"code":4,"error":"Not found: Table not found: my_table"}"#.to_string())
.unwrap()
});
let err = table.checkpoint_lsm().await.unwrap_err();
assert!(
matches!(err, Error::TableNotFound { .. }),
"a missing table must say so: {err:?}"
);
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
1,
"no point re-claiming a table that does not exist"
);
}
/// A lost claim — 421, not 404 — does re-issue from flush, the call that
/// re-claims and replays.
#[tokio::test(start_paused = true)]
async fn test_registry_miss_reissues_from_flush() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = calls.clone();
let table = Table::new_with_handler("my_table", move |request| {
let path = request.url().path().to_string();
let n = seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if path.contains("flush_lsm") {
// First flush lands; the claim is then lost, and the
// re-issued flush succeeds.
return accepted();
}
if path.contains("compact_lsm") {
if n < 4 {
return http::Response::builder()
.status(421)
.body(r#"{"code":19,"error":"table not claimed"}"#.to_string())
.unwrap();
}
return accepted();
}
ok_json(stats_body(if n < 6 { &[1] } else { &[] }, false))
});
table
.checkpoint_lsm()
.await
.expect("a lost claim must be recovered by re-flushing, not surfaced");
}
#[tokio::test]
async fn test_get_lsm_stats_absent_when_wal_off() {
let table = Table::new_with_handler("my_table", |_request| {
http::Response::builder()
.status(200)
.body(serde_json::json!({ "lsm_stats": null }).to_string())
.unwrap()
});
assert!(table.get_lsm_stats(false).await.unwrap().is_none());
}
#[tokio::test]
async fn test_wait_for_index() {
let table = _make_table_with_indices(0);
table
.wait_for_index(&["vector_idx", "my_idx"], Duration::from_secs(1))
.await
.unwrap();
}
#[tokio::test]
async fn test_wait_for_index_timeout() {
let table = _make_table_with_indices(100);
let e = table
.wait_for_index(&["vector_idx", "my_idx"], Duration::from_secs(1))
.await
.unwrap_err();
assert_eq!(
e.to_string(),
"Timeout error: timed out waiting for indices: [\"vector_idx\", \"my_idx\"] after 1s"
);
}
#[tokio::test]
async fn test_wait_for_index_timeout_never_created() {
let table = _make_table_with_indices(0);
let e = table
.wait_for_index(&["doesnt_exist_idx"], Duration::from_secs(1))
.await
.unwrap_err();
assert_eq!(
e.to_string(),
"Timeout error: timed out waiting for indices: [\"doesnt_exist_idx\"] after 1s"
);
}
fn _make_table_with_indices(unindexed_rows: usize) -> Table {
Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
let response_body = match request.url().path() {
"/v1/table/my_table/describe/" => {
let schema = Schema::new(vec![
Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
8,
),
false,
),
Field::new("my_column", DataType::Utf8, false),
]);
serde_json::from_str::<serde_json::Value>(&describe_response(&schema)).unwrap()
}
"/v1/table/my_table/index/list/" => {
serde_json::json!({
"indexes": [
{
"index_name": "vector_idx",
"index_uuid": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
"columns": ["vector"],
"index_status": "done",
},
{
"index_name": "my_idx",
"index_uuid": "34255f64-5717-4562-b3fc-2c963f66afa6",
"columns": ["my_column"],
"index_status": "done",
},
]
})
}
"/v1/table/my_table/index/vector_idx/stats/" => {
serde_json::json!({
"num_indexed_rows": 100000,
"num_unindexed_rows": unindexed_rows,
"index_type": "IVF_PQ",
"distance_type": "l2"
})
}
"/v1/table/my_table/index/my_idx/stats/" => {
serde_json::json!({
"num_indexed_rows": 100000,
"num_unindexed_rows": unindexed_rows,
"index_type": "LABEL_LIST"
})
}
_path => {
serde_json::json!(None::<String>)
}
};
let body = serde_json::to_string(&response_body).unwrap();
let status = if body == "null" { 404 } else { 200 };
http::Response::builder().status(status).body(body).unwrap()
})
}
rust/lancedb/src/remote/table/tests/mod.rs
+254
View File
@@ -0,0 +1,254 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
use crate::index::vector::{
IvfFlatIndexBuilder, IvfHnswFlatIndexBuilder, IvfHnswSqIndexBuilder, IvfRqIndexBuilder,
IvfSqIndexBuilder,
};
use crate::remote::JSON_CONTENT_TYPE;
use crate::remote::client::{ClientConfig, RetryConfig};
use crate::remote::db::DEFAULT_SERVER_VERSION;
use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken};
use crate::utils::background_cache::clock;
use crate::{
DistanceType, Error, Table,
index::{Index, IndexStatistics, IndexType, vector::IvfPqIndexBuilder},
query::{
AnalyzePlanDistributedMetrics, ColumnOrdering, ExecutableQuery, QueryBase,
QueryExecutionOptions,
},
};
use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type};
use arrow_array::Array;
use arrow_array::builder::LargeBinaryBuilder;
use arrow_array::{
BinaryArray, Int32Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray,
StructArray, record_batch,
};
use arrow_schema::{DataType, Field, Schema};
use chrono::{DateTime, Utc};
use futures::{StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture};
use lance::dataset::ROW_ID;
use lance_index::scalar::inverted::{DocumentGranularity, query::MatchQuery};
use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams};
use reqwest::Body;
use rstest::rstest;
use serde_json::json;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use std::{collections::HashMap, pin::Pin};
fn refresh_done(job_id: &str) -> String {
json!({
"job_id": job_id,
"job_state": "DONE",
"result": {
"rows_assigned": 12,
"rows_failed": 0,
"rows_remaining": 0,
"source_version": 7,
"published_version": 8,
}
})
.to_string()
}
async fn collect_body(body: Body) -> Vec<u8> {
use http_body::Body;
let mut body = body;
let mut data = Vec::new();
let mut body_pin = Pin::new(&mut body);
futures::stream::poll_fn(|cx| body_pin.as_mut().poll_frame(cx))
.for_each(|frame| {
data.extend_from_slice(frame.unwrap().data_ref().unwrap());
futures::future::ready(())
})
.await;
data
}
fn write_ipc_stream(data: &RecordBatch) -> Vec<u8> {
let mut body = Vec::new();
{
let options = arrow_ipc::writer::IpcWriteOptions::default()
.try_with_compression(Some(arrow_ipc::CompressionType::LZ4_FRAME))
.expect("Failed to create IPC write options");
let mut writer = arrow_ipc::writer::StreamWriter::try_new_with_options(
&mut body,
&data.schema(),
options,
)
.expect("Failed to create writer");
writer.write(data).expect("Failed to write data");
writer.finish().expect("Failed to finish");
}
body
}
fn write_ipc_stream_uncompressed(data: &RecordBatch) -> Vec<u8> {
let mut body = Vec::new();
{
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &data.schema())
.expect("Failed to create writer");
writer.write(data).expect("Failed to write data");
writer.finish().expect("Failed to finish");
}
body
}
fn write_ipc_file(data: &RecordBatch) -> Vec<u8> {
let mut body = Vec::new();
{
let mut writer = arrow_ipc::writer::FileWriter::try_new(&mut body, &data.schema())
.expect("Failed to create writer");
writer.write(data).expect("Failed to write data");
writer.finish().expect("Failed to finish");
}
body
}
/// Build a JSON describe response for the given schema.
fn describe_response(schema: &Schema) -> String {
let json_schema = JsonSchema::try_from(schema).unwrap();
serde_json::to_string(&json!({
"version": 1,
"schema": json_schema,
}))
.unwrap()
}
fn nested_index_schema() -> Schema {
let vector_type =
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8);
Schema::new(vec![
Field::new("rowId", DataType::Int32, false),
Field::new("row-id", DataType::Int32, false),
Field::new("userId", DataType::Int32, false),
Field::new(
"metadata",
DataType::Struct(vec![Field::new("user_id", DataType::Int32, false)].into()),
false,
),
Field::new(
"MetaData",
DataType::Struct(vec![Field::new("userId", DataType::Int32, false)].into()),
false,
),
Field::new(
"image",
DataType::Struct(vec![Field::new("embedding", vector_type, false)].into()),
false,
),
Field::new(
"payload",
DataType::Struct(vec![Field::new("text", DataType::Utf8, false)].into()),
false,
),
Field::new(
"docs",
DataType::List(Arc::new(Field::new(
"item",
DataType::Struct(vec![Field::new("content", DataType::Utf8, true)].into()),
true,
))),
true,
),
Field::new(
"meta-data",
DataType::Struct(vec![Field::new("user-id", DataType::Int32, false)].into()),
false,
),
Field::new(
"literal",
DataType::Struct(vec![Field::new("a.b", DataType::Int32, false)].into()),
false,
),
])
}
/// Build a `get_lsm_stats` body for one bucket holding `generations`.
fn stats_body(generations: &[u64], compacting: bool) -> String {
serde_json::json!({
"lsm_stats": {
"buckets": [{
"shard_id": "b0",
"status": "Active",
"writer_epoch": 1,
"manifest_version": 1,
"current_generation": generations.iter().max().copied().unwrap_or(0) + 1,
"replay_after_wal_entry_position": 0,
"wal_entry_position_last_seen": 0,
"generations": generations.iter()
.map(|g| serde_json::json!({ "generation": g, "bytes": 1 }))
.collect::<Vec<_>>(),
"compacting": compacting,
"memtables": [],
}],
}
})
.to_string()
}
fn schema_json() -> &'static str {
r#"{"fields": [{"name": "id", "type": {"type": "int32"}, "nullable": true}]}"#
}
fn simple_describe_response() -> http::Response<String> {
http::Response::builder()
.status(200)
.body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json()))
.unwrap()
}
// ----- Branch support -----
/// Parse a request's in-memory JSON body. Only valid for JSON-body ops
/// (not Arrow-stream inserts, whose body is a stream).
fn request_body_json(request: &reqwest::Request) -> serde_json::Value {
let bytes = request
.body()
.expect("request has a body")
.as_bytes()
.expect("body is in-memory");
serde_json::from_slice(bytes).expect("body is valid JSON")
}
/// One leg's worth of results, keyed so the two legs share a row.
/// `row_ids` mirrors a server answering a query that asked for `_rowid`.
fn leg(score_column: &str, ids: Vec<&str>, row_ids: bool) -> RecordBatch {
use arrow_array::{Float32Array, UInt64Array};
let mut fields = vec![
Field::new("id", DataType::Utf8, false),
Field::new("text", DataType::Utf8, true),
Field::new(score_column, DataType::Float32, false),
];
let scores = Float32Array::from(vec![1.0_f32; ids.len()]);
let texts = StringArray::from(ids.clone());
let mut columns: Vec<arrow_array::ArrayRef> = vec![
Arc::new(StringArray::from(ids.clone())),
Arc::new(texts),
Arc::new(scores),
];
if row_ids {
fields.push(Field::new(ROW_ID, DataType::UInt64, false));
columns.push(Arc::new(UInt64Array::from_iter_values(
ids.iter().map(|id| id.as_bytes()[0] as u64),
)));
}
RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
}
mod blobs;
mod branches;
mod columns;
mod freshness;
mod hybrid;
mod index;
mod lsm;
mod multipart;
mod query;
mod schema_cache;
mod versions;
mod write;
rust/lancedb/src/remote/table/tests/multipart.rs
+767
View File
@@ -0,0 +1,767 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_add_insert_fails() {
// Verify that an HTTP error from the insert endpoint is properly
// surfaced with the status code intact. Use 400 (non-retryable).
let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
let describe_body = describe_response(&batch.schema());
let table = Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => http::Response::builder()
.status(400)
.body("bad request".to_string())
.unwrap(),
path => panic!("Unexpected request path: {}", path),
});
let result = table.add(batch).execute().await;
let err = result.unwrap_err();
match &err {
Error::Http { status_code, .. } => {
assert_eq!(*status_code, Some(reqwest::StatusCode::BAD_REQUEST));
}
other => panic!("Expected Http error, got: {:?}", other),
}
}
#[tokio::test]
async fn test_add_retries_on_retryable_status() {
// Verify that rescannable data retries on retryable status codes (e.g. 502)
// and eventually succeeds.
let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
let describe_body = describe_response(&batch.schema());
let attempt = Arc::new(AtomicUsize::new(0));
let attempt_clone = attempt.clone();
let table = Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
let n = attempt_clone.fetch_add(1, Ordering::SeqCst);
if n < 2 {
http::Response::builder()
.status(502)
.body("bad gateway".to_string())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 3}"#.to_string())
.unwrap()
}
}
path => panic!("Unexpected request path: {}", path),
});
let result = table.add(batch).execute().await.unwrap();
assert_eq!(result.version, 3);
assert_eq!(attempt.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_query_with_datafusion_filter() {
use datafusion_expr::{col, lit};
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
// The Datafusion expression should be serialized to SQL
let filter = body.get("filter").expect("filter should be present");
let filter_str = filter.as_str().expect("filter should be a string");
// col("x") > lit(10) AND col("status") = lit("active")
assert!(
filter_str.contains("x") && filter_str.contains("10"),
"Filter should contain 'x' and '10', got: {}",
filter_str
);
assert!(
filter_str.contains("status") && filter_str.contains("active"),
"Filter should contain 'status' and 'active', got: {}",
filter_str
);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
// Use only_if_expr with a Datafusion expression
let expr = col("x").gt(lit(10)).and(col("status").eq(lit("active")));
let data = table
.query()
.only_if_expr(expr)
.execute()
.await
.unwrap()
.collect::<Vec<_>>()
.await;
assert_eq!(data.len(), 1);
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_multipart_write_happy_path() {
use std::sync::Mutex;
let create_count = Arc::new(AtomicUsize::new(0));
let insert_count = Arc::new(AtomicUsize::new(0));
let complete_count = Arc::new(AtomicUsize::new(0));
let abort_count = Arc::new(AtomicUsize::new(0));
let upload_ids = Arc::new(Mutex::new(Vec::<String>::new()));
let create_count_c = create_count.clone();
let insert_count_c = insert_count.clone();
let complete_count_c = complete_count.clone();
let abort_count_c = abort_count.clone();
let upload_ids_c = upload_ids.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
let query = request.url().query().unwrap_or("");
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
create_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"upload_id": "test-upload-123"}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
insert_count_c.fetch_add(1, Ordering::SeqCst);
let uid = url::form_urlencoded::parse(query.as_bytes())
.find(|(k, _)| k == "upload_id")
.map(|(_, v)| v.to_string());
upload_ids_c
.lock()
.unwrap()
.push(uid.expect("missing upload_id on insert"));
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
complete_count_c.fetch_add(1, Ordering::SeqCst);
let uid = url::form_urlencoded::parse(query.as_bytes())
.find(|(k, _)| k == "upload_id")
.map(|(_, v)| v.to_string());
upload_ids_c
.lock()
.unwrap()
.push(uid.expect("missing upload_id on complete"));
return http::Response::builder()
.status(200)
.body(r#"{"version": 5}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/abort" {
abort_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(String::new())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table
.add(vec![batch])
.write_parallelism(2)
.execute()
.await
.unwrap();
assert_eq!(result.version, 5);
assert_eq!(create_count.load(Ordering::SeqCst), 1);
assert!(
insert_count.load(Ordering::SeqCst) > 1,
"Expected multiple insert calls, got {}",
insert_count.load(Ordering::SeqCst)
);
assert_eq!(complete_count.load(Ordering::SeqCst), 1);
assert_eq!(abort_count.load(Ordering::SeqCst), 0);
let ids = upload_ids.lock().unwrap();
assert!(
ids.iter().all(|id| id == "test-upload-123"),
"All requests should use the same upload_id, got: {:?}",
*ids
);
}
#[tokio::test]
async fn test_multipart_write_progress() {
let callback_count = Arc::new(AtomicUsize::new(0));
let max_active = Arc::new(AtomicUsize::new(0));
let last_total_tasks = Arc::new(AtomicUsize::new(0));
let seen_done = Arc::new(std::sync::Mutex::new(false));
let cb_count = callback_count.clone();
let cb_active = max_active.clone();
let cb_total = last_total_tasks.clone();
let cb_done = seen_done.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
return http::Response::builder()
.status(200)
.body(r#"{"upload_id": "prog-upload"}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
return http::Response::builder()
.status(200)
.body(r#"{"version": 3}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
table
.add(vec![batch])
.write_parallelism(2)
.progress(move |p| {
cb_count.fetch_add(1, Ordering::SeqCst);
cb_active.fetch_max(p.active_tasks(), Ordering::SeqCst);
cb_total.store(p.total_tasks(), Ordering::SeqCst);
if p.done() {
*cb_done.lock().unwrap() = true;
}
})
.execute()
.await
.unwrap();
assert!(
callback_count.load(Ordering::SeqCst) >= 1,
"expected at least one progress callback"
);
assert!(*seen_done.lock().unwrap(), "must see done=true");
assert_eq!(last_total_tasks.load(Ordering::SeqCst), 2);
assert!(
max_active.load(Ordering::SeqCst) >= 1,
"expected at least one active task"
);
}
#[tokio::test]
async fn test_multipart_write_fallback_old_server() {
let insert_count = Arc::new(AtomicUsize::new(0));
let create_count = Arc::new(AtomicUsize::new(0));
let insert_count_c = insert_count.clone();
let create_count_c = create_count.clone();
// Server version 0.3.0 does not support multipart writes
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 3, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path.contains("multipart_write") {
create_count_c.fetch_add(1, Ordering::SeqCst);
panic!("Should not call multipart write endpoints on old server");
}
if path == "/v1/table/my_table/insert/" {
let query = request.url().query().unwrap_or("");
assert!(
!query.contains("upload_id"),
"Should not have upload_id for old server"
);
insert_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table
.add(vec![batch])
.write_parallelism(2)
.execute()
.await
.unwrap();
assert_eq!(result.version, 2);
assert_eq!(create_count.load(Ordering::SeqCst), 0);
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_write_small_data_single_partition() {
let insert_count = Arc::new(AtomicUsize::new(0));
let create_count = Arc::new(AtomicUsize::new(0));
let insert_count_c = insert_count.clone();
let create_count_c = create_count.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path.contains("multipart_write") {
create_count_c.fetch_add(1, Ordering::SeqCst);
panic!("Should not call multipart write endpoints for small data");
}
if path == "/v1/table/my_table/insert/" {
let query = request.url().query().unwrap_or("");
assert!(
!query.contains("upload_id"),
"Should not have upload_id for small data"
);
insert_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
// Small data: only 3 rows
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table.add(vec![batch]).execute().await.unwrap();
assert_eq!(result.version, 2);
assert_eq!(create_count.load(Ordering::SeqCst), 0);
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_write_empty_overwrite_uses_single_partition() {
// A multipart write creates its upload session before any partition
// executes. If the input has no batches at all, every partition would
// stage nothing (see `send_multipart_chunked`), so completing the write
// would have nothing to commit and `mode=overwrite` would be silently
// dropped. An explicit `write_parallelism` must not force the multipart
// path for empty input; it should fall back to the single-request path,
// which always sends one schema-only request and carries `mode=overwrite`.
let insert_count = Arc::new(AtomicUsize::new(0));
let multipart_count = Arc::new(AtomicUsize::new(0));
let insert_count_c = insert_count.clone();
let multipart_count_c = multipart_count.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path.contains("multipart_write") {
multipart_count_c.fetch_add(1, Ordering::SeqCst);
panic!("Should not use multipart write endpoints for empty input");
}
if path == "/v1/table/my_table/insert/" {
let query = request.url().query().unwrap_or("");
assert!(
!query.contains("upload_id"),
"Should not have upload_id for empty input"
);
assert!(
query.contains("mode=overwrite"),
"Should carry mode=overwrite, got query: {}",
query
);
insert_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)]));
let empty_batches: Vec<std::result::Result<RecordBatch, arrow_schema::ArrowError>> = Vec::new();
let data: Box<dyn RecordBatchReader + Send> =
Box::new(RecordBatchIterator::new(empty_batches, schema));
let result = table
.add(data)
.mode(AddDataMode::Overwrite)
.write_parallelism(4)
.execute()
.await
.unwrap();
assert_eq!(result.version, 2);
assert_eq!(multipart_count.load(Ordering::SeqCst), 0);
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_write_abort_on_insert_failure() {
let create_count = Arc::new(AtomicUsize::new(0));
let insert_count = Arc::new(AtomicUsize::new(0));
let complete_count = Arc::new(AtomicUsize::new(0));
let abort_count = Arc::new(AtomicUsize::new(0));
let create_count_c = create_count.clone();
let insert_count_c = insert_count.clone();
let complete_count_c = complete_count.clone();
let abort_count_c = abort_count.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
create_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"upload_id": "test-upload-456"}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
let count = insert_count_c.fetch_add(1, Ordering::SeqCst);
// Fail on the first insert with non-retryable status
if count == 0 {
return http::Response::builder()
.status(400)
.body("Bad Request".to_string())
.unwrap();
}
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
complete_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 5}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/abort" {
abort_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(String::new())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table.add(vec![batch]).write_parallelism(2).execute().await;
assert!(result.is_err());
assert_eq!(create_count.load(Ordering::SeqCst), 1);
assert_eq!(complete_count.load(Ordering::SeqCst), 0);
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_write_abort_on_complete_failure() {
let abort_count = Arc::new(AtomicUsize::new(0));
let abort_count_c = abort_count.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
return http::Response::builder()
.status(200)
.body(r#"{"upload_id": "test-upload-789"}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/insert/" {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
return http::Response::builder()
.status(400)
.body("Bad Request".to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/abort" {
abort_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(String::new())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table.add(vec![batch]).write_parallelism(2).execute().await;
assert!(result.is_err());
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
}
fn retry_config_no_backoff() -> ClientConfig {
ClientConfig {
retry_config: RetryConfig {
retries: Some(3),
connect_retries: Some(3),
read_retries: Some(3),
backoff_factor: Some(0.0),
backoff_jitter: Some(0.0),
statuses: Some(vec![502, 503]),
},
..Default::default()
}
}
#[tokio::test]
async fn test_multipart_write_retry_on_partition_failure() {
// All inserts for the first upload session return 503 (retryable).
// After exhausting internal retries, the outer loop retries with a
// new session and succeeds.
let create_count = Arc::new(AtomicUsize::new(0));
let complete_count = Arc::new(AtomicUsize::new(0));
let abort_count = Arc::new(AtomicUsize::new(0));
let create_count_c = create_count.clone();
let complete_count_c = complete_count.clone();
let abort_count_c = abort_count.clone();
let table = Table::new_with_handler_version_and_config(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
let query = request.url().query().unwrap_or("");
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
let n = create_count_c.fetch_add(1, Ordering::SeqCst);
let body = format!(r#"{{"upload_id": "upload-{}"}}"#, n + 1);
return http::Response::builder().status(200).body(body).unwrap();
}
if path == "/v1/table/my_table/insert/" {
// Fail all inserts for the first session
if query.contains("upload_id=upload-1") {
return http::Response::builder()
.status(503)
.body("Service Unavailable".to_string())
.unwrap();
}
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
complete_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(r#"{"version": 7}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/abort" {
abort_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(String::new())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
retry_config_no_backoff(),
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table
.add(vec![batch])
.write_parallelism(2)
.execute()
.await
.unwrap();
assert_eq!(result.version, 7);
assert_eq!(create_count.load(Ordering::SeqCst), 2);
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
assert_eq!(complete_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multipart_write_retry_on_complete_failure() {
// Complete returns 503 for the first session, succeeds for the second.
let create_count = Arc::new(AtomicUsize::new(0));
let abort_count = Arc::new(AtomicUsize::new(0));
let create_count_c = create_count.clone();
let abort_count_c = abort_count.clone();
let table = Table::new_with_handler_version_and_config(
"my_table",
semver::Version::new(0, 4, 0),
move |request| {
let path = request.url().path();
let query = request.url().query().unwrap_or("");
if path == "/v1/table/my_table/describe/" {
return simple_describe_response();
}
if path == "/v1/table/my_table/multipart_write/create" {
let n = create_count_c.fetch_add(1, Ordering::SeqCst);
let body = format!(r#"{{"upload_id": "upload-{}"}}"#, n + 1);
return http::Response::builder().status(200).body(body).unwrap();
}
if path == "/v1/table/my_table/insert/" {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/complete" {
// Fail complete for first session
if query.contains("upload_id=upload-1") {
return http::Response::builder()
.status(503)
.body("Service Unavailable".to_string())
.unwrap();
}
return http::Response::builder()
.status(200)
.body(r#"{"version": 9}"#.to_string())
.unwrap();
}
if path == "/v1/table/my_table/multipart_write/abort" {
abort_count_c.fetch_add(1, Ordering::SeqCst);
return http::Response::builder()
.status(200)
.body(String::new())
.unwrap();
}
panic!("Unexpected request path: {}", path);
},
retry_config_no_backoff(),
);
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let result = table
.add(vec![batch])
.write_parallelism(2)
.execute()
.await
.unwrap();
assert_eq!(result.version, 9);
assert_eq!(create_count.load(Ordering::SeqCst), 2);
assert_eq!(abort_count.load(Ordering::SeqCst), 1);
}
rust/lancedb/src/remote/table/tests/query.rs
+806
View File
@@ -0,0 +1,806 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[rstest]
#[case(DEFAULT_SERVER_VERSION.clone())]
#[case(semver::Version::new(0, 2, 0))]
#[tokio::test]
async fn test_query_vector_non_finite(#[case] version: semver::Version) {
let table =
Table::new_with_handler_version("my_table", version, |_| -> http::Response<String> {
panic!("non-finite vectors must be rejected before sending a request")
});
for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
// Requests built without the query builder must also return an
// error, including a non-finite vector later in a batch.
for batched in [false, true] {
let mut request = table
.query()
.nearest_to(&[0.1, 0.2])
.unwrap()
.into_request();
if !batched {
request.query_vector.clear();
}
request
.query_vector
.push(Arc::new(arrow_array::Float32Array::from(vec![0.1, value])));
let result = table
.base_table()
.query(
&AnyQuery::VectorQuery(request),
QueryExecutionOptions::default(),
)
.await;
let Err(err) = result else {
panic!("non-finite query vector unexpectedly succeeded")
};
assert!(matches!(err, Error::InvalidInput { .. }));
assert!(err.to_string().contains("only finite values"));
}
}
}
#[tokio::test]
async fn test_query_vector_default_values() {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let mut expected_body = serde_json::json!({
"prefilter": true,
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"k": 10,
"ef": Option::<usize>::None,
"refine_factor": null,
"version": null,
});
// Pass vector separately to make sure it matches f32 precision.
expected_body["vector"] = vec![0.1f32, 0.2, 0.3].into();
assert_eq!(body, expected_body);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let data = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.execute()
.await;
let data = data.unwrap().collect::<Vec<_>>().await;
assert_eq!(data.len(), 1);
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_query_vector_approx_mode_sent_when_set() {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let mut expected_body = serde_json::json!({
"prefilter": true,
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
"approx_mode": "accurate",
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"k": 10,
"ef": Option::<usize>::None,
"refine_factor": null,
"version": null,
});
expected_body["vector"] = vec![0.1f32, 0.2, 0.3].into();
assert_eq!(body, expected_body);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let data = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.approx_mode(crate::ApproxMode::Accurate)
.execute()
.await;
let data = data.unwrap().collect::<Vec<_>>().await;
assert_eq!(data.len(), 1);
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_query_fts_default_values() {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let expected_body = serde_json::json!({
"full_text_query": {
"columns": [],
"query": "test",
},
"prefilter": true,
"version": null,
"k": 10,
"vector": [],
});
assert_eq!(body, expected_body);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let data = table
.query()
.full_text_search(FullTextSearchQuery::new("test".to_owned()))
.execute()
.await;
let data = data.unwrap().collect::<Vec<_>>().await;
assert_eq!(data.len(), 1);
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_query_vector_nprobes_zero() {
let table = Table::new_with_handler::<&str>("my_table", |_| {
panic!("invalid nprobes must be rejected before sending a request")
});
let result = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.nprobes(0);
assert!(matches!(
result,
Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0"
));
}
#[tokio::test]
async fn test_query_vector_all_params() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let mut expected_body = serde_json::json!({
"vector_column": "my_vector",
"prefilter": false,
"k": 42,
"offset": 10,
"distance_type": "cosine",
"bypass_vector_index": true,
"columns": ["a", "b"],
"order_by": [
{
"column_name": "score",
"ascending": false,
"nulls_first": true,
},
{
"column_name": "id",
"ascending": true,
"nulls_first": false,
}
],
"nprobes": 12,
"minimum_nprobes": 12,
"maximum_nprobes": 12,
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"ef": Option::<usize>::None,
"refine_factor": 2,
"version": null,
});
// Pass vector separately to make sure it matches f32 precision.
expected_body["vector"] = vec![0.1f32, 0.2, 0.3].into();
assert_eq!(body, expected_body);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let response_body = write_ipc_file(&data);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let _ = table
.query()
.limit(42)
.offset(10)
.select(Select::columns(&["a", "b"]))
.order_by(Some(vec![
ColumnOrdering::desc_nulls_first("score".to_string()),
ColumnOrdering::asc_nulls_last("id".to_string()),
]))
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.column("my_vector")
.postfilter()
.distance_type(crate::DistanceType::Cosine)
.nprobes(12)
.unwrap()
.refine_factor(2)
.bypass_vector_index()
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_query_vector_nested_field_path() {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let mut expected_body = serde_json::json!({
"vector_column": "image.embedding",
"prefilter": true,
"k": 10,
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"ef": Option::<usize>::None,
"refine_factor": Option::<u32>::None,
"version": null,
});
expected_body["vector"] = vec![0.1f32, 0.2, 0.3].into();
assert_eq!(body, expected_body);
let response_body = write_ipc_file(&expected_data_ref);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let _ = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.column("image.embedding")
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_query_fts() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let expected_body = serde_json::json!({
"full_text_query": {
"columns": ["a", "b"],
"query": "hello world",
},
"k": 10,
"vector": [],
"with_row_id": true,
"prefilter": true,
"version": null
});
let expected_body_2 = serde_json::json!({
"full_text_query": {
"columns": ["b","a"],
"query": "hello world",
},
"k": 10,
"vector": [],
"with_row_id": true,
"prefilter": true,
"version": null
});
assert!(body == expected_body || body == expected_body_2);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let response_body = write_ipc_file(&data);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let _ = table
.query()
.full_text_search(
FullTextSearchQuery::new("hello world".into())
.with_columns(&["a".into(), "b".into()])
.unwrap(),
)
.with_row_id()
.limit(10)
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_analyze_plan_distributed_metrics_query_param() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/");
assert_eq!(
request
.url()
.query_pairs()
.find(|(k, _)| k == "distributed_metrics"),
Some(("distributed_metrics".into(), "per_worker".into()))
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(body["k"], serde_json::json!(1));
http::Response::builder()
.status(200)
.body(r#""analyzed plan""#)
.unwrap()
});
let result = table
.query()
.limit(1)
.analyze_plan_with_options(QueryExecutionOptions {
analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker,
..Default::default()
})
.await
.unwrap();
assert_eq!(result, "analyzed plan");
}
#[tokio::test]
async fn test_take_offsets_explain_plan_does_not_execute_query() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/explain_plan/");
http::Response::builder()
.status(200)
.body(r#""RemoteLookupExec""#)
.unwrap()
});
let explained = table
.take_offsets(vec![0, 1, 0, 2])
.select(crate::query::Select::columns(&["id"]))
.limit(3)
.explain_plan(false)
.await
.unwrap();
assert!(explained.contains("GlobalLimitExec"));
assert!(explained.contains("TakeRestoreExec"));
assert!(!explained.contains("CoalescePartitionsExec"));
assert!(explained.contains("RemoteLookupExec"));
}
#[tokio::test]
async fn test_converted_take_request_restores_remote_occurrences() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["columns"], json!(["id", "_rowoffset"]));
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("_rowoffset", DataType::UInt64, false),
])),
vec![
Arc::new(Int32Array::from(vec![5])),
Arc::new(arrow_array::UInt64Array::from(vec![5])),
],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&data))
.unwrap()
});
let request = table
.take_offsets(vec![5, 5])
.select(crate::query::Select::columns(&["id"]))
.into_request();
let batches = table
.base_table()
.query(&AnyQuery::Query(request), QueryExecutionOptions::default())
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
assert!(
batches
.iter()
.all(|batch| batch.schema().fields().len() == 1)
);
}
#[tokio::test]
async fn test_take_offsets_analyze_plan_delegates_to_remote() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/");
assert_eq!(
request
.url()
.query_pairs()
.find(|(key, _)| key == "distributed_metrics"),
Some(("distributed_metrics".into(), "per_worker".into()))
);
http::Response::builder()
.status(200)
.body(r#""Remote analyzed plan: worker metrics""#)
.unwrap()
});
let analyzed = table
.take_offsets(vec![0, 1, 0, 2])
.select(crate::query::Select::columns(&["id"]))
.limit(3)
.analyze_plan_with_options(QueryExecutionOptions {
analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker,
..Default::default()
})
.await
.unwrap();
assert_eq!(analyzed, "Remote analyzed plan: worker metrics");
}
#[tokio::test]
async fn test_query_structured_fts() {
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 6, 0), |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let expected_body = serde_json::json!({
"full_text_query": {
"query": {
"match": {
"terms": "hello world",
"column": "payload.text",
"boost": 1.0,
"fuzziness": 0,
"max_expansions": 50,
"operator": "Or",
"prefix_length": 0,
"document_granularity": "list_element",
},
}
},
"k": 10,
"vector": [],
"with_row_id": true,
"prefilter": true,
"version": null
});
assert_eq!(body, expected_body);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let response_body = write_ipc_file(&data);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let _ = table
.query()
.full_text_search(FullTextSearchQuery::new_query(
MatchQuery::new("hello world".to_owned())
.with_column(Some("payload.text".to_owned()))
.with_document_granularity(DocumentGranularity::ListElement)
.into(),
))
.with_row_id()
.limit(10)
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_query_combined_fields_uses_structured_fts() {
use lance_index::scalar::inverted::query::CombinedFieldsQuery;
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 3, 0), |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["full_text_query"]["query"],
serde_json::json!({
"combined_fields": {
"query": "hello world",
"columns": ["title", "text"],
"boost": [1.0, 1.0],
"operator": "Or"
}
})
);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&data))
.unwrap()
});
table
.query()
.full_text_search(FullTextSearchQuery::new_query(
CombinedFieldsQuery::try_new("hello world".into(), vec!["title".into(), "text".into()])
.unwrap()
.into(),
))
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_query_row_document_granularity_uses_structured_fts() {
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 3, 0), |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["full_text_query"]["query"]["match"]["document_granularity"],
"row"
);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&data))
.unwrap()
});
table
.query()
.full_text_search(FullTextSearchQuery::new_query(
MatchQuery::new("hello world".to_owned())
.with_column(Some("payload.text".to_owned()))
.with_document_granularity(DocumentGranularity::Row)
.into(),
))
.execute()
.await
.unwrap();
}
#[rstest]
#[case(DEFAULT_SERVER_VERSION.clone())]
#[case(semver::Version::new(0, 3, 0))]
#[case(semver::Version::new(0, 5, 0))]
#[tokio::test]
async fn test_query_document_granularity_requires_server_support(#[case] version: semver::Version) {
let table =
Table::new_with_handler_version("my_table", version, |_| -> http::Response<String> {
panic!("unsupported remote query must fail before sending a request")
});
let result = table
.query()
.full_text_search(FullTextSearchQuery::new_query(
MatchQuery::new("hello world".to_owned())
.with_column(Some("payload.text".to_owned()))
.with_document_granularity(DocumentGranularity::ListElement)
.into(),
))
.execute()
.await;
let Err(err) = result else {
panic!("legacy remote query unexpectedly succeeded")
};
assert!(
err.to_string()
.contains("document granularity requires remote server version 0.6.0 or later")
);
}
#[rstest]
#[case(DEFAULT_SERVER_VERSION.clone())]
#[case(semver::Version::new(0, 2, 0))]
#[tokio::test]
async fn test_batch_queries(#[case] version: semver::Version) {
let table = Table::new_with_handler_version("my_table", version.clone(), move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
let query_vectors = body["vector"].as_array().unwrap();
let version = ServerVersion(version.clone());
let data = if version.support_multivector() {
assert_eq!(query_vectors.len(), 2);
assert_eq!(query_vectors[0].as_array().unwrap().len(), 3);
assert_eq!(query_vectors[1].as_array().unwrap().len(), 3);
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("query_index", DataType::Int32, false),
])),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
Arc::new(Int32Array::from(vec![0, 0, 0, 1, 1, 1])),
],
)
.unwrap()
} else {
// it's single flat vector, so here the length is dim
assert_eq!(query_vectors.len(), 3);
RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap()
};
let response_body = write_ipc_file(&data);
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap()
});
let query = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.add_query_vector(vec![0.4, 0.5, 0.6])
.unwrap();
let results = query
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let results = concat_batches(&results[0].schema(), &results).unwrap();
let query_index = results["query_index"].as_primitive::<Int32Type>();
// We don't guarantee order.
assert!(query_index.values().contains(&0));
assert!(query_index.values().contains(&1));
}
rust/lancedb/src/remote/table/tests/schema_cache.rs
+1056
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/table/tests/versions.rs
+369
View File
@@ -0,0 +1,369 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_list_versions() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/version/list/");
let version1 = lance::dataset::Version {
version: 1,
timestamp: "2024-01-01T00:00:00Z".parse().unwrap(),
metadata: Default::default(),
};
let version2 = lance::dataset::Version {
version: 2,
timestamp: "2024-02-01T00:00:00Z".parse().unwrap(),
metadata: Default::default(),
};
let response_body = serde_json::json!({
"versions": [
version1,
version2,
]
});
let response_body = serde_json::to_string(&response_body).unwrap();
http::Response::builder()
.status(200)
.body(response_body)
.unwrap()
});
let versions = table.list_versions().await.unwrap();
assert_eq!(versions.len(), 2);
assert_eq!(versions[0].version, 1);
assert_eq!(
versions[0].timestamp,
"2024-01-01T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
);
assert_eq!(versions[1].version, 2);
assert_eq!(
versions[1].timestamp,
"2024-02-01T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
);
// assert_eq!(versions, expected);
}
/// Namespace-backed servers report `timestamp_millis` instead of
/// `timestamp`, and may omit `metadata` entirely.
#[tokio::test]
async fn test_list_versions_timestamp_millis() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/version/list/");
let response_body = serde_json::json!({
"versions": [
{
"version": 1,
"manifest_path": "path/to/_versions/1.manifest",
"timestamp_millis": 1704067200000i64,
},
{
"version": 2,
"manifest_path": "path/to/_versions/2.manifest",
"timestamp_millis": 1706745600000i64,
"metadata": {"key": "value"},
},
]
});
let response_body = serde_json::to_string(&response_body).unwrap();
http::Response::builder()
.status(200)
.body(response_body)
.unwrap()
});
let versions = table.list_versions().await.unwrap();
assert_eq!(versions.len(), 2);
assert_eq!(versions[0].version, 1);
assert_eq!(
versions[0].timestamp,
"2024-01-01T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
);
assert!(versions[0].metadata.is_empty());
assert_eq!(versions[1].version, 2);
assert_eq!(
versions[1].timestamp,
"2024-02-01T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
);
assert_eq!(
versions[1].metadata.get("key").map(String::as_str),
Some("value")
);
}
#[tokio::test]
async fn test_index_stats() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/stats/"
);
let response_body = serde_json::json!({
"num_indexed_rows": 100000,
"num_unindexed_rows": 0,
"index_type": "IVF_PQ",
"distance_type": "l2"
});
let response_body = serde_json::to_string(&response_body).unwrap();
http::Response::builder()
.status(200)
.body(response_body)
.unwrap()
});
let indices = table.index_stats("my_index").await.unwrap().unwrap();
let expected = IndexStatistics {
num_indexed_rows: 100000,
num_unindexed_rows: 0,
index_type: IndexType::IvfPq,
distance_type: Some(DistanceType::L2),
num_indices: None,
};
assert_eq!(indices, expected);
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/table/my_table/index/my_index/stats/"
);
http::Response::builder().status(404).body("").unwrap()
});
let indices = table.index_stats("my_index").await.unwrap();
assert!(indices.is_none());
}
#[tokio::test]
async fn test_passes_version() {
let table = Table::new_with_handler("my_table", |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let version = body
.as_object()
.unwrap()
.get("version")
.unwrap()
.as_u64()
.unwrap();
assert_eq!(version, 42);
let response_body = match request.url().path() {
"/v1/table/my_table/describe/" => {
serde_json::json!({
"version": 42,
"schema": { "fields": [] }
})
}
"/v1/table/my_table/index/list/" => {
serde_json::json!({
"indexes": []
})
}
"/v1/table/my_table/index/my_idx/stats/" => {
serde_json::json!({
"num_indexed_rows": 100000,
"num_unindexed_rows": 0,
"index_type": "IVF_PQ",
"distance_type": "l2"
})
}
"/v1/table/my_table/count_rows/" => {
serde_json::json!(1000)
}
"/v1/table/my_table/query/" => {
let expected_data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let expected_data_ref = expected_data.clone();
let response_body = write_ipc_file(&expected_data_ref);
return http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(response_body)
.unwrap();
}
path => panic!("Unexpected path: {}", path),
};
http::Response::builder()
.status(200)
.body(
serde_json::to_string(&response_body)
.unwrap()
.as_bytes()
.to_vec(),
)
.unwrap()
});
table.checkout(42).await.unwrap();
// ensure that version is passed to the /describe endpoint
let version = table.version().await.unwrap();
assert_eq!(version, 42);
// ensure it's passed to other read API calls
table.list_indices().await.unwrap();
table.index_stats("my_idx").await.unwrap();
table.count_rows(None).await.unwrap();
table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_restore_requires_checkout() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
request_count_clone.fetch_add(1, Ordering::SeqCst);
assert_eq!(request_body_json(&request)["version"], 42);
let body = match request.url().path() {
"/v1/table/my_table/describe/" => r#"{"version":42,"schema":{"fields":[]}}"#,
"/v1/table/my_table/restore/" => r#"{"version":43}"#,
path => panic!("unexpected request path: {path}"),
};
http::Response::builder().status(200).body(body).unwrap()
});
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 0);
table.checkout(42).await.unwrap();
table.checkout_latest().await.unwrap();
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 1);
table.checkout(42).await.unwrap();
table.restore().await.unwrap();
assert_eq!(request_count.load(Ordering::SeqCst), 3);
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_fails_if_checkout_version_doesnt_exist() {
let table = Table::new_with_handler("my_table", |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let version = body
.as_object()
.unwrap()
.get("version")
.unwrap()
.as_u64()
.unwrap();
if version != 42 {
return http::Response::builder()
.status(404)
.body(format!("Table my_table (version: {}) not found", version))
.unwrap();
}
let response_body = match request.url().path() {
"/v1/table/my_table/describe/" => {
serde_json::json!({
"version": 42,
"schema": { "fields": [] }
})
}
_ => panic!("Unexpected path"),
};
http::Response::builder()
.status(200)
.body(serde_json::to_string(&response_body).unwrap())
.unwrap()
});
let res = table.checkout(43).await;
println!("{:?}", res);
assert!(
matches!(res, Err(Error::TableNotFound { name, .. }) if name == "my_table (version: 43)")
);
}
#[tokio::test]
async fn test_timetravel_immutable() {
let table = Table::new_with_handler::<String>("my_table", |request| {
let response_body = match request.url().path() {
"/v1/table/my_table/describe/" => {
serde_json::json!({
"version": 42,
"schema": { "fields": [] }
})
}
_ => panic!("Should not have made a request: {:?}", request),
};
http::Response::builder()
.status(200)
.body(serde_json::to_string(&response_body).unwrap())
.unwrap()
});
table.checkout(42).await.unwrap();
// Ensure that all mutable operations fail.
let res = table
.update()
.column("a", "a + 1")
.column("b", "b - 1")
.only_if("b > 10")
.execute()
.await;
assert!(matches!(res, Err(Error::NotSupported { .. })));
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
));
let res = table.merge_insert(&["some_col"]).execute(data).await;
assert!(matches!(res, Err(Error::NotSupported { .. })));
let res = table.delete("id in (1, 2, 3)").await;
assert!(matches!(res, Err(Error::NotSupported { .. })));
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let res = table.add(data.clone()).execute().await;
assert!(matches!(res, Err(Error::NotSupported { .. })));
let res = table
.create_index(&["a"], Index::IvfPq(Default::default()))
.execute()
.await;
assert!(matches!(res, Err(Error::NotSupported { .. })));
}
rust/lancedb/src/remote/table/tests/write.rs
+994
View File
@@ -0,0 +1,994 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
#[tokio::test]
async fn test_not_found() {
let table = Table::new_with_handler("my_table", |_| {
http::Response::builder()
.status(404)
.body("table my_table not found")
.unwrap()
});
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let example_data_for_add = || batch.clone();
let example_data_for_merge = || -> Box<dyn RecordBatchReader + Send> {
Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
))
};
// These table operations should translate 404 to TableNotFound.
let results: Vec<BoxFuture<'_, Result<()>>> = vec![
Box::pin(table.version().map_ok(|_| ())),
Box::pin(table.schema().map_ok(|_| ())),
Box::pin(table.count_rows(None).map_ok(|_| ())),
Box::pin(table.update().column("a", "a + 1").execute().map_ok(|_| ())),
Box::pin(table.add(example_data_for_add()).execute().map_ok(|_| ())),
Box::pin(
table
.merge_insert(&["test"])
.execute(example_data_for_merge())
.map_ok(|_| ()),
),
Box::pin(table.delete("false").map_ok(|_| ())),
Box::pin(
table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"x".into(),
"y".into(),
)]))
.execute()
.map_ok(|_| ()),
),
Box::pin(async {
let alterations = vec![ColumnAlteration::new("x".into()).rename("y".into())];
table.alter_columns(&alterations).await.map(|_| ())
}),
Box::pin(table.drop_columns(&["a"]).map_ok(|_| ())),
// TODO: other endpoints.
];
for result in results {
let result = result.await;
assert!(result.is_err());
assert!(
matches!(&result, &Err(Error::TableNotFound { ref name, .. }) if name == "my_table")
);
let full_error_report = snafu::Report::from_error(result.unwrap_err()).to_string();
assert!(full_error_report.contains("table my_table not found"));
}
}
#[tokio::test]
async fn test_version() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/describe/");
http::Response::builder()
.status(200)
.body(r#"{"version": 42, "schema": { "fields": [] }}"#)
.unwrap()
});
let version = table.version().await.unwrap();
assert_eq!(version, 42);
}
#[tokio::test]
async fn test_schema() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/describe/");
http::Response::builder()
.status(200)
.body(
r#"{"version": 42, "schema": {"fields": [
{"name": "a", "type": { "type": "int32" }, "nullable": false},
{"name": "b", "type": { "type": "string" }, "nullable": true}
], "metadata": {"key": "value"}}}"#,
)
.unwrap()
});
let expected = Arc::new(
Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Utf8, true),
])
.with_metadata([("key".into(), "value".into())].into()),
);
let schema = table.schema().await.unwrap();
assert_eq!(schema, expected);
}
#[tokio::test]
async fn test_count_rows() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/count_rows/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
assert_eq!(
request.body().unwrap().as_bytes().unwrap(),
br#"{"version":null}"#
);
http::Response::builder().status(200).body("42").unwrap()
});
let count = table.count_rows(None).await.unwrap();
assert_eq!(count, 42);
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/count_rows/");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
assert_eq!(
request.body().unwrap().as_bytes().unwrap(),
br#"{"predicate":"`A` > 10","version":null}"#
);
http::Response::builder().status(200).body("42").unwrap()
});
let count = table
.base_table()
.count_rows(Some(Filter::Sql(r#""A" > 10"#.into())))
.await
.unwrap();
assert_eq!(count, 42);
}
#[rstest]
#[case("", 0)]
#[case("{}", 0)]
#[case(r#"{"request_id": "test-request-id"}"#, 0)]
#[case(r#"{"version": 43}"#, 43)]
#[tokio::test]
async fn test_add_append(#[case] response_body: &str, #[case] expected_version: u64) {
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
// Clone response_body to give it 'static lifetime for the closure
let response_body = response_body.to_string();
let describe_body = describe_response(&data.schema());
let (sender, receiver) = std::sync::mpsc::channel();
let table =
Table::new_with_handler("my_table", move |mut request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
assert_eq!(request.method(), "POST");
assert!(
request
.url()
.query_pairs()
.filter(|(k, _)| k == "mode")
.all(|(_, v)| v == "append")
);
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
let mut body_out = reqwest::Body::from(Vec::new());
std::mem::swap(request.body_mut().as_mut().unwrap(), &mut body_out);
sender.send(body_out).unwrap();
http::Response::builder()
.status(200)
.body(response_body.clone())
.unwrap()
}
path => panic!("Unexpected request path: {}", path),
});
let result = table.add(data.clone()).execute().await.unwrap();
// Check version matches expected value
assert_eq!(result.version, expected_version);
let body = receiver.recv().unwrap();
let body = collect_body(body).await;
let expected_body = write_ipc_stream(&data);
assert_eq!(&body, &expected_body);
}
#[tokio::test]
async fn add_rejects_external_blob_flag_before_any_request() {
let table = Table::new_with_handler::<String>("my_table", |request| {
panic!("Unexpected request: {}", request.url().path())
});
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
let err = table
.add(data)
.allow_external_blob_outside_bases(true)
.execute()
.await
.unwrap_err();
assert!(matches!(err, Error::NotSupported { .. }), "got {err:?}");
assert!(err.to_string().contains("local tables"));
}
#[tokio::test]
async fn add_string_blob_becomes_uri_struct_without_the_local_flag() {
let table_schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
crate::blob("image", true),
]);
let describe_body = describe_response(&table_schema);
let input = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("image", DataType::Utf8, true),
])),
vec![
Arc::new(Int64Array::from(vec![1])),
Arc::new(StringArray::from(vec![Some("s3://bucket/key")])),
],
)
.unwrap();
let (sender, receiver) = std::sync::mpsc::channel();
let table =
Table::new_with_handler("my_table", move |mut request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
let mut body_out = reqwest::Body::from(Vec::new());
std::mem::swap(request.body_mut().as_mut().unwrap(), &mut body_out);
sender.send(body_out).unwrap();
http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap()
}
path => panic!("Unexpected path: {path}"),
});
table.add(input).execute().await.unwrap();
let body = collect_body(receiver.recv().unwrap()).await;
let mut reader =
arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None).unwrap();
let batch = reader.next().unwrap().unwrap();
let image = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<StructArray>()
.expect("remote add should send the coerced blob struct");
let uri: &StringArray = image
.column_by_name("uri")
.unwrap()
.as_any()
.downcast_ref()
.unwrap();
assert_eq!(uri.value(0), "s3://bucket/key");
assert!(image.column_by_name("data").unwrap().is_null(0));
}
#[rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_add_overwrite(#[case] old_server: bool) {
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let describe_body = describe_response(&data.schema());
let (sender, receiver) = std::sync::mpsc::channel();
let table =
Table::new_with_handler("my_table", move |mut request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
assert_eq!(request.method(), "POST");
assert_eq!(
request
.url()
.query_pairs()
.find(|(k, _)| k == "mode")
.map(|kv| kv.1)
.as_deref(),
Some("overwrite"),
"Expected mode=overwrite"
);
assert_eq!(
request.headers().get("Content-Type").unwrap(),
ARROW_STREAM_CONTENT_TYPE
);
let mut body_out = reqwest::Body::from(Vec::new());
std::mem::swap(request.body_mut().as_mut().unwrap(), &mut body_out);
sender.send(body_out).unwrap();
if old_server {
http::Response::builder()
.status(200)
.body("".to_string())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 43}"#.to_string())
.unwrap()
}
}
path => panic!("Unexpected request path: {}", path),
});
let result = table
.add(data.clone())
.mode(AddDataMode::Overwrite)
.execute()
.await
.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
let body = receiver.recv().unwrap();
let body = collect_body(body).await;
let expected_body = write_ipc_stream(&data);
assert_eq!(&body, &expected_body);
}
#[tokio::test]
async fn test_add_preprocessing() {
use crate::table::NaNVectorBehavior;
use arrow_array::{FixedSizeListArray, Float32Array, Int64Array};
// The table schema: {id: Int64, vec: FixedSizeList<Float32>[3]}
let table_schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new(
"vec",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 3),
false,
),
]);
let json_schema = JsonSchema::try_from(&table_schema).unwrap();
let describe_body = serde_json::to_string(&json!({
"version": 1,
"schema": json_schema,
}))
.unwrap();
// ---- Part 1: NaN vectors should be rejected by default ----
let nan_data = RecordBatch::try_new(
Arc::new(table_schema.clone()),
vec![
Arc::new(Int64Array::from(vec![1])),
Arc::new(
FixedSizeListArray::try_new(
Arc::new(Field::new("item", DataType::Float32, true)),
3,
Arc::new(Float32Array::from(vec![1.0, f32::NAN, 3.0])),
None,
)
.unwrap(),
),
],
)
.unwrap();
let describe_body_clone = describe_body.clone();
let table = Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body_clone.clone())
.unwrap(),
"/v1/table/my_table/insert/" => http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap(),
path => panic!("Unexpected path: {path}"),
});
let result = table.add(nan_data).execute().await;
assert!(result.is_err(), "NaN vectors should be rejected by default");
assert!(
result.unwrap_err().to_string().contains("NaN"),
"error should mention NaN"
);
// ---- Part 2: With Keep, should handle casting and missing columns ----
// Input: {id: Int32 (needs cast to Int64), vec: FixedSizeList<Float32>[3] with NaN}
// Table expects Int64 for id; NaN should be kept.
let input_schema = Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new(
"vec",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 3),
false,
),
]);
let cast_data = RecordBatch::try_new(
Arc::new(input_schema),
vec![
Arc::new(Int32Array::from(vec![42])),
Arc::new(
FixedSizeListArray::try_new(
Arc::new(Field::new("item", DataType::Float32, true)),
3,
Arc::new(Float32Array::from(vec![1.0, f32::NAN, 3.0])),
None,
)
.unwrap(),
),
],
)
.unwrap();
let (sender, receiver) = std::sync::mpsc::channel();
let table =
Table::new_with_handler("my_table", move |mut request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
let mut body_out = reqwest::Body::from(Vec::new());
std::mem::swap(request.body_mut().as_mut().unwrap(), &mut body_out);
sender.send(body_out).unwrap();
http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap()
}
path => panic!("Unexpected path: {path}"),
});
table
.add(cast_data)
.on_nan_vectors(NaNVectorBehavior::Keep)
.execute()
.await
.unwrap();
// Verify the data sent to the server was cast to the table schema.
let body = receiver.recv().unwrap();
let body = collect_body(body).await;
let cursor = std::io::Cursor::new(body);
let mut reader = arrow_ipc::reader::StreamReader::try_new(cursor, None).unwrap();
let batch = reader.next().unwrap().unwrap();
assert_eq!(batch.schema().field(0).data_type(), &DataType::Int64);
let ids: &Int64Array = batch.column(0).as_any().downcast_ref().unwrap();
assert_eq!(ids.value(0), 42);
}
#[rstest]
#[case::old_server("", 0, 0)]
#[case::rows_updated(r#"{"rows_updated": 5, "version": 43}"#, 5, 43)]
#[case::updated_rows(r#"{"updated_rows": 5, "version": 43}"#, 5, 43)]
#[case::zero_updated_rows(r#"{"updated_rows": 0, "version": 43}"#, 0, 43)]
#[case::missing_row_count(r#"{"version": 43}"#, 0, 43)]
#[tokio::test]
async fn test_update(
#[case] response_body: &'static str,
#[case] expected_rows_updated: u64,
#[case] expected_version: u64,
#[values(true, false)] filtered: bool,
) {
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/update/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
if let Some(body) = request.body().unwrap().as_bytes() {
let body = std::str::from_utf8(body).unwrap();
let value: serde_json::Value = serde_json::from_str(body).unwrap();
let updates = value.get("updates").unwrap().as_array().unwrap();
assert!(updates.len() == 2);
let col_name = updates[0][0].as_str().unwrap();
let expression = updates[0][1].as_str().unwrap();
assert_eq!(col_name, "a");
assert_eq!(expression, "a + 1");
let col_name = updates[1][0].as_str().unwrap();
let expression = updates[1][1].as_str().unwrap();
assert_eq!(col_name, "b");
assert_eq!(expression, "b - 1");
assert_eq!(
value.get("predicate").unwrap(),
&serde_json::json!(if filtered { Some("`B` > 10") } else { None })
);
}
http::Response::builder()
.status(200)
.body(response_body)
.unwrap()
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let mut update = table.update().column("a", "a + 1").column("b", "b - 1");
if filtered {
update = update.only_if(r#""B" > 10"#);
}
let result = table.base_table().update(update).await.unwrap();
assert_eq!(result.version, expected_version);
assert_eq!(result.rows_updated, expected_rows_updated);
}
#[tokio::test]
async fn test_alter_columns_rejects_missing_changes_before_request() {
let table = Table::new_with_handler::<String>("my_table", |request| {
panic!("Unexpected request: {}", request.url().path())
});
for alterations in [
vec![ColumnAlteration::new("id".into())],
vec![
ColumnAlteration::new("id".into()).rename("new_id".into()),
ColumnAlteration::new("id".into()),
],
] {
let err = table.alter_columns(&alterations).await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}");
assert!(err.to_string().contains("path 'id'"));
}
}
#[rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_alter_columns(#[case] old_server: bool) {
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/alter_columns/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body = std::str::from_utf8(body).unwrap();
let value: serde_json::Value = serde_json::from_str(body).unwrap();
let alterations = value.get("alterations").unwrap().as_array().unwrap();
assert!(alterations.len() == 2);
let path = alterations[0]["path"].as_str().unwrap();
let data_type = alterations[0]["data_type"]["type"].as_str().unwrap();
assert_eq!(path, "b.c");
assert_eq!(data_type, "int32");
let path = alterations[1]["path"].as_str().unwrap();
let nullable = alterations[1]["nullable"].as_bool().unwrap();
let rename = alterations[1]["rename"].as_str().unwrap();
assert_eq!(path, "x");
assert!(nullable);
assert_eq!(rename, "y");
if old_server {
http::Response::builder()
.status(200)
.body("{}".to_string())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 43}"#.to_string())
.unwrap()
}
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let result = table
.alter_columns(&[
ColumnAlteration::new("b.c".into()).cast_to(DataType::Int32),
ColumnAlteration::new("x".into())
.rename("y".into())
.set_nullable(true),
])
.await
.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
}
#[rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_merge_insert(#[case] old_server: bool) {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
));
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/merge_insert/" {
assert_eq!(request.method(), "POST");
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(params["on"], "some_col");
assert_eq!(params["when_matched_update_all"], "true");
assert_eq!(params["when_not_matched_insert_all"], "false");
assert_eq!(params["when_not_matched_by_source_delete"], "false");
assert_eq!(params["when_matched_update_all_filt"], "target.`A` > 0");
assert!(!params.contains_key("when_not_matched_by_source_delete_filt"));
assert!(!params.contains_key("use_index"));
if old_server {
http::Response::builder().status(200).body("{}").unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 43, "num_deleted_rows": 0, "num_inserted_rows": 3, "num_updated_rows": 0}"#)
.unwrap()
}
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let mut merge = table.merge_insert(&["some_col"]);
merge.when_matched_update_all(Some(r#"target."A" > 0"#.into()));
let result = table.base_table().merge_insert(merge, data).await.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
if !old_server {
assert_eq!(result.num_deleted_rows, 0);
assert_eq!(result.num_inserted_rows, 3);
assert_eq!(result.num_updated_rows, 0);
}
}
#[rstest]
#[case::no_batches_insert(false, false)]
#[case::empty_batch_insert(true, false)]
#[case::no_batches_delete(false, true)]
#[case::empty_batch_delete(true, true)]
#[tokio::test]
async fn test_merge_insert_empty_source(#[case] has_batch: bool, #[case] delete_unmatched: bool) {
let mut fields = vec![Field::new("id", DataType::Int64, false)];
if !delete_unmatched {
fields.extend([
Field::new("k", DataType::Int64, false),
Field::new(
"vector",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
true,
),
Field::new("s", DataType::Utf8, true),
]);
}
let schema = Arc::new(Schema::new_with_metadata(
fields,
HashMap::from([("source".to_string(), "empty".to_string())]),
));
let batches = if has_batch {
vec![Ok(RecordBatch::new_empty(schema.clone()))]
} else {
vec![]
};
let data: Box<dyn RecordBatchReader + Send> =
Box::new(RecordBatchIterator::new(batches, schema.clone()));
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_ref = attempts.clone();
let num_deleted_rows = if delete_unmatched { 3 } else { 0 };
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE);
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(params["on"], "id");
assert_eq!(
params["when_not_matched_insert_all"],
(!delete_unmatched).to_string()
);
assert_eq!(
params["when_not_matched_by_source_delete"],
delete_unmatched.to_string()
);
let body = request.body().unwrap().as_bytes().unwrap();
let reader = StreamReader::try_new(Cursor::new(body), None).unwrap();
assert_eq!(reader.schema(), schema);
for batch in reader {
assert_eq!(batch.unwrap().num_rows(), 0);
}
// The empty source must retain its schema when replayed after a conflict.
if attempts_ref.fetch_add(1, Ordering::SeqCst) == 0 {
http::Response::builder()
.status(409)
.body(String::new())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(
json!({
"version": 43,
"num_deleted_rows": num_deleted_rows,
"num_inserted_rows": 0,
"num_updated_rows": 0,
})
.to_string(),
)
.unwrap()
}
});
let mut merge = table.merge_insert(&["id"]);
if delete_unmatched {
merge.when_not_matched_by_source_delete(None);
} else {
merge.when_not_matched_insert_all();
}
let result = merge.execute(data).await.unwrap();
assert_eq!(attempts.load(Ordering::SeqCst), 2);
assert_eq!(result.version, 43);
assert_eq!(result.num_deleted_rows, num_deleted_rows);
assert_eq!(result.num_inserted_rows, 0);
assert_eq!(result.num_updated_rows, 0);
}
#[tokio::test]
async fn test_merge_insert_composite_key() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
));
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
// One repeated `on` per column, in the order the caller gave them.
let on = request
.url()
.query_pairs()
.filter(|(key, _)| key == "on")
.map(|(_, value)| value.into_owned())
.collect::<Vec<_>>();
assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]);
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(params["when_matched_update_all"], "true");
assert_eq!(params["when_not_matched_insert_all"], "true");
http::Response::builder()
.status(200)
.body(r#"{"version": 43, "num_deleted_rows": 0, "num_inserted_rows": 3, "num_updated_rows": 0}"#)
.unwrap()
});
let mut merge = table.merge_insert(&["shard_key", "id"]);
merge.when_matched_update_all(None);
merge.when_not_matched_insert_all();
let result = table.base_table().merge_insert(merge, data).await.unwrap();
assert_eq!(result.num_inserted_rows, 3);
}
#[tokio::test]
async fn test_merge_insert_rejects_repeated_on_column() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
));
let table = Table::new_with_handler::<&str>("my_table", |request| {
panic!("Unexpected request: {}", request.url());
});
let merge = table.merge_insert(&["id", "id"]);
let err = table
.base_table()
.merge_insert(merge, data)
.await
.unwrap_err();
assert!(
matches!(&err, Error::InvalidInput { message } if message.contains("'id' is repeated")),
"unexpected error: {err}"
);
}
#[tokio::test]
async fn test_merge_insert_retries_on_409() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
[Ok(batch.clone())],
batch.schema(),
));
// Default parameters
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(params["on"], "some_col");
assert_eq!(params["when_matched_update_all"], "false");
assert_eq!(params["when_not_matched_insert_all"], "false");
assert_eq!(params["when_not_matched_by_source_delete"], "false");
assert!(!params.contains_key("when_matched_update_all_filt"));
assert!(!params.contains_key("when_not_matched_by_source_delete_filt"));
http::Response::builder().status(409).body("").unwrap()
});
let e = table
.merge_insert(&["some_col"])
.execute(data)
.await
.unwrap_err();
assert!(e.to_string().contains("Hit retry limit"));
}
#[rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_delete(#[case] old_server: bool) {
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/delete/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
let predicate = body.get("predicate").unwrap().as_str().unwrap();
assert_eq!(predicate, "`ID` in (1, 2, 3)");
if old_server {
http::Response::builder()
.status(200)
.body("{}".to_string())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 43}"#.to_string())
.unwrap()
}
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let result = table
.base_table()
.delete(Predicate::String(r#""ID" in (1, 2, 3)"#))
.await
.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
}
#[tokio::test]
async fn test_delete_expr() {
use datafusion_expr::{col, lit};
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/delete/" {
assert_eq!(request.method(), "POST");
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert!(body.get("predicate").unwrap().is_string());
http::Response::builder()
.status(200)
.body(r#"{"num_deleted_rows": 4, "version": 2}"#)
.unwrap()
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let expr = col("id").gt(lit(5));
let result = table.delete(&expr).await.unwrap();
assert_eq!(result.num_deleted_rows, 4);
assert_eq!(result.version, 2);
}
#[rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_drop_columns(#[case] old_server: bool) {
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/drop_columns/" {
assert_eq!(request.method(), "POST");
assert_eq!(
request.headers().get("Content-Type").unwrap(),
JSON_CONTENT_TYPE
);
let body = request.body().unwrap().as_bytes().unwrap();
let body = std::str::from_utf8(body).unwrap();
let value: serde_json::Value = serde_json::from_str(body).unwrap();
let columns = value.get("columns").unwrap().as_array().unwrap();
assert!(columns.len() == 2);
let col1 = columns[0].as_str().unwrap();
let col2 = columns[1].as_str().unwrap();
assert_eq!(col1, "a");
assert_eq!(col2, "b");
if old_server {
http::Response::builder().status(200).body("{}").unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"version": 43}"#)
.unwrap()
}
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let result = table.drop_columns(&["a", "b"]).await.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
}
rust/lancedb/src/remote/token_cache.rs
+1 -915
View File
@@ -979,918 +979,4 @@ pub fn token_cache_for_config(config: &OAuthConfig) -> Result<Option<Arc<TokenCa
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::remote::HeaderProvider;
use crate::remote::oauth::OAuthHeaderProvider;
use oauth2::basic::BasicTokenType;
use oauth2::{AccessToken, RefreshToken};
use serial_test::serial;
/// Temp directory that satisfies the cache hardening checks. CI runners
/// can create temp directories with group/other bits set, which the
/// private-directory validation correctly rejects.
fn cache_tempdir() -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
}
dir
}
fn device_config(cache_dir: &Path) -> OAuthConfig {
OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: None,
client_auth_method: None,
scopes: vec!["openid".to_string()],
flow: OAuthFlow::DeviceCode,
refresh_buffer_secs: None,
resource: None,
audience: None,
token_cache: Some(TokenCacheOptions::new().cache_dir(cache_dir)),
}
}
/// Stateful mock IdP covering discovery, device authorization, device
/// polling, client credentials, and refresh with strict rotation: a
/// refresh token that is not the currently issued one is rejected with
/// `invalid_grant`, which is exactly what real providers do on rotation.
struct MockIdp {
issuer_url: String,
requests: Arc<std::sync::Mutex<Vec<String>>>,
device_authorizations: Arc<AtomicUsize>,
refresh_attempts: Arc<AtomicUsize>,
invalid_grant_rejections: Arc<AtomicUsize>,
access_tokens_issued: Arc<AtomicUsize>,
current_refresh: Arc<std::sync::Mutex<Option<String>>>,
fail_refreshes: Arc<AtomicBool>,
issue_refresh_tokens: Arc<AtomicBool>,
}
impl MockIdp {
async fn start() -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = Self {
issuer_url: issuer_url.clone(),
requests: Arc::new(std::sync::Mutex::new(Vec::new())),
device_authorizations: Arc::new(AtomicUsize::new(0)),
refresh_attempts: Arc::new(AtomicUsize::new(0)),
invalid_grant_rejections: Arc::new(AtomicUsize::new(0)),
access_tokens_issued: Arc::new(AtomicUsize::new(0)),
current_refresh: Arc::new(std::sync::Mutex::new(None)),
fail_refreshes: Arc::new(AtomicBool::new(false)),
issue_refresh_tokens: Arc::new(AtomicBool::new(true)),
};
let requests = Arc::clone(&server.requests);
let device_authorizations = Arc::clone(&server.device_authorizations);
let refresh_attempts = Arc::clone(&server.refresh_attempts);
let invalid_grant_rejections = Arc::clone(&server.invalid_grant_rejections);
let access_tokens_issued = Arc::clone(&server.access_tokens_issued);
let current_refresh = Arc::clone(&server.current_refresh);
let fail_refreshes = Arc::clone(&server.fail_refreshes);
let issue_refresh_tokens = Arc::clone(&server.issue_refresh_tokens);
tokio::spawn(async move {
loop {
let Ok((mut stream, _)) = listener.accept().await else {
return;
};
let requests = Arc::clone(&requests);
let device_authorizations = Arc::clone(&device_authorizations);
let refresh_attempts = Arc::clone(&refresh_attempts);
let invalid_grant_rejections = Arc::clone(&invalid_grant_rejections);
let access_tokens_issued = Arc::clone(&access_tokens_issued);
let current_refresh = Arc::clone(&current_refresh);
let fail_refreshes = Arc::clone(&fail_refreshes);
let issue_refresh_tokens = Arc::clone(&issue_refresh_tokens);
tokio::spawn(async move {
let (request_line, body) = read_http_request(&mut stream).await;
if request_line.starts_with("POST ") {
requests.lock().unwrap().push(body.clone());
}
if request_line.starts_with("GET /.well-known/openid-configuration ") {
let discovery = format!(
r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"#
);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request_line.starts_with("POST /device ") {
device_authorizations.fetch_add(1, Ordering::SeqCst);
let device = format!(
r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","expires_in":60,"interval":1}}"#
);
write_json_response(&mut stream, "200 OK", &device).await;
} else if request_line.starts_with("POST /token ") {
if body.contains("grant_type=refresh_token") {
refresh_attempts.fetch_add(1, Ordering::SeqCst);
if fail_refreshes.load(Ordering::SeqCst) {
write_json_response(
&mut stream,
"503 Service Unavailable",
r#"{"error":"temporarily_unavailable"}"#,
)
.await;
return;
}
let expected = current_refresh.lock().unwrap().clone();
let matched = body
.split('&')
.find_map(|pair| pair.strip_prefix("refresh_token="))
.map(|token| token.to_string())
.zip(expected)
.is_some_and(|(offered, expected)| offered == expected);
if !matched {
invalid_grant_rejections.fetch_add(1, Ordering::SeqCst);
write_json_response(
&mut stream,
"400 Bad Request",
r#"{"error":"invalid_grant"}"#,
)
.await;
return;
}
issue_access_token(
&mut stream,
&access_tokens_issued,
&current_refresh,
issue_refresh_tokens.load(Ordering::SeqCst),
)
.await;
} else {
// Device polling or client credentials: issue
// a token and, for interactive grants, a fresh
// refresh token with strict rotation.
let grant_device =
body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code");
if grant_device {
issue_access_token(
&mut stream,
&access_tokens_issued,
&current_refresh,
issue_refresh_tokens.load(Ordering::SeqCst),
)
.await;
} else {
let token = format!(
r#"{{"access_token":"access-{}","expires_in":3600}}"#,
access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1
);
write_json_response(&mut stream, "200 OK", &token).await;
}
}
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
});
}
});
server
}
fn config(&self, cache_dir: &Path) -> OAuthConfig {
let mut config = device_config(cache_dir);
config.issuer_url = self.issuer_url.clone();
config
}
}
async fn issue_access_token(
stream: &mut TcpStream,
access_tokens_issued: &AtomicUsize,
current_refresh: &std::sync::Mutex<Option<String>>,
with_refresh: bool,
) {
let number = access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1;
let token = if with_refresh {
*current_refresh.lock().unwrap() = Some(format!("refresh-{number}"));
format!(
r#"{{"access_token":"access-{number}","refresh_token":"refresh-{number}","expires_in":3600}}"#
)
} else {
format!(r#"{{"access_token":"access-{number}","expires_in":3600}}"#)
};
write_json_response(stream, "200 OK", &token).await;
}
async fn read_http_request(stream: &mut TcpStream) -> (String, String) {
let mut buffer = Vec::new();
let mut header_end = None;
while header_end.is_none() {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request headers");
buffer.extend_from_slice(&chunk[..read]);
header_end = find_subsequence(&buffer, b"\r\n\r\n").map(|pos| pos + 4);
}
let header_end = header_end.unwrap();
let headers = String::from_utf8_lossy(&buffer[..header_end]).to_string();
let request_line = headers.lines().next().unwrap_or_default().to_string();
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while buffer.len() < header_end + content_length {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request body");
buffer.extend_from_slice(&chunk[..read]);
}
let body =
String::from_utf8_lossy(&buffer[header_end..header_end + content_length]).to_string();
(request_line, body)
}
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack
.windows(needle.len())
.position(|window| window == needle)
}
async fn write_json_response(stream: &mut TcpStream, status: &str, body: &str) {
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).await.unwrap();
}
fn suppress_browser() {
// Point the OAuth browser helper at a no-op so device-flow tests never
// open a real browser window.
unsafe { std::env::set_var("LANCEDB_OAUTH_BROWSER", "/usr/bin/true") };
}
#[tokio::test]
#[serial]
async fn test_provider_reuses_cached_session_across_instances() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let first = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = first.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
// A second, independent provider (simulating a second process) must
// refresh silently instead of starting another device flow.
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = second.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2");
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 1);
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-2");
}
#[tokio::test]
#[serial]
async fn test_concurrent_providers_serialize_rotation() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
let provider_a = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let provider_b = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let (headers_a, headers_b) =
tokio::join!(provider_a.get_headers(), provider_b.get_headers());
let token_a = headers_a.unwrap().remove("authorization").unwrap();
let token_b = headers_b.unwrap().remove("authorization").unwrap();
assert!(
{
let mut tokens = [token_a, token_b];
tokens.sort();
tokens
} == ["Bearer access-2".to_string(), "Bearer access-3".to_string()],
"each provider must observe its own refreshed token"
);
// Rotation raced would produce an invalid_grant and a second device
// flow; the lock prevents both.
assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 0);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 2);
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-3");
}
#[tokio::test]
#[serial]
async fn test_transient_refresh_failure_retains_record() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
idp.fail_refreshes.store(true, Ordering::SeqCst);
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let error = second.get_headers().await.unwrap_err();
assert!(error.to_string().contains("503"));
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(
status.refreshable,
"transient failures must keep the record"
);
}
#[tokio::test]
#[serial]
async fn test_invalid_grant_deletes_record_and_reauthenticates() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
// Simulate a revoked refresh token by replacing the record with one
// the provider never issued.
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let mut record = cache.load().await.unwrap().unwrap();
record.refresh_token = "revoked-refresh".to_string();
cache.store(&record).await.unwrap();
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = second.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2");
assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 1);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 2);
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-2");
}
#[tokio::test]
#[serial]
async fn test_session_login_status_logout_lifecycle() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!status.refreshable);
let status = session.login().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.issuer_url, idp.issuer_url);
assert_eq!(status.client_id, "client-id");
assert_eq!(status.scopes, vec!["openid".to_string()]);
assert_eq!(status.flow, "device_code");
assert!(status.obtained_at.is_some());
// An independent session manager sees the same cached login.
let other = OAuthSession::new(idp.config(dir.path())).unwrap();
assert!(other.status().await.unwrap().refreshable);
assert!(session.logout().await.unwrap().removed);
assert!(!session.logout().await.unwrap().removed);
assert!(!other.status().await.unwrap().refreshable);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
}
#[tokio::test]
#[serial]
async fn test_login_without_refresh_token_clears_prior_record() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
session.login().await.unwrap();
assert!(session.status().await.unwrap().refreshable);
// A provider that stops issuing refresh tokens (for example a login
// without offline_access) must not leave the earlier account behind.
idp.issue_refresh_tokens.store(false, Ordering::SeqCst);
let status = session.login().await.unwrap();
assert!(!status.refreshable);
assert!(!session.status().await.unwrap().refreshable);
}
#[tokio::test]
async fn test_client_credentials_with_cache_stays_memory_only() {
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let mut config = idp.config(dir.path());
config.flow = OAuthFlow::ClientCredentials;
config.client_secret = Some("secret".to_string());
let provider = OAuthHeaderProvider::new(config).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
assert_eq!(dir.path().read_dir().unwrap().count(), 0);
}
#[test]
fn test_managed_identity_with_cache_is_rejected() {
let dir = cache_tempdir();
let mut config = device_config(dir.path());
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
let err = OAuthHeaderProvider::new(config).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity"))
);
}
#[tokio::test]
#[serial]
async fn test_provider_debug_and_status_reveal_no_secrets() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let provider = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
let debug = format!("{provider:?}");
assert!(!debug.contains("access-1"));
assert!(!debug.contains("refresh-1"));
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!format!("{status:?}").contains("refresh-"));
}
#[tokio::test]
async fn test_targeted_cache_refresh_and_logout_isolation() {
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let mut config = idp.config(dir.path());
let options = config.token_cache.clone().unwrap();
let untargeted_key = CacheKey::new(&config).unwrap().file_stem;
assert_eq!(
untargeted_key,
hex_sha256(
format!(
"v1\n{}\nclient-id\nopenid\ndevice_code\npublic",
idp.issuer_url
)
.as_bytes()
)
);
let mut sessions = Vec::new();
let mut keys = std::collections::HashSet::new();
for (resource, audience) in [
(None, None),
(Some("urn:one"), None),
(None, Some("audience & ü")),
(Some("urn:one"), Some("audience & ü")),
(Some("urn:two"), Some("audience & ü")),
(Some("urn:one"), Some("other")),
] {
config.resource = resource.map(str::to_owned);
config.audience = audience.map(str::to_owned);
let cache = TokenCache::new(&config, &options).unwrap();
assert!(keys.insert(cache.key.file_stem.clone()));
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("unused".into()),
refresh_token: Some(RefreshToken::new("seed-refresh".into())),
expires_in: Some(3600),
token_type: None,
})
.unwrap();
cache.store(&record).await.unwrap();
*idp.current_refresh.lock().unwrap() = Some("seed-refresh".into());
let provider = OAuthHeaderProvider::new(config.clone()).unwrap();
provider.get_headers().await.unwrap();
let request = idp.requests.lock().unwrap().last().unwrap().clone();
let params: std::collections::HashMap<_, _> =
url::form_urlencoded::parse(request.as_bytes()).collect();
assert_eq!(params.get("resource").map(|s| s.as_ref()), resource);
assert_eq!(params.get("audience").map(|s| s.as_ref()), audience);
assert_eq!(params.get("grant_type").unwrap(), "refresh_token");
let session = OAuthSession::new(config.clone()).unwrap();
let status = session.status().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.resource, config.resource);
assert_eq!(status.audience, config.audience);
sessions.push(session);
}
assert!(sessions.pop().unwrap().logout().await.unwrap().removed);
for session in sessions {
assert!(session.status().await.unwrap().refreshable);
}
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_legacy_cache_record_without_target_fields() {
let dir = cache_tempdir();
let config = device_config(dir.path());
let cache = TokenCache::new(&config, config.token_cache.as_ref().unwrap()).unwrap();
let legacy = br#"{"version":1,"issuer_url":"https://issuer.example.com","client_id":"client-id","scopes":["openid"],"flow":"device_code","client_auth":"public","refresh_token":"legacy","obtained_at":1}"#;
write_record(dir.path(), &cache.record_path(), legacy).unwrap();
let session = OAuthSession::new(config).unwrap();
let status = session.status().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.resource, None);
assert_eq!(status.audience, None);
assert!(session.logout().await.unwrap().removed);
}
#[test]
fn test_cache_key_canonicalizes_scopes_and_issuer() {
let mut config = device_config(Path::new("/tmp/cache"));
config.issuer_url = "https://issuer.example.com/".to_string();
config.scopes = vec!["b".to_string(), " a ".to_string(), "a".to_string()];
let key = CacheKey::new(&config).unwrap();
assert_eq!(key.issuer_url, "https://issuer.example.com");
assert_eq!(key.scopes, vec!["a".to_string(), "b".to_string()]);
config.issuer_url = "https://issuer.example.com".to_string();
let canonical = CacheKey::new(&config).unwrap();
assert_eq!(canonical.file_stem, key.file_stem);
}
#[test]
fn test_cache_key_separates_identity_dimensions() {
let base = device_config(Path::new("/tmp/cache"));
let base_key = CacheKey::new(&base).unwrap();
let mut other = base.clone();
other.client_id = "other-client".to_string();
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.issuer_url = "https://other.example.com".to_string();
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.scopes = vec!["profile".to_string()];
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.flow = OAuthFlow::AuthorizationCode(Default::default());
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.client_secret = Some("secret".to_string());
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
}
#[test]
fn test_cache_key_contains_no_secret_material() {
let mut config = device_config(Path::new("/tmp/cache"));
config.client_secret = Some("super-secret-value".to_string());
let key = CacheKey::new(&config).unwrap();
assert!(!key.file_stem.contains("super-secret-value"));
assert_eq!(key.file_stem.len(), 64);
}
#[test]
fn test_flow_key_rejects_non_persistent_flows() {
let mut config = device_config(Path::new("/tmp/cache"));
config.flow = OAuthFlow::ClientCredentials;
let err = CacheKey::new(&config).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("not supported"))
);
let mut config = device_config(Path::new("/tmp/cache"));
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
assert!(CacheKey::new(&config).is_err());
}
#[test]
fn test_token_cache_options_defaults() {
let options = TokenCacheOptions::new();
assert!(options.cache_dir.is_none());
assert!(options.lock_timeout_secs.is_none());
assert_eq!(options.lock_timeout(), Duration::from_secs(30));
let options = options.cache_dir("/tmp/x").lock_timeout_secs(5);
assert_eq!(options.cache_dir.as_deref(), Some(Path::new("/tmp/x")));
assert_eq!(options.lock_timeout(), Duration::from_secs(5));
}
#[test]
fn test_token_cache_options_rejects_empty_dir() {
let options = TokenCacheOptions::new().cache_dir("");
assert!(matches!(
options.resolved_dir(),
Err(Error::InvalidInput { message }) if message.contains("must not be empty")
));
}
#[tokio::test]
async fn test_session_lifecycle_without_cache_entry() {
let dir = cache_tempdir();
let session = OAuthSession::new(device_config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!status.refreshable);
assert_eq!(status.issuer_url, "https://issuer.example.com");
assert_eq!(status.client_id, "client-id");
assert_eq!(status.scopes, vec!["openid".to_string()]);
assert_eq!(status.flow, "device_code");
assert_eq!(status.obtained_at, None);
let logout = session.logout().await.unwrap();
assert!(!logout.removed);
}
#[tokio::test]
async fn test_record_round_trip_and_redaction() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let response = TokenResponse {
access_token: AccessToken::new("access-token".to_string()),
refresh_token: Some(RefreshToken::new("refresh-token".to_string())),
expires_in: Some(3600),
token_type: Some(BasicTokenType::Bearer),
};
let record = cache.record_from_response(&response).unwrap();
let debug = format!("{record:?}");
assert!(!debug.contains("refresh-token"));
assert!(debug.contains("<redacted>"));
// No access-token material is persisted.
let json = serde_json::to_string(&record).unwrap();
assert!(!json.contains("access-token"));
cache.store(&record).await.unwrap();
let loaded = cache.load().await.unwrap().unwrap();
assert_eq!(loaded.refresh_token, "refresh-token");
assert_eq!(loaded.version, CACHE_RECORD_VERSION);
// The on-disk file must not be group/other readable and must not be a symlink target.
let path = cache.record_path();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o077, 0);
}
assert!(std::fs::symlink_metadata(&path).unwrap().is_file());
assert!(cache.delete().await.unwrap());
assert!(cache.load().await.unwrap().is_none());
assert!(!cache.delete().await.unwrap());
}
#[cfg(unix)]
#[tokio::test]
async fn test_record_rejects_symlink() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
.unwrap();
cache.store(&record).await.unwrap();
let path = cache.record_path();
let target = dir.path().join("evil.json");
std::fs::write(&target, "{}").unwrap();
std::fs::remove_file(&path).unwrap();
std::os::unix::fs::symlink(&target, &path).unwrap();
let err = cache.load().await.unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("not a regular file"))
);
}
#[cfg(unix)]
#[tokio::test]
async fn test_record_rejects_world_readable_file() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let path = cache.record_path();
std::fs::write(&path, "{}").unwrap();
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
let err = cache.load().await.unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("group or other"))
);
}
/// Write a record file that passes the permission hardening checks.
fn write_record_file(path: &Path, contents: &str) {
std::fs::write(path, contents).unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).unwrap();
}
}
#[tokio::test]
async fn test_record_rejects_unknown_version_and_corruption() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let path = cache.record_path();
// A complete, well-formed record with an unknown schema version.
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
.unwrap();
let mut json = serde_json::to_value(&record).unwrap();
json["version"] = serde_json::json!(99);
write_record_file(&path, &json.to_string());
let err = cache.load().await.unwrap_err();
assert!(
matches!(err, Error::Runtime { message } if message.contains("unsupported version"))
);
write_record_file(&path, r#"{"version":1,"issuer_url":"x""#);
let err = cache.load().await.unwrap_err();
assert!(matches!(err, Error::Runtime { message } if message.contains("corrupt")));
write_record_file(&path, "");
assert!(
cache
.load()
.await
.unwrap_err()
.to_string()
.contains("corrupt")
);
}
#[cfg(unix)]
#[tokio::test]
async fn test_cache_dir_rejects_open_permissions() {
let dir = cache_tempdir();
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap();
let err = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("chmod 700")));
}
#[tokio::test]
async fn test_lock_serializes_and_releases() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let guard = cache.acquire_lock().await.unwrap();
let contender = {
let dir = dir.path().to_path_buf();
let cache2 = TokenCache::new(
&device_config(&dir),
&TokenCacheOptions::new()
.cache_dir(&dir)
.lock_timeout_secs(1),
)
.unwrap();
tokio::time::timeout(Duration::from_millis(300), cache2.acquire_lock()).await
};
assert!(contender.is_err(), "second acquire must block while held");
drop(guard);
let cache3 = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
tokio::time::timeout(Duration::from_secs(5), cache3.acquire_lock())
.await
.expect("lock re-acquirable after release")
.unwrap();
}
#[test]
fn test_token_cache_for_config_gating() {
let dir = cache_tempdir();
assert!(
token_cache_for_config(&device_config(dir.path()))
.unwrap()
.is_some()
);
let mut config = device_config(dir.path());
config.token_cache = None;
assert!(token_cache_for_config(&config).unwrap().is_none());
let mut config = device_config(dir.path());
config.flow = OAuthFlow::ClientCredentials;
assert!(token_cache_for_config(&config).unwrap().is_none());
let mut config = device_config(dir.path());
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
let err = token_cache_for_config(&config).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity"))
);
}
#[test]
fn test_oauth_session_requires_cache_options() {
let mut config = device_config(Path::new("/tmp/cache"));
config.token_cache = None;
let err = OAuthSession::new(config).unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("token_cache")));
}
#[tokio::test]
async fn test_store_skips_responses_without_refresh_token() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let response = TokenResponse {
access_token: AccessToken::new("access-token".to_string()),
refresh_token: None,
expires_in: Some(3600),
token_type: None,
};
assert!(cache.record_from_response(&response).is_none());
cache.store_if_refreshable(&response).await.unwrap();
assert!(cache.load().await.unwrap().is_none());
}
#[test]
fn test_session_status_debug_has_no_secrets() {
let status = SessionStatus {
refreshable: true,
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["openid".to_string()],
resource: None,
audience: None,
flow: "device_code".to_string(),
obtained_at: Some(100),
};
let debug = format!("{status:?}");
assert!(!debug.contains("refresh-token"));
}
}
mod tests;
rust/lancedb/src/remote/token_cache/tests.rs
+910
View File
@@ -0,0 +1,910 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::remote::HeaderProvider;
use crate::remote::oauth::OAuthHeaderProvider;
use oauth2::basic::BasicTokenType;
use oauth2::{AccessToken, RefreshToken};
use serial_test::serial;
/// Temp directory that satisfies the cache hardening checks. CI runners
/// can create temp directories with group/other bits set, which the
/// private-directory validation correctly rejects.
fn cache_tempdir() -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
}
dir
}
fn device_config(cache_dir: &Path) -> OAuthConfig {
OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: None,
client_auth_method: None,
scopes: vec!["openid".to_string()],
flow: OAuthFlow::DeviceCode,
refresh_buffer_secs: None,
resource: None,
audience: None,
token_cache: Some(TokenCacheOptions::new().cache_dir(cache_dir)),
}
}
/// Stateful mock IdP covering discovery, device authorization, device
/// polling, client credentials, and refresh with strict rotation: a
/// refresh token that is not the currently issued one is rejected with
/// `invalid_grant`, which is exactly what real providers do on rotation.
struct MockIdp {
issuer_url: String,
requests: Arc<std::sync::Mutex<Vec<String>>>,
device_authorizations: Arc<AtomicUsize>,
refresh_attempts: Arc<AtomicUsize>,
invalid_grant_rejections: Arc<AtomicUsize>,
access_tokens_issued: Arc<AtomicUsize>,
current_refresh: Arc<std::sync::Mutex<Option<String>>>,
fail_refreshes: Arc<AtomicBool>,
issue_refresh_tokens: Arc<AtomicBool>,
}
impl MockIdp {
async fn start() -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = Self {
issuer_url: issuer_url.clone(),
requests: Arc::new(std::sync::Mutex::new(Vec::new())),
device_authorizations: Arc::new(AtomicUsize::new(0)),
refresh_attempts: Arc::new(AtomicUsize::new(0)),
invalid_grant_rejections: Arc::new(AtomicUsize::new(0)),
access_tokens_issued: Arc::new(AtomicUsize::new(0)),
current_refresh: Arc::new(std::sync::Mutex::new(None)),
fail_refreshes: Arc::new(AtomicBool::new(false)),
issue_refresh_tokens: Arc::new(AtomicBool::new(true)),
};
let requests = Arc::clone(&server.requests);
let device_authorizations = Arc::clone(&server.device_authorizations);
let refresh_attempts = Arc::clone(&server.refresh_attempts);
let invalid_grant_rejections = Arc::clone(&server.invalid_grant_rejections);
let access_tokens_issued = Arc::clone(&server.access_tokens_issued);
let current_refresh = Arc::clone(&server.current_refresh);
let fail_refreshes = Arc::clone(&server.fail_refreshes);
let issue_refresh_tokens = Arc::clone(&server.issue_refresh_tokens);
tokio::spawn(async move {
loop {
let Ok((mut stream, _)) = listener.accept().await else {
return;
};
let requests = Arc::clone(&requests);
let device_authorizations = Arc::clone(&device_authorizations);
let refresh_attempts = Arc::clone(&refresh_attempts);
let invalid_grant_rejections = Arc::clone(&invalid_grant_rejections);
let access_tokens_issued = Arc::clone(&access_tokens_issued);
let current_refresh = Arc::clone(&current_refresh);
let fail_refreshes = Arc::clone(&fail_refreshes);
let issue_refresh_tokens = Arc::clone(&issue_refresh_tokens);
tokio::spawn(async move {
let (request_line, body) = read_http_request(&mut stream).await;
if request_line.starts_with("POST ") {
requests.lock().unwrap().push(body.clone());
}
if request_line.starts_with("GET /.well-known/openid-configuration ") {
let discovery = format!(
r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"#
);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request_line.starts_with("POST /device ") {
device_authorizations.fetch_add(1, Ordering::SeqCst);
let device = format!(
r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","expires_in":60,"interval":1}}"#
);
write_json_response(&mut stream, "200 OK", &device).await;
} else if request_line.starts_with("POST /token ") {
if body.contains("grant_type=refresh_token") {
refresh_attempts.fetch_add(1, Ordering::SeqCst);
if fail_refreshes.load(Ordering::SeqCst) {
write_json_response(
&mut stream,
"503 Service Unavailable",
r#"{"error":"temporarily_unavailable"}"#,
)
.await;
return;
}
let expected = current_refresh.lock().unwrap().clone();
let matched = body
.split('&')
.find_map(|pair| pair.strip_prefix("refresh_token="))
.map(|token| token.to_string())
.zip(expected)
.is_some_and(|(offered, expected)| offered == expected);
if !matched {
invalid_grant_rejections.fetch_add(1, Ordering::SeqCst);
write_json_response(
&mut stream,
"400 Bad Request",
r#"{"error":"invalid_grant"}"#,
)
.await;
return;
}
issue_access_token(
&mut stream,
&access_tokens_issued,
&current_refresh,
issue_refresh_tokens.load(Ordering::SeqCst),
)
.await;
} else {
// Device polling or client credentials: issue
// a token and, for interactive grants, a fresh
// refresh token with strict rotation.
let grant_device = body.contains(
"grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code",
);
if grant_device {
issue_access_token(
&mut stream,
&access_tokens_issued,
&current_refresh,
issue_refresh_tokens.load(Ordering::SeqCst),
)
.await;
} else {
let token = format!(
r#"{{"access_token":"access-{}","expires_in":3600}}"#,
access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1
);
write_json_response(&mut stream, "200 OK", &token).await;
}
}
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
});
}
});
server
}
fn config(&self, cache_dir: &Path) -> OAuthConfig {
let mut config = device_config(cache_dir);
config.issuer_url = self.issuer_url.clone();
config
}
}
async fn issue_access_token(
stream: &mut TcpStream,
access_tokens_issued: &AtomicUsize,
current_refresh: &std::sync::Mutex<Option<String>>,
with_refresh: bool,
) {
let number = access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1;
let token = if with_refresh {
*current_refresh.lock().unwrap() = Some(format!("refresh-{number}"));
format!(
r#"{{"access_token":"access-{number}","refresh_token":"refresh-{number}","expires_in":3600}}"#
)
} else {
format!(r#"{{"access_token":"access-{number}","expires_in":3600}}"#)
};
write_json_response(stream, "200 OK", &token).await;
}
async fn read_http_request(stream: &mut TcpStream) -> (String, String) {
let mut buffer = Vec::new();
let mut header_end = None;
while header_end.is_none() {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request headers");
buffer.extend_from_slice(&chunk[..read]);
header_end = find_subsequence(&buffer, b"\r\n\r\n").map(|pos| pos + 4);
}
let header_end = header_end.unwrap();
let headers = String::from_utf8_lossy(&buffer[..header_end]).to_string();
let request_line = headers.lines().next().unwrap_or_default().to_string();
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while buffer.len() < header_end + content_length {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request body");
buffer.extend_from_slice(&chunk[..read]);
}
let body =
String::from_utf8_lossy(&buffer[header_end..header_end + content_length]).to_string();
(request_line, body)
}
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack
.windows(needle.len())
.position(|window| window == needle)
}
async fn write_json_response(stream: &mut TcpStream, status: &str, body: &str) {
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).await.unwrap();
}
fn suppress_browser() {
// Point the OAuth browser helper at a no-op so device-flow tests never
// open a real browser window.
unsafe { std::env::set_var("LANCEDB_OAUTH_BROWSER", "/usr/bin/true") };
}
#[tokio::test]
#[serial]
async fn test_provider_reuses_cached_session_across_instances() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let first = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = first.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
// A second, independent provider (simulating a second process) must
// refresh silently instead of starting another device flow.
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = second.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2");
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 1);
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-2");
}
#[tokio::test]
#[serial]
async fn test_concurrent_providers_serialize_rotation() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
let provider_a = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let provider_b = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let (headers_a, headers_b) = tokio::join!(provider_a.get_headers(), provider_b.get_headers());
let token_a = headers_a.unwrap().remove("authorization").unwrap();
let token_b = headers_b.unwrap().remove("authorization").unwrap();
assert!(
{
let mut tokens = [token_a, token_b];
tokens.sort();
tokens
} == ["Bearer access-2".to_string(), "Bearer access-3".to_string()],
"each provider must observe its own refreshed token"
);
// Rotation raced would produce an invalid_grant and a second device
// flow; the lock prevents both.
assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 0);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 2);
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-3");
}
#[tokio::test]
#[serial]
async fn test_transient_refresh_failure_retains_record() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
idp.fail_refreshes.store(true, Ordering::SeqCst);
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let error = second.get_headers().await.unwrap_err();
assert!(error.to_string().contains("503"));
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(
status.refreshable,
"transient failures must keep the record"
);
}
#[tokio::test]
#[serial]
async fn test_invalid_grant_deletes_record_and_reauthenticates() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
priming.get_headers().await.unwrap();
// Simulate a revoked refresh token by replacing the record with one
// the provider never issued.
let cache = crate::remote::token_cache::TokenCache::new(
&idp.config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let mut record = cache.load().await.unwrap().unwrap();
record.refresh_token = "revoked-refresh".to_string();
cache.store(&record).await.unwrap();
let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = second.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2");
assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 1);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 2);
let record = cache.load().await.unwrap().unwrap();
assert_eq!(record.refresh_token, "refresh-2");
}
#[tokio::test]
#[serial]
async fn test_session_login_status_logout_lifecycle() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!status.refreshable);
let status = session.login().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.issuer_url, idp.issuer_url);
assert_eq!(status.client_id, "client-id");
assert_eq!(status.scopes, vec!["openid".to_string()]);
assert_eq!(status.flow, "device_code");
assert!(status.obtained_at.is_some());
// An independent session manager sees the same cached login.
let other = OAuthSession::new(idp.config(dir.path())).unwrap();
assert!(other.status().await.unwrap().refreshable);
assert!(session.logout().await.unwrap().removed);
assert!(!session.logout().await.unwrap().removed);
assert!(!other.status().await.unwrap().refreshable);
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1);
}
#[tokio::test]
#[serial]
async fn test_login_without_refresh_token_clears_prior_record() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
session.login().await.unwrap();
assert!(session.status().await.unwrap().refreshable);
// A provider that stops issuing refresh tokens (for example a login
// without offline_access) must not leave the earlier account behind.
idp.issue_refresh_tokens.store(false, Ordering::SeqCst);
let status = session.login().await.unwrap();
assert!(!status.refreshable);
assert!(!session.status().await.unwrap().refreshable);
}
#[tokio::test]
async fn test_client_credentials_with_cache_stays_memory_only() {
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let mut config = idp.config(dir.path());
config.flow = OAuthFlow::ClientCredentials;
config.client_secret = Some("secret".to_string());
let provider = OAuthHeaderProvider::new(config).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
assert_eq!(dir.path().read_dir().unwrap().count(), 0);
}
#[test]
fn test_managed_identity_with_cache_is_rejected() {
let dir = cache_tempdir();
let mut config = device_config(dir.path());
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
let err = OAuthHeaderProvider::new(config).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity"))
);
}
#[tokio::test]
#[serial]
async fn test_provider_debug_and_status_reveal_no_secrets() {
suppress_browser();
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let provider = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1");
let debug = format!("{provider:?}");
assert!(!debug.contains("access-1"));
assert!(!debug.contains("refresh-1"));
let session = OAuthSession::new(idp.config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!format!("{status:?}").contains("refresh-"));
}
#[tokio::test]
async fn test_targeted_cache_refresh_and_logout_isolation() {
let dir = cache_tempdir();
let idp = MockIdp::start().await;
let mut config = idp.config(dir.path());
let options = config.token_cache.clone().unwrap();
let untargeted_key = CacheKey::new(&config).unwrap().file_stem;
assert_eq!(
untargeted_key,
hex_sha256(
format!(
"v1\n{}\nclient-id\nopenid\ndevice_code\npublic",
idp.issuer_url
)
.as_bytes()
)
);
let mut sessions = Vec::new();
let mut keys = std::collections::HashSet::new();
for (resource, audience) in [
(None, None),
(Some("urn:one"), None),
(None, Some("audience & ü")),
(Some("urn:one"), Some("audience & ü")),
(Some("urn:two"), Some("audience & ü")),
(Some("urn:one"), Some("other")),
] {
config.resource = resource.map(str::to_owned);
config.audience = audience.map(str::to_owned);
let cache = TokenCache::new(&config, &options).unwrap();
assert!(keys.insert(cache.key.file_stem.clone()));
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("unused".into()),
refresh_token: Some(RefreshToken::new("seed-refresh".into())),
expires_in: Some(3600),
token_type: None,
})
.unwrap();
cache.store(&record).await.unwrap();
*idp.current_refresh.lock().unwrap() = Some("seed-refresh".into());
let provider = OAuthHeaderProvider::new(config.clone()).unwrap();
provider.get_headers().await.unwrap();
let request = idp.requests.lock().unwrap().last().unwrap().clone();
let params: std::collections::HashMap<_, _> =
url::form_urlencoded::parse(request.as_bytes()).collect();
assert_eq!(params.get("resource").map(|s| s.as_ref()), resource);
assert_eq!(params.get("audience").map(|s| s.as_ref()), audience);
assert_eq!(params.get("grant_type").unwrap(), "refresh_token");
let session = OAuthSession::new(config.clone()).unwrap();
let status = session.status().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.resource, config.resource);
assert_eq!(status.audience, config.audience);
sessions.push(session);
}
assert!(sessions.pop().unwrap().logout().await.unwrap().removed);
for session in sessions {
assert!(session.status().await.unwrap().refreshable);
}
assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_legacy_cache_record_without_target_fields() {
let dir = cache_tempdir();
let config = device_config(dir.path());
let cache = TokenCache::new(&config, config.token_cache.as_ref().unwrap()).unwrap();
let legacy = br#"{"version":1,"issuer_url":"https://issuer.example.com","client_id":"client-id","scopes":["openid"],"flow":"device_code","client_auth":"public","refresh_token":"legacy","obtained_at":1}"#;
write_record(dir.path(), &cache.record_path(), legacy).unwrap();
let session = OAuthSession::new(config).unwrap();
let status = session.status().await.unwrap();
assert!(status.refreshable);
assert_eq!(status.resource, None);
assert_eq!(status.audience, None);
assert!(session.logout().await.unwrap().removed);
}
#[test]
fn test_cache_key_canonicalizes_scopes_and_issuer() {
let mut config = device_config(Path::new("/tmp/cache"));
config.issuer_url = "https://issuer.example.com/".to_string();
config.scopes = vec!["b".to_string(), " a ".to_string(), "a".to_string()];
let key = CacheKey::new(&config).unwrap();
assert_eq!(key.issuer_url, "https://issuer.example.com");
assert_eq!(key.scopes, vec!["a".to_string(), "b".to_string()]);
config.issuer_url = "https://issuer.example.com".to_string();
let canonical = CacheKey::new(&config).unwrap();
assert_eq!(canonical.file_stem, key.file_stem);
}
#[test]
fn test_cache_key_separates_identity_dimensions() {
let base = device_config(Path::new("/tmp/cache"));
let base_key = CacheKey::new(&base).unwrap();
let mut other = base.clone();
other.client_id = "other-client".to_string();
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.issuer_url = "https://other.example.com".to_string();
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.scopes = vec!["profile".to_string()];
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.flow = OAuthFlow::AuthorizationCode(Default::default());
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
let mut other = base.clone();
other.client_secret = Some("secret".to_string());
assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem);
}
#[test]
fn test_cache_key_contains_no_secret_material() {
let mut config = device_config(Path::new("/tmp/cache"));
config.client_secret = Some("super-secret-value".to_string());
let key = CacheKey::new(&config).unwrap();
assert!(!key.file_stem.contains("super-secret-value"));
assert_eq!(key.file_stem.len(), 64);
}
#[test]
fn test_flow_key_rejects_non_persistent_flows() {
let mut config = device_config(Path::new("/tmp/cache"));
config.flow = OAuthFlow::ClientCredentials;
let err = CacheKey::new(&config).unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("not supported")));
let mut config = device_config(Path::new("/tmp/cache"));
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
assert!(CacheKey::new(&config).is_err());
}
#[test]
fn test_token_cache_options_defaults() {
let options = TokenCacheOptions::new();
assert!(options.cache_dir.is_none());
assert!(options.lock_timeout_secs.is_none());
assert_eq!(options.lock_timeout(), Duration::from_secs(30));
let options = options.cache_dir("/tmp/x").lock_timeout_secs(5);
assert_eq!(options.cache_dir.as_deref(), Some(Path::new("/tmp/x")));
assert_eq!(options.lock_timeout(), Duration::from_secs(5));
}
#[test]
fn test_token_cache_options_rejects_empty_dir() {
let options = TokenCacheOptions::new().cache_dir("");
assert!(matches!(
options.resolved_dir(),
Err(Error::InvalidInput { message }) if message.contains("must not be empty")
));
}
#[tokio::test]
async fn test_session_lifecycle_without_cache_entry() {
let dir = cache_tempdir();
let session = OAuthSession::new(device_config(dir.path())).unwrap();
let status = session.status().await.unwrap();
assert!(!status.refreshable);
assert_eq!(status.issuer_url, "https://issuer.example.com");
assert_eq!(status.client_id, "client-id");
assert_eq!(status.scopes, vec!["openid".to_string()]);
assert_eq!(status.flow, "device_code");
assert_eq!(status.obtained_at, None);
let logout = session.logout().await.unwrap();
assert!(!logout.removed);
}
#[tokio::test]
async fn test_record_round_trip_and_redaction() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let response = TokenResponse {
access_token: AccessToken::new("access-token".to_string()),
refresh_token: Some(RefreshToken::new("refresh-token".to_string())),
expires_in: Some(3600),
token_type: Some(BasicTokenType::Bearer),
};
let record = cache.record_from_response(&response).unwrap();
let debug = format!("{record:?}");
assert!(!debug.contains("refresh-token"));
assert!(debug.contains("<redacted>"));
// No access-token material is persisted.
let json = serde_json::to_string(&record).unwrap();
assert!(!json.contains("access-token"));
cache.store(&record).await.unwrap();
let loaded = cache.load().await.unwrap().unwrap();
assert_eq!(loaded.refresh_token, "refresh-token");
assert_eq!(loaded.version, CACHE_RECORD_VERSION);
// The on-disk file must not be group/other readable and must not be a symlink target.
let path = cache.record_path();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o077, 0);
}
assert!(std::fs::symlink_metadata(&path).unwrap().is_file());
assert!(cache.delete().await.unwrap());
assert!(cache.load().await.unwrap().is_none());
assert!(!cache.delete().await.unwrap());
}
#[cfg(unix)]
#[tokio::test]
async fn test_record_rejects_symlink() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
.unwrap();
cache.store(&record).await.unwrap();
let path = cache.record_path();
let target = dir.path().join("evil.json");
std::fs::write(&target, "{}").unwrap();
std::fs::remove_file(&path).unwrap();
std::os::unix::fs::symlink(&target, &path).unwrap();
let err = cache.load().await.unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("not a regular file"))
);
}
#[cfg(unix)]
#[tokio::test]
async fn test_record_rejects_world_readable_file() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let path = cache.record_path();
std::fs::write(&path, "{}").unwrap();
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
let err = cache.load().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("group or other")));
}
/// Write a record file that passes the permission hardening checks.
fn write_record_file(path: &Path, contents: &str) {
std::fs::write(path, contents).unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).unwrap();
}
}
#[tokio::test]
async fn test_record_rejects_unknown_version_and_corruption() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let path = cache.record_path();
// A complete, well-formed record with an unknown schema version.
let record = cache
.record_from_response(&TokenResponse {
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
.unwrap();
let mut json = serde_json::to_value(&record).unwrap();
json["version"] = serde_json::json!(99);
write_record_file(&path, &json.to_string());
let err = cache.load().await.unwrap_err();
assert!(matches!(err, Error::Runtime { message } if message.contains("unsupported version")));
write_record_file(&path, r#"{"version":1,"issuer_url":"x""#);
let err = cache.load().await.unwrap_err();
assert!(matches!(err, Error::Runtime { message } if message.contains("corrupt")));
write_record_file(&path, "");
assert!(
cache
.load()
.await
.unwrap_err()
.to_string()
.contains("corrupt")
);
}
#[cfg(unix)]
#[tokio::test]
async fn test_cache_dir_rejects_open_permissions() {
let dir = cache_tempdir();
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap();
let err = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("chmod 700")));
}
#[tokio::test]
async fn test_lock_serializes_and_releases() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let guard = cache.acquire_lock().await.unwrap();
let contender = {
let dir = dir.path().to_path_buf();
let cache2 = TokenCache::new(
&device_config(&dir),
&TokenCacheOptions::new()
.cache_dir(&dir)
.lock_timeout_secs(1),
)
.unwrap();
tokio::time::timeout(Duration::from_millis(300), cache2.acquire_lock()).await
};
assert!(contender.is_err(), "second acquire must block while held");
drop(guard);
let cache3 = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
tokio::time::timeout(Duration::from_secs(5), cache3.acquire_lock())
.await
.expect("lock re-acquirable after release")
.unwrap();
}
#[test]
fn test_token_cache_for_config_gating() {
let dir = cache_tempdir();
assert!(
token_cache_for_config(&device_config(dir.path()))
.unwrap()
.is_some()
);
let mut config = device_config(dir.path());
config.token_cache = None;
assert!(token_cache_for_config(&config).unwrap().is_none());
let mut config = device_config(dir.path());
config.flow = OAuthFlow::ClientCredentials;
assert!(token_cache_for_config(&config).unwrap().is_none());
let mut config = device_config(dir.path());
config.flow = OAuthFlow::AzureManagedIdentity { client_id: None };
let err = token_cache_for_config(&config).unwrap_err();
assert!(
matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity"))
);
}
#[test]
fn test_oauth_session_requires_cache_options() {
let mut config = device_config(Path::new("/tmp/cache"));
config.token_cache = None;
let err = OAuthSession::new(config).unwrap_err();
assert!(matches!(err, Error::InvalidInput { message } if message.contains("token_cache")));
}
#[tokio::test]
async fn test_store_skips_responses_without_refresh_token() {
let dir = cache_tempdir();
let cache = TokenCache::new(
&device_config(dir.path()),
&TokenCacheOptions::new().cache_dir(dir.path()),
)
.unwrap();
let response = TokenResponse {
access_token: AccessToken::new("access-token".to_string()),
refresh_token: None,
expires_in: Some(3600),
token_type: None,
};
assert!(cache.record_from_response(&response).is_none());
cache.store_if_refreshable(&response).await.unwrap();
assert!(cache.load().await.unwrap().is_none());
}
#[test]
fn test_session_status_debug_has_no_secrets() {
let status = SessionStatus {
refreshable: true,
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["openid".to_string()],
resource: None,
audience: None,
flow: "device_code".to_string(),
obtained_at: Some(100),
};
let debug = format!("{status:?}");
assert!(!debug.contains("refresh-token"));
}