From 6dde01de337d0356c1b67f229970830f3155aad0 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 06:23:00 +0000 Subject: [PATCH] refactor(rust): move remote module unit tests into separate files Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BgkiPYS2eTmAWVyqn6KBnN --- rust/lancedb/src/remote/client.rs | 753 +- rust/lancedb/src/remote/client/tests.rs | 752 ++ rust/lancedb/src/remote/db.rs | 2749 +--- rust/lancedb/src/remote/db/tests/catalog.rs | 634 + rust/lancedb/src/remote/db/tests/jobs.rs | 489 + rust/lancedb/src/remote/db/tests/mod.rs | 25 + rust/lancedb/src/remote/db/tests/namespace.rs | 561 + rust/lancedb/src/remote/db/tests/tables.rs | 1049 ++ rust/lancedb/src/remote/oauth.rs | 1865 +-- .../src/remote/oauth/tests/authorization.rs | 478 + rust/lancedb/src/remote/oauth/tests/config.rs | 255 + rust/lancedb/src/remote/oauth/tests/device.rs | 300 + rust/lancedb/src/remote/oauth/tests/flows.rs | 361 + rust/lancedb/src/remote/oauth/tests/mod.rs | 485 + rust/lancedb/src/remote/table.rs | 10429 +--------------- rust/lancedb/src/remote/table/blobs.rs | 1239 +- rust/lancedb/src/remote/table/blobs/tests.rs | 1234 ++ rust/lancedb/src/remote/table/insert.rs | 728 +- rust/lancedb/src/remote/table/insert/tests.rs | 722 ++ rust/lancedb/src/remote/table/tests/blobs.rs | 753 ++ .../src/remote/table/tests/branches.rs | 1223 ++ .../lancedb/src/remote/table/tests/columns.rs | 1041 ++ .../src/remote/table/tests/freshness.rs | 717 ++ rust/lancedb/src/remote/table/tests/hybrid.rs | 250 + rust/lancedb/src/remote/table/tests/index.rs | 1280 ++ rust/lancedb/src/remote/table/tests/lsm.rs | 901 ++ rust/lancedb/src/remote/table/tests/mod.rs | 254 + .../src/remote/table/tests/multipart.rs | 767 ++ rust/lancedb/src/remote/table/tests/query.rs | 806 ++ .../src/remote/table/tests/schema_cache.rs | 1056 ++ .../src/remote/table/tests/versions.rs | 369 + rust/lancedb/src/remote/table/tests/write.rs | 994 ++ rust/lancedb/src/remote/token_cache.rs | 916 +- rust/lancedb/src/remote/token_cache/tests.rs | 910 ++ 34 files changed, 18673 insertions(+), 18672 deletions(-) create mode 100644 rust/lancedb/src/remote/client/tests.rs create mode 100644 rust/lancedb/src/remote/db/tests/catalog.rs create mode 100644 rust/lancedb/src/remote/db/tests/jobs.rs create mode 100644 rust/lancedb/src/remote/db/tests/mod.rs create mode 100644 rust/lancedb/src/remote/db/tests/namespace.rs create mode 100644 rust/lancedb/src/remote/db/tests/tables.rs create mode 100644 rust/lancedb/src/remote/oauth/tests/authorization.rs create mode 100644 rust/lancedb/src/remote/oauth/tests/config.rs create mode 100644 rust/lancedb/src/remote/oauth/tests/device.rs create mode 100644 rust/lancedb/src/remote/oauth/tests/flows.rs create mode 100644 rust/lancedb/src/remote/oauth/tests/mod.rs create mode 100644 rust/lancedb/src/remote/table/blobs/tests.rs create mode 100644 rust/lancedb/src/remote/table/insert/tests.rs create mode 100644 rust/lancedb/src/remote/table/tests/blobs.rs create mode 100644 rust/lancedb/src/remote/table/tests/branches.rs create mode 100644 rust/lancedb/src/remote/table/tests/columns.rs create mode 100644 rust/lancedb/src/remote/table/tests/freshness.rs create mode 100644 rust/lancedb/src/remote/table/tests/hybrid.rs create mode 100644 rust/lancedb/src/remote/table/tests/index.rs create mode 100644 rust/lancedb/src/remote/table/tests/lsm.rs create mode 100644 rust/lancedb/src/remote/table/tests/mod.rs create mode 100644 rust/lancedb/src/remote/table/tests/multipart.rs create mode 100644 rust/lancedb/src/remote/table/tests/query.rs create mode 100644 rust/lancedb/src/remote/table/tests/schema_cache.rs create mode 100644 rust/lancedb/src/remote/table/tests/versions.rs create mode 100644 rust/lancedb/src/remote/table/tests/write.rs create mode 100644 rust/lancedb/src/remote/token_cache/tests.rs diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 705506705..358611dca 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -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::::default_headers( - "", - "us-east-1", - "db-name", - false, - &RemoteOptions::default(), - None, - &ClientConfig::default(), - ) - .unwrap(); - assert!(!headers.contains_key("x-api-key")); - - let headers = RestfulLanceDbClient::::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::::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::::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::::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, - } - - impl TestHeaderProvider { - fn new(headers: HashMap) -> Self { - Self { headers } - } - } - - #[async_trait::async_trait] - impl HeaderProvider for TestHeaderProvider { - async fn get_headers(&self) -> Result> { - 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> { - 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), - ..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), - 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), - 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), - 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::::resolve_max_bytes_per_request(Some(1234)).unwrap(); - assert_eq!(resolved, Some(1234)); - } - - #[test] - fn test_resolve_max_bytes_zero_disables() { - let resolved = - RestfulLanceDbClient::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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; diff --git a/rust/lancedb/src/remote/client/tests.rs b/rust/lancedb/src/remote/client/tests.rs new file mode 100644 index 000000000..a9c8f9c13 --- /dev/null +++ b/rust/lancedb/src/remote/client/tests.rs @@ -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::::default_headers( + "", + "us-east-1", + "db-name", + false, + &RemoteOptions::default(), + None, + &ClientConfig::default(), + ) + .unwrap(); + assert!(!headers.contains_key("x-api-key")); + + let headers = RestfulLanceDbClient::::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::::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::::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::::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, +} + +impl TestHeaderProvider { + fn new(headers: HashMap) -> Self { + Self { headers } + } +} + +#[async_trait::async_trait] +impl HeaderProvider for TestHeaderProvider { + async fn get_headers(&self) -> Result> { + 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> { + 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), + ..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), + 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), + 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), + 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::::resolve_max_bytes_per_request(Some(1234)).unwrap(); + assert_eq!(resolved, Some(1234)); +} + +#[test] +fn test_resolve_max_bytes_zero_disables() { + let resolved = RestfulLanceDbClient::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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::::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:?}"); +} diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 1453d8d73..e0e4c27b3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1903,2751 +1903,4 @@ impl From for RemoteOptions { } #[cfg(test)] -mod tests { - use super::{NamespaceHeaderProviderContext, build_cache_key}; - use std::collections::HashMap; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, OnceLock}; - - use arrow_array::{Int32Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use lance_namespace_impls::{DynamicContextProvider, OperationInfo}; - use rstest::rstest; - - 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}, - }; - - #[test] - fn test_cache_key_security() { - // Test that cache keys are unique regardless of delimiter manipulation - - // Case 1: Different delimiters should not affect cache key - let key1 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]); - let key2 = build_cache_key("table1", &["ns1$ns2".to_string()]); - assert_ne!( - key1, key2, - "Cache keys should differ for different namespace structures" - ); - - // Case 2: Table name containing delimiter should not cause collision - let key3 = build_cache_key("ns2$table1", &["ns1".to_string()]); - assert_ne!( - key1, key3, - "Cache key should be different when table name contains delimiter" - ); - - // Case 3: Empty namespace vs namespace with empty string - let key4 = build_cache_key("table1", &[]); - let key5 = build_cache_key("table1", &["".to_string()]); - assert_ne!( - key4, key5, - "Empty namespace should differ from namespace with empty string" - ); - - // Case 4: Verify same inputs produce same key (consistency) - let key6 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]); - assert_eq!(key1, key6, "Same inputs should produce same cache key"); - } - - #[tokio::test] - async fn test_create_materialized_view_requires_sql_client() { - let db = super::RemoteDatabase::new_mock(|request| -> http::Response { - panic!("unexpected REST request: {}", request.url().path()) - }); - let error = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10" - .into(), - with_no_data: false, - }) - .await - .unwrap_err(); - assert!(error.to_string().contains("SQL is unavailable")); - } - - #[tokio::test] - async fn test_drop_materialized_view_uses_item_route_and_job() { - let db = super::RemoteDatabase::new_mock(|request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/drop" - ); - assert!(request.body().is_none()); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-drop"}).to_string()) - .unwrap() - }); - let job = db - .drop_materialized_view_async("adults", &["analytics".into()]) - .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-drop")); - } - - #[tokio::test] - async fn test_drop_function_async_returns_job() { - let db = super::RemoteDatabase::new_mock(|request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/function/embed/drop"); - http::Response::builder() - .status(202) - .body(serde_json::json!({"dropped": true, "job_id": "j1-fn-drop"}).to_string()) - .unwrap() - }); - let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); - assert!(dropped); - assert_eq!(job.id(), Some("j1-fn-drop")); - } - - /// An unbound name and an inline deletion both answer `200`: the name is gone and - /// nothing is left to wait for, so the job is already finished. - #[tokio::test] - async fn test_drop_function_async_completed_inline() { - let db = super::RemoteDatabase::new_mock(|_| { - http::Response::builder() - .status(200) - .body(serde_json::json!({"dropped": false}).to_string()) - .unwrap() - }); - let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); - assert!(!dropped); - assert_eq!(job.id(), None); - assert_eq!(job.status().await.unwrap(), "finished"); - job.wait().await.unwrap(); - } - - #[tokio::test] - async fn test_drop_function_async_rejects_incomplete_acceptance() { - for body in [ - r#"{"dropped":true}"#, - r#"{"dropped":true,"job_id":""}"#, - r#"{"dropped":true,"job_id":null}"#, - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(202).body(body).unwrap() - }); - let error = db - .drop_function_async("embed", "1", &[]) - .await - .err() - .unwrap(); - assert!(error.to_string().contains("valid job_id")); - } - } - - #[tokio::test] - async fn test_drop_materialized_view_completed_inline() { - let db = super::RemoteDatabase::new_mock(|request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/materialized_view/adults/drop"); - http::Response::builder().status(200).body("{}").unwrap() - }); - let job = db - .drop_materialized_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_materialized_view_rejects_incomplete_acceptance() { - for body in ["{}", r#"{"job_id":""}"#, r#"{"job_id":null}"#] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(202).body(body).unwrap() - }); - let error = db - .drop_materialized_view_async("adults", &[]) - .await - .err() - .unwrap(); - assert!(error.to_string().contains("valid job_id")); - } - } - - #[tokio::test] - async fn test_drop_materialized_view_rejects_unexpected_success_status() { - let db = super::RemoteDatabase::new_mock(|_| { - http::Response::builder().status(204).body("").unwrap() - }); - let error = db - .drop_materialized_view_async("adults", &[]) - .await - .err() - .unwrap(); - assert!(error.to_string().contains("200 OK or 202 Accepted")); - } - - #[tokio::test] - async fn test_list_materialized_views_follows_empty_pages() { - let page = Arc::new(AtomicUsize::new(0)); - let db = super::RemoteDatabase::new_mock({ - let page = page.clone(); - move |request| { - assert_eq!(request.method(), "GET"); - assert_eq!( - request.url().path(), - "/v1/namespace/analytics/materialized_view/list" - ); - match page.fetch_add(1, Ordering::SeqCst) { - 0 => { - assert!(request.url().query().is_none()); - http::Response::builder() - .status(200) - .body( - serde_json::json!({"views": [], "page_token": "next"}).to_string(), - ) - .unwrap() - } - 1 => { - assert_eq!(request.url().query(), Some("page_token=next")); - http::Response::builder() - .status(200) - .body(serde_json::json!({"views": ["adults"]}).to_string()) - .unwrap() - } - _ => panic!("listing requested too many pages"), - } - } - }); - assert_eq!( - db.list_materialized_views(&["analytics".into()]) - .await - .unwrap(), - ["adults"] - ); - } - - #[tokio::test] - async fn test_retries() { - // We'll record the request_id here, to check it matches the one in the error. - let seen_request_id = Arc::new(OnceLock::new()); - let seen_request_id_ref = seen_request_id.clone(); - let conn = Connection::new_with_handler(move |request| { - // Request id should be the same on each retry. - let request_id = request.headers()["x-request-id"] - .to_str() - .unwrap() - .to_string(); - let seen_id = seen_request_id_ref.get_or_init(|| request_id.clone()); - assert_eq!(&request_id, seen_id); - - http::Response::builder() - .status(500) - .body("internal server error") - .unwrap() - }); - let result = conn.table_names().execute().await; - if let Err(Error::Retry { - request_id, - request_failures, - max_request_failures, - source, - .. - }) = result - { - let expected_id = seen_request_id.get().unwrap(); - assert_eq!(&request_id, expected_id); - assert_eq!(request_failures, max_request_failures); - assert!( - source.to_string().contains("internal server error"), - "source: {:?}", - source - ); - } else { - panic!("unexpected result: {:?}", result); - }; - } - - #[tokio::test] - async fn test_table_names() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/table/"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - .body(r#"{"tables": ["table1", "table2"]}"#) - .unwrap() - }); - let names = conn.table_names().execute().await.unwrap(); - assert_eq!(names, vec!["table1", "table2"]); - } - - #[tokio::test] - async fn test_table_names_in_a_namespace_never_invents_a_page_token() { - // The namespace route's token belongs to the store, so `table_names` cannot build one - // from `start_after`. It walks the namespace on the server's own tokens and applies the - // name semantics itself, which is what keeps it working either side of the change. - let page = Arc::new(AtomicUsize::new(0)); - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.url().path(), "/v1/namespace/ns/table/list"); - let query = request.url().query().unwrap_or(""); - assert!( - !query.contains("page_token=users"), - "a table name must never be sent as a page token: {query}" - ); - match page.fetch_add(1, Ordering::SeqCst) { - 0 => { - assert!( - !query.contains("page_token"), - "the walk starts with no token" - ); - http::Response::builder() - .status(200) - .body(r#"{"tables": ["users", "orders"], "page_token": "opaque-1"}"#) - .unwrap() - } - _ => { - assert!(query.contains("page_token=opaque-1")); - http::Response::builder() - .status(200) - .body(r#"{"tables": ["widgets"]}"#) - .unwrap() - } - } - }); - - let names = conn - .table_names() - .namespace(vec!["ns".to_string()]) - .start_after("users") - .execute() - .await - .unwrap(); - // Name order, resumed after "users": "orders" sorts before it and is dropped. - assert_eq!(names, vec!["widgets"]); - } - - #[tokio::test] - async fn test_table_names_in_a_namespace_stops_on_a_repeated_token() { - // A server that handed back the token it was given would never finish the walk. - let conn = Connection::new_with_handler(|_request| { - http::Response::builder() - .status(200) - .body(r#"{"tables": ["a"], "page_token": "same"}"#) - .unwrap() - }); - - let names = conn - .table_names() - .namespace(vec!["ns".to_string()]) - .execute() - .await - .unwrap(); - // The guard bounds the walk instead of letting it run forever. The repeat is the - // server breaking the token contract and is not papered over here. - assert_eq!(names, vec!["a", "a"]); - } - - #[tokio::test] - async fn test_table_names_in_a_namespace_stops_on_an_empty_token() { - // An empty token ends the listing. Sending it back would ask a server that reads it - // as "start from the beginning" for the first page a second time, and every name on - // that page would be collected twice. - 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!( - !request.url().query().unwrap_or("").contains("page_token"), - "an empty token must never be sent back" - ); - http::Response::builder() - .status(200) - .body(r#"{"tables": ["a"], "page_token": ""}"#) - .unwrap() - }); - - let names = conn - .table_names() - .namespace(vec!["ns".to_string()]) - .execute() - .await - .unwrap(); - assert_eq!(names, vec!["a"]); - assert_eq!(requests.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_table_names_pagination() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/table/"); - assert!(request.url().query().unwrap().contains("limit=2")); - assert!(request.url().query().unwrap().contains("page_token=table2")); - - http::Response::builder() - .status(200) - .body(r#"{"tables": ["table3", "table4"], "page_token": "token"}"#) - .unwrap() - }); - let names = conn - .table_names() - .start_after("table2") - .limit(2) - .execute() - .await - .unwrap(); - assert_eq!(names, vec!["table3", "table4"]); - } - - #[tokio::test] - async fn test_open_table() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/describe/"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - .body(r#"{"table": "table1"}"#) - .unwrap() - }); - let table = conn.open_table("table1").execute().await.unwrap(); - assert_eq!(table.name(), "table1"); - - // Storage options should be ignored. - let table = conn - .open_table("table1") - .storage_option("key", "value") - .execute() - .await - .unwrap(); - assert_eq!(table.name(), "table1"); - } - - #[tokio::test] - async fn test_open_table_checkout_is_independent_per_handle() { - let latest = Arc::new(AtomicUsize::new(2)); - let current = latest.clone(); - let mut db = super::RemoteDatabase::new_mock(move |request| { - let body = match request.url().path() { - "/v1/table/table1/describe/" => { - let requested_version = request - .body() - .and_then(|body| body.as_bytes()) - .and_then(|body| serde_json::from_slice::(body).ok()) - .and_then(|body| body["version"].as_u64()); - let version = - requested_version.unwrap_or_else(|| current.load(Ordering::SeqCst) as u64); - serde_json::json!({ - "version": version, - "schema": {"fields": [ - {"name": "id", "type": {"type": "int32"}, "nullable": false} - ]} - }) - .to_string() - } - "/v1/table/table1/insert/" => { - let version = current.fetch_add(1, Ordering::SeqCst) + 1; - serde_json::json!({"version": version}).to_string() - } - path => panic!("unexpected request path: {path}"), - }; - http::Response::builder().status(200).body(body).unwrap() - }); - db.table_cache = moka::future::Cache::new(10); - let conn = Connection::new( - Arc::new(db), - Arc::new(crate::embeddings::MemoryRegistry::new()), - ); - - let writer = conn.open_table("table1").execute().await.unwrap(); - let reader = conn.open_table("table1").execute().await.unwrap(); - reader.checkout(1).await.unwrap(); - - assert_eq!(reader.version().await.unwrap(), 1); - assert_eq!(writer.version().await.unwrap(), 2); - - let data = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), - vec![Arc::new(Int32Array::from(vec![3]))], - ) - .unwrap(); - assert_eq!(writer.add(data).execute().await.unwrap().version, 3); - assert_eq!(writer.version().await.unwrap(), 3); - assert_eq!(reader.version().await.unwrap(), 1); - } - - #[tokio::test] - async fn test_open_table_seeds_the_schema_from_its_describe() { - let describe_calls = Arc::new(AtomicUsize::new(0)); - let counted = describe_calls.clone(); - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.url().path(), "/v1/table/table1/describe/"); - counted.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "id", "type": {"type": "int64"}, "nullable": false} - ]}}"# - .to_string(), - ) - .unwrap() - }); - - let table = conn.open_table("table1").execute().await.unwrap(); - let schema = table.schema().await.unwrap(); - - assert_eq!(schema.field(0).name(), "id"); - assert_eq!(describe_calls.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_open_table_survives_a_describe_body_it_cannot_parse() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.url().path(), "/v1/table/table1/describe/"); - http::Response::builder() - .status(200) - .body(r#"{"table": "table1"}"#.to_string()) - .unwrap() - }); - - let table = conn.open_table("table1").execute().await.unwrap(); - - assert_eq!(table.name(), "table1"); - } - - #[tokio::test] - async fn test_open_table_branch_and_version() { - let conn = Connection::new_with_handler(|request| { - let body = if request.url().path() == "/v1/table/t/branches/list/" { - // checkout_branch validates the branch exists via list_branches. - r#"{"branches":{"exp":{"parentVersion":1,"createAt":1,"manifestSize":1}}}"# - } else { - // describe (table open + version/branch validation) - r#"{"table": "t", "version": 2, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"# - }; - http::Response::builder().status(200).body(body).unwrap() - }); - - // version-only (and "main" + version) time-travel the main chain - let v2 = conn.open_table("t").version(2).execute().await.unwrap(); - assert_eq!(v2.current_branch(), None); - let main_v2 = conn - .open_table("t") - .branch("main") - .version(2) - .execute() - .await - .unwrap(); - assert_eq!(main_v2.current_branch(), None); - - // a non-main branch opens a handle scoped to that branch - let exp = conn.open_table("t").branch("exp").execute().await.unwrap(); - assert_eq!(exp.current_branch(), Some("exp".to_string())); - let exp_v2 = conn - .open_table("t") - .branch("exp") - .version(2) - .execute() - .await - .unwrap(); - assert_eq!(exp_v2.current_branch(), Some("exp".to_string())); - } - - #[tokio::test] - async fn test_open_table_not_found() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder() - .status(404) - .body("table not found") - .unwrap() - }); - let result = conn.open_table("table1").execute().await; - assert!(result.is_err()); - assert!(matches!(result, Err(crate::Error::TableNotFound { .. }))); - } - - #[tokio::test] - async fn test_create_table() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/create/"); - assert_eq!( - request - .headers() - .get(reqwest::header::CONTENT_TYPE) - .unwrap(), - ARROW_STREAM_CONTENT_TYPE.as_bytes() - ); - - http::Response::builder().status(200).body("").unwrap() - }); - 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 table = conn.create_table("table1", data).execute().await.unwrap(); - assert_eq!(table.name(), "table1"); - } - - #[tokio::test] - async fn test_create_table_already_exists() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder() - .status(400) - .body("table table1 already exists") - .unwrap() - }); - 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 result = conn.create_table("table1", data).execute().await; - assert!(result.is_err()); - assert!( - matches!(result, Err(crate::Error::TableAlreadyExists { name }) if name == "table1") - ); - } - - #[rstest] - #[case::matching(Field::new("a", DataType::Int32, false), true)] - #[case::different_name(Field::new("x", DataType::Int32, false), false)] - #[case::different_type(Field::new("a", DataType::Int64, false), false)] - #[case::different_nullability(Field::new("a", DataType::Int32, true), false)] - #[tokio::test] - async fn test_create_table_exist_ok_validates_schema( - #[case] existing_field: Field, - #[case] matches: bool, - #[values(false, true)] empty: bool, - #[values(false, true)] cached: bool, - ) { - let existing_schema = Schema::new(vec![existing_field]); - let description = serde_json::json!({ - "version": 1, - "schema": lance::arrow::json::JsonSchema::try_from(&existing_schema).unwrap(), - }) - .to_string(); - let mut db = super::RemoteDatabase::new_mock(move |request| { - assert_eq!(request.method(), &reqwest::Method::POST); - match request.url().path() { - "/v1/table/table1/create/" => { - assert_eq!(request.url().query(), Some("mode=exist_ok")); - http::Response::builder() - .status(400) - .body("Table table1 already exists".to_string()) - .unwrap() - } - "/v1/table/table1/describe/" => http::Response::builder() - .status(200) - .body(description.clone()) - .unwrap(), - path => panic!("unexpected path: {path}"), - } - }); - db.table_cache = moka::future::Cache::new(10); - let conn = Connection::new( - Arc::new(db), - Arc::new(crate::embeddings::MemoryRegistry::new()), - ); - if cached { - conn.open_table("table1").execute().await.unwrap(); - } - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let builder = if empty { - conn.create_empty_table("table1", schema.clone()) - } else { - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - conn.create_table("table1", data) - }; - let result = builder - .mode(CreateTableMode::exist_ok(|b| b)) - .execute() - .await; - if matches { - let table = result.unwrap(); - assert_eq!(table.name(), "table1"); - assert_eq!(table.schema().await.unwrap(), schema); - } else { - assert!(matches!(result, Err(Error::Schema { message }) - if message == "Provided schema does not match existing table schema")); - } - } - - #[tokio::test] - async fn test_create_table_modes() { - let test_cases = [ - (None, "mode=create"), - (Some(CreateTableMode::Create), "mode=create"), - (Some(CreateTableMode::Overwrite), "mode=overwrite"), - ( - Some(CreateTableMode::ExistOk(Box::new(|b| b))), - "mode=exist_ok", - ), - ]; - - for (mode, expected_query_string) in test_cases { - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/create/"); - assert_eq!(request.url().query(), Some(expected_query_string)); - - http::Response::builder().status(200).body("").unwrap() - }); - - 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 mut builder = conn.create_table("table1", data.clone()); - if let Some(mode) = mode { - builder = builder.mode(mode); - } - builder.execute().await.unwrap(); - } - - // check that the open table callback is called with exist_ok - let conn = Connection::new_with_handler(|request| match request.url().path() { - "/v1/table/table1/create/" => http::Response::builder() - .status(400) - .body("Table table1 already exists") - .unwrap(), - "/v1/table/table1/describe/" => http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": {"type": "int32"}, "nullable": false} - ]}}"#, - ) - .unwrap(), - _ => { - panic!("unexpected path: {:?}", 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, 2, 3]))], - ) - .unwrap(); - - let called: Arc> = Arc::new(OnceLock::new()); - let called_in_cb = called.clone(); - conn.create_table("table1", data) - .mode(CreateTableMode::ExistOk(Box::new(move |b| { - called_in_cb.clone().set(true).unwrap(); - b - }))) - .execute() - .await - .unwrap(); - - let called = *called.get().unwrap_or(&false); - assert!(called); - } - - #[tokio::test] - async fn test_create_table_empty() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/create/"); - assert_eq!( - request - .headers() - .get(reqwest::header::CONTENT_TYPE) - .unwrap(), - ARROW_STREAM_CONTENT_TYPE.as_bytes() - ); - - http::Response::builder().status(200).body("").unwrap() - }); - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - conn.create_empty_table("table1", schema) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_drop_table() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/drop/"); - assert_eq!(request.url().query(), None); - assert!(request.body().is_none()); - - http::Response::builder().status(200).body("").unwrap() - }); - conn.drop_table("table1", &[]).await.unwrap(); - // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404. - } - - #[tokio::test] - async fn test_drop_table_does_not_read_response_body() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder() - .status(200) - .body(vec![0xff]) - .unwrap() - }); - - conn.drop_table("table1", &[]).await.unwrap(); - } - - #[tokio::test] - async fn test_drop_table_async_returns_job() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/drop/"); - http::Response::builder() - .status(202) - .body(r#"{"job_id":"drop-job-123"}"#) - .unwrap() - }); - - let job = conn.drop_table_async("table1", &[]).await.unwrap(); - assert_eq!(job.id(), Some("drop-job-123")); - } - - #[tokio::test] - async fn test_drop_table_async_old_server_returns_done_job() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder().status(200).body("").unwrap() - }); - - let job = conn.drop_table_async("table1", &[]).await.unwrap(); - assert_eq!(job.id(), None); - assert_eq!(job.status().await.unwrap(), "finished"); - } - - #[tokio::test] - async fn test_drop_table_async_rejects_accepted_response_without_job_id() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder().status(202).body("{}").unwrap() - }); - - let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); - assert!(error.to_string().contains("valid job_id")); - } - - #[tokio::test] - async fn test_drop_table_async_rejects_empty_job_id() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder() - .status(202) - .body(r#"{"job_id":""}"#) - .unwrap() - }); - - let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); - assert!(error.to_string().contains("valid job_id")); - } - - #[tokio::test] - async fn test_rename_table() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/table1/rename/"); - 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["new_table_name"], "table2"); - - http::Response::builder().status(200).body("").unwrap() - }); - conn.rename_table("table1", "table2", &[], &[]) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_connect_remote_options() { - let db_uri = "db://my-container/my-prefix"; - let _ = ConnectBuilder::new(db_uri) - .region("us-east-1") - .api_key("my-api-key") - .storage_options(vec![("azure_storage_account_name", "my-storage-account")]) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_table_names_with_root_namespace() { - // When namespace is empty (root namespace), should use /v1/table/ for backwards compatibility - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/table/"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - .body(r#"{"tables": ["table1", "table2"]}"#) - .unwrap() - }); - let names = conn - .table_names() - .namespace(vec![]) - .execute() - .await - .unwrap(); - assert_eq!(names, vec!["table1", "table2"]); - } - - #[tokio::test] - async fn test_table_names_with_namespace() { - // When namespace is non-empty, should use /v1/namespace/{id}/table/list - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/test/table/list"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - .body(r#"{"tables": ["table1", "table2"]}"#) - .unwrap() - }); - let names = conn - .table_names() - .namespace(vec!["test".to_string()]) - .execute() - .await - .unwrap(); - assert_eq!(names, vec!["table1", "table2"]); - } - - #[tokio::test] - async fn test_table_names_with_nested_namespace() { - // When namespace is vec!["ns1", "ns2"], should use /v1/namespace/ns1$ns2/table/list - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/ns1$ns2/table/list"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - // A namespace listing names the tables in that namespace; the - // namespace is the route, not part of each name. The client - // joins the two itself to build each table's identifier, so a - // listing that repeated the namespace would be joined twice. - .body(r#"{"tables": ["table1", "table2"]}"#) - .unwrap() - }); - let names = conn - .table_names() - .namespace(vec!["ns1".to_string(), "ns2".to_string()]) - .execute() - .await - .unwrap(); - assert_eq!(names, vec!["table1", "table2"]); - } - - #[tokio::test] - async fn test_open_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$table1/describe/"); - assert_eq!(request.url().query(), None); - - http::Response::builder() - .status(200) - .body(r#"{"table": "table1"}"#) - .unwrap() - }); - let table = conn - .open_table("table1") - .namespace(vec!["ns1".to_string(), "ns2".to_string()]) - .execute() - .await - .unwrap(); - assert_eq!(table.name(), "table1"); - } - - #[tokio::test] - async fn test_create_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$table1/create/"); - assert_eq!( - request - .headers() - .get(reqwest::header::CONTENT_TYPE) - .unwrap(), - ARROW_STREAM_CONTENT_TYPE.as_bytes() - ); - - http::Response::builder().status(200).body("").unwrap() - }); - 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 table = conn - .create_table("table1", data) - .namespace(vec!["ns1".to_string()]) - .execute() - .await - .unwrap(); - assert_eq!(table.name(), "table1"); - } - - #[tokio::test] - async fn test_drop_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$table1/drop/"); - assert_eq!(request.url().query(), None); - assert!(request.body().is_none()); - - http::Response::builder().status(200).body("").unwrap() - }); - conn.drop_table("table1", &["ns1".to_string(), "ns2".to_string()]) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_rename_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$table1/rename/"); - 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["new_table_name"], "table2"); - assert_eq!(body["new_namespace"], serde_json::json!(["ns2"])); - - http::Response::builder().status(200).body("").unwrap() - }); - conn.rename_table( - "table1", - "table2", - &["ns1".to_string()], - &["ns2".to_string()], - ) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_empty_table_with_namespace() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/table/prod$data$metrics/create/"); - assert_eq!( - request - .headers() - .get(reqwest::header::CONTENT_TYPE) - .unwrap(), - ARROW_STREAM_CONTENT_TYPE.as_bytes() - ); - - http::Response::builder().status(200).body("").unwrap() - }); - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - conn.create_empty_table("metrics", schema) - .namespace(vec!["prod".to_string(), "data".to_string()]) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_header_provider_in_request() { - // Test HeaderProvider implementation that adds custom headers - #[derive(Debug, Clone)] - struct TestHeaderProvider { - headers: HashMap, - } - - #[async_trait::async_trait] - impl HeaderProvider for TestHeaderProvider { - async fn get_headers(&self) -> crate::Result> { - 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; - - // 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> { - Err(crate::Error::Runtime { - message: "Failed to fetch auth token".to_string(), - }) - } - } - - let provider = Arc::new(ErrorHeaderProvider) as Arc; - 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> { - Ok(HashMap::from([( - "authorization".to_string(), - "Bearer token".to_string(), - )])) - } - } - - let context_provider = NamespaceHeaderProviderContext { - header_provider: Arc::new(TestHeaderProvider) as Arc, - }; - - 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> { - Ok(HashMap::from([( - "authorization".to_string(), - "Bearer token".to_string(), - )])) - } - } - - let client_config = ClientConfig { - header_provider: Some(Arc::new(TestHeaderProvider) as Arc), - ..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())); - } - } - - #[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" - ); - } - } - - #[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::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::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::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::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 { - 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::>(); - 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::>(); - 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::>(); - 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::(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 { - 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); - } -} +mod tests; diff --git a/rust/lancedb/src/remote/db/tests/catalog.rs b/rust/lancedb/src/remote/db/tests/catalog.rs new file mode 100644 index 000000000..6e5f31657 --- /dev/null +++ b/rust/lancedb/src/remote/db/tests/catalog.rs @@ -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 { + 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::>(); + 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::>(); + 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::>(); + 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::(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 { + 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); +} diff --git a/rust/lancedb/src/remote/db/tests/jobs.rs b/rust/lancedb/src/remote/db/tests/jobs.rs new file mode 100644 index 000000000..a609f487e --- /dev/null +++ b/rust/lancedb/src/remote/db/tests/jobs.rs @@ -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" + ); + } +} diff --git a/rust/lancedb/src/remote/db/tests/mod.rs b/rust/lancedb/src/remote/db/tests/mod.rs new file mode 100644 index 000000000..67ef365f1 --- /dev/null +++ b/rust/lancedb/src/remote/db/tests/mod.rs @@ -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; diff --git a/rust/lancedb/src/remote/db/tests/namespace.rs b/rust/lancedb/src/remote/db/tests/namespace.rs new file mode 100644 index 000000000..cb66b15aa --- /dev/null +++ b/rust/lancedb/src/remote/db/tests/namespace.rs @@ -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, + } + + #[async_trait::async_trait] + impl HeaderProvider for TestHeaderProvider { + async fn get_headers(&self) -> crate::Result> { + 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; + + // 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> { + Err(crate::Error::Runtime { + message: "Failed to fetch auth token".to_string(), + }) + } + } + + let provider = Arc::new(ErrorHeaderProvider) as Arc; + 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> { + Ok(HashMap::from([( + "authorization".to_string(), + "Bearer token".to_string(), + )])) + } + } + + let context_provider = NamespaceHeaderProviderContext { + header_provider: Arc::new(TestHeaderProvider) as Arc, + }; + + 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> { + Ok(HashMap::from([( + "authorization".to_string(), + "Bearer token".to_string(), + )])) + } + } + + let client_config = ClientConfig { + header_provider: Some(Arc::new(TestHeaderProvider) as Arc), + ..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())); + } +} diff --git a/rust/lancedb/src/remote/db/tests/tables.rs b/rust/lancedb/src/remote/db/tests/tables.rs new file mode 100644 index 000000000..47211851f --- /dev/null +++ b/rust/lancedb/src/remote/db/tests/tables.rs @@ -0,0 +1,1049 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[test] +fn test_cache_key_security() { + // Test that cache keys are unique regardless of delimiter manipulation + + // Case 1: Different delimiters should not affect cache key + let key1 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]); + let key2 = build_cache_key("table1", &["ns1$ns2".to_string()]); + assert_ne!( + key1, key2, + "Cache keys should differ for different namespace structures" + ); + + // Case 2: Table name containing delimiter should not cause collision + let key3 = build_cache_key("ns2$table1", &["ns1".to_string()]); + assert_ne!( + key1, key3, + "Cache key should be different when table name contains delimiter" + ); + + // Case 3: Empty namespace vs namespace with empty string + let key4 = build_cache_key("table1", &[]); + let key5 = build_cache_key("table1", &["".to_string()]); + assert_ne!( + key4, key5, + "Empty namespace should differ from namespace with empty string" + ); + + // Case 4: Verify same inputs produce same key (consistency) + let key6 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]); + assert_eq!(key1, key6, "Same inputs should produce same cache key"); +} + +#[tokio::test] +async fn test_create_materialized_view_requires_sql_client() { + let db = super::super::RemoteDatabase::new_mock(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()) + }); + let error = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["analytics".into()], + query: "SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10".into(), + with_no_data: false, + }) + .await + .unwrap_err(); + assert!(error.to_string().contains("SQL is unavailable")); +} + +#[tokio::test] +async fn test_drop_materialized_view_uses_item_route_and_job() { + let db = super::super::RemoteDatabase::new_mock(|request| { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.url().path(), + "/v1/materialized_view/analytics$adults/drop" + ); + assert!(request.body().is_none()); + http::Response::builder() + .status(202) + .body(serde_json::json!({"job_id": "j1-mv-drop"}).to_string()) + .unwrap() + }); + let job = db + .drop_materialized_view_async("adults", &["analytics".into()]) + .await + .unwrap(); + assert_eq!(job.id(), Some("j1-mv-drop")); +} + +#[tokio::test] +async fn test_drop_function_async_returns_job() { + let db = super::super::RemoteDatabase::new_mock(|request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/function/embed/drop"); + http::Response::builder() + .status(202) + .body(serde_json::json!({"dropped": true, "job_id": "j1-fn-drop"}).to_string()) + .unwrap() + }); + let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); + assert!(dropped); + assert_eq!(job.id(), Some("j1-fn-drop")); +} + +/// An unbound name and an inline deletion both answer `200`: the name is gone and +/// nothing is left to wait for, so the job is already finished. +#[tokio::test] +async fn test_drop_function_async_completed_inline() { + let db = super::super::RemoteDatabase::new_mock(|_| { + http::Response::builder() + .status(200) + .body(serde_json::json!({"dropped": false}).to_string()) + .unwrap() + }); + let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); + assert!(!dropped); + assert_eq!(job.id(), None); + assert_eq!(job.status().await.unwrap(), "finished"); + job.wait().await.unwrap(); +} + +#[tokio::test] +async fn test_drop_function_async_rejects_incomplete_acceptance() { + for body in [ + r#"{"dropped":true}"#, + r#"{"dropped":true,"job_id":""}"#, + r#"{"dropped":true,"job_id":null}"#, + ] { + let db = super::super::RemoteDatabase::new_mock(move |_| { + http::Response::builder().status(202).body(body).unwrap() + }); + let error = db + .drop_function_async("embed", "1", &[]) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("valid job_id")); + } +} + +#[tokio::test] +async fn test_drop_materialized_view_completed_inline() { + let db = super::super::RemoteDatabase::new_mock(|request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/materialized_view/adults/drop"); + http::Response::builder().status(200).body("{}").unwrap() + }); + let job = db + .drop_materialized_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_materialized_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_materialized_view_async("adults", &[]) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("valid job_id")); + } +} + +#[tokio::test] +async fn test_drop_materialized_view_rejects_unexpected_success_status() { + let db = super::super::RemoteDatabase::new_mock(|_| { + http::Response::builder().status(204).body("").unwrap() + }); + let error = db + .drop_materialized_view_async("adults", &[]) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("200 OK or 202 Accepted")); +} + +#[tokio::test] +async fn test_list_materialized_views_follows_empty_pages() { + let page = Arc::new(AtomicUsize::new(0)); + let db = super::super::RemoteDatabase::new_mock({ + let page = page.clone(); + move |request| { + assert_eq!(request.method(), "GET"); + assert_eq!( + request.url().path(), + "/v1/namespace/analytics/materialized_view/list" + ); + match page.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert!(request.url().query().is_none()); + http::Response::builder() + .status(200) + .body(serde_json::json!({"views": [], "page_token": "next"}).to_string()) + .unwrap() + } + 1 => { + assert_eq!(request.url().query(), Some("page_token=next")); + http::Response::builder() + .status(200) + .body(serde_json::json!({"views": ["adults"]}).to_string()) + .unwrap() + } + _ => panic!("listing requested too many pages"), + } + } + }); + assert_eq!( + db.list_materialized_views(&["analytics".into()]) + .await + .unwrap(), + ["adults"] + ); +} + +#[tokio::test] +async fn test_retries() { + // We'll record the request_id here, to check it matches the one in the error. + let seen_request_id = Arc::new(OnceLock::new()); + let seen_request_id_ref = seen_request_id.clone(); + let conn = Connection::new_with_handler(move |request| { + // Request id should be the same on each retry. + let request_id = request.headers()["x-request-id"] + .to_str() + .unwrap() + .to_string(); + let seen_id = seen_request_id_ref.get_or_init(|| request_id.clone()); + assert_eq!(&request_id, seen_id); + + http::Response::builder() + .status(500) + .body("internal server error") + .unwrap() + }); + let result = conn.table_names().execute().await; + if let Err(Error::Retry { + request_id, + request_failures, + max_request_failures, + source, + .. + }) = result + { + let expected_id = seen_request_id.get().unwrap(); + assert_eq!(&request_id, expected_id); + assert_eq!(request_failures, max_request_failures); + assert!( + source.to_string().contains("internal server error"), + "source: {:?}", + source + ); + } else { + panic!("unexpected result: {:?}", result); + }; +} + +#[tokio::test] +async fn test_table_names() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + .body(r#"{"tables": ["table1", "table2"]}"#) + .unwrap() + }); + let names = conn.table_names().execute().await.unwrap(); + assert_eq!(names, vec!["table1", "table2"]); +} + +#[tokio::test] +async fn test_table_names_in_a_namespace_never_invents_a_page_token() { + // The namespace route's token belongs to the store, so `table_names` cannot build one + // from `start_after`. It walks the namespace on the server's own tokens and applies the + // name semantics itself, which is what keeps it working either side of the change. + let page = Arc::new(AtomicUsize::new(0)); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/namespace/ns/table/list"); + let query = request.url().query().unwrap_or(""); + assert!( + !query.contains("page_token=users"), + "a table name must never be sent as a page token: {query}" + ); + match page.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert!( + !query.contains("page_token"), + "the walk starts with no token" + ); + http::Response::builder() + .status(200) + .body(r#"{"tables": ["users", "orders"], "page_token": "opaque-1"}"#) + .unwrap() + } + _ => { + assert!(query.contains("page_token=opaque-1")); + http::Response::builder() + .status(200) + .body(r#"{"tables": ["widgets"]}"#) + .unwrap() + } + } + }); + + let names = conn + .table_names() + .namespace(vec!["ns".to_string()]) + .start_after("users") + .execute() + .await + .unwrap(); + // Name order, resumed after "users": "orders" sorts before it and is dropped. + assert_eq!(names, vec!["widgets"]); +} + +#[tokio::test] +async fn test_table_names_in_a_namespace_stops_on_a_repeated_token() { + // A server that handed back the token it was given would never finish the walk. + let conn = Connection::new_with_handler(|_request| { + http::Response::builder() + .status(200) + .body(r#"{"tables": ["a"], "page_token": "same"}"#) + .unwrap() + }); + + let names = conn + .table_names() + .namespace(vec!["ns".to_string()]) + .execute() + .await + .unwrap(); + // The guard bounds the walk instead of letting it run forever. The repeat is the + // server breaking the token contract and is not papered over here. + assert_eq!(names, vec!["a", "a"]); +} + +#[tokio::test] +async fn test_table_names_in_a_namespace_stops_on_an_empty_token() { + // An empty token ends the listing. Sending it back would ask a server that reads it + // as "start from the beginning" for the first page a second time, and every name on + // that page would be collected twice. + 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!( + !request.url().query().unwrap_or("").contains("page_token"), + "an empty token must never be sent back" + ); + http::Response::builder() + .status(200) + .body(r#"{"tables": ["a"], "page_token": ""}"#) + .unwrap() + }); + + let names = conn + .table_names() + .namespace(vec!["ns".to_string()]) + .execute() + .await + .unwrap(); + assert_eq!(names, vec!["a"]); + assert_eq!(requests.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_table_names_pagination() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + assert!(request.url().query().unwrap().contains("limit=2")); + assert!(request.url().query().unwrap().contains("page_token=table2")); + + http::Response::builder() + .status(200) + .body(r#"{"tables": ["table3", "table4"], "page_token": "token"}"#) + .unwrap() + }); + let names = conn + .table_names() + .start_after("table2") + .limit(2) + .execute() + .await + .unwrap(); + assert_eq!(names, vec!["table3", "table4"]); +} + +#[tokio::test] +async fn test_open_table() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + .body(r#"{"table": "table1"}"#) + .unwrap() + }); + let table = conn.open_table("table1").execute().await.unwrap(); + assert_eq!(table.name(), "table1"); + + // Storage options should be ignored. + let table = conn + .open_table("table1") + .storage_option("key", "value") + .execute() + .await + .unwrap(); + assert_eq!(table.name(), "table1"); +} + +#[tokio::test] +async fn test_open_table_checkout_is_independent_per_handle() { + let latest = Arc::new(AtomicUsize::new(2)); + let current = latest.clone(); + let mut db = super::super::RemoteDatabase::new_mock(move |request| { + let body = match request.url().path() { + "/v1/table/table1/describe/" => { + let requested_version = request + .body() + .and_then(|body| body.as_bytes()) + .and_then(|body| serde_json::from_slice::(body).ok()) + .and_then(|body| body["version"].as_u64()); + let version = + requested_version.unwrap_or_else(|| current.load(Ordering::SeqCst) as u64); + serde_json::json!({ + "version": version, + "schema": {"fields": [ + {"name": "id", "type": {"type": "int32"}, "nullable": false} + ]} + }) + .to_string() + } + "/v1/table/table1/insert/" => { + let version = current.fetch_add(1, Ordering::SeqCst) + 1; + serde_json::json!({"version": version}).to_string() + } + path => panic!("unexpected request path: {path}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + + let writer = conn.open_table("table1").execute().await.unwrap(); + let reader = conn.open_table("table1").execute().await.unwrap(); + reader.checkout(1).await.unwrap(); + + assert_eq!(reader.version().await.unwrap(), 1); + assert_eq!(writer.version().await.unwrap(), 2); + + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![3]))], + ) + .unwrap(); + assert_eq!(writer.add(data).execute().await.unwrap().version, 3); + assert_eq!(writer.version().await.unwrap(), 3); + assert_eq!(reader.version().await.unwrap(), 1); +} + +#[tokio::test] +async fn test_open_table_seeds_the_schema_from_its_describe() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let counted = describe_calls.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + counted.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "id", "type": {"type": "int64"}, "nullable": false} + ]}}"# + .to_string(), + ) + .unwrap() + }); + + let table = conn.open_table("table1").execute().await.unwrap(); + let schema = table.schema().await.unwrap(); + + assert_eq!(schema.field(0).name(), "id"); + assert_eq!(describe_calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_open_table_survives_a_describe_body_it_cannot_parse() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + http::Response::builder() + .status(200) + .body(r#"{"table": "table1"}"#.to_string()) + .unwrap() + }); + + let table = conn.open_table("table1").execute().await.unwrap(); + + assert_eq!(table.name(), "table1"); +} + +#[tokio::test] +async fn test_open_table_branch_and_version() { + let conn = Connection::new_with_handler(|request| { + let body = if request.url().path() == "/v1/table/t/branches/list/" { + // checkout_branch validates the branch exists via list_branches. + r#"{"branches":{"exp":{"parentVersion":1,"createAt":1,"manifestSize":1}}}"# + } else { + // describe (table open + version/branch validation) + r#"{"table": "t", "version": 2, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"# + }; + http::Response::builder().status(200).body(body).unwrap() + }); + + // version-only (and "main" + version) time-travel the main chain + let v2 = conn.open_table("t").version(2).execute().await.unwrap(); + assert_eq!(v2.current_branch(), None); + let main_v2 = conn + .open_table("t") + .branch("main") + .version(2) + .execute() + .await + .unwrap(); + assert_eq!(main_v2.current_branch(), None); + + // a non-main branch opens a handle scoped to that branch + let exp = conn.open_table("t").branch("exp").execute().await.unwrap(); + assert_eq!(exp.current_branch(), Some("exp".to_string())); + let exp_v2 = conn + .open_table("t") + .branch("exp") + .version(2) + .execute() + .await + .unwrap(); + assert_eq!(exp_v2.current_branch(), Some("exp".to_string())); +} + +#[tokio::test] +async fn test_open_table_not_found() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(404) + .body("table not found") + .unwrap() + }); + let result = conn.open_table("table1").execute().await; + assert!(result.is_err()); + assert!(matches!(result, Err(crate::Error::TableNotFound { .. }))); +} + +#[tokio::test] +async fn test_create_table() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/create/"); + assert_eq!( + request + .headers() + .get(reqwest::header::CONTENT_TYPE) + .unwrap(), + ARROW_STREAM_CONTENT_TYPE.as_bytes() + ); + + http::Response::builder().status(200).body("").unwrap() + }); + 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 table = conn.create_table("table1", data).execute().await.unwrap(); + assert_eq!(table.name(), "table1"); +} + +#[tokio::test] +async fn test_create_table_already_exists() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(400) + .body("table table1 already exists") + .unwrap() + }); + 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 result = conn.create_table("table1", data).execute().await; + assert!(result.is_err()); + assert!(matches!(result, Err(crate::Error::TableAlreadyExists { name }) if name == "table1")); +} + +#[rstest] +#[case::matching(Field::new("a", DataType::Int32, false), true)] +#[case::different_name(Field::new("x", DataType::Int32, false), false)] +#[case::different_type(Field::new("a", DataType::Int64, false), false)] +#[case::different_nullability(Field::new("a", DataType::Int32, true), false)] +#[tokio::test] +async fn test_create_table_exist_ok_validates_schema( + #[case] existing_field: Field, + #[case] matches: bool, + #[values(false, true)] empty: bool, + #[values(false, true)] cached: bool, +) { + let existing_schema = Schema::new(vec![existing_field]); + let description = serde_json::json!({ + "version": 1, + "schema": lance::arrow::json::JsonSchema::try_from(&existing_schema).unwrap(), + }) + .to_string(); + let mut db = super::super::RemoteDatabase::new_mock(move |request| { + assert_eq!(request.method(), &reqwest::Method::POST); + match request.url().path() { + "/v1/table/table1/create/" => { + assert_eq!(request.url().query(), Some("mode=exist_ok")); + http::Response::builder() + .status(400) + .body("Table table1 already exists".to_string()) + .unwrap() + } + "/v1/table/table1/describe/" => http::Response::builder() + .status(200) + .body(description.clone()) + .unwrap(), + path => panic!("unexpected path: {path}"), + } + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + if cached { + conn.open_table("table1").execute().await.unwrap(); + } + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let builder = if empty { + conn.create_empty_table("table1", schema.clone()) + } else { + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + conn.create_table("table1", data) + }; + let result = builder + .mode(CreateTableMode::exist_ok(|b| b)) + .execute() + .await; + if matches { + let table = result.unwrap(); + assert_eq!(table.name(), "table1"); + assert_eq!(table.schema().await.unwrap(), schema); + } else { + assert!(matches!(result, Err(Error::Schema { message }) + if message == "Provided schema does not match existing table schema")); + } +} + +#[tokio::test] +async fn test_create_table_modes() { + let test_cases = [ + (None, "mode=create"), + (Some(CreateTableMode::Create), "mode=create"), + (Some(CreateTableMode::Overwrite), "mode=overwrite"), + ( + Some(CreateTableMode::ExistOk(Box::new(|b| b))), + "mode=exist_ok", + ), + ]; + + for (mode, expected_query_string) in test_cases { + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/create/"); + assert_eq!(request.url().query(), Some(expected_query_string)); + + http::Response::builder().status(200).body("").unwrap() + }); + + 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 mut builder = conn.create_table("table1", data.clone()); + if let Some(mode) = mode { + builder = builder.mode(mode); + } + builder.execute().await.unwrap(); + } + + // check that the open table callback is called with exist_ok + let conn = Connection::new_with_handler(|request| match request.url().path() { + "/v1/table/table1/create/" => http::Response::builder() + .status(400) + .body("Table table1 already exists") + .unwrap(), + "/v1/table/table1/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": {"type": "int32"}, "nullable": false} + ]}}"#, + ) + .unwrap(), + _ => { + panic!("unexpected path: {:?}", 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, 2, 3]))], + ) + .unwrap(); + + let called: Arc> = Arc::new(OnceLock::new()); + let called_in_cb = called.clone(); + conn.create_table("table1", data) + .mode(CreateTableMode::ExistOk(Box::new(move |b| { + called_in_cb.clone().set(true).unwrap(); + b + }))) + .execute() + .await + .unwrap(); + + let called = *called.get().unwrap_or(&false); + assert!(called); +} + +#[tokio::test] +async fn test_create_table_empty() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/create/"); + assert_eq!( + request + .headers() + .get(reqwest::header::CONTENT_TYPE) + .unwrap(), + ARROW_STREAM_CONTENT_TYPE.as_bytes() + ); + + http::Response::builder().status(200).body("").unwrap() + }); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + conn.create_empty_table("table1", schema) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_drop_table() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/drop/"); + assert_eq!(request.url().query(), None); + assert!(request.body().is_none()); + + http::Response::builder().status(200).body("").unwrap() + }); + conn.drop_table("table1", &[]).await.unwrap(); + // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404. +} + +#[tokio::test] +async fn test_drop_table_does_not_read_response_body() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(200) + .body(vec![0xff]) + .unwrap() + }); + + conn.drop_table("table1", &[]).await.unwrap(); +} + +#[tokio::test] +async fn test_drop_table_async_returns_job() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/drop/"); + http::Response::builder() + .status(202) + .body(r#"{"job_id":"drop-job-123"}"#) + .unwrap() + }); + + let job = conn.drop_table_async("table1", &[]).await.unwrap(); + assert_eq!(job.id(), Some("drop-job-123")); +} + +#[tokio::test] +async fn test_drop_table_async_old_server_returns_done_job() { + let conn = + Connection::new_with_handler(|_| http::Response::builder().status(200).body("").unwrap()); + + let job = conn.drop_table_async("table1", &[]).await.unwrap(); + assert_eq!(job.id(), None); + assert_eq!(job.status().await.unwrap(), "finished"); +} + +#[tokio::test] +async fn test_drop_table_async_rejects_accepted_response_without_job_id() { + let conn = + Connection::new_with_handler(|_| http::Response::builder().status(202).body("{}").unwrap()); + + let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); + assert!(error.to_string().contains("valid job_id")); +} + +#[tokio::test] +async fn test_drop_table_async_rejects_empty_job_id() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(202) + .body(r#"{"job_id":""}"#) + .unwrap() + }); + + let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); + assert!(error.to_string().contains("valid job_id")); +} + +#[tokio::test] +async fn test_rename_table() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/rename/"); + 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["new_table_name"], "table2"); + + http::Response::builder().status(200).body("").unwrap() + }); + conn.rename_table("table1", "table2", &[], &[]) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_connect_remote_options() { + let db_uri = "db://my-container/my-prefix"; + let _ = ConnectBuilder::new(db_uri) + .region("us-east-1") + .api_key("my-api-key") + .storage_options(vec![("azure_storage_account_name", "my-storage-account")]) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_table_names_with_root_namespace() { + // When namespace is empty (root namespace), should use /v1/table/ for backwards compatibility + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + .body(r#"{"tables": ["table1", "table2"]}"#) + .unwrap() + }); + let names = conn + .table_names() + .namespace(vec![]) + .execute() + .await + .unwrap(); + assert_eq!(names, vec!["table1", "table2"]); +} + +#[tokio::test] +async fn test_table_names_with_namespace() { + // When namespace is non-empty, should use /v1/namespace/{id}/table/list + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/namespace/test/table/list"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + .body(r#"{"tables": ["table1", "table2"]}"#) + .unwrap() + }); + let names = conn + .table_names() + .namespace(vec!["test".to_string()]) + .execute() + .await + .unwrap(); + assert_eq!(names, vec!["table1", "table2"]); +} + +#[tokio::test] +async fn test_table_names_with_nested_namespace() { + // When namespace is vec!["ns1", "ns2"], should use /v1/namespace/ns1$ns2/table/list + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/namespace/ns1$ns2/table/list"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + // A namespace listing names the tables in that namespace; the + // namespace is the route, not part of each name. The client + // joins the two itself to build each table's identifier, so a + // listing that repeated the namespace would be joined twice. + .body(r#"{"tables": ["table1", "table2"]}"#) + .unwrap() + }); + let names = conn + .table_names() + .namespace(vec!["ns1".to_string(), "ns2".to_string()]) + .execute() + .await + .unwrap(); + assert_eq!(names, vec!["table1", "table2"]); +} + +#[tokio::test] +async fn test_open_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$table1/describe/"); + assert_eq!(request.url().query(), None); + + http::Response::builder() + .status(200) + .body(r#"{"table": "table1"}"#) + .unwrap() + }); + let table = conn + .open_table("table1") + .namespace(vec!["ns1".to_string(), "ns2".to_string()]) + .execute() + .await + .unwrap(); + assert_eq!(table.name(), "table1"); +} + +#[tokio::test] +async fn test_create_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$table1/create/"); + assert_eq!( + request + .headers() + .get(reqwest::header::CONTENT_TYPE) + .unwrap(), + ARROW_STREAM_CONTENT_TYPE.as_bytes() + ); + + http::Response::builder().status(200).body("").unwrap() + }); + 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 table = conn + .create_table("table1", data) + .namespace(vec!["ns1".to_string()]) + .execute() + .await + .unwrap(); + assert_eq!(table.name(), "table1"); +} + +#[tokio::test] +async fn test_drop_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$table1/drop/"); + assert_eq!(request.url().query(), None); + assert!(request.body().is_none()); + + http::Response::builder().status(200).body("").unwrap() + }); + conn.drop_table("table1", &["ns1".to_string(), "ns2".to_string()]) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_rename_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$table1/rename/"); + 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["new_table_name"], "table2"); + assert_eq!(body["new_namespace"], serde_json::json!(["ns2"])); + + http::Response::builder().status(200).body("").unwrap() + }); + conn.rename_table( + "table1", + "table2", + &["ns1".to_string()], + &["ns2".to_string()], + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_empty_table_with_namespace() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/prod$data$metrics/create/"); + assert_eq!( + request + .headers() + .get(reqwest::header::CONTENT_TYPE) + .unwrap(), + ARROW_STREAM_CONTENT_TYPE.as_bytes() + ); + + http::Response::builder().status(200).body("").unwrap() + }); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + conn.create_empty_table("metrics", schema) + .namespace(vec!["prod".to_string(), "data".to_string()]) + .execute() + .await + .unwrap(); +} diff --git a/rust/lancedb/src/remote/oauth.rs b/rust/lancedb/src/remote/oauth.rs index c7cffa327..e0c12b5c1 100644 --- a/rust/lancedb/src/remote/oauth.rs +++ b/rust/lancedb/src/remote/oauth.rs @@ -1743,1867 +1743,4 @@ impl HeaderProvider for OAuthHeaderProvider { } #[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - use base64::Engine; - use oauth2::TokenResponse as _; - 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, - ) -> 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 { - self.headers.lines().find_map(|line| { - let (key, value) = line.split_once(':')?; - key.trim() - .eq_ignore_ascii_case(name) - .then(|| value.trim().to_string()) - }) - } - } - - #[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::>()); - } - 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::>() - ); - } - 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::(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: \"\"")); - } - - #[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::().unwrap() - ); - } - - #[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(); - } - - #[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, - refreshes: Arc, - } - - #[async_trait] - impl TokenSource for RefreshingTokenSource { - async fn fetch_token(&self) -> Result { - self.fetches.fetch_add(1, Ordering::SeqCst); - Ok(token_response("initial", Some("refresh"), Some(3600))) - } - - async fn refresh_token(&self, refresh_token: &str) -> Result { - 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, - refreshes: Arc, - } - - #[async_trait] - impl TokenSource for FailedRefreshTokenSource { - async fn fetch_token(&self) -> Result { - 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 { - 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, - } - - #[async_trait] - impl TokenSource for TransientRefreshFailureSource { - async fn fetch_token(&self) -> Result { - 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 { - 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") - ); - } - - fn test_config(flow: OAuthFlow, client_secret: Option) -> 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(\"\")")); - 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(\"\")")); - } - - #[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_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>>, - 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, 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, 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 spawn_oauth_server() -> (String, Arc, 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) - } - - 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::().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 { - 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 tests; diff --git a/rust/lancedb/src/remote/oauth/tests/authorization.rs b/rust/lancedb/src/remote/oauth/tests/authorization.rs new file mode 100644 index 000000000..75befc6f4 --- /dev/null +++ b/rust/lancedb/src/remote/oauth/tests/authorization.rs @@ -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(); +} diff --git a/rust/lancedb/src/remote/oauth/tests/config.rs b/rust/lancedb/src/remote/oauth/tests/config.rs new file mode 100644 index 000000000..21891d530 --- /dev/null +++ b/rust/lancedb/src/remote/oauth/tests/config.rs @@ -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) -> 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(\"\")")); + 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(\"\")")); +} + +#[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, 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) +} diff --git a/rust/lancedb/src/remote/oauth/tests/device.rs b/rust/lancedb/src/remote/oauth/tests/device.rs new file mode 100644 index 000000000..78ff2cfb8 --- /dev/null +++ b/rust/lancedb/src/remote/oauth/tests/device.rs @@ -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, + refreshes: Arc, +} + +#[async_trait] +impl TokenSource for RefreshingTokenSource { + async fn fetch_token(&self) -> Result { + self.fetches.fetch_add(1, Ordering::SeqCst); + Ok(token_response("initial", Some("refresh"), Some(3600))) + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + 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, + refreshes: Arc, +} + +#[async_trait] +impl TokenSource for FailedRefreshTokenSource { + async fn fetch_token(&self) -> Result { + 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 { + 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, +} + +#[async_trait] +impl TokenSource for TransientRefreshFailureSource { + async fn fetch_token(&self) -> Result { + 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 { + 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") + ); +} diff --git a/rust/lancedb/src/remote/oauth/tests/flows.rs b/rust/lancedb/src/remote/oauth/tests/flows.rs new file mode 100644 index 000000000..51d8800ad --- /dev/null +++ b/rust/lancedb/src/remote/oauth/tests/flows.rs @@ -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::>()); + } + 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::>() + ); + } + 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::(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: \"\"")); +} + +#[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::().unwrap() + ); +} diff --git a/rust/lancedb/src/remote/oauth/tests/mod.rs b/rust/lancedb/src/remote/oauth/tests/mod.rs new file mode 100644 index 000000000..0539d2aaa --- /dev/null +++ b/rust/lancedb/src/remote/oauth/tests/mod.rs @@ -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, +) -> 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 { + 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>>, + 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, 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, 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::().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 { + 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; diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index ca8fe644f..f3ca7593b 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3965,10431 +3965,4 @@ impl TryFrom for MergeInsertRequest { } #[cfg(test)] -mod tests { - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - use std::{collections::HashMap, pin::Pin}; - - use super::*; - - use crate::remote::client::{ClientConfig, RetryConfig}; - use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken}; - - 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_index::scalar::inverted::{DocumentGranularity, query::MatchQuery}; - use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams}; - use reqwest::Body; - use rstest::rstest; - use serde_json::json; - - use crate::index::vector::{ - IvfFlatIndexBuilder, IvfHnswFlatIndexBuilder, IvfHnswSqIndexBuilder, IvfRqIndexBuilder, - IvfSqIndexBuilder, - }; - use crate::remote::JSON_CONTENT_TYPE; - use crate::remote::db::DEFAULT_SERVER_VERSION; - use crate::utils::background_cache::clock; - use crate::{ - DistanceType, Error, Table, - index::{Index, IndexStatistics, IndexType, vector::IvfPqIndexBuilder}, - query::{ - AnalyzePlanDistributedMetrics, ColumnOrdering, ExecutableQuery, QueryBase, - QueryExecutionOptions, - }, - }; - - 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() - } - - #[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 { - Box::new(RecordBatchIterator::new( - [Ok(batch.clone())], - batch.schema(), - )) - }; - - // These table operations should translate 404 to TableNotFound. - let results: Vec>> = 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); - } - - async fn collect_body(body: Body) -> Vec { - 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 { - 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 { - 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 { - 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, - ), - ]) - } - - #[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::("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::() - .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[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[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::("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 = 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::>(); - 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 = - 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::>(); - 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 = 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::>(); - assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]); - - let params = request.url().query_pairs().collect::>(); - 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 = 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 = 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::>(); - 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 }); - } - - #[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::>() - .await; - assert_eq!(data.len(), 1); - assert_eq!(data[0].as_ref().unwrap(), &expected_data); - } - - fn blob_describe_response() -> http::Response { - 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 = (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) -> 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, - ) -> 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 { - 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 { - 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 { - 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 { - 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, - ) { - 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"); - } - - #[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 { - 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::::None, - "upper_bound": Option::::None, - "k": 10, - "ef": Option::::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::>().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::::None, - "upper_bound": Option::::None, - "k": 10, - "ef": Option::::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::>().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::>().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::::None, - "upper_bound": Option::::None, - "ef": Option::::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::::None, - "upper_bound": Option::::None, - "ef": Option::::None, - "refine_factor": Option::::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::>() - .await - .unwrap(); - - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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 { - 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::>() - .await - .unwrap(); - let results = concat_batches(&results[0].schema(), &results).unwrap(); - - let query_index = results["query_index"].as_primitive::(); - // We don't guarantee order. - assert!(query_index.values().contains(&0)); - assert!(query_index.values().contains(&1)); - } - - #[tokio::test] - async fn test_create_index() { - let cases = [ - ( - "IVF_FLAT", - json!({ - "metric_type": "hamming", - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfFlat(IvfFlatIndexBuilder::default().distance_type(DistanceType::Hamming)), - ), - ( - "IVF_FLAT", - json!({ - "metric_type": "hamming", - "num_partitions": 128, - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .distance_type(DistanceType::Hamming) - .num_partitions(128), - ), - ), - ( - "IVF_PQ", - json!({ - "metric_type": "l2", - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfPq(Default::default()), - ), - ( - "IVF_PQ", - json!({ - "metric_type": "cosine", - "num_partitions": 128, - "num_bits": 4, - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfPq( - IvfPqIndexBuilder::default() - .distance_type(DistanceType::Cosine) - .num_partitions(128) - .num_bits(4), - ), - ), - ( - "IVF_PQ", - json!({ - "metric_type": "l2", - "num_sub_vectors": 16, - "sample_rate": 512, - "max_iterations": 100, - }), - Index::IvfPq( - IvfPqIndexBuilder::default() - .num_sub_vectors(16) - .sample_rate(512) - .max_iterations(100), - ), - ), - ( - "IVF_HNSW_SQ", - json!({ - "metric_type": "l2", - "sample_rate": 256, - "max_iterations": 50, - "m": 20, - "ef_construction": 300, - }), - Index::IvfHnswSq(Default::default()), - ), - ( - "IVF_HNSW_SQ", - json!({ - "metric_type": "l2", - "num_partitions": 128, - "sample_rate": 256, - "max_iterations": 50, - "m": 40, - "ef_construction": 500, - }), - Index::IvfHnswSq( - IvfHnswSqIndexBuilder::default() - .distance_type(DistanceType::L2) - .num_partitions(128) - .num_edges(40) - .ef_construction(500), - ), - ), - ( - "IVF_HNSW_FLAT", - json!({ - "metric_type": "l2", - "sample_rate": 256, - "max_iterations": 50, - "m": 20, - "ef_construction": 300, - }), - Index::IvfHnswFlat(Default::default()), - ), - ( - "IVF_HNSW_FLAT", - json!({ - "metric_type": "cosine", - "num_partitions": 64, - "sample_rate": 256, - "max_iterations": 50, - "m": 40, - "ef_construction": 500, - }), - Index::IvfHnswFlat( - IvfHnswFlatIndexBuilder::default() - .distance_type(DistanceType::Cosine) - .num_partitions(64) - .num_edges(40) - .ef_construction(500), - ), - ), - ( - "IVF_SQ", - json!({ - "metric_type": "l2", - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfSq(Default::default()), - ), - ( - "IVF_SQ", - json!({ - "metric_type": "cosine", - "num_partitions": 64, - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfSq( - IvfSqIndexBuilder::default() - .distance_type(DistanceType::Cosine) - .num_partitions(64), - ), - ), - ( - "IVF_RQ", - json!({ - "metric_type": "l2", - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfRq(Default::default()), - ), - ( - "IVF_RQ", - json!({ - "metric_type": "cosine", - "num_partitions": 64, - "num_bits": 8, - "sample_rate": 256, - "max_iterations": 50, - }), - Index::IvfRq( - IvfRqIndexBuilder::default() - .distance_type(DistanceType::Cosine) - .num_partitions(64) - .num_bits(8), - ), - ), - // HNSW_PQ isn't yet supported on SaaS - ("BTREE", json!({}), Index::BTree(Default::default())), - ("BITMAP", json!({}), Index::Bitmap(Default::default())), - ("ZONEMAP", json!({}), Index::ZoneMap(Default::default())), - ("NGRAM", json!({}), Index::NGram(Default::default())), - ( - "BLOOM_FILTER", - json!({}), - Index::BloomFilter(Default::default()), - ), - ("RTREE", json!({}), Index::RTree(Default::default())), - ( - "BLOOM_FILTER", - json!({"number_of_items": 4096, "probability": 0.01}), - Index::BloomFilter( - crate::index::scalar::BloomFilterIndexBuilder::default() - .number_of_items(4096) - .unwrap() - .probability(0.01) - .unwrap(), - ), - ), - ( - "RTREE", - json!({"page_size": 1024}), - Index::RTree( - crate::index::scalar::RTreeIndexBuilder::default() - .page_size(1024) - .unwrap(), - ), - ), - ( - "LABEL_LIST", - json!({}), - Index::LabelList(Default::default()), - ), - ( - "FTS", - serde_json::to_value(InvertedIndexParams::default()).unwrap(), - Index::FTS(Default::default()), - ), - ( - "FTS", - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["block_size"] = 256.into(); - body - }, - Index::FTS(InvertedIndexParams::default().block_size(256).unwrap()), - ), - ( - "FTS", - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["custom_stop_words"] = json!(["cat", " cat ", "CAT"]); - body - }, - Index::FTS(InvertedIndexParams::default().custom_stop_words(Some(vec![ - "cat".to_string(), - " cat ".to_string(), - "CAT".to_string(), - ]))), - ), - ( - "FTS", - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["document_granularity"] = "list_element".into(); - body - }, - Index::FTS( - InvertedIndexParams::default() - .document_granularity(DocumentGranularity::ListElement), - ), - ), - ]; - - for (index_type, expected_body, index) in cases { - let table = Table::new_with_handler_version( - "my_table", - semver::Version::new(0, 6, 0), - move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => { - 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 = expected_body.clone(); - expected_body["column"] = "a".into(); - expected_body[INDEX_TYPE_KEY] = index_type.into(); - - assert_eq!(body, expected_body); - - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }, - ); - - table.create_index(&["a"], index).execute().await.unwrap(); - } - } - - #[tokio::test] - async fn test_create_index_forwards_replace_false_on_existing_route() { - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => { - let body = request.body().unwrap().as_bytes().unwrap(); - let body: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(body["replace"], json!(false)); - - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }); - - table - .create_index(&["a"], Index::BTree(Default::default())) - .replace(false) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_index_returns_job() { - let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let describe_calls_in_handler = describe_calls.clone(); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => http::Response::builder() - .status(200) - .body(r#"{"job_id": "job-123"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => { - let body = request.body().unwrap().as_bytes().unwrap(); - let body: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(body["job_id"], "job-123"); - let state = if describe_calls_in_handler - .fetch_add(1, std::sync::atomic::Ordering::SeqCst) - == 0 - { - "IN_PROGRESS" - } else { - "DONE" - }; - http::Response::builder() - .status(200) - .body(format!( - r#"{{"job_id": "job-123", "job_state": "{state}"}}"# - )) - .unwrap() - } - "/v1/jobs/cancel" => { - let body = request.body().unwrap().as_bytes().unwrap(); - let body: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(body["job_id"], "job-123"); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }); - - let job = table - .create_index(&["a"], Index::BTree(Default::default())) - .execute_async() - .await - .unwrap(); - assert_eq!(job.id(), Some("job-123")); - job.wait().await.unwrap(); - assert_eq!(describe_calls.load(std::sync::atomic::Ordering::SeqCst), 2); - job.cancel().await.unwrap(); - } - - /// An unrecognized state is treated as still running, so the client keeps - /// polling rather than reporting a wrong terminal outcome. - #[tokio::test] - async fn test_job_wait_treats_unknown_state_as_running() { - let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let calls_in_handler = calls.clone(); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => http::Response::builder() - .status(200) - .body(r#"{"job_id": "job-unknown"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => { - let state = if calls_in_handler - .fetch_add(1, std::sync::atomic::Ordering::SeqCst) - == 0 - { - "SOMETHING_NEW" - } else { - "DONE" - }; - http::Response::builder() - .status(200) - .body(format!( - r#"{{"job_id": "job-unknown", "job_state": "{state}"}}"# - )) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - }); - - let job = table - .create_index(&["a"], Index::BTree(Default::default())) - .execute_async() - .await - .unwrap(); - job.wait().await.unwrap(); - assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn test_job_wait_surfaces_failure() { - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => http::Response::builder() - .status(200) - .body(r#"{"job_id": "job-err"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(r#"{"job_id": "job-err", "job_state": "FAILED"}"#.to_string()) - .unwrap(), - path => panic!("Unexpected path: {}", path), - }); - - let job = table - .create_index(&["a"], Index::BTree(Default::default())) - .execute_async() - .await - .unwrap(); - let err = job.wait().await.unwrap_err(); - let crate::Error::JobFailed { failure, .. } = &err else { - panic!("expected JobFailed, got {err:?}"); - }; - // The server said only that it failed, so nothing may be invented. - assert!(failure.message.is_none(), "{failure:?}"); - assert!(failure.phase.is_none(), "{failure:?}"); - assert!(failure.retryable.is_none(), "{failure:?}"); - assert_eq!(err.to_string(), "Job job-err failed"); - } - - /// A server that reports why the job failed has that reason surfaced - /// verbatim rather than replaced with a generic message. - #[tokio::test] - async fn test_job_wait_reports_the_server_failure_reason() { - let table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => http::Response::builder() - .status(200) - .body(r#"{"job_id": "job-err"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body( - r#"{"job_id": "job-err", "job_state": "FAILED", "failure": {"phase": "commit", "message": "preempted", "retryable": true}}"# - .to_string(), - ) - .unwrap(), - path => panic!("Unexpected path: {}", path), - } - }); - - let job = table - .create_index(&["a"], Index::BTree(Default::default())) - .execute_async() - .await - .unwrap(); - let err = job.wait().await.unwrap_err(); - let crate::Error::JobFailed { failure, .. } = &err else { - panic!("expected JobFailed, got {err:?}"); - }; - assert_eq!(failure.message.as_deref(), Some("preempted")); - assert_eq!(failure.phase.as_deref(), Some("commit")); - assert_eq!(failure.retryable, Some(true)); - assert_eq!(err.to_string(), "Job job-err failed: preempted (in commit)"); - } - - /// Servers that return no job id (e.g. an empty create-index response) - /// yield an already-done job. - #[tokio::test] - async fn test_create_index_without_job_id_is_done() { - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/create_index/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - path => panic!("Unexpected path: {}", path), - }); - - let job = table - .create_index(&["a"], Index::BTree(Default::default())) - .execute_async() - .await - .unwrap(); - job.wait().await.unwrap(); - job.cancel().await.unwrap(); - } - - #[tokio::test] - async fn test_create_index_nested_field_paths() { - let schema = nested_index_schema(); - let expected_requests = Arc::new(vec![ - json!({ - "column": "rowId", - "index_type": "BTREE", - }), - json!({ - "column": "`row-id`", - "index_type": "BTREE", - }), - json!({ - "column": "userId", - "index_type": "BTREE", - }), - json!({ - "column": "MetaData.userId", - "index_type": "BTREE", - }), - json!({ - "column": "metadata.user_id", - "index_type": "BTREE", - }), - json!({ - "column": "image.embedding", - "index_type": "IVF_PQ", - "metric_type": "l2", - }), - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["column"] = "payload.text".into(); - body["index_type"] = "FTS".into(); - body - }, - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["column"] = "docs.content".into(); - body["index_type"] = "FTS".into(); - body - }, - { - let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); - body["column"] = "docs.content".into(); - body["index_type"] = "FTS".into(); - body["document_granularity"] = "list_element".into(); - body - }, - json!({ - "column": "`meta-data`.`user-id`", - "index_type": "BTREE", - }), - json!({ - "column": "literal.`a.b`", - "index_type": "BTREE", - }), - ]); - let request_idx = Arc::new(AtomicUsize::new(0)); - let table = Table::new_with_handler_version("my_table", semver::Version::new(0, 6, 0), { - let schema = schema.clone(); - let expected_requests = expected_requests.clone(); - let request_idx = request_idx.clone(); - move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(), - "/v1/table/my_table/create_index/" => { - assert_eq!( - request.headers().get("Content-Type").unwrap(), - JSON_CONTENT_TYPE - ); - let idx = request_idx.fetch_add(1, Ordering::SeqCst); - let body = request.body().unwrap().as_bytes().unwrap(); - let body: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(body, expected_requests[idx]); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - } - }); - - table - .create_index(&["rowId"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["`ROW-ID`"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["userId"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["MetaData.userId"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["Metadata.USER_ID"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["Image.Embedding"], Index::Auto) - .execute() - .await - .unwrap(); - table - .create_index(&["Payload.Text"], Index::FTS(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["Docs.Content"], Index::FTS(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index( - &["Docs.Content"], - Index::FTS( - InvertedIndexParams::default() - .document_granularity(DocumentGranularity::ListElement), - ), - ) - .execute() - .await - .unwrap(); - table - .create_index(&["`META-DATA`.`USER-ID`"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - table - .create_index(&["literal.`A.B`"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - - assert_eq!(request_idx.load(Ordering::SeqCst), expected_requests.len()); - } - - #[tokio::test] - async fn test_create_list_element_fts_requires_server_support() { - let table = Table::new_with_handler_version( - "my_table", - semver::Version::new(0, 5, 0), - |_| -> http::Response { - panic!("unsupported index creation must fail before sending a request") - }, - ); - - let result = table - .create_index( - &["docs.content"], - Index::FTS( - InvertedIndexParams::default() - .document_granularity(DocumentGranularity::ListElement), - ), - ) - .execute() - .await; - - assert!( - result - .unwrap_err() - .to_string() - .contains("document granularity requires remote server version 0.6.0 or later") - ); - } - - #[tokio::test] - async fn test_list_indices() { - let schema = Schema::new(vec![ - Field::new( - "vector", - DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8), - false, - ), - Field::new( - "metadata", - DataType::Struct(vec![Field::new("my.column", DataType::Utf8, true)].into()), - false, - ), - ]); - let 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/" => { - return http::Response::builder() - .status(200) - .body(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": ["metadata.`my.column`"], - "index_status": "done", - }, - ] - }) - } - "/v1/table/my_table/index/vector_idx/stats/" => { - serde_json::json!({ - "num_indexed_rows": 100000, - "num_unindexed_rows": 0, - "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": 0, - "index_type": "LABEL_LIST" - }) - } - path => panic!("Unexpected path: {}", path), - }; - http::Response::builder() - .status(200) - .body(serde_json::to_string(&response_body).unwrap()) - .unwrap() - }); - - let indices = table.list_indices().await.unwrap(); - let expected = vec![ - IndexConfig { - name: "vector_idx".into(), - index_type: IndexType::IvfPq, - columns: vec!["vector".into()], - index_uuid: None, - type_url: None, - created_at: None, - num_indexed_rows: None, - num_unindexed_rows: None, - size_bytes: None, - num_segments: None, - index_version: None, - index_details: None, - }, - IndexConfig { - name: "my_idx".into(), - index_type: IndexType::LabelList, - columns: vec!["metadata.`my.column`".into()], - index_uuid: None, - type_url: None, - created_at: None, - num_indexed_rows: None, - num_unindexed_rows: None, - size_bytes: None, - num_segments: None, - index_version: None, - index_details: None, - }, - ]; - assert_eq!(indices, expected); - } - - #[rstest] - #[case::legacy(false)] - #[case::enriched(true)] - #[tokio::test] - async fn test_list_indices_fts_public_list_path(#[case] enriched: bool) { - let schema = nested_index_schema(); - let table = Table::new_with_handler("my_table", move |request| { - let body = match request.url().path() { - "/v1/table/my_table/describe/" => describe_response(&schema), - "/v1/table/my_table/index/list/" => serde_json::json!({ - "indexes": [{ - "index_name": "docs_idx", - "columns": ["docs.item.content"], - "index_type": enriched.then_some("FTS"), - }], - }) - .to_string(), - "/v1/table/my_table/index/docs_idx/stats/" => { - assert!(!enriched, "enriched responses must not fetch index stats"); - serde_json::json!({ - "num_indexed_rows": 1, - "num_unindexed_rows": 0, - "index_type": "FTS", - }) - .to_string() - } - path => panic!("Unexpected path: {path}"), - }; - http::Response::builder().status(200).body(body).unwrap() - }); - - let indices = table.list_indices().await.unwrap(); - assert_eq!(indices.len(), 1); - assert_eq!(indices[0].index_type, IndexType::FTS); - assert_eq!(indices[0].columns, vec!["docs.content"]); - } - - #[tokio::test] - async fn test_list_indices_nested_field_paths() { - let schema = nested_index_schema(); - let 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/" => { - return http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(); - } - "/v1/table/my_table/index/list/" => { - serde_json::json!({ - "indexes": [ - { - "index_name": "row_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000001", - "columns": ["rowId"], - "index_status": "done", - }, - { - "index_name": "row_dash_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000002", - "columns": ["`ROW-ID`"], - "index_status": "done", - }, - { - "index_name": "user_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000003", - "columns": ["userId"], - "index_status": "done", - }, - { - "index_name": "mixed_case_metadata_user_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000004", - "columns": ["MetaData.userId"], - "index_status": "done", - }, - { - "index_name": "metadata_user_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000005", - "columns": ["Metadata.USER_ID"], - "index_status": "done", - }, - { - "index_name": "image_embedding_idx", - "index_uuid": "00000000-0000-0000-0000-000000000006", - "columns": ["Image.Embedding"], - "index_status": "done", - }, - { - "index_name": "payload_text_idx", - "index_uuid": "00000000-0000-0000-0000-000000000007", - "columns": ["Payload.Text"], - "index_status": "done", - }, - { - "index_name": "meta_data_user_id_idx", - "index_uuid": "00000000-0000-0000-0000-000000000008", - "columns": ["`META-DATA`.`USER-ID`"], - "index_status": "done", - }, - { - "index_name": "literal_dot_idx", - "index_uuid": "00000000-0000-0000-0000-000000000009", - "columns": ["literal.`A.B`"], - "index_status": "done", - }, - ] - }) - } - "/v1/table/my_table/index/row_id_idx/stats/" - | "/v1/table/my_table/index/row_dash_id_idx/stats/" - | "/v1/table/my_table/index/user_id_idx/stats/" - | "/v1/table/my_table/index/mixed_case_metadata_user_id_idx/stats/" - | "/v1/table/my_table/index/metadata_user_id_idx/stats/" - | "/v1/table/my_table/index/meta_data_user_id_idx/stats/" - | "/v1/table/my_table/index/literal_dot_idx/stats/" => { - serde_json::json!({ - "num_indexed_rows": 100000, - "num_unindexed_rows": 0, - "index_type": "BTREE" - }) - } - "/v1/table/my_table/index/image_embedding_idx/stats/" => { - serde_json::json!({ - "num_indexed_rows": 100000, - "num_unindexed_rows": 0, - "index_type": "IVF_PQ", - "distance_type": "l2" - }) - } - "/v1/table/my_table/index/payload_text_idx/stats/" => { - serde_json::json!({ - "num_indexed_rows": 100000, - "num_unindexed_rows": 0, - "index_type": "FTS" - }) - } - path => panic!("Unexpected path: {}", path), - }; - http::Response::builder() - .status(200) - .body(serde_json::to_string(&response_body).unwrap()) - .unwrap() - }); - - let indices = table.list_indices().await.unwrap(); - // The remote path leaves the rich metadata fields None until the server - // wires them through. See https://github.com/lancedb/lancedb/issues/3494 - let expected: Vec = [ - ("row_id_idx", IndexType::BTree, "rowId"), - ("row_dash_id_idx", IndexType::BTree, "`row-id`"), - ("user_id_idx", IndexType::BTree, "userId"), - ( - "mixed_case_metadata_user_id_idx", - IndexType::BTree, - "MetaData.userId", - ), - ("metadata_user_id_idx", IndexType::BTree, "metadata.user_id"), - ("image_embedding_idx", IndexType::IvfPq, "image.embedding"), - ("payload_text_idx", IndexType::FTS, "payload.text"), - ( - "meta_data_user_id_idx", - IndexType::BTree, - "`meta-data`.`user-id`", - ), - ("literal_dot_idx", IndexType::BTree, "literal.`a.b`"), - ] - .into_iter() - .map(|(name, index_type, column)| IndexConfig { - name: name.into(), - index_type, - columns: vec![column.into()], - index_uuid: None, - type_url: None, - created_at: None, - num_indexed_rows: None, - num_unindexed_rows: None, - size_bytes: None, - num_segments: None, - index_version: None, - index_details: None, - }) - .collect(); - assert_eq!(indices, expected); - } - - /// Verifies that when the server returns `index_type` in the list response, - /// `list_indices` uses all enriched fields directly and does **not** make a - /// per-index `/index/{name}/stats/` call. - #[tokio::test] - async fn test_list_indices_enriched() { - let schema = Schema::new(vec![ - Field::new( - "vector", - DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8), - false, - ), - Field::new("text", DataType::Utf8, false), - ]); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(), - "/v1/table/my_table/index/list/" => { - let body = serde_json::json!({ - "indexes": [ - { - "index_name": "vector_idx", - "index_uuid": "3fa85f64-5717-4562-b3fc-2c963f66afa6", - "columns": ["vector"], - "index_type": "IVF_PQ", - "index_status": "done", - "num_indexed_rows": 1000, - "num_unindexed_rows": 50, - "size_bytes": 204800, - "num_segments": 2, - "index_version": 1, - "index_details": "{\"num_partitions\":16}", - "created_at": "2026-06-18T21:37:36.637Z", - "type_url": "type.googleapis.com/lance.index.vector.IvfPq", - }, - { - "index_name": "text_idx", - "index_uuid": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", - "columns": ["text"], - "index_type": "FTS", - "index_status": "done", - "num_indexed_rows": 1000, - "num_unindexed_rows": 0, - "size_bytes": 8192, - "num_segments": 1, - }, - ] - }); - http::Response::builder() - .status(200) - .body(serde_json::to_string(&body).unwrap()) - .unwrap() - } - // stats endpoint must NOT be called for enriched responses - path => panic!("Unexpected path (stats should not be called): {}", path), - } - }); - - let indices = table.list_indices().await.unwrap(); - assert_eq!(indices.len(), 2); - - let vec_idx = &indices[0]; - assert_eq!(vec_idx.name, "vector_idx"); - assert_eq!(vec_idx.index_type, IndexType::IvfPq); - assert_eq!(vec_idx.columns, vec!["vector".to_string()]); - assert_eq!( - vec_idx.index_uuid, - Some("3fa85f64-5717-4562-b3fc-2c963f66afa6".to_string()) - ); - assert_eq!(vec_idx.num_indexed_rows, Some(1000)); - assert_eq!(vec_idx.num_unindexed_rows, Some(50)); - assert_eq!(vec_idx.size_bytes, Some(204800)); - assert_eq!(vec_idx.num_segments, Some(2)); - assert_eq!(vec_idx.index_version, Some(1)); - assert_eq!( - vec_idx.index_details, - Some("{\"num_partitions\":16}".to_string()) - ); - assert_eq!( - vec_idx.type_url, - Some("type.googleapis.com/lance.index.vector.IvfPq".to_string()) - ); - assert_eq!( - vec_idx.created_at, - Some("2026-06-18T21:37:36.637Z".parse::>().unwrap()) - ); - - let text_idx = &indices[1]; - assert_eq!(text_idx.name, "text_idx"); - assert_eq!(text_idx.index_type, IndexType::FTS); - assert_eq!(text_idx.columns, vec!["text".to_string()]); - assert_eq!( - text_idx.index_uuid, - Some("aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee".to_string()) - ); - assert_eq!(text_idx.num_indexed_rows, Some(1000)); - assert_eq!(text_idx.num_unindexed_rows, Some(0)); - assert_eq!(text_idx.size_bytes, Some(8192)); - assert_eq!(text_idx.num_segments, Some(1)); - // optional fields not present in the response are None - assert_eq!(text_idx.index_version, None); - assert_eq!(text_idx.index_details, None); - assert_eq!(text_idx.type_url, None); - assert_eq!(text_idx.created_at, None); - } - - #[tokio::test] - async fn test_tokenize_uses_remote_index_details() { - let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); - let index_details = serde_json::json!({ - "base_tokenizer": "icu", - "language": "English", - "with_position": false, - "max_token_length": 40, - "lower_case": true, - "stem": false, - "remove_stop_words": true, - "ascii_folding": true, - "custom_stop_words": ["hello"], - }) - .to_string(); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(), - "/v1/table/my_table/index/list/" => { - let body = serde_json::json!({ - "indexes": [ - { - "index_name": "text_idx", - "columns": ["text"], - "index_type": "FTS", - "index_details": index_details, - }, - ] - }); - http::Response::builder() - .status(200) - .body(serde_json::to_string(&body).unwrap()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }); - - let tokens = table - .tokenize("Hello, こんにちは世界!", "text_idx") - .await - .unwrap(); - - // Positions are relative to the first retained token, so dropping the - // leading "hello" stop word does not shift the remaining tokens. - assert_eq!( - tokens, - vec![ - FtsToken { - text: "こんにちは".to_string(), - position: 0, - }, - FtsToken { - text: "世界".to_string(), - position: 1, - }, - ] - ); - } - - #[tokio::test] - async fn test_tokenize_requires_existing_index_name() { - let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); - let table = Table::new_with_handler("my_table", move |request| -> http::Response { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(), - "/v1/table/my_table/index/list/" => { - let body = serde_json::json!({ "indexes": [] }); - http::Response::builder() - .status(200) - .body(serde_json::to_string(&body).unwrap()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }); - - let err = table.tokenize("hello", "text_idx").await.unwrap_err(); - assert!(matches!( - err, - Error::InvalidInput { message } - if message.contains("No index named 'text_idx'") - )); - } - - #[tokio::test] - async fn test_tokenize_with_column_remote_requires_index_details() { - let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.method(), "POST"); - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap(), - "/v1/table/my_table/index/list/" => { - let body = serde_json::json!({ - "indexes": [ - { - "index_name": "text_idx", - "columns": ["text"], - "index_type": "FTS", - }, - ] - }); - http::Response::builder() - .status(200) - .body(serde_json::to_string(&body).unwrap()) - .unwrap() - } - path => panic!("Unexpected path: {}", path), - } - }); - - let err = table - .tokenize_with_column("hello", "text") - .await - .unwrap_err(); - - assert!(matches!( - err, - Error::InvalidInput { message } - if message.contains("does not include tokenizer details") - )); - } - - #[test] - fn test_deserialize_created_at() { - #[derive(Deserialize)] - struct Wrapper { - #[serde(default, deserialize_with = "deserialize_created_at")] - created_at: Option>, - } - - // RFC 3339 string (current server format). - let w: Wrapper = - serde_json::from_str(r#"{"created_at": "2026-06-18T21:37:36.637Z"}"#).unwrap(); - assert_eq!( - w.created_at, - Some("2026-06-18T21:37:36.637Z".parse::>().unwrap()) - ); - - // Unix milliseconds (legacy server format). - let w: Wrapper = serde_json::from_str(r#"{"created_at": 1700000000000}"#).unwrap(); - assert_eq!(w.created_at, DateTime::from_timestamp_millis(1700000000000)); - - // Null and missing both yield None. - let w: Wrapper = serde_json::from_str(r#"{"created_at": null}"#).unwrap(); - assert_eq!(w.created_at, None); - let w: Wrapper = serde_json::from_str(r#"{}"#).unwrap(); - assert_eq!(w.created_at, None); - - // A malformed string is rejected rather than silently dropped to None. - assert!(serde_json::from_str::(r#"{"created_at": "not-a-date"}"#).is_err()); - } - - #[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::>().unwrap() - ); - assert_eq!(versions[1].version, 2); - assert_eq!( - versions[1].timestamp, - "2024-02-01T00:00:00Z".parse::>().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::>().unwrap() - ); - assert!(versions[0].metadata.is_empty()); - assert_eq!(versions[1].version, 2); - assert_eq!( - versions[1].timestamp, - "2024-02-01T00:00:00Z".parse::>().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::("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 = 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 { .. }))); - } - - #[rstest] - #[case(true)] - #[case(false)] - #[tokio::test] - async fn test_add_columns(#[case] old_server: bool) { - let table = Table::new_with_handler("my_table", move |request| { - if request.url().path() == "/v1/table/my_table/describe/" { - simple_describe_response() - } else if request.url().path() == "/v1/table/my_table/add_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 new_columns = value.get("new_columns").unwrap().as_array().unwrap(); - assert!(new_columns.len() == 2); - - let col_name = new_columns[0]["name"].as_str().unwrap(); - let expression = new_columns[0]["expression"].as_str().unwrap(); - assert_eq!(col_name, "b"); - assert_eq!(expression, "a + 1"); - - let col_name = new_columns[1]["name"].as_str().unwrap(); - let expression = new_columns[1]["expression"].as_str().unwrap(); - assert_eq!(col_name, "x"); - assert_eq!(expression, "cast(NULL as int32)"); - - 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 - .add_columns() - .transform(NewColumnTransform::SqlExpressions(vec![ - ("b".into(), "a + 1".into()), - ("x".into(), "cast(NULL as int32)".into()), - ])) - .execute() - .await - .unwrap(); - - assert_eq!(result.version, if old_server { 0 } else { 43 }); - } - - /// A pyarrow schema goes over the wire as an Arrow IPC schema under an - /// Arrow content type, not as a JSON description of the types, so a - /// field arrives exactly as sent. The branch has no JSON envelope to - /// ride in and goes on the query string. - #[tokio::test] - async fn test_add_columns_sends_an_arrow_schema_for_all_nulls() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/add_columns/" => { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.headers().get("Content-Type").unwrap(), - ARROW_STREAM_CONTENT_TYPE - ); - assert_eq!( - request.url().query_pairs().find(|(k, _)| k == "branch"), - None, - "the main branch sends no branch parameter" - ); - - // The payload decodes back to the exact schema, types and - // field metadata included. - let body = request.body().unwrap().as_bytes().unwrap(); - let reader = - arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None) - .unwrap(); - let sent = reader.schema(); - assert_eq!(sent.fields().len(), 2); - assert_eq!(sent.field(0).name(), "ts"); - assert_eq!( - sent.field(0).data_type(), - &DataType::Timestamp( - arrow_schema::TimeUnit::Nanosecond, - Some("America/New_York".into()) - ) - ); - assert!(sent.field(0).is_nullable()); - assert_eq!(sent.field(1).name(), "price"); - assert_eq!(sent.field(1).data_type(), &DataType::Decimal128(38, 10)); - assert_eq!( - sent.field(1).metadata().get("unit").map(String::as_str), - Some("cents") - ); - - http::Response::builder() - .status(200) - .body(r#"{"version": 7}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected request path: {path}"), - }); - - let schema = Arc::new(Schema::new(vec![ - Field::new( - "ts", - DataType::Timestamp( - arrow_schema::TimeUnit::Nanosecond, - Some("America/New_York".into()), - ), - true, - ), - Field::new("price", DataType::Decimal128(38, 10), true).with_metadata( - std::collections::HashMap::from([("unit".to_string(), "cents".to_string())]), - ), - ])); - - let result = table - .add_columns() - .transform(NewColumnTransform::AllNulls(schema)) - .execute() - .await - .unwrap(); - - assert_eq!(result.version, 7); - } - - /// A declaration is sent as `{name, computed}` for the server to plan; the - /// client never types the expression itself. - #[tokio::test] - async fn test_add_computed_columns_sends_the_expression() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/add_columns/" => { - assert_eq!(request.method(), "POST"); - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!( - value["new_columns"], - serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) - ); - http::Response::builder() - .status(200) - .body(r#"{"version": 7}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - }); - - let result = table - .add_columns() - .computed("doubled", "x * 2") - .execute() - .await - .unwrap(); - assert_eq!(result.version, 7); - } - - #[tokio::test] - async fn test_add_scalar_function_column_sends_atomic_null_declaration() { - let table = Table::new_with_handler("my_table", |request| { - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body( - r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#, - ) - .unwrap(), - "/v1/table/my_table/add_columns/" => { - let actual: serde_json::Value = serde_json::from_slice( - request.body().unwrap().as_bytes().unwrap(), - ) - .unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json" - )) - .unwrap(); - assert_eq!(actual, expected); - http::Response::builder() - .status(200) - .body(r#"{"version":8}"#) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - } - }); - let application = crate::function::FunctionApplication::from_json( - r#"{ - "function":{"name":"embed","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs":[{"parameter":"text","kind":"column","value":{"path":"description"}}], - "output":{"kind":"scalar","arrow_type":"list","nullable":false} - }"#, - ) - .unwrap(); - - let result = table - .add_columns() - .function_as("embedding", application) - .execute() - .await - .unwrap(); - assert_eq!(result.version, 8); - } - - /// The fixture binding's table: `title` and `body` bound as inputs, its - /// two outputs declared, plus an unbound `spare`. - fn fixture_bound_schema() -> Schema { - let binding = crate::function::FunctionBinding::from_json(include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" - )) - .unwrap(); - let binding_metadata = crate::table::computed_columns::function_bindings_metadata( - std::slice::from_ref(&binding), - ) - .unwrap(); - let mut fields = vec![ - Field::new("title", DataType::Utf8, true), - Field::new("body", DataType::Utf8, true), - ]; - fields.extend(binding.outputs().iter().map(|output| { - let data_type = match output.arrow_type.as_str() { - "utf8" => DataType::Utf8, - "int64" => DataType::Int64, - other => panic!("unexpected fixture output type {other}"), - }; - Field::new(&output.output_name, data_type, true).with_metadata( - crate::table::computed_columns::function_computed_column_metadata( - binding.binding_id(), - output.output_ordinal, - &["title".into(), "body".into()], - ), - ) - })); - fields.push(Field::new("spare", DataType::Int32, true)); - Schema::new_with_metadata( - fields, - HashMap::from([( - crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(), - binding_metadata, - )]), - ) - } - - /// Only a column the binding uses is refused, and it is refused before - /// any request goes out; the rest reach the server as usual. - #[tokio::test] - async fn test_add_columns_scopes_to_the_columns_a_binding_uses() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&fixture_bound_schema())) - .unwrap(), - "/v1/table/my_table/add_columns/" => http::Response::builder() - .status(200) - .body(r#"{"version":10}"#.to_string()) - .unwrap(), - path => panic!("Unexpected path: {path}"), - }); - table - .add_columns() - .computed("doubled", "spare * 2") - .execute() - .await - .unwrap(); - table - .add_columns() - .transform(NewColumnTransform::SqlExpressions(vec![( - "eager".into(), - "spare + 1".into(), - )])) - .execute() - .await - .unwrap(); - - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&fixture_bound_schema())) - .unwrap(), - path => panic!("mutation request must not be sent: {path}"), - }); - for name in ["title", "search_text"] { - let err = table - .add_columns() - .computed(name, "1") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); - let err = table - .add_columns() - .transform(NewColumnTransform::SqlExpressions(vec![( - name.into(), - "1".into(), - )])) - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); - } - } - - #[tokio::test] - async fn test_add_function_column_allows_an_existing_binding() { - let schema = fixture_bound_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_response(&schema)) - .unwrap(), - "/v1/table/my_table/add_columns/" => { - let actual: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) - .unwrap(); - assert_eq!( - actual["new_columns"], - serde_json::json!([ - {"name":"secondary_text","all_null":true}, - {"name":"secondary_token_count","all_null":true} - ]) - ); - http::Response::builder() - .status(200) - .body(r#"{"version":10}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - }); - let application = crate::function::FunctionApplication::from_json( - r#"{ - "function":{"name":"text_features","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs":[ - {"parameter":"title","kind":"column","value":{"path":"title"}}, - {"parameter":"body","kind":"column","value":{"path":"body"}} - ], - "output":{"kind":"named_struct","fields":[ - {"name":"normalized_text","arrow_type":"utf8","nullable":false}, - {"name":"token_count","arrow_type":"int64","nullable":false} - ]}, - "columns":{ - "normalized_text":"secondary_text", - "token_count":"secondary_token_count" - } - }"#, - ) - .unwrap(); - - let result = table - .add_columns() - .function(application) - .execute() - .await - .unwrap(); - assert_eq!(result.version, 10); - } - - #[tokio::test] - async fn test_add_fixed_size_list_function_column_declares_the_vector_type() { - let table = Table::new_with_handler("my_table", |request| { - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body( - r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#, - ) - .unwrap(), - "/v1/table/my_table/add_columns/" => { - let actual: serde_json::Value = serde_json::from_slice( - request.body().unwrap().as_bytes().unwrap(), - ) - .unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" - )) - .unwrap(); - assert_eq!(actual, expected); - http::Response::builder() - .status(200) - .body(r#"{"version":8}"#) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - } - }); - let application = crate::function::FunctionApplication::from_json( - r#"{ - "function":{"name":"embed","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs":[{"parameter":"text","kind":"column","value":{"path":"description"}}], - "output":{"kind":"scalar","arrow_type":"fixed_size_list","nullable":false} - }"#, - ) - .unwrap(); - - let result = table - .add_columns() - .function_as("embedding", application) - .execute() - .await - .unwrap(); - assert_eq!(result.version, 8); - } - - #[tokio::test] - async fn test_add_named_struct_function_expands_one_atomic_binding() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body( - r#"{"version":1,"schema":{"fields":[ - {"name":"title","nullable":true,"type":{"type":"string"}}, - {"name":"body","nullable":true,"type":{"type":"string"}} - ]}}"#, - ) - .unwrap(), - "/v1/table/my_table/add_columns/" => { - let actual: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_multi_output_declaration_request.json" - )) - .unwrap(); - assert_eq!(actual, expected); - http::Response::builder() - .status(200) - .body(r#"{"version":9}"#) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - }); - let application = crate::function::FunctionApplication::from_json( - r#"{ - "function":{"name":"text_features","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs":[ - {"parameter":"title","kind":"column","value":{"path":"title"}}, - {"parameter":"body","kind":"column","value":{"path":"body"}} - ], - "output":{"kind":"named_struct","fields":[ - {"name":"normalized_text","arrow_type":"utf8","nullable":false}, - {"name":"token_count","arrow_type":"int64","nullable":false} - ]}, - "columns":{"normalized_text":"search_text"} - }"#, - ) - .unwrap(); - - let result = table - .add_columns() - .function(application) - .execute() - .await - .unwrap(); - assert_eq!(result.version, 9); - } - - #[tokio::test] - async fn test_add_columns_fails_closed_on_newer_function_binding_metadata() { - let table = Table::new_with_handler("my_table", |request| { - match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body( - r#"{"version":1,"schema":{"fields":[{"name":"x","nullable":true,"type":{"type":"int32"}}],"metadata":{"lancedb::function_bindings":"{\"version\":2,\"bindings\":[]}"}}}"#, - ) - .unwrap(), - path => panic!("mutation request must not be sent: {path}"), - } - }); - - let err = table - .add_columns() - .computed("doubled", "x * 2") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::NotSupported { .. })); - } - - /// A remote refresh is a server job: the async form returns its handle, - /// and the blocking form refuses rather than invent a fill count. - #[tokio::test] - async fn test_refresh_column_async_submits_a_backfill_job() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column"); - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(value["column"], "doubled"); - http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-42"}"#) - .unwrap() - }); - - let job = table.refresh_column_async("doubled").await.unwrap(); - assert_eq!(job.id(), Some("j-42")); - - let err = table.refresh_column("doubled").await.unwrap_err(); - assert!( - matches!(&err, Error::NotSupported { message } - if message.contains("refresh_column_async")), - "{err:?}" - ); - } - - #[rstest] - #[case(false, "Function is unavailable")] - #[case(true, "Function name was not found")] - #[tokio::test] - async fn test_function_column_operations_preserve_dependency_not_found( - #[case] declare_column: bool, - #[case] message: &'static str, - ) { - let schema_requests = Arc::new(AtomicUsize::new(0)); - let requests = schema_requests.clone(); - let table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/describe/" => { - requests.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#.to_string(), - ) - .unwrap() - } - "/v1/table/my_table/backfill_column" | "/v1/table/my_table/add_columns/" => { - http::Response::builder() - .status(404) - .body(json!({"code": 4, "error": message}).to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - } - }); - table.schema().await.unwrap(); - assert_eq!(schema_requests.load(Ordering::SeqCst), 1); - - let error = if declare_column { - let fixture: serde_json::Value = serde_json::from_str(include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" - )) - .unwrap(); - let application = crate::function::FunctionApplication::from_json( - &fixture["function"]["application"].to_string(), - ) - .unwrap(); - table - .add_columns() - .function_as("embedding", application) - .execute() - .await - .unwrap_err() - } else { - table.refresh_column_async("embedding").await.unwrap_err() - }; - let Error::Http { - source, - status_code, - request_id, - } = error - else { - panic!("dependency 404 was misclassified: {error:?}"); - }; - assert_eq!(status_code, Some(StatusCode::NOT_FOUND)); - assert!(source.to_string().contains(message)); - assert!(!request_id.is_empty()); - - table.schema().await.unwrap(); - assert_eq!(schema_requests.load(Ordering::SeqCst), 2); - } - - /// The error listing is table-addressed with optional job and column - /// filters, mirroring the server's SQL surface, and the two non-record - /// signals come back as their own fields rather than as rows. - #[tokio::test] - async fn test_function_errors_lists_the_rows_a_refresh_skipped() { - use crate::function::{FunctionErrorFragment, FunctionErrorRecord, FunctionErrorsRequest}; - - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/errors"); - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!( - value, - serde_json::json!({"job_id": "j-7", "column": "embedding", "limit": 2}) - ); - http::Response::builder() - .status(200) - .body( - r#"{"records": [{"job_id": "j-7", "fragment_id": 3, "row_offset": 9, - "column": "embedding", "function": "embed", "function_version": "2", - "table_version": 11, "error_type": "ValueError", - "error_message": "bad input 'x'", "created_at_millis": 1700000000000}], - "fragments": [{"job_id": "j-7", "fragment_id": 4, "rows_skipped": 500, - "rows_recorded": 100}], - "truncated": true}"#, - ) - .unwrap() - }); - - let errors = table - .function_errors( - FunctionErrorsRequest::new() - .job_id("j-7") - .column("embedding") - .limit(2), - ) - .await - .unwrap(); - assert_eq!( - errors.records, - [FunctionErrorRecord { - job_id: "j-7".into(), - fragment_id: 3, - row_offset: Some(9), - column: "embedding".into(), - function: "embed".into(), - function_version: "2".into(), - table_version: 11, - error_type: "ValueError".into(), - error_message: "bad input 'x'".into(), - created_at_millis: 1_700_000_000_000, - }] - ); - assert_eq!( - errors.fragments, - [FunctionErrorFragment { - job_id: "j-7".into(), - fragment_id: 4, - rows_skipped: 500, - rows_recorded: 100, - }] - ); - assert!(errors.truncated); - - // No filter sends no filter, and an empty listing reads as such. - let table = Table::new_with_handler("my_table", |request| { - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(value, serde_json::json!({})); - http::Response::builder() - .status(200) - .body(r#"{"records": []}"#) - .unwrap() - }); - let errors = table - .function_errors(FunctionErrorsRequest::new()) - .await - .unwrap(); - assert_eq!(errors, crate::function::FunctionErrors::default()); - } - - /// The refresh handle is wrapped for read-freshness tracking, so it has to - /// forward the detail APIs too -- this is the job an operator is holding - /// when a backfill goes quiet. - #[tokio::test] - async fn test_refresh_job_handle_reports_detail_and_events() { - 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 table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/backfill_column" => http::Response::builder() - .status(202) - .body(br#"{"job_id": "j-42"}"#.to_vec()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body( - r#"{"job_id": "j-42", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 7, "spec": {"column": "doubled"}}"# - .as_bytes() - .to_vec(), - ) - .unwrap(), - "/v1/jobs/query_events" => { - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) - .unwrap(); - assert_eq!(body["job_id"], "j-42"); - http::Response::builder() - .status(200) - .body(events.clone()) - .unwrap() - } - other => panic!("unexpected path {other}"), - } - }); - - let job = table.refresh_column_async("doubled").await.unwrap(); - job.refresh().await.unwrap(); - assert_eq!(job.state().as_deref(), Some("running")); - assert_eq!(job.job_type().as_deref(), Some("refresh_column")); - assert_eq!(job.creation_ms(), Some(7)); - assert_eq!(job.spec().unwrap()["column"], "doubled"); - - let batches = job - .events(crate::job::JobEventsRequest::default()) - .await - .unwrap(); - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); - } - - #[tokio::test] - async fn test_refresh_submission_uses_add_columns_version_fence() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/add_columns/" => http::Response::builder() - .status(200) - .body(r#"{"version": 7}"#.to_string()) - .unwrap(), - "/v1/table/my_table/backfill_column" => { - let min_version = request - .headers() - .get("x-lancedb-min-version") - .and_then(|value| value.to_str().ok()); - if min_version != Some("7") { - return http::Response::builder() - .status(400) - .body(r#"{"error":"Column not found: doubled"}"#.to_string()) - .unwrap(); - } - http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-43"}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - }); - - let result = table - .add_columns() - .computed("doubled", "a * 2") - .execute() - .await - .unwrap(); - assert_eq!(result.version, 7); - - let job = table.refresh_column_async("doubled").await.unwrap(); - assert_eq!(job.id(), Some("j-43")); - } - - /// The gate's reproducer: after a successful wait, a same-handle read - /// must carry the exact published version so a stale server cache cannot - /// serve the pre-backfill snapshot. - #[tokio::test] - async fn test_backfill_wait_establishes_read_freshness() { - let saw_published_version = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let saw = saw_published_version.clone(); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/backfill_column" => http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-7"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(refresh_done("j-7")) - .unwrap(), - "/v1/table/my_table/count_rows/" => { - saw.store( - request - .headers() - .get("x-lancedb-min-read-version") - .and_then(|value| value.to_str().ok()) - == Some("8"), - std::sync::atomic::Ordering::SeqCst, - ); - http::Response::builder() - .status(200) - .body("1".to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - }); - - let job = table.refresh_column_async("doubled").await.unwrap(); - let result = job.wait().await.unwrap(); - assert_eq!(result.rows_assigned, 12); - assert_eq!(result.published_version, Some(8)); - table.count_rows(None).await.unwrap(); - assert!( - saw_published_version.load(std::sync::atomic::Ordering::SeqCst), - "read after wait did not carry the published version" - ); - } - - /// A checkout after submission wins over the completion fence: the - /// pinned view must not regain a timestamp floor from the job. - #[tokio::test] - async fn test_checkout_after_submit_beats_the_completion_fence() { - let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let saw = saw_min_timestamp.clone(); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/backfill_column" => http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-8"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(refresh_done("j-8")) - .unwrap(), - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/count_rows/" => { - saw.store( - request.headers().contains_key("x-lancedb-min-timestamp"), - std::sync::atomic::Ordering::SeqCst, - ); - http::Response::builder() - .status(200) - .body("1".to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - }); - - let job = table.refresh_column_async("doubled").await.unwrap(); - table.checkout(3).await.unwrap(); - job.wait().await.unwrap(); - table.count_rows(None).await.unwrap(); - assert!( - !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), - "completion fence overrode an explicit checkout" - ); - } - - /// Tag checkout resets freshness state wholesale; the fence must not - /// survive it. - #[tokio::test] - async fn test_tag_checkout_after_submit_beats_the_completion_fence() { - let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let saw = saw_min_timestamp.clone(); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/backfill_column" => http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-9"}"#.to_string()) - .unwrap(), - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(refresh_done("j-9")) - .unwrap(), - "/v1/table/my_table/tags/version/" => http::Response::builder() - .status(200) - .body(r#"{"version": 5}"#.to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/count_rows/" => { - saw.store( - request.headers().contains_key("x-lancedb-min-timestamp"), - std::sync::atomic::Ordering::SeqCst, - ); - http::Response::builder() - .status(200) - .body("1".to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - }); - - let job = table.refresh_column_async("doubled").await.unwrap(); - table.checkout_tag("v1").await.unwrap(); - job.wait().await.unwrap(); - table.count_rows(None).await.unwrap(); - assert!( - !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), - "completion fence overrode a tag checkout" - ); - } - - /// A checkout landing while the submission request is in flight advances - /// the epoch past the token captured at submit. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_checkout_during_submission_beats_the_completion_fence() { - let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let saw = saw_min_timestamp.clone(); - 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 table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/backfill_column" => { - // Signal arrival, then hold the response until the - // test's checkout completes. - arrived_tx.lock().unwrap().send(()).unwrap(); - release_rx - .lock() - .unwrap() - .recv_timeout(std::time::Duration::from_secs(10)) - .unwrap(); - http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-10"}"#.to_string()) - .unwrap() - } - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(refresh_done("j-10")) - .unwrap(), - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - "/v1/table/my_table/count_rows/" => { - saw.store( - request.headers().contains_key("x-lancedb-min-timestamp"), - std::sync::atomic::Ordering::SeqCst, - ); - http::Response::builder() - .status(200) - .body("1".to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - } - }); - - let submit = tokio::spawn({ - let table = table.clone(); - async move { table.refresh_column_async("doubled").await } - }); - tokio::task::spawn_blocking(move || { - arrived_rx - .recv_timeout(std::time::Duration::from_secs(10)) - .unwrap() - }) - .await - .unwrap(); - table.checkout(7).await.unwrap(); - release_tx.send(()).unwrap(); - - let job = submit.await.unwrap().unwrap(); - job.wait().await.unwrap(); - table.count_rows(None).await.unwrap(); - assert!( - !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), - "completion fence overrode a checkout that landed mid-submission" - ); - } - - /// checkout_latest keeps the handle on latest, so a completed backfill - /// must retain the checkout timestamp and add its exact published version. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_checkout_latest_during_submission_keeps_the_fence() { - let seen_headers = Arc::new(std::sync::Mutex::new(None::)); - let saw = seen_headers.clone(); - 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 table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/backfill_column" => { - arrived_tx.lock().unwrap().send(()).unwrap(); - release_rx - .lock() - .unwrap() - .recv_timeout(std::time::Duration::from_secs(10)) - .unwrap(); - http::Response::builder() - .status(202) - .body(r#"{"job_id": "j-11"}"#.to_string()) - .unwrap() - } - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(refresh_done("j-11")) - .unwrap(), - "/v1/table/my_table/count_rows/" => { - *saw.lock().unwrap() = Some(request.headers().clone()); - http::Response::builder() - .status(200) - .body("1".to_string()) - .unwrap() - } - path => panic!("unexpected request: {path}"), - }); - - let submit = tokio::spawn({ - let table = table.clone(); - async move { table.refresh_column_async("doubled").await } - }); - tokio::task::spawn_blocking(move || { - arrived_rx - .recv_timeout(std::time::Duration::from_secs(10)) - .unwrap() - }) - .await - .unwrap(); - table.checkout_latest().await.unwrap(); - release_tx.send(()).unwrap(); - - let job = submit.await.unwrap().unwrap(); - job.wait().await.unwrap(); - table.count_rows(None).await.unwrap(); - let headers = seen_headers.lock().unwrap().clone().expect("no request"); - assert!(headers.contains_key("x-lancedb-min-timestamp")); - assert_eq!( - headers - .get("x-lancedb-min-read-version") - .and_then(|value| value.to_str().ok()), - Some("8") - ); - } - - #[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>, - ) { - 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()); - } - - /// 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::>(), - "compacting": compacting, - "memtables": [], - }], - } - }) - .to_string() - } - - /// `flush_lsm` / `compact_lsm` answer 202 with no body at all. - fn accepted() -> http::Response { - http::Response::builder() - .status(202) - .body(String::new()) - .unwrap() - } - - fn ok_json(body: String) -> http::Response { - 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 = [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 = 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::(&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::) - } - }; - 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() - }) - } - - #[tokio::test] - async fn test_table_with_namespace_identifier() { - // Test that a table created with namespace uses the correct identifier in API calls - let table = Table::new_with_handler("ns1$ns2$table1", |request| { - assert_eq!(request.method(), "POST"); - // All API calls should use the full identifier in the path - assert_eq!(request.url().path(), "/v1/table/ns1$ns2$table1/describe/"); - - http::Response::builder() - .status(200) - .body(r#"{"version": 1, "schema": { "fields": [] }}"#) - .unwrap() - }); - - // The name() method should return just the base name, not the full identifier - assert_eq!(table.name(), "ns1$ns2$table1"); - - // API operations should work correctly - let version = table.version().await.unwrap(); - assert_eq!(version, 1); - } - - #[tokio::test] - async fn test_query_with_namespace() { - let table = Table::new_with_handler("analytics$events", |request| { - match request.url().path() { - "/v1/table/analytics$events/query/" => { - assert_eq!(request.method(), "POST"); - - // Return empty arrow stream - let data = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - let body = write_ipc_file(&data); - - http::Response::builder() - .status(200) - .header("Content-Type", ARROW_FILE_CONTENT_TYPE) - .body(body) - .unwrap() - } - _ => { - panic!("Unexpected path: {}", request.url().path()); - } - } - }); - - let results = table.query().execute().await.unwrap(); - let batches = results.try_collect::>().await.unwrap(); - assert_eq!(batches.len(), 1); - assert_eq!(batches[0].num_rows(), 3); - } - - #[tokio::test] - async fn test_add_data_with_namespace() { - 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("prod$metrics", move |mut request| { - match request.url().path() { - "/v1/table/prod$metrics/describe/" => http::Response::builder() - .status(200) - .body(describe_body.clone()) - .unwrap(), - "/v1/table/prod$metrics/insert/" => { - assert_eq!(request.method(), "POST"); - 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(r#"{"version": 2}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected request path: {}", path), - } - }); - - let result = table.add(data.clone()).execute().await.unwrap(); - - assert_eq!(result.version, 2); - - 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_create_index_with_namespace() { - let table = Table::new_with_handler("dev$users", |request| { - match request.url().path() { - "/v1/table/dev$users/create_index/" => { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.headers().get("Content-Type").unwrap(), - JSON_CONTENT_TYPE - ); - - // Verify the request body contains the column name - 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(); - assert_eq!(value["column"], "embedding"); - assert_eq!(value["index_type"], "IVF_PQ"); - } - - http::Response::builder() - .status(200) - .body("".to_string()) - .unwrap() - } - "/v1/table/dev$users/describe/" => { - let schema = Schema::new(vec![Field::new( - "embedding", - DataType::FixedSizeList( - Arc::new(Field::new("item", DataType::Float32, true)), - 8, - ), - false, - )]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } - _ => { - panic!("Unexpected path: {}", request.url().path()); - } - } - }); - - table - .create_index(&["embedding"], Index::IvfPq(IvfPqIndexBuilder::default())) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_drop_columns_with_namespace() { - let table = Table::new_with_handler("test$schema_ops", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/table/test$schema_ops/drop_columns/" - ); - 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 columns = value["columns"].as_array().unwrap(); - assert_eq!(columns.len(), 2); - assert_eq!(columns[0], "old_col1"); - assert_eq!(columns[1], "old_col2"); - } - - http::Response::builder() - .status(200) - .body(r#"{"version": 5}"#) - .unwrap() - }); - - let result = table.drop_columns(&["old_col1", "old_col2"]).await.unwrap(); - assert_eq!(result.version, 5); - } - - #[tokio::test] - async fn test_uri() { - 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": 1, "schema": {"fields": []}, "location": "s3://bucket/path/to/table"}"#) - .unwrap() - }); - - let uri = table.uri().await.unwrap(); - assert_eq!(uri, "s3://bucket/path/to/table"); - } - - #[tokio::test] - async fn test_uri_missing_location() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); - - // Server returns response without location field - http::Response::builder() - .status(200) - .body(r#"{"version": 1, "schema": {"fields": []}}"#) - .unwrap() - }); - - let result = table.uri().await; - assert!(result.is_err()); - assert!(matches!(&result, Err(Error::NotSupported { .. }))); - } - - #[tokio::test] - async fn test_uri_caching() { - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); - call_count_clone.fetch_add(1, Ordering::SeqCst); - - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": []}, "location": "gs://bucket/table"}"#, - ) - .unwrap() - }); - - // First call should fetch from server - let uri1 = table.uri().await.unwrap(); - assert_eq!(uri1, "gs://bucket/table"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second call should use cached value - let uri2 = table.uri().await.unwrap(); - assert_eq!(uri2, "gs://bucket/table"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call - } - - /// Test that schema is fetched once and cached for subsequent calls - #[tokio::test] - async fn test_schema_caching() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); - call_count_clone.fetch_add(1, Ordering::SeqCst); - - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - }); - - // First call should fetch from server - let schema1 = table.schema().await.unwrap(); - assert_eq!(schema1.fields().len(), 1); - assert_eq!(schema1.field(0).name(), "a"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second call should use cached value - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call - - // Third call should still use cached value - let schema3 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call - } - - /// Test that schema cache expires after 30 seconds TTL - #[tokio::test] - async fn test_schema_cache_invalidation_after_ttl() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); - call_count_clone.fetch_add(1, Ordering::SeqCst); - - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - }); - - // First call should fetch from server - let _schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second call should use cached value (within TTL) - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&_schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Advance mock time past TTL (no real wait) - clock::advance_by(Duration::from_secs(31)); - - // Third call should re-fetch from server (TTL expired) - let schema3 = table.schema().await.unwrap(); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&_schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - } - - /// Test that schema cache is invalidated after schema-changing operations - #[rstest] - #[case("overwrite")] - #[case("add_columns")] - #[case("drop_columns")] - #[case("alter_columns")] - #[tokio::test] - async fn test_schema_cache_invalidation_after_operation(#[case] operation: &str) { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false}, - {"name": "b", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else if path == "/v1/table/my_table/insert/" - || path == "/v1/table/my_table/add_columns/" - || path == "/v1/table/my_table/drop_columns/" - || path == "/v1/table/my_table/alter_columns/" - { - http::Response::builder() - .status(200) - .body(r#"{"version": 2}"#) - .unwrap() - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call should fetch from server - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second schema call should use cached value - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Perform the schema-changing operation - match operation { - "overwrite" => { - let data = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); - let _ = table.add(data).mode(AddDataMode::Overwrite).execute().await; - } - "add_columns" => { - let _ = table - .add_columns() - .transform(NewColumnTransform::SqlExpressions(vec![( - "c".into(), - "a + 1".into(), - )])) - .execute() - .await; - } - "drop_columns" => { - let _ = table.drop_columns(&["b"]).await; - } - "alter_columns" => { - let alterations = vec![ColumnAlteration::new("a".into()).rename("new_a".into())]; - let _ = table.alter_columns(&alterations).await; - } - _ => panic!("Unknown operation: {}", operation), - } - - // Schema call after operation should re-fetch from server (cache invalidated) - let schema3 = table.schema().await.unwrap(); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - } - - /// Test that schema cache is invalidated when server returns certain error codes - #[rstest] - #[case(400, true)] // 400 Bad Request should invalidate cache - #[case(401, false)] // 401 Unauthorized should NOT invalidate cache - #[case(403, false)] // 403 Forbidden should NOT invalidate cache - #[case(404, true)] // 404 Not Found should invalidate (table might be recreated) - #[case(500, true)] // 500 Internal Server Error should invalidate cache - #[case(503, false)] // 503 Service Unavailable should NOT invalidate cache - #[tokio::test] - async fn test_schema_cache_invalidation_on_errors( - #[case] error_status: u16, - #[case] should_invalidate: bool, - ) { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - let current_count = call_count_clone.load(Ordering::SeqCst); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else if path == "/v1/table/my_table/count_rows/" { - // Return error on first count_rows call - if current_count == 1 { - http::Response::builder() - .status(error_status) - .body("error") - .unwrap() - } else { - http::Response::builder().status(200).body("10").unwrap() - } - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call should fetch from server - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second schema call should use cached value - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Perform operation that returns error - let result = table.count_rows(None).await; - assert!(result.is_err()); - - // Schema call after error - check if cache was invalidated - let schema3 = table.schema().await.unwrap(); - if should_invalidate { - assert_eq!( - call_count.load(Ordering::SeqCst), - 2, - "Cache should be invalidated for {} error", - error_status - ); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - } else { - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "Cache should NOT be invalidated for {} error", - error_status - ); - assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - } - } - - /// A pinned snapshot should reuse the version and schema returned by its - /// initial describe instead of issuing two more describe requests. - #[tokio::test] - async fn test_checkout_current_seeds_schema_from_single_describe() { - let describe_calls = Arc::new(AtomicUsize::new(0)); - let calls = describe_calls.clone(); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); - calls.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version":42,"schema":{"fields":[{"name":"a","type":{"type":"int32"},"nullable":false}]}}"#, - ) - .unwrap() - }); - - let snapshot = table.checkout_current().await.unwrap(); - assert_eq!(snapshot.schema().await.unwrap().fields().len(), 1); - assert_eq!(describe_calls.load(Ordering::SeqCst), 1); - } - - /// Test that schema cache is invalidated after checkout - #[tokio::test] - async fn test_schema_cache_invalidation_on_checkout() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - call_count_clone.fetch_add(1, Ordering::SeqCst); - let count = call_count_clone.load(Ordering::SeqCst); - - if path == "/v1/table/my_table/describe/" { - // Return different schemas for different calls - if count <= 2 { - // First schema call and checkout validation - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - // After checkout - http::Response::builder() - .status(200) - .body( - r#"{"version": 2, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false}, - {"name": "b", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call - let schema1 = table.schema().await.unwrap(); - assert_eq!(schema1.fields().len(), 1); - - // Second schema call should use cached value (no new call) - let call_count_before = call_count.load(Ordering::SeqCst); - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), call_count_before); - - // Checkout to version 2 (makes a describe call to validate) - let _ = table.checkout(2).await; - - // Schema call after checkout should re-fetch (cache was invalidated) - let schema3 = table.schema().await.unwrap(); - assert_eq!(schema3.fields().len(), 2); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_schema_fetch_does_not_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 table = Table::new_with_handler("my_table", move |request| { - let body = request_body_json(&request); - let pinned = body["version"].as_u64() == Some(5); - if !pinned { - arrived_tx.lock().unwrap().send(()).unwrap(); - release_rx - .lock() - .unwrap() - .recv_timeout(Duration::from_secs(10)) - .unwrap(); - } - let field = if pinned { "pinned" } else { "latest" }; - http::Response::builder() - .status(200) - .body(format!( - r#"{{"version":5,"schema":{{"fields":[{{"name":"{field}","type":{{"type":"int32"}},"nullable":false}}]}}}}"# - )) - .unwrap() - }); - - let schema_fetch = 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(); - - let schema = schema_fetch.await.unwrap().unwrap(); - assert!(schema.field_with_name("pinned").is_ok()); - let cached = table.schema().await.unwrap(); - assert!(cached.field_with_name("pinned").is_ok()); - } - - /// Test that schema cache is invalidated after checkout_latest - #[tokio::test] - async fn test_schema_cache_invalidation_on_checkout_latest() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second schema call should use cached value - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Checkout latest - let _ = table.checkout_latest().await; - - // Schema call after checkout_latest should re-fetch (cache was invalidated) - let schema3 = table.schema().await.unwrap(); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - } - - /// Test that schema cache is invalidated after checkout_tag - #[tokio::test] - async fn test_schema_cache_invalidation_on_checkout_tag() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else if path == "/v1/table/my_table/tags/list/" { - http::Response::builder() - .status(200) - .body(r#"{"tags": {"v2": {"version": 2}}}"#) - .unwrap() - } else if path == "/v1/table/my_table/tags/version/" { - http::Response::builder() - .status(200) - .body(r#"{"version": 2}"#) - .unwrap() - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second schema call should use cached value - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Checkout tag - table - .checkout_tag("v2") - .await - .expect("checkout_tag should succeed"); - - // Schema call after checkout_tag should re-fetch (cache was invalidated) - let schema3 = table.schema().await.unwrap(); - assert_eq!( - call_count.load(Ordering::SeqCst), - 2, - "Cache should have been invalidated and re-fetched" - ); - assert_ne!( - Arc::as_ptr(&schema3), - Arc::as_ptr(&schema1), - "Should be different Arc instances" - ); - } - - /// Test that restore invalidates cache (via checkout_latest) - #[tokio::test] - async fn test_schema_cache_invalidation_on_restore() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else if path == "/v1/table/my_table/restore/" { - http::Response::builder() - .status(200) - .body(r#"{"version": 1}"#) - .unwrap() - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - table.checkout(1).await.unwrap(); - - // First schema call - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - - // Second schema call uses cache - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - - // Restore operation - table.restore().await.unwrap(); - - // Schema call after restore should re-fetch (cache invalidated) - let schema3 = table.schema().await.unwrap(); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 3); - } - - /// Test that centralized error handling invalidates cache on query errors - #[tokio::test] - async fn test_centralized_error_invalidation_on_query() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - let current_count = call_count_clone.load(Ordering::SeqCst); - - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else if path == "/v1/table/my_table/query/" { - // Return 400 error on first query (could be schema mismatch) - if current_count == 1 { - http::Response::builder() - .status(400) - .body("Bad request") - .unwrap() - } else { - // Return empty result for successful query - http::Response::builder() - .status(200) - .header("content-type", "application/vnd.apache.arrow.stream") - .body("") - .unwrap() - } - } else { - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - } - }); - - // First schema call - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Second schema call uses cache - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Query that returns 400 error - let result = table.query().execute().await; - assert!(result.is_err()); - - // Schema call after error should re-fetch (cache invalidated by centralized handler) - let schema3 = table.schema().await.unwrap(); - assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - } - - /// Test that concurrent schema() calls with an empty cache only trigger one fetch. - #[tokio::test] - async fn test_concurrent_schema_calls_single_fetch() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Arc::new(Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - panic!("Unexpected request: {}", path); - } - })); - - let mut handles = Vec::new(); - for _ in 0..10 { - let table = table.clone(); - handles.push(tokio::spawn(async move { table.schema().await.unwrap() })); - } - - let schemas: Vec = futures::future::try_join_all(handles).await.unwrap(); - - // All callers should get the same Arc - for schema in &schemas { - assert_eq!(Arc::as_ptr(schema), Arc::as_ptr(&schemas[0])); - } - // Only one describe call should have been made - assert_eq!(call_count.load(Ordering::SeqCst), 1); - } - - /// Test that a background refresh is triggered in the refresh window and - /// returns the cached value immediately. - #[tokio::test] - async fn test_background_refresh_triggers_in_window() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - if path == "/v1/table/my_table/describe/" { - let count = call_count_clone.fetch_add(1, Ordering::SeqCst); - if count == 0 { - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - http::Response::builder() - .status(200) - .body( - r#"{"version": 2, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false}, - {"name": "b", "type": { "type": "string" }, "nullable": true} - ]}}"#, - ) - .unwrap() - } - } else { - panic!("Unexpected request: {}", path); - } - }); - - // Populate cache and trigger peek transition to Current state - let schema1 = table.schema().await.unwrap(); - assert_eq!(schema1.fields().len(), 1); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - // Second call transitions cache from Refreshing to Current via peek() - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - - // Advance into refresh window (TTL=30s, window=5s, so 26s is in window) - clock::advance_by(Duration::from_secs(26)); - - // This call enters the refresh window: returns cached value and creates - // a background shared future (Refreshing state with previous). - let schema3 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - // Only the initial fetch so far - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Advance past TTL so the previous value expires. This forces the next - // schema() to Wait on the in-flight shared future, driving it to completion. - clock::advance_by(Duration::from_secs(30)); - - let schema4 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 2); - assert_eq!(schema4.fields().len(), 2); - assert_ne!(Arc::as_ptr(&schema4), Arc::as_ptr(&schema1)); - } - - /// Test that multiple calls during the refresh window don't trigger - /// duplicate background refreshes. - #[tokio::test] - async fn test_no_duplicate_background_refreshes() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - if path == "/v1/table/my_table/describe/" { - call_count_clone.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - panic!("Unexpected request: {}", path); - } - }); - - // Populate cache and transition to Current state - let schema1 = table.schema().await.unwrap(); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - let _ = table.schema().await.unwrap(); // peek transition - - // Advance into refresh window - clock::advance_by(Duration::from_secs(26)); - - // Multiple rapid calls should all return cached. The first one enters - // the refresh window and starts a background fetch (Refreshing state). - // Subsequent calls see Refreshing with a valid previous and return it. - for _ in 0..5 { - let schema = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema), Arc::as_ptr(&schema1)); - } - - // Advance past TTL and drive the shared future to completion - clock::advance_by(Duration::from_secs(30)); - let _ = table.schema().await.unwrap(); - - // Only one additional describe call (the background refresh), - // not five separate ones - assert_eq!(call_count.load(Ordering::SeqCst), 2); - } - - /// Test that if a background refresh fails, the previously cached value - /// is preserved and still returned. - #[tokio::test] - async fn test_background_refresh_error_preserves_cache() { - let call_count = Arc::new(AtomicUsize::new(0)); - let call_count_clone = call_count.clone(); - - let table = Table::new_with_handler("my_table", move |request| { - let path = request.url().path(); - if path == "/v1/table/my_table/describe/" { - let count = call_count_clone.fetch_add(1, Ordering::SeqCst); - if count == 0 { - // First call succeeds - http::Response::builder() - .status(200) - .body( - r#"{"version": 1, "schema": {"fields": [ - {"name": "a", "type": { "type": "int32" }, "nullable": false} - ]}}"#, - ) - .unwrap() - } else { - // Subsequent calls fail (422 is not retried) - http::Response::builder() - .status(422) - .body("Unprocessable Entity") - .unwrap() - } - } else { - panic!("Unexpected request: {}", path); - } - }); - - // Populate cache and transition to Current state - let schema1 = table.schema().await.unwrap(); - assert_eq!(schema1.fields().len(), 1); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - let _ = table.schema().await.unwrap(); // peek transition - - // Advance into refresh window - clock::advance_by(Duration::from_secs(26)); - - // Trigger background refresh (returns cached value). The background - // fetch will fail but the previous value should be preserved. - let schema2 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); - - // Still in the refresh window: the previous value is valid, - // so calling schema() should still return it. - let schema3 = table.schema().await.unwrap(); - assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); - - // Advance past TTL. The shared future will be driven and fail. - // The peek() error path should revert to the previous cached value. - clock::advance_by(Duration::from_secs(30)); - - // After the error, the previous is restored but its timestamp is old, - // so the next call triggers a new fetch which also fails. - let result = table.schema().await; - assert_eq!(call_count.load(Ordering::SeqCst), 2); - // The error from the failed fetch should be propagated - assert!(result.is_err()); - } - - #[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::>() - .await; - - assert_eq!(data.len(), 1); - assert_eq!(data[0].as_ref().unwrap(), &expected_data); - } - - fn schema_json() -> &'static str { - r#"{"fields": [{"name": "id", "type": {"type": "int32"}, "nullable": true}]}"# - } - - fn simple_describe_response() -> http::Response { - http::Response::builder() - .status(200) - .body(format!(r#"{{"version": 1, "schema": {}}}"#, schema_json())) - .unwrap() - } - - #[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::::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> = - Vec::new(); - let data: Box = - 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); - } - - // ---- 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( - body_for: F, - ) -> ( - impl Fn(reqwest::Request) -> http::Response + Clone + Send + Sync + 'static, - Arc>>, - ) - 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 + Clone + Send + Sync + 'static, - Arc>>, - ) { - 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( - body_for: F, - ) -> ( - impl Fn(reqwest::Request) -> http::Response + Clone + Send + Sync + 'static, - Arc>>, - ) - where - F: Fn(&str) -> String + Clone + Send + Sync + 'static, - { - let captured: Arc>> = - 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")); - } - - #[tokio::test] - async fn test_update_field_metadata() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/table/my_table/update_field_metadata/" - ); - http::Response::builder() - .status(200) - .body(r#"{"version": 7, "fields": {"category": {"unit": "label"}}}"#) - .unwrap() - }); - - let result = table - .update_field_metadata(&[FieldMetadataUpdate::new("category").set("unit", "label")]) - .await - .unwrap(); - assert_eq!(result.version, 7); - } - - // ----- 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") - } - - #[tokio::test] - async fn test_create_branch_default_source() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/branches/create/"); - let body = request_body_json(&request); - assert_eq!(body["name"], "exp"); - assert!( - body.get("from_branch").is_none(), - "a main source omits from_branch" - ); - assert!( - body.get("from_version").is_none(), - "a latest source omits from_version" - ); - http::Response::builder().status(200).body("{}").unwrap() - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - assert_eq!(branch.current_branch(), Some("exp".to_string())); - assert_eq!(table.current_branch(), None); - } - - #[tokio::test] - async fn test_create_branch_from_branch_and_version() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| { - let body = request_body_json(&request); - assert_eq!(body["name"], "exp"); - assert_eq!(body["from_branch"], "base"); - assert_eq!(body["from_version"], 3); - http::Response::builder().status(200).body("{}").unwrap() - }); - table - .create_branch("exp", Ref::Version(Some("base".into()), Some(3))) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_branch_from_main_normalizes_to_none() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| { - let body = request_body_json(&request); - assert!( - body.get("from_branch").is_none(), - "\"main\" normalizes to an absent from_branch" - ); - assert_eq!(body["from_version"], 7); - http::Response::builder().status(200).body("{}").unwrap() - }); - table - .create_branch("exp", Ref::Version(Some("main".into()), Some(7))) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_branch_from_version_number_on_main() { - use lance::dataset::refs::Ref; - // A bare version number on a main handle resolves to (main, version). - let table = Table::new_with_handler("my_table", |request| { - let body = request_body_json(&request); - assert!(body.get("from_branch").is_none()); - assert_eq!(body["from_version"], 5); - http::Response::builder().status(200).body("{}").unwrap() - }); - table - .create_branch("exp", Ref::VersionNumber(5)) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_branch_from_tag_resolves_via_tags_endpoint() { - use lance::dataset::refs::Ref; - // A tag source has no from_tag in the create contract; it is resolved to - // its (branch, version) via the tags/version endpoint first. - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/tags/version/" => { - assert_eq!(request_body_json(&request)["tag"], "t"); - http::Response::builder() - .status(200) - .body(r#"{"version":3,"branch":"base"}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/branches/create/" => { - let body = request_body_json(&request); - assert_eq!(body["name"], "exp"); - assert_eq!(body["from_branch"], "base"); - assert_eq!(body["from_version"], 3); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - table - .create_branch("exp", Ref::Tag("t".into())) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_branch_from_tag_on_main_normalizes() { - use lance::dataset::refs::Ref; - // A tag resolving to the main branch collapses from_branch to absent. - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/tags/version/" => http::Response::builder() - .status(200) - .body(r#"{"version":4,"branch":"main"}"#.to_string()) - .unwrap(), - "/v1/table/my_table/branches/create/" => { - let body = request_body_json(&request); - assert!( - body.get("from_branch").is_none(), - "a resolved \"main\" normalizes to an absent from_branch" - ); - assert_eq!(body["from_version"], 4); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - table - .create_branch("exp", Ref::Tag("t".into())) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_create_branch_invalid_request_maps_to_invalid_input() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |_| { - http::Response::builder() - .status(400) - .body("unsafe branch name") - .unwrap() - }); - let err = table - .create_branch("../evil", Ref::Version(None, None)) - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); - } - - #[tokio::test] - async fn test_create_branch_conflict_maps_to_already_exists() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |_| { - http::Response::builder() - .status(409) - .body("branch already exists") - .unwrap() - }); - let err = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap_err(); - assert!( - matches!(err, Error::TableAlreadyExists { .. }), - "409 should map to AlreadyExists, got {err:?}" - ); - } - - #[tokio::test] - async fn test_materialized_view_describe_and_refresh_requires_sql() { - const QUERY: &str = - "SELECT s.x, d.label FROM analytics.source s JOIN analytics.dim d ON s.id = d.id"; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/materialized_view/my_table/describe" => http::Response::builder() - .status(200) - .body( - json!({ - "name": "my_table", - "query": QUERY - }) - .to_string(), - ) - .unwrap(), - path => panic!("unexpected request: {path}"), - }); - let view = crate::MaterializedView::from_table(table).await.unwrap(); - assert_eq!(view.definition_sql(), QUERY); - assert!(view.definition().is_err()); - let error = view.refresh().execute().await.unwrap_err(); - assert!(error.to_string().contains("SQL is unavailable")); - } - - #[tokio::test] - async fn test_create_branch_empty_name_rejected_client_side() { - use lance::dataset::refs::Ref; - // The empty name is rejected before any request is sent. - let table = Table::new_with_handler("my_table", |request| -> http::Response { - panic!("unexpected request: {}", request.url().path()) - }); - let err = table - .create_branch("", Ref::Version(None, None)) - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); - } - - #[tokio::test] - async fn test_list_branches() { - use lance::dataset::refs::BranchIdentifier; - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/branches/list/"); - // A branch forked off main: the server omits `parentBranch` entirely - // (skip_serializing_if), not `null`, so this mirrors the real wire. - http::Response::builder() - .status(200) - .body( - r#"{"branches":{"exp":{"parentVersion":2,"createAt":1234,"manifestSize":4096}}}"#, - ) - .unwrap() - }); - let branches = table.list_branches().await.unwrap(); - let exp = branches.get("exp").expect("exp present"); - assert_eq!(exp.parent_version, 2); - assert_eq!(exp.create_at, 1234); - assert_eq!(exp.manifest_size, 4096); - assert_eq!(exp.parent_branch, None); - assert!(exp.metadata.is_empty()); - // The server omits the internal lineage token; it defaults to the sentinel. - assert_eq!( - exp.identifier, - BranchIdentifier::missing_identifier_sentinel() - ); - } - - #[tokio::test] - async fn test_main_only_metadata_is_unfenced_from_branch_timeline() { - let requests = Arc::new(std::sync::Mutex::new(HashMap::new())); - let captured = requests.clone(); - let saw_delete_response_floor = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let saw_high_floor = saw_delete_response_floor.clone(); - let table = RemoteTable::new_mock( - "my_table".to_string(), - move |request| { - let path = request.url().path().to_string(); - captured - .lock() - .unwrap() - .insert(path.clone(), request.headers().clone()); - match path.as_str() { - "/v1/table/my_table/count_rows/" => { - saw_high_floor.store( - request - .headers() - .get("x-lancedb-min-read-version") - .and_then(|value| value.to_str().ok()) - == Some("100"), - std::sync::atomic::Ordering::SeqCst, - ); - http::Response::builder() - .status(200) - .header("x-lancedb-version", "2") - .body("1".to_string()) - .unwrap() - } - "/v1/table/my_table/tags/list/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/tags/version/" => http::Response::builder() - .status(200) - .body(r#"{"version":1}"#.to_string()) - .unwrap(), - "/v1/table/my_table/tags/delete/" => http::Response::builder() - .status(200) - .header("x-lancedb-version", "100") - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/branches/list/" => http::Response::builder() - .status(200) - .body(r#"{"branches":{}}"#.to_string()) - .unwrap(), - path => panic!("unexpected path: {path}"), - } - }, - None, - ); - let branch = table.with_branch(Some("exp".to_string())); - - branch.count_rows(None).await.unwrap(); - let mut tags = branch.tags().await.unwrap(); - tags.list().await.unwrap(); - tags.get_version("v1").await.unwrap(); - tags.delete("v1").await.unwrap(); - branch.list_branches().await.unwrap(); - branch.count_rows(None).await.unwrap(); - - let requests = requests.lock().unwrap(); - for path in [ - "/v1/table/my_table/tags/list/", - "/v1/table/my_table/tags/version/", - "/v1/table/my_table/tags/delete/", - "/v1/table/my_table/branches/list/", - ] { - assert!( - !requests[path].contains_key("x-lancedb-min-read-version"), - "{path} inherited the branch timeline" - ); - } - assert!( - !saw_delete_response_floor.load(std::sync::atomic::Ordering::SeqCst), - "tag deletion contaminated the branch timeline" - ); - } - - #[tokio::test] - async fn test_delete_branch() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/branches/delete/"); - let body = request_body_json(&request); - assert_eq!(body["name"], "exp"); - http::Response::builder().status(200).body("{}").unwrap() - }); - table.delete_branch("exp").await.unwrap(); - } - - #[tokio::test] - async fn test_delete_branch_not_found() { - let table = Table::new_with_handler("my_table", |_| { - http::Response::builder() - .status(404) - .body("no such branch") - .unwrap() - }); - let err = table.delete_branch("ghost").await.unwrap_err(); - assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}"); - } - - fn sample_branch_diff_json() -> &'static str { - r#"{ - "fromBranch":"exp", - "parentVersion":1, - "mainVersion":1, - "branchVersion":2, - "baseMoved":false, - "rowCountMain":3, - "rowCountBranch":3, - "rowSummary":{ - "unchanged":3, - "newOnBase":0, - "newOnBranch":0, - "staleRecompute":0, - "inputsChanged":0, - "deltaAvailable":false - }, - "addedColumns":[{"name":"tag","dataType":"utf8","nullable":true}], - "removedColumns":[], - "changedColumns":[], - "addedIndexes":[], - "removedIndexes":[], - "errors":[] - }"# - } - - #[tokio::test] - async fn test_diff_branch() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/branches/diff/"); - let body = request_body_json(&request); - assert_eq!(body["from_branch"], "exp"); - http::Response::builder() - .status(200) - .body(sample_branch_diff_json()) - .unwrap() - }); - let diff = table.diff_branch("exp").await.unwrap(); - assert_eq!(diff.from_branch, "exp"); - assert!(diff.errors.is_empty()); - assert_eq!(diff.added_columns.len(), 1); - assert_eq!(diff.added_columns[0].name, "tag"); - } - - #[tokio::test] - async fn test_cherry_pick_dry_run() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!( - request.url().path(), - "/v1/table/my_table/branches/cherry_pick/" - ); - let body = request_body_json(&request); - assert_eq!(body["from_branch"], "exp"); - assert_eq!(body["dry_run"], true); - let resp = format!( - r#"{{"status":"ready","diff":{},"preview":{{"promotedColumns":["tag"]}}}}"#, - sample_branch_diff_json() - ); - http::Response::builder().status(200).body(resp).unwrap() - }); - let result = table.cherry_pick("exp", true).await.unwrap(); - assert_eq!(result.status, crate::table::CherryPickStatus::Ready); - assert_eq!(result.preview.promoted_columns, vec!["tag".to_string()]); - assert!(result.main_version_after.is_none()); - } - - #[tokio::test] - async fn test_successful_cherry_pick_advances_main_read_watermark() { - 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/branches/cherry_pick/" => { - let response = serde_json::json!({ - "status": "cherryPicked", - "diff": serde_json::from_str::(sample_branch_diff_json()) - .unwrap(), - "preview": { "promotedColumns": ["tag"] }, - "mainVersionAfter": 2 - }); - http::Response::builder() - .status(200) - .body(response.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 result = table.cherry_pick("exp", false).await.unwrap(); - assert_eq!(result.status, crate::table::CherryPickStatus::CherryPicked); - table.count_rows(None).await.unwrap(); - assert_eq!( - count_headers - .lock() - .unwrap() - .as_ref() - .unwrap() - .get("x-lancedb-min-read-version") - .and_then(|value| value.to_str().ok()), - Some("2") - ); - } - - #[tokio::test] - async fn test_cherry_pick_failed_returns_ok_with_body() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!( - request.url().path(), - "/v1/table/my_table/branches/cherry_pick/" - ); - let body = request_body_json(&request); - assert_eq!(body["dry_run"], false); - let mut diff: serde_json::Value = - serde_json::from_str(sample_branch_diff_json()).unwrap(); - diff["errors"] = serde_json::json!([{ - "code": "baseMoved", - "message": "main has advanced" - }]); - let resp = serde_json::json!({ - "status": "failed", - "diff": diff, - "preview": { "promotedColumns": [] } - }); - http::Response::builder() - .status(409) - .body(resp.to_string()) - .unwrap() - }); - let result = table.cherry_pick("exp", false).await.unwrap(); - assert_eq!(result.status, crate::table::CherryPickStatus::Failed); - assert!(!result.diff.errors.is_empty()); - assert_eq!(result.diff.errors.len(), 1); - } - - #[tokio::test] - async fn test_cherry_pick_unknown_error_code_parses() { - let table = Table::new_with_handler("my_table", |_| { - let mut diff: serde_json::Value = - serde_json::from_str(sample_branch_diff_json()).unwrap(); - diff["errors"] = serde_json::json!([{ - "code": "multipleCommits", - "message": "branch has more than one data commit" - }]); - let resp = serde_json::json!({ - "status": "failed", - "diff": diff, - "preview": { "operation": "append", "rowsAdded": 2 } - }); - http::Response::builder() - .status(409) - .body(resp.to_string()) - .unwrap() - }); - let result = table.cherry_pick("exp", false).await.unwrap(); - assert_eq!(result.status, crate::table::CherryPickStatus::Failed); - assert_eq!( - result.diff.errors[0].code, - crate::table::CherryPickErrorCode::Unknown - ); - assert!(result.preview.promoted_columns.is_empty()); - } - - #[tokio::test] - async fn test_cherry_pick_unexpected_2xx_is_error() { - let table = Table::new_with_handler("my_table", |_| { - http::Response::builder() - .status(204) - .body(String::new()) - .unwrap() - }); - let err = table.cherry_pick("exp", false).await.unwrap_err(); - match err { - Error::Http { - status_code: Some(code), - .. - } => assert_eq!(code, reqwest::StatusCode::NO_CONTENT), - other => panic!("expected Http error, got {other:?}"), - } - } - - #[tokio::test] - async fn test_checkout_branch_validates_via_list() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/branches/list/"); - http::Response::builder() - .status(200) - .body( - r#"{"branches":{"exp":{"parentBranch":null,"parentVersion":1,"createAt":1,"manifestSize":1}}}"#, - ) - .unwrap() - }); - let branch = table.checkout_branch("exp", None).await.unwrap(); - assert_eq!(branch.current_branch(), Some("exp".to_string())); - } - - #[tokio::test] - async fn test_checkout_branch_missing() { - let table = Table::new_with_handler("my_table", |_| { - http::Response::builder() - .status(200) - .body(r#"{"branches":{}}"#) - .unwrap() - }); - let err = table.checkout_branch("ghost", None).await.unwrap_err(); - assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}"); - } - - #[tokio::test] - async fn test_checkout_main_returns_main_handle() { - // "main" yields a main-scoped handle without any validation request. - let table = Table::new_with_handler("my_table", |request| -> http::Response { - panic!("unexpected request: {}", request.url().path()) - }); - let main = table.checkout_branch("main", None).await.unwrap(); - assert_eq!(main.current_branch(), None); - } - - #[tokio::test] - async fn test_branch_count_rows_carries_branch_in_body() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/count_rows/" => { - let body = request_body_json(&request); - assert_eq!(body["branch"], "exp"); - http::Response::builder() - .status(200) - .body("7".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - assert_eq!(branch.count_rows(None).await.unwrap(), 7); - } - - #[tokio::test] - async fn test_main_handle_omits_branch_in_body() { - // A main handle must not send a branch field (byte-compatible with - // pre-branch servers and the existing wire format). - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/count_rows/"); - let body = request_body_json(&request); - assert!( - body.get("branch").is_none(), - "main handle must not send a branch field" - ); - http::Response::builder().status(200).body("0").unwrap() - }); - table.count_rows(None).await.unwrap(); - } - - #[tokio::test] - async fn test_branch_update_and_delete_carry_branch() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/update/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(r#"{"version":5}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/delete/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(r#"{"version":6,"num_deleted_rows":1}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch - .update() - .column("a", "a + 1") - .execute() - .await - .unwrap(); - branch.delete("a > 1").await.unwrap(); - } - - #[tokio::test] - async fn test_branch_list_indices_carries_branch_in_body() { - use lance::dataset::refs::Ref; - // list_indices posts to index/list and then fetches the schema (describe) - // to resolve column names; both must carry the branch. - let describe_body = - describe_response(&Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/index/list/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(r#"{"indexes":[]}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/describe/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(describe_body.clone()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - assert!(branch.list_indices().await.unwrap().is_empty()); - } - - #[tokio::test] - async fn test_branch_list_versions_carries_query_param() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/version/list/" => { - assert_eq!( - request - .url() - .query_pairs() - .find(|(k, _)| k == "branch") - .map(|(_, v)| v.into_owned()), - Some("exp".to_string()), - "version/list must carry ?branch=exp" - ); - http::Response::builder() - .status(200) - .body(r#"{"versions":[]}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - assert!(branch.list_versions().await.unwrap().is_empty()); - } - - #[tokio::test] - async fn test_branch_drop_index_carries_query_param() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/index/my_idx/drop/" => { - assert_eq!( - request - .url() - .query_pairs() - .find(|(k, _)| k == "branch") - .map(|(_, v)| v.into_owned()), - Some("exp".to_string()) - ); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.drop_index("my_idx").await.unwrap(); - } - - #[tokio::test] - async fn test_branch_insert_carries_query_param() { - use lance::dataset::refs::Ref; - 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 table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => { - // schema() fetch on the branch handle carries branch in the body. - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(describe_body.clone()) - .unwrap() - } - "/v1/table/my_table/insert/" => { - assert_eq!( - request - .url() - .query_pairs() - .find(|(k, _)| k == "branch") - .map(|(_, v)| v.into_owned()), - Some("exp".to_string()), - "insert must carry ?branch=exp" - ); - http::Response::builder() - .status(200) - .body(r#"{"version":2}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.add(data.clone()).execute().await.unwrap(); - } - - #[tokio::test] - async fn test_branch_tag_create_carries_branch_in_body() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/tags/create/" => { - let body = request_body_json(&request); - assert_eq!(body["branch"], "exp"); - assert_eq!(body["tag"], "v1"); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.tags().await.unwrap().create("v1", 1).await.unwrap(); - } - - #[tokio::test] - async fn test_checkout_branch_version_forwards_branch() { - // Branch versions overlap main's, so the version must resolve on the - // branch's own chain -- the validating describe and reads carry both. - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - let describe_body = describe_response(&schema); - let table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/branches/list/" => http::Response::builder() - .status(200) - .body( - r#"{"branches":{"exp":{"parentBranch":null,"parentVersion":1,"createAt":1,"manifestSize":1}}}"# - .to_string(), - ) - .unwrap(), - "/v1/table/my_table/describe/" => { - let body = request_body_json(&request); - assert_eq!(body["branch"], "exp", "checkout validate carries branch"); - assert_eq!(body["version"], 2, "checkout validate carries version"); - http::Response::builder().status(200).body(describe_body.clone()).unwrap() - } - "/v1/table/my_table/count_rows/" => { - let body = request_body_json(&request); - assert_eq!(body["branch"], "exp"); - assert_eq!(body["version"], 2, "overlapping version resolves on the branch chain"); - http::Response::builder().status(200).body("3".to_string()).unwrap() - } - path => panic!("unexpected request path: {path}"), - } - }); - let branch = table.checkout_branch("exp", Some(2)).await.unwrap(); - assert_eq!(branch.current_branch(), Some("exp".to_string())); - assert_eq!(branch.count_rows(None).await.unwrap(), 3); - } - - fn branch_query_param(request: &reqwest::Request) -> Option { - request - .url() - .query_pairs() - .find(|(k, _)| k == "branch") - .map(|(_, v)| v.into_owned()) - } - - #[tokio::test] - async fn test_branch_query_carries_branch_in_body() { - use lance::dataset::refs::Ref; - // /query/ is the hot read path; the branch must ride in the JSON 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 data_ref = data.clone(); - let table = - Table::new_with_handler("my_table", 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/query/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) - .body(write_ipc_file(&data_ref)) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - let rows: usize = branch - .query() - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 3); - } - - #[tokio::test] - async fn test_branch_merge_insert_carries_query_param() { - use lance::dataset::refs::Ref; - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - let data: Box = Box::new(RecordBatchIterator::new( - [Ok(batch.clone())], - batch.schema(), - )); - let table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/merge_insert/" => { - assert_eq!( - branch_query_param(&request).as_deref(), - Some("exp"), - "merge_insert must carry ?branch=exp" - ); - http::Response::builder() - .status(200) - .body( - r#"{"version":2,"num_deleted_rows":0,"num_inserted_rows":3,"num_updated_rows":0}"# - .to_string(), - ) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - } - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.merge_insert(&["id"]).execute(data).await.unwrap(); - } - - #[tokio::test] - async fn test_branch_multipart_write_carries_query_param() { - use lance::dataset::refs::Ref; - // The multipart path (create -> insert parts -> complete) must carry - // ?branch= on every leg; an old server version forces it. - let table = Table::new_with_handler_version( - "my_table", - semver::Version::new(0, 4, 0), - move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/multipart_write/create" => { - assert_eq!(branch_query_param(&request).as_deref(), Some("exp")); - http::Response::builder() - .status(200) - .body(r#"{"upload_id": "u1"}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/insert/" => { - assert_eq!( - branch_query_param(&request).as_deref(), - Some("exp"), - "multipart insert must carry ?branch=exp" - ); - http::Response::builder() - .status(200) - .body(r#"{"version": 1}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/multipart_write/complete" => { - assert_eq!(branch_query_param(&request).as_deref(), Some("exp")); - http::Response::builder() - .status(200) - .body(r#"{"version": 5}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }, - ); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); - branch - .add(vec![batch]) - .write_parallelism(2) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_branch_restore_carries_branch_in_body() { - use lance::dataset::refs::Ref; - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - assert_eq!(request_body_json(&request)["version"], 1); - http::Response::builder() - .status(200) - .body(r#"{"version":1,"schema":{"fields":[]}}"#.to_string()) - .unwrap() - } - "/v1/table/my_table/restore/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - assert_eq!(request_body_json(&request)["version"], 1); - http::Response::builder() - .status(200) - .body(r#"{"version":1}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.checkout(1).await.unwrap(); - branch.restore().await.unwrap(); - } - - #[tokio::test] - async fn test_branch_create_index_carries_branch_in_body() { - use lance::dataset::refs::Ref; - // create_index fetches the schema (describe) to resolve the column and - // then posts create_index; both must carry the branch in the body. - let describe_body = - describe_response(&Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(describe_body.clone()) - .unwrap() - } - "/v1/table/my_table/create_index/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch - .create_index(&["a"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - } - - #[tokio::test] - async fn test_branch_column_ops_carry_branch_in_body() { - use lance::dataset::refs::Ref; - // add_columns / alter_columns / drop_columns all stamp the branch into - // the JSON body via apply_branch_body. - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/add_columns/" - | "/v1/table/my_table/alter_columns/" - | "/v1/table/my_table/drop_columns/" => { - assert_eq!( - request_body_json(&request)["branch"], - "exp", - "{} must carry the branch", - request.url().path() - ); - http::Response::builder() - .status(200) - .body(r#"{"version":43}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch - .add_columns() - .transform(NewColumnTransform::SqlExpressions(vec![( - "b".into(), - "a + 1".into(), - )])) - .execute() - .await - .unwrap(); - branch - .alter_columns(&[ColumnAlteration::new("a".into()).rename("b".into())]) - .await - .unwrap(); - branch.drop_columns(&["a"]).await.unwrap(); - } - - #[tokio::test] - async fn test_branch_update_field_metadata_carries_branch_in_body() { - use lance::dataset::refs::Ref; - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/update_field_metadata/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body(r#"{"version":7,"fields":{}}"#.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch - .update_field_metadata(&[FieldMetadataUpdate::new("category").set("unit", "label")]) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_branch_index_stats_carries_branch_in_body() { - use lance::dataset::refs::Ref; - let table = Table::new_with_handler("my_table", move |request| { - match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body("{}".to_string()) - .unwrap(), - "/v1/table/my_table/index/my_index/stats/" => { - assert_eq!(request_body_json(&request)["branch"], "exp"); - http::Response::builder() - .status(200) - .body( - r#"{"num_indexed_rows":1,"num_unindexed_rows":0,"index_type":"IVF_PQ","distance_type":"l2"}"# - .to_string(), - ) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - } - }); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - assert!(branch.index_stats("my_index").await.unwrap().is_some()); - } - - #[tokio::test] - async fn test_branch_stats_attaches_body_while_main_omits_it() { - use lance::dataset::refs::Ref; - // stats has a bespoke conditional body: a main handle stays a bodyless - // POST, while a branch handle attaches {"branch": ...}. - let stats_body = r#"{"total_bytes":1,"num_rows":3,"num_indices":0,"fragment_stats":{"num_fragments":1,"num_small_fragments":0,"lengths":{"min":3,"max":3,"mean":3,"p25":3,"p50":3,"p75":3,"p99":3}}}"#; - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/branches/create/" => http::Response::builder() - .status(200) - .body(stats_body.to_string()) - .unwrap(), - "/v1/table/my_table/stats/" => { - match request.body() { - // main handle: byte-identical to the pre-branch wire format. - None => {} - // branch handle: branch travels in the body. - Some(_) => assert_eq!(request_body_json(&request)["branch"], "exp"), - } - http::Response::builder() - .status(200) - .body(stats_body.to_string()) - .unwrap() - } - path => panic!("unexpected request path: {path}"), - }); - table.stats().await.unwrap(); - let branch = table - .create_branch("exp", Ref::Version(None, None)) - .await - .unwrap(); - branch.stats().await.unwrap(); - } - - // 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)"}"#; - - use lance::dataset::ROW_ID; - - 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), - ]) - } - - /// 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 = 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() - } - - /// 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>>) -> 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::>() - .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::>() - .await - .unwrap(); - - assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 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::>() - .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::>() - .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" - ); - } -} +mod tests; diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index ac37bf7ae..62bd056d2 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -895,1241 +895,4 @@ async fn probe_blob_file(requester: Arc) -> Result http::Response> { - http::Response::builder() - .status(StatusCode::NO_CONTENT) - .body(Vec::new()) - .unwrap() - } - - fn empty_blob_response() -> http::Response> { - http::Response::builder() - .status(StatusCode::RANGE_NOT_SATISFIABLE) - .header(header::CONTENT_RANGE, "bytes */0") - .body(Vec::new()) - .unwrap() - } - - fn range_response(request: &Request, payload: &[u8]) -> http::Response> { - let value = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap(); - let (start, end) = value - .strip_prefix("bytes=") - .unwrap() - .split_once('-') - .unwrap(); - let start = start.parse::().unwrap(); - let end = if end.is_empty() { - payload.len() - 1 - } else { - end.parse::().unwrap() - }; - http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes {start}-{end}/{}", payload.len()), - ) - .header(VERSION_HEADER, "5") - .body(payload[start..=end].to_vec()) - .unwrap() - } - - fn mock_remote_blob_table( - requests: Arc>>, - ) -> RemoteTable { - RemoteTable::new_mock( - "my_table".to_string(), - move |request| { - assert_eq!(request.method(), reqwest::Method::GET); - let path = request.url().path(); - assert!(path.starts_with("/v1/table/my_table/blob/image/")); - requests.lock().unwrap().push( - request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap() - .to_string(), - ); - if path.contains("/20/bytes") { - return null_blob_response(); - } - if path.contains("/30/bytes") { - return empty_blob_response(); - } - range_response(&request, PAYLOAD) - }, - Some(Version::new(0, 5, 0)), - ) - } - - #[tokio::test] - async fn remote_blob_files_probe_sizes_and_preserve_nulls() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - - let mut files = table.probe_files("image", &[10, 20]).await.unwrap(); - - assert_eq!(files.len(), 2); - assert!(files[1].is_none()); - let file = files[0].take().unwrap(); - assert_eq!(file.size(), PAYLOAD.len() as u64); - assert_eq!( - requests.lock().unwrap().as_slice(), - ["bytes=0-0", "bytes=0-0"] - ); - } - - /// Descriptors as a remote query returns them, one `(kind, size)` per row. - fn descriptors(rows: &[Option<(BlobKind, u64)>]) -> StructArray { - let kinds = rows - .iter() - .map(|row| row.map_or(0, |(kind, _)| kind as u8)) - .collect::(); - let sizes = rows - .iter() - .map(|row| row.map_or(0, |(_, size)| size)) - .collect::(); - let len = rows.len(); - StructArray::new( - lance_core::datatypes::BLOB_V2_DESC_FIELDS.clone(), - vec![ - Arc::new(kinds), - Arc::new(UInt64Array::from(vec![0; len])), - Arc::new(sizes), - Arc::new(arrow_array::UInt32Array::from(vec![0; len])), - Arc::new(arrow_array::StringArray::from(vec![""; len])), - ], - Some(rows.iter().map(Option::is_some).collect()), - ) - } - - impl RemoteTable { - /// Open handles through the per-row size probe alone. - async fn probe_files( - &self, - column: &str, - row_ids: &[u64], - ) -> Result>> { - let read_snapshot = self.snapshot_read_state().await; - self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot) - .await - } - } - - /// A request the take mock received. - #[derive(Debug, Clone, PartialEq)] - enum TakeMockRequest { - Query { - row_ids: Vec, - version: Option, - }, - Describe, - Range { - row_id: u64, - range: String, - version: Option, - }, - } - - /// Row ids named by a `_rowid IN (...)` filter. - fn filter_row_ids(filter: &str) -> Vec { - let list = filter.split_once('(').unwrap().1.trim_end_matches(')'); - list.split(',') - .map(|id| id.trim().parse().unwrap()) - .collect() - } - - /// A table whose query route answers `_rowid` takes from `rows` and whose - /// byte routes serve [`PAYLOAD`]. Takes return rows in reverse order, and - /// report `query_version` in `x-lancedb-version` when it is set. - fn mock_take_blob_table( - rows: Vec<(u64, Option<(BlobKind, u64)>)>, - query_version: Option<&'static str>, - ) -> ( - RemoteTable, - Arc>>, - ) { - let requests = Arc::new(StdMutex::new(Vec::new())); - let captured = requests.clone(); - let table = RemoteTable::new_mock( - "my_table".to_string(), - move |request| { - let path = request.url().path().to_string(); - if 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"], serde_json::json!(["image"])); - assert_eq!(body["with_row_id"], serde_json::json!(true)); - assert_eq!(body["use_lsm"], serde_json::json!(false)); - let requested = filter_row_ids(body["filter"].as_str().unwrap()); - captured.lock().unwrap().push(TakeMockRequest::Query { - row_ids: requested.clone(), - version: body["version"].as_u64(), - }); - let taken = rows - .iter() - .rev() - .filter(|(row_id, _)| requested.contains(row_id)) - .collect::>(); - let batch = RecordBatch::try_from_iter([ - ( - "image", - Arc::new(descriptors( - &taken.iter().map(|(_, row)| *row).collect::>(), - )) as Arc, - ), - ( - ROW_ID, - Arc::new(UInt64Array::from_iter_values( - taken.iter().map(|(row_id, _)| *row_id), - )) as Arc, - ), - ]) - .unwrap(); - let mut body = Vec::new(); - let mut writer = - arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); - writer.write(&batch).unwrap(); - writer.finish().unwrap(); - drop(writer); - let mut response = http::Response::builder().status(200); - if let Some(version) = query_version { - response = response.header(VERSION_HEADER, version); - } - return response.body(body).unwrap(); - } - if path == "/v1/table/my_table/describe/" { - captured.lock().unwrap().push(TakeMockRequest::Describe); - return http::Response::builder() - .status(200) - .body(r#"{"version":9,"schema":{"fields":[]}}"#.as_bytes().to_vec()) - .unwrap(); - } - let row_id = path - .strip_prefix("/v1/table/my_table/blob/image/") - .and_then(|rest| rest.strip_suffix("/bytes")) - .unwrap_or_else(|| panic!("unexpected path: {path}")) - .parse() - .unwrap(); - captured.lock().unwrap().push(TakeMockRequest::Range { - row_id, - range: request.headers()[header::RANGE] - .to_str() - .unwrap() - .to_string(), - version: request - .url() - .query_pairs() - .find(|(name, _)| name == "version") - .map(|(_, value)| value.into_owned()), - }); - range_response(&request, PAYLOAD) - }, - Some(Version::new(0, 5, 0)), - ); - (table, requests) - } - - #[tokio::test] - async fn remote_blob_files_are_sized_from_one_descriptor_take() { - let payload_size = PAYLOAD.len() as u64; - let (table, requests) = mock_take_blob_table( - vec![ - (10, Some((BlobKind::Inline, payload_size))), - (20, None), - (30, Some((BlobKind::Inline, 0))), - (40, Some((BlobKind::Packed, payload_size))), - (50, Some((BlobKind::Dedicated, payload_size))), - (60, Some((BlobKind::External, 0))), - ], - Some("7"), - ); - - // Reordered and duplicated row ids come back in request order. - let files = table - .fetch_blob_files_impl("image", &[60, 10, 20, 30, 40, 50, 10]) - .await - .unwrap(); - - let sizes = files - .iter() - .map(|file| file.as_ref().map(BlobFile::size)) - .collect::>(); - assert_eq!( - sizes, - [ - Some(payload_size), - Some(payload_size), - None, - Some(0), - Some(payload_size), - Some(payload_size), - Some(payload_size), - ] - ); - // One take, then a probe for the external blob that recorded no size. - assert_eq!( - requests.lock().unwrap().as_slice(), - [ - TakeMockRequest::Query { - row_ids: vec![10, 20, 30, 40, 50, 60], - version: None, - }, - TakeMockRequest::Range { - row_id: 60, - range: "bytes=0-0".into(), - version: Some("7".into()), - }, - ] - ); - - // Sized handles read the version the take read. - assert_eq!( - files[1].as_ref().unwrap().read_range(5..12).await.unwrap(), - &PAYLOAD[5..12] - ); - assert!(files[3].as_ref().unwrap().read().await.unwrap().is_empty()); - assert_eq!( - requests.lock().unwrap().last().unwrap(), - &TakeMockRequest::Range { - row_id: 10, - range: "bytes=5-11".into(), - version: Some("7".into()), - } - ); - } - - #[tokio::test] - async fn remote_blob_files_take_every_chunk_at_the_first_version() { - let row_ids = (0..MAX_FETCH_BLOBS_ROW_IDS as u64 + 1).collect::>(); - let (table, requests) = mock_take_blob_table( - row_ids.iter().map(|row_id| (*row_id, None)).collect(), - Some("7"), - ); - - let files = table - .fetch_blob_files_impl("image", &row_ids) - .await - .unwrap(); - - assert_eq!(files.len(), row_ids.len()); - let versions = requests - .lock() - .unwrap() - .iter() - .map(|request| match request { - TakeMockRequest::Query { row_ids, version } => (row_ids.len(), *version), - other => panic!("unexpected request: {other:?}"), - }) - .collect::>(); - assert_eq!(versions, [(MAX_FETCH_BLOBS_ROW_IDS, None), (1, Some(7))]); - } - - #[tokio::test] - async fn remote_blob_files_retake_at_the_described_version_when_a_take_does_not_report_it() { - let payload_size = PAYLOAD.len() as u64; - let (table, requests) = - mock_take_blob_table(vec![(10, Some((BlobKind::Inline, payload_size)))], None); - - let file = table - .fetch_blob_files_impl("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - file.read_range(0..4).await.unwrap(); - - // The unversioned take is repeated at the described version. - assert_eq!( - requests.lock().unwrap().as_slice(), - [ - TakeMockRequest::Query { - row_ids: vec![10], - version: None, - }, - TakeMockRequest::Describe, - TakeMockRequest::Query { - row_ids: vec![10], - version: Some(9), - }, - TakeMockRequest::Range { - row_id: 10, - range: "bytes=0-3".into(), - version: Some("9".into()), - }, - ] - ); - } - - #[tokio::test] - async fn remote_blob_files_take_at_the_checked_out_version() { - let payload_size = PAYLOAD.len() as u64; - let (table, requests) = mock_take_blob_table( - vec![(10, Some((BlobKind::Inline, payload_size)))], - Some("7"), - ); - table.checkout(9).await.unwrap(); - requests.lock().unwrap().clear(); - - table.fetch_blob_files_impl("image", &[10]).await.unwrap(); - - assert_eq!( - requests.lock().unwrap().as_slice(), - [TakeMockRequest::Query { - row_ids: vec![10], - version: Some(9), - }] - ); - } - - #[tokio::test] - async fn remote_blob_files_reject_row_ids_missing_from_the_take() { - let (table, requests) = mock_take_blob_table(vec![(10, None)], Some("7")); - - let error = table - .fetch_blob_files_impl("image", &[10, 99]) - .await - .unwrap_err(); - - assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); - assert!(error.to_string().contains("only 1 exist"), "{error}"); - assert_eq!(requests.lock().unwrap().len(), 1); - } - - #[tokio::test] - async fn remote_blob_files_reject_a_non_blob_column() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |request| { - assert_eq!(request.url().path(), "/v1/table/my_table/query/"); - let batch = RecordBatch::try_from_iter([ - ( - "id", - Arc::new(arrow_array::Int64Array::from(vec![1])) as Arc, - ), - ( - ROW_ID, - Arc::new(UInt64Array::from(vec![10])) as Arc, - ), - ]) - .unwrap(); - let mut body = Vec::new(); - let mut writer = - arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); - writer.write(&batch).unwrap(); - writer.finish().unwrap(); - drop(writer); - http::Response::builder().status(200).body(body).unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - - let error = table.fetch_blob_files_impl("id", &[10]).await.unwrap_err(); - - assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); - assert!( - error.to_string().contains("'id' is not a blob column"), - "{error}" - ); - } - - #[tokio::test] - async fn remote_blob_file_reads_the_requested_range() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - assert_eq!(file.read_range(5..12).await.unwrap(), &PAYLOAD[5..12]); - assert!(requests.lock().unwrap().contains(&"bytes=5-11".to_string())); - } - - #[tokio::test] - async fn remote_blob_file_reads_the_probed_version_after_table_changes() { - let changed = Arc::new(AtomicBool::new(false)); - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = RemoteTable::new_mock( - "my_table".to_string(), - { - let changed = changed.clone(); - let requests = requests.clone(); - move |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap(); - let version = request - .url() - .query_pairs() - .find(|(name, _)| name == "version") - .map(|(_, value)| value.into_owned()); - requests - .lock() - .unwrap() - .push((range.to_string(), version.clone())); - if changed.load(Ordering::SeqCst) && version.as_deref() != Some("5") { - return http::Response::builder() - .status(StatusCode::BAD_REQUEST) - .body(b"row id is absent from the latest version".to_vec()) - .unwrap(); - } - range_response(&request, PAYLOAD) - } - }, - Some(Version::new(0, 5, 0)), - ); - - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - changed.store(true, Ordering::SeqCst); - - assert_eq!(file.read_range(5..9).await.unwrap(), &PAYLOAD[5..9]); - assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[..4]); - file.seek(10).await.unwrap(); - assert_eq!(file.read().await.unwrap(), &PAYLOAD[10..]); - - assert_eq!( - requests.lock().unwrap().as_slice(), - [ - ("bytes=0-0".to_string(), None), - ("bytes=5-8".to_string(), Some("5".to_string())), - ("bytes=0-".to_string(), Some("5".to_string())), - ("bytes=10-".to_string(), Some("5".to_string())), - ] - ); - } - - #[tokio::test] - async fn remote_blob_file_keeps_the_open_timeline_after_parent_checkout() { - let range_requests = Arc::new(StdMutex::new(Vec::new())); - let captured = range_requests.clone(); - let table = RemoteTable::new_mock( - "my_table".to_string(), - move |request| match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(r#"{"version":5,"schema":{"fields":[]}}"#.as_bytes().to_vec()) - .unwrap(), - "/v1/table/my_table/blob/image/10/bytes" => { - captured.lock().unwrap().push(( - request.url().query().unwrap_or_default().to_string(), - request.headers().clone(), - )); - range_response(&request, PAYLOAD) - } - path => panic!("unexpected path: {path}"), - }, - Some(Version::new(0, 5, 0)), - ); - - table.checkout(5).await.unwrap(); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - table.checkout_latest().await.unwrap(); - file.read_range(5..12).await.unwrap(); - - let requests = range_requests.lock().unwrap(); - let (query, headers) = requests.last().unwrap(); - assert!(query.contains("version=5")); - assert!(!headers.contains_key("x-lancedb-min-timestamp")); - } - - #[tokio::test] - async fn remote_blob_file_reuses_sequential_response_until_seek() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - assert_eq!(file.read_up_to(4).await.unwrap(), b"0123".as_slice()); - assert_eq!(file.read_up_to(3).await.unwrap(), b"456".as_slice()); - file.seek(20).await.unwrap(); - assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[20..24]); - assert_eq!(file.tell().await.unwrap(), 24); - assert_eq!( - requests.lock().unwrap().as_slice(), - ["bytes=0-0", "bytes=0-", "bytes=20-"] - ); - } - - #[tokio::test] - async fn remote_blob_files_allow_an_empty_blob_without_a_version_header() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - - let mut files = table.probe_files("image", &[10, 20, 30]).await.unwrap(); - - assert_eq!(files.len(), 3); - assert!(files[1].is_none()); - let empty = files[2].take().unwrap(); - assert_eq!(empty.size(), 0); - let probe_requests = requests.lock().unwrap().len(); - assert!(empty.read_range(0..0).await.unwrap().is_empty()); - assert!(empty.read().await.unwrap().is_empty()); - assert_eq!(requests.lock().unwrap().len(), probe_requests); - let nonempty = files[0].take().unwrap(); - assert_eq!(nonempty.size(), PAYLOAD.len() as u64); - } - - #[tokio::test] - async fn remote_blob_files_reject_unsatisfied_probe_for_nonempty_blob() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |_| { - http::Response::builder() - .status(StatusCode::RANGE_NOT_SATISFIABLE) - .header(header::CONTENT_RANGE, "bytes */36") - .body(Vec::new()) - .unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - - let error = table.probe_files("image", &[10]).await.unwrap_err(); - assert!( - error - .to_string() - .contains("blob size probe returned HTTP 416 for a 36-byte blob"), - "got: {error}" - ); - } - - #[tokio::test] - async fn remote_blob_files_reject_a_self_contradictory_probe_response() { - // A zero-length blob must use `416 Content-Range: bytes */0`. - let table = RemoteTable::new_mock( - "my_table".to_string(), - |_| { - http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header(header::CONTENT_RANGE, "bytes 0-0/0") - .body(vec![0u8]) - .unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - - let error = table.probe_files("image", &[10]).await.unwrap_err(); - assert!( - error - .to_string() - .contains("blob size probe returned an invalid Content-Range header"), - "got: {error}" - ); - } - - #[tokio::test] - async fn remote_blob_files_reject_a_probe_without_a_version() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |request| { - let mut response = range_response(&request, PAYLOAD); - response.headers_mut().remove(VERSION_HEADER); - response - }, - Some(Version::new(0, 5, 0)), - ); - - let error = table.probe_files("image", &[10]).await.unwrap_err(); - assert!( - error - .to_string() - .contains("missing or invalid x-lancedb-version"), - "got: {error}" - ); - } - - #[tokio::test] - async fn remote_blob_files_reject_unsupported_server_version() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |_| -> http::Response { - panic!("old servers must be rejected before a range request") - }, - Some(Version::new(0, 4, 9)), - ); - let error = table - .fetch_blob_files_impl("image", &[10]) - .await - .unwrap_err(); - assert!(matches!(error, Error::NotSupported { .. })); - assert!(error.to_string().contains("0.5.0")); - } - - #[tokio::test] - async fn remote_blob_file_rejects_mismatched_content_range() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = RemoteTable::new_mock( - "my_table".to_string(), - { - let requests = requests.clone(); - move |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap() - .to_string(); - requests.lock().unwrap().push(range.clone()); - if range == "bytes=0-0" { - return range_response(&request, PAYLOAD); - } - http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 6-12/{}", PAYLOAD.len()), - ) - .body(PAYLOAD[6..=12].to_vec()) - .unwrap() - } - }, - Some(Version::new(0, 5, 0)), - ); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - let error = file.read_range(5..12).await.unwrap_err(); - assert!( - error.to_string().contains("expected Content-Range"), - "got: {error}" - ); - } - - #[tokio::test] - async fn remote_blob_file_rejects_short_response_body() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - move |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap(); - if range == "bytes=0-0" { - return range_response(&request, PAYLOAD); - } - http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 5-11/{}", PAYLOAD.len()), - ) - .body(PAYLOAD[5..=7].to_vec()) - .unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - let error = file.read_range(5..12).await.unwrap_err(); - assert!( - error.to_string().contains("returned 3 bytes, expected 7"), - "got: {error}" - ); - } - - #[tokio::test] - async fn remote_blob_file_failed_read_preserves_cursor_and_retries_fresh() { - let sequential_requests = Arc::new(AtomicUsize::new(0)); - let table = RemoteTable::new_mock( - "my_table".to_string(), - { - let sequential_requests = sequential_requests.clone(); - move |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap() - .to_string(); - if range == "bytes=0-0" { - return range_response(&request, PAYLOAD); - } - let attempt = sequential_requests.fetch_add(1, Ordering::SeqCst); - if attempt == 0 { - // End the response five bytes early to simulate a truncated - // sequential read. - return http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), - ) - .body(PAYLOAD[..5].to_vec()) - .unwrap(); - } - range_response(&request, PAYLOAD) - } - }, - Some(Version::new(0, 5, 0)), - ); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - let error = file.read_up_to(10).await.unwrap_err(); - assert!( - error - .to_string() - .contains("response ended before the requested blob range"), - "got: {error}" - ); - assert_eq!(file.tell().await.unwrap(), 0); - - // Retry from the last committed cursor with a fresh request. - let retried = file.read_up_to(4).await.unwrap(); - assert_eq!(retried, b"0123".as_slice()); - assert_eq!(sequential_requests.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn remote_blob_file_empty_range_sends_no_request() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - assert!(file.read_range(3..3).await.unwrap().is_empty()); - assert_eq!(requests.lock().unwrap().as_slice(), ["bytes=0-0"]); - } - - #[derive(Debug)] - struct CountingProbeRequester { - index: usize, - in_flight: Arc, - max_in_flight: Arc, - } - - #[async_trait::async_trait] - impl BlobRangeRequester for CountingProbeRequester { - async fn request_range( - &self, - range_header: &str, - _mode: RangeRequestMode, - _version: Option, - ) -> Result<(String, Response)> { - assert_eq!(range_header, "bytes=0-0"); - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.max_in_flight.fetch_max(now, Ordering::SeqCst); - // Earlier probes sleep longer, so later probes finish first and the - // ordered collection has to do real reordering work. - tokio::time::sleep(Duration::from_millis(20 - self.index as u64)).await; - self.in_flight.fetch_sub(1, Ordering::SeqCst); - let response = http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 0-0/{}", 100 + self.index), - ) - .header(VERSION_HEADER, "5") - .body(vec![0u8]) - .unwrap(); - Ok((format!("probe-{}", self.index), Response::from(response))) - } - } - - #[tokio::test] - async fn remote_blob_file_probes_are_bounded_and_preserve_order() { - let in_flight = Arc::new(AtomicUsize::new(0)); - let max_in_flight = Arc::new(AtomicUsize::new(0)); - let probes = (0..16) - .map(|index| { - let requester: Arc = Arc::new(CountingProbeRequester { - index, - in_flight: in_flight.clone(), - max_in_flight: max_in_flight.clone(), - }); - requester - }) - .collect(); - - let files = probe_blob_files(probes).await.unwrap(); - - let sizes: Vec = files.into_iter().map(|file| file.unwrap().size()).collect(); - let expected: Vec = (0..16).map(|index| 100 + index as u64).collect(); - assert_eq!(sizes, expected); - let max = max_in_flight.load(Ordering::SeqCst); - assert!(max > 1, "probes never overlapped"); - assert!( - max <= BLOB_REQUEST_CONCURRENCY, - "{max} probes in flight exceeds the bound" - ); - } - - #[tokio::test] - async fn remote_blob_file_metadata_reports_none() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - - let mut files = table.probe_files("image", &[10]).await.unwrap(); - let file = files.remove(0).unwrap(); - - assert_eq!(file.size(), PAYLOAD.len() as u64); - assert_eq!(file.position(), None); - assert_eq!(file.kind(), None); - assert_eq!(file.data_path(), None); - assert_eq!(file.uri(), None); - } - - #[tokio::test] - async fn closed_remote_blob_file_rejects_every_operation_without_requests() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - - let mut files = table.probe_files("image", &[10]).await.unwrap(); - let file = files.remove(0).unwrap(); - let probe_requests = requests.lock().unwrap().len(); - - file.close().await.unwrap(); - - assert!(file.is_closed().await); - for error in [ - file.read().await.unwrap_err(), - file.read_range(0..1).await.unwrap_err(), - file.read_ranges(&[0..1, 1..2]).await.unwrap_err(), - file.read_up_to(1).await.unwrap_err(), - file.seek(0).await.unwrap_err(), - file.tell().await.unwrap_err(), - ] { - assert!(error.to_string().contains("already closed"), "got: {error}"); - } - assert_eq!(requests.lock().unwrap().len(), probe_requests); - } - - #[tokio::test] - async fn out_of_range_read_fails_without_a_request() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - - let mut files = table.probe_files("image", &[10]).await.unwrap(); - let file = files.remove(0).unwrap(); - let probe_requests = requests.lock().unwrap().len(); - - let past_end = file.read_range(0..file.size() + 1).await.unwrap_err(); - assert!( - past_end.to_string().contains("exceeds blob size"), - "got: {past_end}" - ); - let inverted = file - .read_range(Range { start: 3, end: 1 }) - .await - .unwrap_err(); - assert!( - inverted.to_string().contains("exceeds end"), - "got: {inverted}" - ); - assert_eq!(requests.lock().unwrap().len(), probe_requests); - } - - #[tokio::test] - async fn data_read_rejects_416_response() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap() - .to_string(); - if range == "bytes=0-0" { - return http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 0-0/{}", PAYLOAD.len()), - ) - .header(VERSION_HEADER, "5") - .body(vec![PAYLOAD[0]]) - .unwrap(); - } - http::Response::builder() - .status(StatusCode::RANGE_NOT_SATISFIABLE) - .header(header::CONTENT_RANGE, "bytes */0") - .body(b"stale range".to_vec()) - .unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - - let mut files = table.probe_files("image", &[10]).await.unwrap(); - let file = files.remove(0).unwrap(); - - let error = file.read_range(1..3).await.unwrap_err(); - assert!(error.to_string().contains("416"), "got: {error}"); - } - - #[tokio::test] - async fn empty_row_id_lists_bypass_server_version_gate() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |_| -> http::Response { panic!("an empty selection sends no request") }, - Some(Version::new(0, 4, 9)), - ); - - let files = table.fetch_blob_files_impl("image", &[]).await.unwrap(); - assert!(files.is_empty()); - - let blobs = table.fetch_blobs_impl("image", &[]).await.unwrap(); - assert_eq!(blobs.len(), 0); - } - - #[derive(Debug)] - struct CountingRangeRequester { - in_flight: Arc, - max_in_flight: Arc, - } - - #[async_trait::async_trait] - impl BlobRangeRequester for CountingRangeRequester { - async fn request_range( - &self, - range_header: &str, - _mode: RangeRequestMode, - _version: Option, - ) -> Result<(String, Response)> { - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.max_in_flight.fetch_max(now, Ordering::SeqCst); - tokio::time::sleep(Duration::from_millis(5)).await; - self.in_flight.fetch_sub(1, Ordering::SeqCst); - let (start, end) = range_header - .strip_prefix("bytes=") - .unwrap() - .split_once('-') - .unwrap(); - let start = start.parse::().unwrap(); - let end = end.parse::().unwrap(); - let response = http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes {start}-{end}/{}", PAYLOAD.len()), - ) - .body(PAYLOAD[start..=end].to_vec()) - .unwrap(); - Ok(("range".to_string(), Response::from(response))) - } - } - - #[tokio::test] - async fn read_ranges_run_bounded_and_preserve_order() { - let in_flight = Arc::new(AtomicUsize::new(0)); - let max_in_flight = Arc::new(AtomicUsize::new(0)); - let requester: Arc = Arc::new(CountingRangeRequester { - in_flight, - max_in_flight: max_in_flight.clone(), - }); - let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64, Some(5)); - - let ranges: Vec<_> = (0..16u64).map(|start| start..start + 2).collect(); - let output = file.read_ranges(&ranges).await.unwrap(); - - for (range, bytes) in ranges.iter().zip(&output) { - assert_eq!( - bytes.as_ref(), - &PAYLOAD[range.start as usize..range.end as usize] - ); - } - let max = max_in_flight.load(Ordering::SeqCst); - assert!(max > 1, "range reads never overlapped"); - assert!( - max <= BLOB_REQUEST_CONCURRENCY, - "{max} range reads in flight exceeds the bound" - ); - } - - #[tokio::test] - async fn read_ranges_reject_out_of_bounds_range() { - let requests = Arc::new(StdMutex::new(Vec::new())); - let table = mock_remote_blob_table(requests.clone()); - let file = table - .probe_files("image", &[10]) - .await - .unwrap() - .pop() - .flatten() - .unwrap(); - - let oob = file - .read_ranges(&[0..2, 0..PAYLOAD.len() as u64 + 1]) - .await - .unwrap_err(); - assert!(oob.to_string().contains("exceeds blob size"), "got: {oob}"); - } - - #[tokio::test] - async fn range_read_rejects_200_response() { - let table = RemoteTable::new_mock( - "my_table".to_string(), - |request| { - let range = request - .headers() - .get(header::RANGE) - .unwrap() - .to_str() - .unwrap() - .to_string(); - if range == "bytes=0-0" { - return http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 0-0/{}", PAYLOAD.len()), - ) - .header(VERSION_HEADER, "5") - .body(vec![PAYLOAD[0]]) - .unwrap(); - } - http::Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) - .body(PAYLOAD.to_vec()) - .unwrap() - }, - Some(Version::new(0, 5, 0)), - ); - - let mut files = table.probe_files("image", &[10]).await.unwrap(); - let file = files.remove(0).unwrap(); - let error = file.read_range(1..3).await.unwrap_err(); - assert!(error.to_string().contains("206"), "got: {error}"); - } - - #[derive(Debug)] - struct BlockingRangeRequester { - release: Arc, - started: Arc, - } - - #[async_trait::async_trait] - impl BlobRangeRequester for BlockingRangeRequester { - async fn request_range( - &self, - range_header: &str, - _mode: RangeRequestMode, - _version: Option, - ) -> Result<(String, Response)> { - if range_header.ends_with('-') { - self.started.notify_one(); - self.release.wait().await; - } - let response = http::Response::builder() - .status(StatusCode::PARTIAL_CONTENT) - .header( - header::CONTENT_RANGE, - format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), - ) - .body(PAYLOAD.to_vec()) - .unwrap(); - Ok(("hung".to_string(), Response::from(response))) - } - } - - #[tokio::test] - async fn close_returns_while_a_sequential_read_is_in_flight() { - let release = Arc::new(tokio::sync::Barrier::new(2)); - let started = Arc::new(tokio::sync::Notify::new()); - let requester: Arc = Arc::new(BlockingRangeRequester { - release: release.clone(), - started: started.clone(), - }); - let file = Arc::new(RemoteBlobFile::new( - requester, - PAYLOAD.len() as u64, - Some(5), - )); - - let reader = { - let file = file.clone(); - tokio::spawn(async move { file.read_up_to(4).await }) - }; - started.notified().await; - - tokio::time::timeout(Duration::from_millis(50), file.close()) - .await - .expect("close waited on the hung read") - .unwrap(); - assert!(file.is_closed()); - - release.wait().await; - let error = reader.await.unwrap().unwrap_err(); - assert!(error.to_string().contains("already closed"), "got: {error}"); - assert!( - file.read_range(0..1) - .await - .unwrap_err() - .to_string() - .contains("already closed") - ); - } -} +mod tests; diff --git a/rust/lancedb/src/remote/table/blobs/tests.rs b/rust/lancedb/src/remote/table/blobs/tests.rs new file mode 100644 index 000000000..3cc15d7f6 --- /dev/null +++ b/rust/lancedb/src/remote/table/blobs/tests.rs @@ -0,0 +1,1234 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex as StdMutex}; +use std::time::Duration; + +use reqwest::Request; +use semver::Version; + +use arrow_array::RecordBatch; + +use super::*; + +const PAYLOAD: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyz"; + +fn null_blob_response() -> http::Response> { + http::Response::builder() + .status(StatusCode::NO_CONTENT) + .body(Vec::new()) + .unwrap() +} + +fn empty_blob_response() -> http::Response> { + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */0") + .body(Vec::new()) + .unwrap() +} + +fn range_response(request: &Request, payload: &[u8]) -> http::Response> { + let value = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + let (start, end) = value + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let start = start.parse::().unwrap(); + let end = if end.is_empty() { + payload.len() - 1 + } else { + end.parse::().unwrap() + }; + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes {start}-{end}/{}", payload.len()), + ) + .header(VERSION_HEADER, "5") + .body(payload[start..=end].to_vec()) + .unwrap() +} + +fn mock_remote_blob_table( + requests: Arc>>, +) -> RemoteTable { + RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + assert_eq!(request.method(), reqwest::Method::GET); + let path = request.url().path(); + assert!(path.starts_with("/v1/table/my_table/blob/image/")); + requests.lock().unwrap().push( + request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(), + ); + if path.contains("/20/bytes") { + return null_blob_response(); + } + if path.contains("/30/bytes") { + return empty_blob_response(); + } + range_response(&request, PAYLOAD) + }, + Some(Version::new(0, 5, 0)), + ) +} + +#[tokio::test] +async fn remote_blob_files_probe_sizes_and_preserve_nulls() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.probe_files("image", &[10, 20]).await.unwrap(); + + assert_eq!(files.len(), 2); + assert!(files[1].is_none()); + let file = files[0].take().unwrap(); + assert_eq!(file.size(), PAYLOAD.len() as u64); + assert_eq!( + requests.lock().unwrap().as_slice(), + ["bytes=0-0", "bytes=0-0"] + ); +} + +/// Descriptors as a remote query returns them, one `(kind, size)` per row. +fn descriptors(rows: &[Option<(BlobKind, u64)>]) -> StructArray { + let kinds = rows + .iter() + .map(|row| row.map_or(0, |(kind, _)| kind as u8)) + .collect::(); + let sizes = rows + .iter() + .map(|row| row.map_or(0, |(_, size)| size)) + .collect::(); + let len = rows.len(); + StructArray::new( + lance_core::datatypes::BLOB_V2_DESC_FIELDS.clone(), + vec![ + Arc::new(kinds), + Arc::new(UInt64Array::from(vec![0; len])), + Arc::new(sizes), + Arc::new(arrow_array::UInt32Array::from(vec![0; len])), + Arc::new(arrow_array::StringArray::from(vec![""; len])), + ], + Some(rows.iter().map(Option::is_some).collect()), + ) +} + +impl RemoteTable { + /// Open handles through the per-row size probe alone. + async fn probe_files(&self, column: &str, row_ids: &[u64]) -> Result>> { + let read_snapshot = self.snapshot_read_state().await; + self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot) + .await + } +} + +/// A request the take mock received. +#[derive(Debug, Clone, PartialEq)] +enum TakeMockRequest { + Query { + row_ids: Vec, + version: Option, + }, + Describe, + Range { + row_id: u64, + range: String, + version: Option, + }, +} + +/// Row ids named by a `_rowid IN (...)` filter. +fn filter_row_ids(filter: &str) -> Vec { + let list = filter.split_once('(').unwrap().1.trim_end_matches(')'); + list.split(',') + .map(|id| id.trim().parse().unwrap()) + .collect() +} + +/// A table whose query route answers `_rowid` takes from `rows` and whose +/// byte routes serve [`PAYLOAD`]. Takes return rows in reverse order, and +/// report `query_version` in `x-lancedb-version` when it is set. +fn mock_take_blob_table( + rows: Vec<(u64, Option<(BlobKind, u64)>)>, + query_version: Option<&'static str>, +) -> ( + RemoteTable, + Arc>>, +) { + let requests = Arc::new(StdMutex::new(Vec::new())); + let captured = requests.clone(); + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + let path = request.url().path().to_string(); + if 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"], serde_json::json!(["image"])); + assert_eq!(body["with_row_id"], serde_json::json!(true)); + assert_eq!(body["use_lsm"], serde_json::json!(false)); + let requested = filter_row_ids(body["filter"].as_str().unwrap()); + captured.lock().unwrap().push(TakeMockRequest::Query { + row_ids: requested.clone(), + version: body["version"].as_u64(), + }); + let taken = rows + .iter() + .rev() + .filter(|(row_id, _)| requested.contains(row_id)) + .collect::>(); + let batch = RecordBatch::try_from_iter([ + ( + "image", + Arc::new(descriptors( + &taken.iter().map(|(_, row)| *row).collect::>(), + )) as Arc, + ), + ( + ROW_ID, + Arc::new(UInt64Array::from_iter_values( + taken.iter().map(|(row_id, _)| *row_id), + )) as Arc, + ), + ]) + .unwrap(); + let mut body = Vec::new(); + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + let mut response = http::Response::builder().status(200); + if let Some(version) = query_version { + response = response.header(VERSION_HEADER, version); + } + return response.body(body).unwrap(); + } + if path == "/v1/table/my_table/describe/" { + captured.lock().unwrap().push(TakeMockRequest::Describe); + return http::Response::builder() + .status(200) + .body(r#"{"version":9,"schema":{"fields":[]}}"#.as_bytes().to_vec()) + .unwrap(); + } + let row_id = path + .strip_prefix("/v1/table/my_table/blob/image/") + .and_then(|rest| rest.strip_suffix("/bytes")) + .unwrap_or_else(|| panic!("unexpected path: {path}")) + .parse() + .unwrap(); + captured.lock().unwrap().push(TakeMockRequest::Range { + row_id, + range: request.headers()[header::RANGE] + .to_str() + .unwrap() + .to_string(), + version: request + .url() + .query_pairs() + .find(|(name, _)| name == "version") + .map(|(_, value)| value.into_owned()), + }); + range_response(&request, PAYLOAD) + }, + Some(Version::new(0, 5, 0)), + ); + (table, requests) +} + +#[tokio::test] +async fn remote_blob_files_are_sized_from_one_descriptor_take() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = mock_take_blob_table( + vec![ + (10, Some((BlobKind::Inline, payload_size))), + (20, None), + (30, Some((BlobKind::Inline, 0))), + (40, Some((BlobKind::Packed, payload_size))), + (50, Some((BlobKind::Dedicated, payload_size))), + (60, Some((BlobKind::External, 0))), + ], + Some("7"), + ); + + // Reordered and duplicated row ids come back in request order. + let files = table + .fetch_blob_files_impl("image", &[60, 10, 20, 30, 40, 50, 10]) + .await + .unwrap(); + + let sizes = files + .iter() + .map(|file| file.as_ref().map(BlobFile::size)) + .collect::>(); + assert_eq!( + sizes, + [ + Some(payload_size), + Some(payload_size), + None, + Some(0), + Some(payload_size), + Some(payload_size), + Some(payload_size), + ] + ); + // One take, then a probe for the external blob that recorded no size. + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + TakeMockRequest::Query { + row_ids: vec![10, 20, 30, 40, 50, 60], + version: None, + }, + TakeMockRequest::Range { + row_id: 60, + range: "bytes=0-0".into(), + version: Some("7".into()), + }, + ] + ); + + // Sized handles read the version the take read. + assert_eq!( + files[1].as_ref().unwrap().read_range(5..12).await.unwrap(), + &PAYLOAD[5..12] + ); + assert!(files[3].as_ref().unwrap().read().await.unwrap().is_empty()); + assert_eq!( + requests.lock().unwrap().last().unwrap(), + &TakeMockRequest::Range { + row_id: 10, + range: "bytes=5-11".into(), + version: Some("7".into()), + } + ); +} + +#[tokio::test] +async fn remote_blob_files_take_every_chunk_at_the_first_version() { + let row_ids = (0..MAX_FETCH_BLOBS_ROW_IDS as u64 + 1).collect::>(); + let (table, requests) = mock_take_blob_table( + row_ids.iter().map(|row_id| (*row_id, None)).collect(), + Some("7"), + ); + + let files = table + .fetch_blob_files_impl("image", &row_ids) + .await + .unwrap(); + + assert_eq!(files.len(), row_ids.len()); + let versions = requests + .lock() + .unwrap() + .iter() + .map(|request| match request { + TakeMockRequest::Query { row_ids, version } => (row_ids.len(), *version), + other => panic!("unexpected request: {other:?}"), + }) + .collect::>(); + assert_eq!(versions, [(MAX_FETCH_BLOBS_ROW_IDS, None), (1, Some(7))]); +} + +#[tokio::test] +async fn remote_blob_files_retake_at_the_described_version_when_a_take_does_not_report_it() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = + mock_take_blob_table(vec![(10, Some((BlobKind::Inline, payload_size)))], None); + + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + file.read_range(0..4).await.unwrap(); + + // The unversioned take is repeated at the described version. + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + TakeMockRequest::Query { + row_ids: vec![10], + version: None, + }, + TakeMockRequest::Describe, + TakeMockRequest::Query { + row_ids: vec![10], + version: Some(9), + }, + TakeMockRequest::Range { + row_id: 10, + range: "bytes=0-3".into(), + version: Some("9".into()), + }, + ] + ); +} + +#[tokio::test] +async fn remote_blob_files_take_at_the_checked_out_version() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = mock_take_blob_table( + vec![(10, Some((BlobKind::Inline, payload_size)))], + Some("7"), + ); + table.checkout(9).await.unwrap(); + requests.lock().unwrap().clear(); + + table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + + assert_eq!( + requests.lock().unwrap().as_slice(), + [TakeMockRequest::Query { + row_ids: vec![10], + version: Some(9), + }] + ); +} + +#[tokio::test] +async fn remote_blob_files_reject_row_ids_missing_from_the_take() { + let (table, requests) = mock_take_blob_table(vec![(10, None)], Some("7")); + + let error = table + .fetch_blob_files_impl("image", &[10, 99]) + .await + .unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains("only 1 exist"), "{error}"); + assert_eq!(requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn remote_blob_files_reject_a_non_blob_column() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/query/"); + let batch = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(arrow_array::Int64Array::from(vec![1])) as Arc, + ), + ( + ROW_ID, + Arc::new(UInt64Array::from(vec![10])) as Arc, + ), + ]) + .unwrap(); + let mut body = Vec::new(); + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + http::Response::builder().status(200).body(body).unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table.fetch_blob_files_impl("id", &[10]).await.unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!( + error.to_string().contains("'id' is not a blob column"), + "{error}" + ); +} + +#[tokio::test] +async fn remote_blob_file_reads_the_requested_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert_eq!(file.read_range(5..12).await.unwrap(), &PAYLOAD[5..12]); + assert!(requests.lock().unwrap().contains(&"bytes=5-11".to_string())); +} + +#[tokio::test] +async fn remote_blob_file_reads_the_probed_version_after_table_changes() { + let changed = Arc::new(AtomicBool::new(false)); + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = RemoteTable::new_mock( + "my_table".to_string(), + { + let changed = changed.clone(); + let requests = requests.clone(); + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + let version = request + .url() + .query_pairs() + .find(|(name, _)| name == "version") + .map(|(_, value)| value.into_owned()); + requests + .lock() + .unwrap() + .push((range.to_string(), version.clone())); + if changed.load(Ordering::SeqCst) && version.as_deref() != Some("5") { + return http::Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(b"row id is absent from the latest version".to_vec()) + .unwrap(); + } + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + changed.store(true, Ordering::SeqCst); + + assert_eq!(file.read_range(5..9).await.unwrap(), &PAYLOAD[5..9]); + assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[..4]); + file.seek(10).await.unwrap(); + assert_eq!(file.read().await.unwrap(), &PAYLOAD[10..]); + + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + ("bytes=0-0".to_string(), None), + ("bytes=5-8".to_string(), Some("5".to_string())), + ("bytes=0-".to_string(), Some("5".to_string())), + ("bytes=10-".to_string(), Some("5".to_string())), + ] + ); +} + +#[tokio::test] +async fn remote_blob_file_keeps_the_open_timeline_after_parent_checkout() { + let range_requests = Arc::new(StdMutex::new(Vec::new())); + let captured = range_requests.clone(); + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(r#"{"version":5,"schema":{"fields":[]}}"#.as_bytes().to_vec()) + .unwrap(), + "/v1/table/my_table/blob/image/10/bytes" => { + captured.lock().unwrap().push(( + request.url().query().unwrap_or_default().to_string(), + request.headers().clone(), + )); + range_response(&request, PAYLOAD) + } + path => panic!("unexpected path: {path}"), + }, + Some(Version::new(0, 5, 0)), + ); + + table.checkout(5).await.unwrap(); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + table.checkout_latest().await.unwrap(); + file.read_range(5..12).await.unwrap(); + + let requests = range_requests.lock().unwrap(); + let (query, headers) = requests.last().unwrap(); + assert!(query.contains("version=5")); + assert!(!headers.contains_key("x-lancedb-min-timestamp")); +} + +#[tokio::test] +async fn remote_blob_file_reuses_sequential_response_until_seek() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert_eq!(file.read_up_to(4).await.unwrap(), b"0123".as_slice()); + assert_eq!(file.read_up_to(3).await.unwrap(), b"456".as_slice()); + file.seek(20).await.unwrap(); + assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[20..24]); + assert_eq!(file.tell().await.unwrap(), 24); + assert_eq!( + requests.lock().unwrap().as_slice(), + ["bytes=0-0", "bytes=0-", "bytes=20-"] + ); +} + +#[tokio::test] +async fn remote_blob_files_allow_an_empty_blob_without_a_version_header() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.probe_files("image", &[10, 20, 30]).await.unwrap(); + + assert_eq!(files.len(), 3); + assert!(files[1].is_none()); + let empty = files[2].take().unwrap(); + assert_eq!(empty.size(), 0); + let probe_requests = requests.lock().unwrap().len(); + assert!(empty.read_range(0..0).await.unwrap().is_empty()); + assert!(empty.read().await.unwrap().is_empty()); + assert_eq!(requests.lock().unwrap().len(), probe_requests); + let nonempty = files[0].take().unwrap(); + assert_eq!(nonempty.size(), PAYLOAD.len() as u64); +} + +#[tokio::test] +async fn remote_blob_files_reject_unsatisfied_probe_for_nonempty_blob() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| { + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */36") + .body(Vec::new()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table.probe_files("image", &[10]).await.unwrap_err(); + assert!( + error + .to_string() + .contains("blob size probe returned HTTP 416 for a 36-byte blob"), + "got: {error}" + ); +} + +#[tokio::test] +async fn remote_blob_files_reject_a_self_contradictory_probe_response() { + // A zero-length blob must use `416 Content-Range: bytes */0`. + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| { + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header(header::CONTENT_RANGE, "bytes 0-0/0") + .body(vec![0u8]) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table.probe_files("image", &[10]).await.unwrap_err(); + assert!( + error + .to_string() + .contains("blob size probe returned an invalid Content-Range header"), + "got: {error}" + ); +} + +#[tokio::test] +async fn remote_blob_files_reject_a_probe_without_a_version() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + let mut response = range_response(&request, PAYLOAD); + response.headers_mut().remove(VERSION_HEADER); + response + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table.probe_files("image", &[10]).await.unwrap_err(); + assert!( + error + .to_string() + .contains("missing or invalid x-lancedb-version"), + "got: {error}" + ); +} + +#[tokio::test] +async fn remote_blob_files_reject_unsupported_server_version() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| -> http::Response { + panic!("old servers must be rejected before a range request") + }, + Some(Version::new(0, 4, 9)), + ); + let error = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + assert!(error.to_string().contains("0.5.0")); +} + +#[tokio::test] +async fn remote_blob_file_rejects_mismatched_content_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = RemoteTable::new_mock( + "my_table".to_string(), + { + let requests = requests.clone(); + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + requests.lock().unwrap().push(range.clone()); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 6-12/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[6..=12].to_vec()) + .unwrap() + } + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_range(5..12).await.unwrap_err(); + assert!( + error.to_string().contains("expected Content-Range"), + "got: {error}" + ); +} + +#[tokio::test] +async fn remote_blob_file_rejects_short_response_body() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 5-11/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[5..=7].to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_range(5..12).await.unwrap_err(); + assert!( + error.to_string().contains("returned 3 bytes, expected 7"), + "got: {error}" + ); +} + +#[tokio::test] +async fn remote_blob_file_failed_read_preserves_cursor_and_retries_fresh() { + let sequential_requests = Arc::new(AtomicUsize::new(0)); + let table = RemoteTable::new_mock( + "my_table".to_string(), + { + let sequential_requests = sequential_requests.clone(); + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + let attempt = sequential_requests.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + // End the response five bytes early to simulate a truncated + // sequential read. + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + ) + .body(PAYLOAD[..5].to_vec()) + .unwrap(); + } + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_up_to(10).await.unwrap_err(); + assert!( + error + .to_string() + .contains("response ended before the requested blob range"), + "got: {error}" + ); + assert_eq!(file.tell().await.unwrap(), 0); + + // Retry from the last committed cursor with a fresh request. + let retried = file.read_up_to(4).await.unwrap(); + assert_eq!(retried, b"0123".as_slice()); + assert_eq!(sequential_requests.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn remote_blob_file_empty_range_sends_no_request() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert!(file.read_range(3..3).await.unwrap().is_empty()); + assert_eq!(requests.lock().unwrap().as_slice(), ["bytes=0-0"]); +} + +#[derive(Debug)] +struct CountingProbeRequester { + index: usize, + in_flight: Arc, + max_in_flight: Arc, +} + +#[async_trait::async_trait] +impl BlobRangeRequester for CountingProbeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + _version: Option, + ) -> Result<(String, Response)> { + assert_eq!(range_header, "bytes=0-0"); + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + // Earlier probes sleep longer, so later probes finish first and the + // ordered collection has to do real reordering work. + tokio::time::sleep(Duration::from_millis(20 - self.index as u64)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", 100 + self.index), + ) + .header(VERSION_HEADER, "5") + .body(vec![0u8]) + .unwrap(); + Ok((format!("probe-{}", self.index), Response::from(response))) + } +} + +#[tokio::test] +async fn remote_blob_file_probes_are_bounded_and_preserve_order() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let probes = (0..16) + .map(|index| { + let requester: Arc = Arc::new(CountingProbeRequester { + index, + in_flight: in_flight.clone(), + max_in_flight: max_in_flight.clone(), + }); + requester + }) + .collect(); + + let files = probe_blob_files(probes).await.unwrap(); + + let sizes: Vec = files.into_iter().map(|file| file.unwrap().size()).collect(); + let expected: Vec = (0..16).map(|index| 100 + index as u64).collect(); + assert_eq!(sizes, expected); + let max = max_in_flight.load(Ordering::SeqCst); + assert!(max > 1, "probes never overlapped"); + assert!( + max <= BLOB_REQUEST_CONCURRENCY, + "{max} probes in flight exceeds the bound" + ); +} + +#[tokio::test] +async fn remote_blob_file_metadata_reports_none() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.probe_files("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + + assert_eq!(file.size(), PAYLOAD.len() as u64); + assert_eq!(file.position(), None); + assert_eq!(file.kind(), None); + assert_eq!(file.data_path(), None); + assert_eq!(file.uri(), None); +} + +#[tokio::test] +async fn closed_remote_blob_file_rejects_every_operation_without_requests() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.probe_files("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let probe_requests = requests.lock().unwrap().len(); + + file.close().await.unwrap(); + + assert!(file.is_closed().await); + for error in [ + file.read().await.unwrap_err(), + file.read_range(0..1).await.unwrap_err(), + file.read_ranges(&[0..1, 1..2]).await.unwrap_err(), + file.read_up_to(1).await.unwrap_err(), + file.seek(0).await.unwrap_err(), + file.tell().await.unwrap_err(), + ] { + assert!(error.to_string().contains("already closed"), "got: {error}"); + } + assert_eq!(requests.lock().unwrap().len(), probe_requests); +} + +#[tokio::test] +async fn out_of_range_read_fails_without_a_request() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.probe_files("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let probe_requests = requests.lock().unwrap().len(); + + let past_end = file.read_range(0..file.size() + 1).await.unwrap_err(); + assert!( + past_end.to_string().contains("exceeds blob size"), + "got: {past_end}" + ); + let inverted = file + .read_range(Range { start: 3, end: 1 }) + .await + .unwrap_err(); + assert!( + inverted.to_string().contains("exceeds end"), + "got: {inverted}" + ); + assert_eq!(requests.lock().unwrap().len(), probe_requests); +} + +#[tokio::test] +async fn data_read_rejects_416_response() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", PAYLOAD.len()), + ) + .header(VERSION_HEADER, "5") + .body(vec![PAYLOAD[0]]) + .unwrap(); + } + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */0") + .body(b"stale range".to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let mut files = table.probe_files("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + + let error = file.read_range(1..3).await.unwrap_err(); + assert!(error.to_string().contains("416"), "got: {error}"); +} + +#[tokio::test] +async fn empty_row_id_lists_bypass_server_version_gate() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| -> http::Response { panic!("an empty selection sends no request") }, + Some(Version::new(0, 4, 9)), + ); + + let files = table.fetch_blob_files_impl("image", &[]).await.unwrap(); + assert!(files.is_empty()); + + let blobs = table.fetch_blobs_impl("image", &[]).await.unwrap(); + assert_eq!(blobs.len(), 0); +} + +#[derive(Debug)] +struct CountingRangeRequester { + in_flight: Arc, + max_in_flight: Arc, +} + +#[async_trait::async_trait] +impl BlobRangeRequester for CountingRangeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + _version: Option, + ) -> Result<(String, Response)> { + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(5)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + let (start, end) = range_header + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let start = start.parse::().unwrap(); + let end = end.parse::().unwrap(); + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes {start}-{end}/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[start..=end].to_vec()) + .unwrap(); + Ok(("range".to_string(), Response::from(response))) + } +} + +#[tokio::test] +async fn read_ranges_run_bounded_and_preserve_order() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let requester: Arc = Arc::new(CountingRangeRequester { + in_flight, + max_in_flight: max_in_flight.clone(), + }); + let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64, Some(5)); + + let ranges: Vec<_> = (0..16u64).map(|start| start..start + 2).collect(); + let output = file.read_ranges(&ranges).await.unwrap(); + + for (range, bytes) in ranges.iter().zip(&output) { + assert_eq!( + bytes.as_ref(), + &PAYLOAD[range.start as usize..range.end as usize] + ); + } + let max = max_in_flight.load(Ordering::SeqCst); + assert!(max > 1, "range reads never overlapped"); + assert!( + max <= BLOB_REQUEST_CONCURRENCY, + "{max} range reads in flight exceeds the bound" + ); +} + +#[tokio::test] +async fn read_ranges_reject_out_of_bounds_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let oob = file + .read_ranges(&[0..2, 0..PAYLOAD.len() as u64 + 1]) + .await + .unwrap_err(); + assert!(oob.to_string().contains("exceeds blob size"), "got: {oob}"); +} + +#[tokio::test] +async fn range_read_rejects_200_response() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", PAYLOAD.len()), + ) + .header(VERSION_HEADER, "5") + .body(vec![PAYLOAD[0]]) + .unwrap(); + } + http::Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) + .body(PAYLOAD.to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let mut files = table.probe_files("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let error = file.read_range(1..3).await.unwrap_err(); + assert!(error.to_string().contains("206"), "got: {error}"); +} + +#[derive(Debug)] +struct BlockingRangeRequester { + release: Arc, + started: Arc, +} + +#[async_trait::async_trait] +impl BlobRangeRequester for BlockingRangeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + _version: Option, + ) -> Result<(String, Response)> { + if range_header.ends_with('-') { + self.started.notify_one(); + self.release.wait().await; + } + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + ) + .body(PAYLOAD.to_vec()) + .unwrap(); + Ok(("hung".to_string(), Response::from(response))) + } +} + +#[tokio::test] +async fn close_returns_while_a_sequential_read_is_in_flight() { + let release = Arc::new(tokio::sync::Barrier::new(2)); + let started = Arc::new(tokio::sync::Notify::new()); + let requester: Arc = Arc::new(BlockingRangeRequester { + release: release.clone(), + started: started.clone(), + }); + let file = Arc::new(RemoteBlobFile::new( + requester, + PAYLOAD.len() as u64, + Some(5), + )); + + let reader = { + let file = file.clone(); + tokio::spawn(async move { file.read_up_to(4).await }) + }; + started.notified().await; + + tokio::time::timeout(Duration::from_millis(50), file.close()) + .await + .expect("close waited on the hung read") + .unwrap(); + assert!(file.is_closed()); + + release.wait().await; + let error = reader.await.unwrap().unwrap_err(); + assert!(error.to_string().contains("already closed"), "got: {error}"); + assert!( + file.read_range(0..1) + .await + .unwrap_err() + .to_string() + .contains("already closed") + ); +} diff --git a/rust/lancedb/src/remote/table/insert.rs b/rust/lancedb/src/remote/table/insert.rs index eef1d8e42..2acb6366e 100644 --- a/rust/lancedb/src/remote/table/insert.rs +++ b/rust/lancedb/src/remote/table/insert.rs @@ -869,730 +869,4 @@ impl ExecutionPlan for RemoteWriteExec { } #[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, - batches: Vec, - ) -> Arc { - 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, - partitions: Vec>, - ) -> Arc { - 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, - ) -> crate::remote::client::RestfulLanceDbClient - { - 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>>, - ) -> crate::remote::client::RestfulLanceDbClient - { - 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, - properties: Arc, - } - - 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 { - &self.properties - } - fn children(&self) -> Vec<&Arc> { - vec![] - } - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> DataFusionResult> { - Ok(self) - } - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> DataFusionResult { - 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 = 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 = 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::::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; diff --git a/rust/lancedb/src/remote/table/insert/tests.rs b/rust/lancedb/src/remote/table/insert/tests.rs new file mode 100644 index 000000000..d08ffa4b1 --- /dev/null +++ b/rust/lancedb/src/remote/table/insert/tests.rs @@ -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, + batches: Vec, +) -> Arc { + 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, + partitions: Vec>, +) -> Arc { + 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, +) -> crate::remote::client::RestfulLanceDbClient { + 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>>, +) -> crate::remote::client::RestfulLanceDbClient { + 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, + properties: Arc, +} + +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 { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + 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 = 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 = 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::::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:?}" + ); +} diff --git a/rust/lancedb/src/remote/table/tests/blobs.rs b/rust/lancedb/src/remote/table/tests/blobs.rs new file mode 100644 index 000000000..0a19db04b --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/blobs.rs @@ -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::>() + .await; + assert_eq!(data.len(), 1); + assert_eq!(data[0].as_ref().unwrap(), &expected_data); +} + +fn blob_describe_response() -> http::Response { + 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 = (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) -> 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) -> 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 { + 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 { + 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 { + 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 { + 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, +) { + 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"); +} diff --git a/rust/lancedb/src/remote/table/tests/branches.rs b/rust/lancedb/src/remote/table/tests/branches.rs new file mode 100644 index 000000000..2437abbc3 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/branches.rs @@ -0,0 +1,1223 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[tokio::test] +async fn test_update_field_metadata() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.url().path(), + "/v1/table/my_table/update_field_metadata/" + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 7, "fields": {"category": {"unit": "label"}}}"#) + .unwrap() + }); + + let result = table + .update_field_metadata(&[FieldMetadataUpdate::new("category").set("unit", "label")]) + .await + .unwrap(); + assert_eq!(result.version, 7); +} + +#[tokio::test] +async fn test_create_branch_default_source() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/branches/create/"); + let body = request_body_json(&request); + assert_eq!(body["name"], "exp"); + assert!( + body.get("from_branch").is_none(), + "a main source omits from_branch" + ); + assert!( + body.get("from_version").is_none(), + "a latest source omits from_version" + ); + http::Response::builder().status(200).body("{}").unwrap() + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert_eq!(branch.current_branch(), Some("exp".to_string())); + assert_eq!(table.current_branch(), None); +} + +#[tokio::test] +async fn test_create_branch_from_branch_and_version() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| { + let body = request_body_json(&request); + assert_eq!(body["name"], "exp"); + assert_eq!(body["from_branch"], "base"); + assert_eq!(body["from_version"], 3); + http::Response::builder().status(200).body("{}").unwrap() + }); + table + .create_branch("exp", Ref::Version(Some("base".into()), Some(3))) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_branch_from_main_normalizes_to_none() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| { + let body = request_body_json(&request); + assert!( + body.get("from_branch").is_none(), + "\"main\" normalizes to an absent from_branch" + ); + assert_eq!(body["from_version"], 7); + http::Response::builder().status(200).body("{}").unwrap() + }); + table + .create_branch("exp", Ref::Version(Some("main".into()), Some(7))) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_branch_from_version_number_on_main() { + use lance::dataset::refs::Ref; + // A bare version number on a main handle resolves to (main, version). + let table = Table::new_with_handler("my_table", |request| { + let body = request_body_json(&request); + assert!(body.get("from_branch").is_none()); + assert_eq!(body["from_version"], 5); + http::Response::builder().status(200).body("{}").unwrap() + }); + table + .create_branch("exp", Ref::VersionNumber(5)) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_branch_from_tag_resolves_via_tags_endpoint() { + use lance::dataset::refs::Ref; + // A tag source has no from_tag in the create contract; it is resolved to + // its (branch, version) via the tags/version endpoint first. + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/tags/version/" => { + assert_eq!(request_body_json(&request)["tag"], "t"); + http::Response::builder() + .status(200) + .body(r#"{"version":3,"branch":"base"}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/branches/create/" => { + let body = request_body_json(&request); + assert_eq!(body["name"], "exp"); + assert_eq!(body["from_branch"], "base"); + assert_eq!(body["from_version"], 3); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + table + .create_branch("exp", Ref::Tag("t".into())) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_branch_from_tag_on_main_normalizes() { + use lance::dataset::refs::Ref; + // A tag resolving to the main branch collapses from_branch to absent. + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/tags/version/" => http::Response::builder() + .status(200) + .body(r#"{"version":4,"branch":"main"}"#.to_string()) + .unwrap(), + "/v1/table/my_table/branches/create/" => { + let body = request_body_json(&request); + assert!( + body.get("from_branch").is_none(), + "a resolved \"main\" normalizes to an absent from_branch" + ); + assert_eq!(body["from_version"], 4); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + table + .create_branch("exp", Ref::Tag("t".into())) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_branch_invalid_request_maps_to_invalid_input() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(400) + .body("unsafe branch name") + .unwrap() + }); + let err = table + .create_branch("../evil", Ref::Version(None, None)) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); +} + +#[tokio::test] +async fn test_create_branch_conflict_maps_to_already_exists() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(409) + .body("branch already exists") + .unwrap() + }); + let err = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap_err(); + assert!( + matches!(err, Error::TableAlreadyExists { .. }), + "409 should map to AlreadyExists, got {err:?}" + ); +} + +#[tokio::test] +async fn test_materialized_view_describe_and_refresh_requires_sql() { + const QUERY: &str = + "SELECT s.x, d.label FROM analytics.source s JOIN analytics.dim d ON s.id = d.id"; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/materialized_view/my_table/describe" => http::Response::builder() + .status(200) + .body( + json!({ + "name": "my_table", + "query": QUERY + }) + .to_string(), + ) + .unwrap(), + path => panic!("unexpected request: {path}"), + }); + let view = crate::MaterializedView::from_table(table).await.unwrap(); + assert_eq!(view.definition_sql(), QUERY); + assert!(view.definition().is_err()); + let error = view.refresh().execute().await.unwrap_err(); + assert!(error.to_string().contains("SQL is unavailable")); +} + +#[tokio::test] +async fn test_create_branch_empty_name_rejected_client_side() { + use lance::dataset::refs::Ref; + // The empty name is rejected before any request is sent. + let table = Table::new_with_handler("my_table", |request| -> http::Response { + panic!("unexpected request: {}", request.url().path()) + }); + let err = table + .create_branch("", Ref::Version(None, None)) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); +} + +#[tokio::test] +async fn test_list_branches() { + use lance::dataset::refs::BranchIdentifier; + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/branches/list/"); + // A branch forked off main: the server omits `parentBranch` entirely + // (skip_serializing_if), not `null`, so this mirrors the real wire. + http::Response::builder() + .status(200) + .body(r#"{"branches":{"exp":{"parentVersion":2,"createAt":1234,"manifestSize":4096}}}"#) + .unwrap() + }); + let branches = table.list_branches().await.unwrap(); + let exp = branches.get("exp").expect("exp present"); + assert_eq!(exp.parent_version, 2); + assert_eq!(exp.create_at, 1234); + assert_eq!(exp.manifest_size, 4096); + assert_eq!(exp.parent_branch, None); + assert!(exp.metadata.is_empty()); + // The server omits the internal lineage token; it defaults to the sentinel. + assert_eq!( + exp.identifier, + BranchIdentifier::missing_identifier_sentinel() + ); +} + +#[tokio::test] +async fn test_main_only_metadata_is_unfenced_from_branch_timeline() { + let requests = Arc::new(std::sync::Mutex::new(HashMap::new())); + let captured = requests.clone(); + let saw_delete_response_floor = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw_high_floor = saw_delete_response_floor.clone(); + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + let path = request.url().path().to_string(); + captured + .lock() + .unwrap() + .insert(path.clone(), request.headers().clone()); + match path.as_str() { + "/v1/table/my_table/count_rows/" => { + saw_high_floor.store( + request + .headers() + .get("x-lancedb-min-read-version") + .and_then(|value| value.to_str().ok()) + == Some("100"), + std::sync::atomic::Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .header("x-lancedb-version", "2") + .body("1".to_string()) + .unwrap() + } + "/v1/table/my_table/tags/list/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/tags/version/" => http::Response::builder() + .status(200) + .body(r#"{"version":1}"#.to_string()) + .unwrap(), + "/v1/table/my_table/tags/delete/" => http::Response::builder() + .status(200) + .header("x-lancedb-version", "100") + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/branches/list/" => http::Response::builder() + .status(200) + .body(r#"{"branches":{}}"#.to_string()) + .unwrap(), + path => panic!("unexpected path: {path}"), + } + }, + None, + ); + let branch = table.with_branch(Some("exp".to_string())); + + branch.count_rows(None).await.unwrap(); + let mut tags = branch.tags().await.unwrap(); + tags.list().await.unwrap(); + tags.get_version("v1").await.unwrap(); + tags.delete("v1").await.unwrap(); + branch.list_branches().await.unwrap(); + branch.count_rows(None).await.unwrap(); + + let requests = requests.lock().unwrap(); + for path in [ + "/v1/table/my_table/tags/list/", + "/v1/table/my_table/tags/version/", + "/v1/table/my_table/tags/delete/", + "/v1/table/my_table/branches/list/", + ] { + assert!( + !requests[path].contains_key("x-lancedb-min-read-version"), + "{path} inherited the branch timeline" + ); + } + assert!( + !saw_delete_response_floor.load(std::sync::atomic::Ordering::SeqCst), + "tag deletion contaminated the branch timeline" + ); +} + +#[tokio::test] +async fn test_delete_branch() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/branches/delete/"); + let body = request_body_json(&request); + assert_eq!(body["name"], "exp"); + http::Response::builder().status(200).body("{}").unwrap() + }); + table.delete_branch("exp").await.unwrap(); +} + +#[tokio::test] +async fn test_delete_branch_not_found() { + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(404) + .body("no such branch") + .unwrap() + }); + let err = table.delete_branch("ghost").await.unwrap_err(); + assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}"); +} + +fn sample_branch_diff_json() -> &'static str { + r#"{ + "fromBranch":"exp", + "parentVersion":1, + "mainVersion":1, + "branchVersion":2, + "baseMoved":false, + "rowCountMain":3, + "rowCountBranch":3, + "rowSummary":{ + "unchanged":3, + "newOnBase":0, + "newOnBranch":0, + "staleRecompute":0, + "inputsChanged":0, + "deltaAvailable":false + }, + "addedColumns":[{"name":"tag","dataType":"utf8","nullable":true}], + "removedColumns":[], + "changedColumns":[], + "addedIndexes":[], + "removedIndexes":[], + "errors":[] + }"# +} + +#[tokio::test] +async fn test_diff_branch() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/branches/diff/"); + let body = request_body_json(&request); + assert_eq!(body["from_branch"], "exp"); + http::Response::builder() + .status(200) + .body(sample_branch_diff_json()) + .unwrap() + }); + let diff = table.diff_branch("exp").await.unwrap(); + assert_eq!(diff.from_branch, "exp"); + assert!(diff.errors.is_empty()); + assert_eq!(diff.added_columns.len(), 1); + assert_eq!(diff.added_columns[0].name, "tag"); +} + +#[tokio::test] +async fn test_cherry_pick_dry_run() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!( + request.url().path(), + "/v1/table/my_table/branches/cherry_pick/" + ); + let body = request_body_json(&request); + assert_eq!(body["from_branch"], "exp"); + assert_eq!(body["dry_run"], true); + let resp = format!( + r#"{{"status":"ready","diff":{},"preview":{{"promotedColumns":["tag"]}}}}"#, + sample_branch_diff_json() + ); + http::Response::builder().status(200).body(resp).unwrap() + }); + let result = table.cherry_pick("exp", true).await.unwrap(); + assert_eq!(result.status, crate::table::CherryPickStatus::Ready); + assert_eq!(result.preview.promoted_columns, vec!["tag".to_string()]); + assert!(result.main_version_after.is_none()); +} + +#[tokio::test] +async fn test_successful_cherry_pick_advances_main_read_watermark() { + 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/branches/cherry_pick/" => { + let response = serde_json::json!({ + "status": "cherryPicked", + "diff": serde_json::from_str::(sample_branch_diff_json()) + .unwrap(), + "preview": { "promotedColumns": ["tag"] }, + "mainVersionAfter": 2 + }); + http::Response::builder() + .status(200) + .body(response.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 result = table.cherry_pick("exp", false).await.unwrap(); + assert_eq!(result.status, crate::table::CherryPickStatus::CherryPicked); + table.count_rows(None).await.unwrap(); + assert_eq!( + count_headers + .lock() + .unwrap() + .as_ref() + .unwrap() + .get("x-lancedb-min-read-version") + .and_then(|value| value.to_str().ok()), + Some("2") + ); +} + +#[tokio::test] +async fn test_cherry_pick_failed_returns_ok_with_body() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!( + request.url().path(), + "/v1/table/my_table/branches/cherry_pick/" + ); + let body = request_body_json(&request); + assert_eq!(body["dry_run"], false); + let mut diff: serde_json::Value = serde_json::from_str(sample_branch_diff_json()).unwrap(); + diff["errors"] = serde_json::json!([{ + "code": "baseMoved", + "message": "main has advanced" + }]); + let resp = serde_json::json!({ + "status": "failed", + "diff": diff, + "preview": { "promotedColumns": [] } + }); + http::Response::builder() + .status(409) + .body(resp.to_string()) + .unwrap() + }); + let result = table.cherry_pick("exp", false).await.unwrap(); + assert_eq!(result.status, crate::table::CherryPickStatus::Failed); + assert!(!result.diff.errors.is_empty()); + assert_eq!(result.diff.errors.len(), 1); +} + +#[tokio::test] +async fn test_cherry_pick_unknown_error_code_parses() { + let table = Table::new_with_handler("my_table", |_| { + let mut diff: serde_json::Value = serde_json::from_str(sample_branch_diff_json()).unwrap(); + diff["errors"] = serde_json::json!([{ + "code": "multipleCommits", + "message": "branch has more than one data commit" + }]); + let resp = serde_json::json!({ + "status": "failed", + "diff": diff, + "preview": { "operation": "append", "rowsAdded": 2 } + }); + http::Response::builder() + .status(409) + .body(resp.to_string()) + .unwrap() + }); + let result = table.cherry_pick("exp", false).await.unwrap(); + assert_eq!(result.status, crate::table::CherryPickStatus::Failed); + assert_eq!( + result.diff.errors[0].code, + crate::table::CherryPickErrorCode::Unknown + ); + assert!(result.preview.promoted_columns.is_empty()); +} + +#[tokio::test] +async fn test_cherry_pick_unexpected_2xx_is_error() { + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(204) + .body(String::new()) + .unwrap() + }); + let err = table.cherry_pick("exp", false).await.unwrap_err(); + match err { + Error::Http { + status_code: Some(code), + .. + } => assert_eq!(code, reqwest::StatusCode::NO_CONTENT), + other => panic!("expected Http error, got {other:?}"), + } +} + +#[tokio::test] +async fn test_checkout_branch_validates_via_list() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/branches/list/"); + http::Response::builder() + .status(200) + .body( + r#"{"branches":{"exp":{"parentBranch":null,"parentVersion":1,"createAt":1,"manifestSize":1}}}"#, + ) + .unwrap() + }); + let branch = table.checkout_branch("exp", None).await.unwrap(); + assert_eq!(branch.current_branch(), Some("exp".to_string())); +} + +#[tokio::test] +async fn test_checkout_branch_missing() { + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(200) + .body(r#"{"branches":{}}"#) + .unwrap() + }); + let err = table.checkout_branch("ghost", None).await.unwrap_err(); + assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}"); +} + +#[tokio::test] +async fn test_checkout_main_returns_main_handle() { + // "main" yields a main-scoped handle without any validation request. + let table = Table::new_with_handler("my_table", |request| -> http::Response { + panic!("unexpected request: {}", request.url().path()) + }); + let main = table.checkout_branch("main", None).await.unwrap(); + assert_eq!(main.current_branch(), None); +} + +#[tokio::test] +async fn test_branch_count_rows_carries_branch_in_body() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/count_rows/" => { + let body = request_body_json(&request); + assert_eq!(body["branch"], "exp"); + http::Response::builder() + .status(200) + .body("7".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert_eq!(branch.count_rows(None).await.unwrap(), 7); +} + +#[tokio::test] +async fn test_main_handle_omits_branch_in_body() { + // A main handle must not send a branch field (byte-compatible with + // pre-branch servers and the existing wire format). + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/count_rows/"); + let body = request_body_json(&request); + assert!( + body.get("branch").is_none(), + "main handle must not send a branch field" + ); + http::Response::builder().status(200).body("0").unwrap() + }); + table.count_rows(None).await.unwrap(); +} + +#[tokio::test] +async fn test_branch_update_and_delete_carry_branch() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/update/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(r#"{"version":5}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/delete/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(r#"{"version":6,"num_deleted_rows":1}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch + .update() + .column("a", "a + 1") + .execute() + .await + .unwrap(); + branch.delete("a > 1").await.unwrap(); +} + +#[tokio::test] +async fn test_branch_list_indices_carries_branch_in_body() { + use lance::dataset::refs::Ref; + // list_indices posts to index/list and then fetches the schema (describe) + // to resolve column names; both must carry the branch. + let describe_body = + describe_response(&Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/index/list/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(r#"{"indexes":[]}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/describe/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert!(branch.list_indices().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_branch_list_versions_carries_query_param() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/version/list/" => { + assert_eq!( + request + .url() + .query_pairs() + .find(|(k, _)| k == "branch") + .map(|(_, v)| v.into_owned()), + Some("exp".to_string()), + "version/list must carry ?branch=exp" + ); + http::Response::builder() + .status(200) + .body(r#"{"versions":[]}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert!(branch.list_versions().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_branch_drop_index_carries_query_param() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/index/my_idx/drop/" => { + assert_eq!( + request + .url() + .query_pairs() + .find(|(k, _)| k == "branch") + .map(|(_, v)| v.into_owned()), + Some("exp".to_string()) + ); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.drop_index("my_idx").await.unwrap(); +} + +#[tokio::test] +async fn test_branch_insert_carries_query_param() { + use lance::dataset::refs::Ref; + 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 table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => { + // schema() fetch on the branch handle carries branch in the body. + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap() + } + "/v1/table/my_table/insert/" => { + assert_eq!( + request + .url() + .query_pairs() + .find(|(k, _)| k == "branch") + .map(|(_, v)| v.into_owned()), + Some("exp".to_string()), + "insert must carry ?branch=exp" + ); + http::Response::builder() + .status(200) + .body(r#"{"version":2}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.add(data.clone()).execute().await.unwrap(); +} + +#[tokio::test] +async fn test_branch_tag_create_carries_branch_in_body() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/tags/create/" => { + let body = request_body_json(&request); + assert_eq!(body["branch"], "exp"); + assert_eq!(body["tag"], "v1"); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.tags().await.unwrap().create("v1", 1).await.unwrap(); +} + +#[tokio::test] +async fn test_checkout_branch_version_forwards_branch() { + // Branch versions overlap main's, so the version must resolve on the + // branch's own chain -- the validating describe and reads carry both. + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let describe_body = describe_response(&schema); + let table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/branches/list/" => http::Response::builder() + .status(200) + .body( + r#"{"branches":{"exp":{"parentBranch":null,"parentVersion":1,"createAt":1,"manifestSize":1}}}"# + .to_string(), + ) + .unwrap(), + "/v1/table/my_table/describe/" => { + let body = request_body_json(&request); + assert_eq!(body["branch"], "exp", "checkout validate carries branch"); + assert_eq!(body["version"], 2, "checkout validate carries version"); + http::Response::builder().status(200).body(describe_body.clone()).unwrap() + } + "/v1/table/my_table/count_rows/" => { + let body = request_body_json(&request); + assert_eq!(body["branch"], "exp"); + assert_eq!(body["version"], 2, "overlapping version resolves on the branch chain"); + http::Response::builder().status(200).body("3".to_string()).unwrap() + } + path => panic!("unexpected request path: {path}"), + } + }); + let branch = table.checkout_branch("exp", Some(2)).await.unwrap(); + assert_eq!(branch.current_branch(), Some("exp".to_string())); + assert_eq!(branch.count_rows(None).await.unwrap(), 3); +} + +fn branch_query_param(request: &reqwest::Request) -> Option { + request + .url() + .query_pairs() + .find(|(k, _)| k == "branch") + .map(|(_, v)| v.into_owned()) +} + +#[tokio::test] +async fn test_branch_query_carries_branch_in_body() { + use lance::dataset::refs::Ref; + // /query/ is the hot read path; the branch must ride in the JSON 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 data_ref = data.clone(); + let table = Table::new_with_handler("my_table", 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/query/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&data_ref)) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + let rows: usize = branch + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum(); + assert_eq!(rows, 3); +} + +#[tokio::test] +async fn test_branch_merge_insert_carries_query_param() { + use lance::dataset::refs::Ref; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let data: Box = Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/merge_insert/" => { + assert_eq!( + branch_query_param(&request).as_deref(), + Some("exp"), + "merge_insert must carry ?branch=exp" + ); + http::Response::builder() + .status(200) + .body( + r#"{"version":2,"num_deleted_rows":0,"num_inserted_rows":3,"num_updated_rows":0}"# + .to_string(), + ) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.merge_insert(&["id"]).execute(data).await.unwrap(); +} + +#[tokio::test] +async fn test_branch_multipart_write_carries_query_param() { + use lance::dataset::refs::Ref; + // The multipart path (create -> insert parts -> complete) must carry + // ?branch= on every leg; an old server version forces it. + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 4, 0), + move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/multipart_write/create" => { + assert_eq!(branch_query_param(&request).as_deref(), Some("exp")); + http::Response::builder() + .status(200) + .body(r#"{"upload_id": "u1"}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/insert/" => { + assert_eq!( + branch_query_param(&request).as_deref(), + Some("exp"), + "multipart insert must carry ?branch=exp" + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 1}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/multipart_write/complete" => { + assert_eq!(branch_query_param(&request).as_deref(), Some("exp")); + http::Response::builder() + .status(200) + .body(r#"{"version": 5}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }, + ); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); + branch + .add(vec![batch]) + .write_parallelism(2) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_branch_restore_carries_branch_in_body() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + assert_eq!(request_body_json(&request)["version"], 1); + http::Response::builder() + .status(200) + .body(r#"{"version":1,"schema":{"fields":[]}}"#.to_string()) + .unwrap() + } + "/v1/table/my_table/restore/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + assert_eq!(request_body_json(&request)["version"], 1); + http::Response::builder() + .status(200) + .body(r#"{"version":1}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.checkout(1).await.unwrap(); + branch.restore().await.unwrap(); +} + +#[tokio::test] +async fn test_branch_create_index_carries_branch_in_body() { + use lance::dataset::refs::Ref; + // create_index fetches the schema (describe) to resolve the column and + // then posts create_index; both must carry the branch in the body. + let describe_body = + describe_response(&Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap() + } + "/v1/table/my_table/create_index/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch + .create_index(&["a"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_branch_column_ops_carry_branch_in_body() { + use lance::dataset::refs::Ref; + // add_columns / alter_columns / drop_columns all stamp the branch into + // the JSON body via apply_branch_body. + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" + | "/v1/table/my_table/alter_columns/" + | "/v1/table/my_table/drop_columns/" => { + assert_eq!( + request_body_json(&request)["branch"], + "exp", + "{} must carry the branch", + request.url().path() + ); + http::Response::builder() + .status(200) + .body(r#"{"version":43}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "b".into(), + "a + 1".into(), + )])) + .execute() + .await + .unwrap(); + branch + .alter_columns(&[ColumnAlteration::new("a".into()).rename("b".into())]) + .await + .unwrap(); + branch.drop_columns(&["a"]).await.unwrap(); +} + +#[tokio::test] +async fn test_branch_update_field_metadata_carries_branch_in_body() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/update_field_metadata/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body(r#"{"version":7,"fields":{}}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch + .update_field_metadata(&[FieldMetadataUpdate::new("category").set("unit", "label")]) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_branch_index_stats_carries_branch_in_body() { + use lance::dataset::refs::Ref; + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/index/my_index/stats/" => { + assert_eq!(request_body_json(&request)["branch"], "exp"); + http::Response::builder() + .status(200) + .body( + r#"{"num_indexed_rows":1,"num_unindexed_rows":0,"index_type":"IVF_PQ","distance_type":"l2"}"# + .to_string(), + ) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert!(branch.index_stats("my_index").await.unwrap().is_some()); +} + +#[tokio::test] +async fn test_branch_stats_attaches_body_while_main_omits_it() { + use lance::dataset::refs::Ref; + // stats has a bespoke conditional body: a main handle stays a bodyless + // POST, while a branch handle attaches {"branch": ...}. + let stats_body = r#"{"total_bytes":1,"num_rows":3,"num_indices":0,"fragment_stats":{"num_fragments":1,"num_small_fragments":0,"lengths":{"min":3,"max":3,"mean":3,"p25":3,"p50":3,"p75":3,"p99":3}}}"#; + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body(stats_body.to_string()) + .unwrap(), + "/v1/table/my_table/stats/" => { + match request.body() { + // main handle: byte-identical to the pre-branch wire format. + None => {} + // branch handle: branch travels in the body. + Some(_) => assert_eq!(request_body_json(&request)["branch"], "exp"), + } + http::Response::builder() + .status(200) + .body(stats_body.to_string()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + table.stats().await.unwrap(); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + branch.stats().await.unwrap(); +} diff --git a/rust/lancedb/src/remote/table/tests/columns.rs b/rust/lancedb/src/remote/table/tests/columns.rs new file mode 100644 index 000000000..bb01f2800 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/columns.rs @@ -0,0 +1,1041 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[rstest] +#[case(true)] +#[case(false)] +#[tokio::test] +async fn test_add_columns(#[case] old_server: bool) { + let table = Table::new_with_handler("my_table", move |request| { + if request.url().path() == "/v1/table/my_table/describe/" { + simple_describe_response() + } else if request.url().path() == "/v1/table/my_table/add_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 new_columns = value.get("new_columns").unwrap().as_array().unwrap(); + assert!(new_columns.len() == 2); + + let col_name = new_columns[0]["name"].as_str().unwrap(); + let expression = new_columns[0]["expression"].as_str().unwrap(); + assert_eq!(col_name, "b"); + assert_eq!(expression, "a + 1"); + + let col_name = new_columns[1]["name"].as_str().unwrap(); + let expression = new_columns[1]["expression"].as_str().unwrap(); + assert_eq!(col_name, "x"); + assert_eq!(expression, "cast(NULL as int32)"); + + 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 + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![ + ("b".into(), "a + 1".into()), + ("x".into(), "cast(NULL as int32)".into()), + ])) + .execute() + .await + .unwrap(); + + assert_eq!(result.version, if old_server { 0 } else { 43 }); +} + +/// A pyarrow schema goes over the wire as an Arrow IPC schema under an +/// Arrow content type, not as a JSON description of the types, so a +/// field arrives exactly as sent. The branch has no JSON envelope to +/// ride in and goes on the query string. +#[tokio::test] +async fn test_add_columns_sends_an_arrow_schema_for_all_nulls() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" => { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.headers().get("Content-Type").unwrap(), + ARROW_STREAM_CONTENT_TYPE + ); + assert_eq!( + request.url().query_pairs().find(|(k, _)| k == "branch"), + None, + "the main branch sends no branch parameter" + ); + + // The payload decodes back to the exact schema, types and + // field metadata included. + let body = request.body().unwrap().as_bytes().unwrap(); + let reader = + arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None).unwrap(); + let sent = reader.schema(); + assert_eq!(sent.fields().len(), 2); + assert_eq!(sent.field(0).name(), "ts"); + assert_eq!( + sent.field(0).data_type(), + &DataType::Timestamp( + arrow_schema::TimeUnit::Nanosecond, + Some("America/New_York".into()) + ) + ); + assert!(sent.field(0).is_nullable()); + assert_eq!(sent.field(1).name(), "price"); + assert_eq!(sent.field(1).data_type(), &DataType::Decimal128(38, 10)); + assert_eq!( + sent.field(1).metadata().get("unit").map(String::as_str), + Some("cents") + ); + + http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected request path: {path}"), + }); + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp( + arrow_schema::TimeUnit::Nanosecond, + Some("America/New_York".into()), + ), + true, + ), + Field::new("price", DataType::Decimal128(38, 10), true).with_metadata( + std::collections::HashMap::from([("unit".to_string(), "cents".to_string())]), + ), + ])); + + let result = table + .add_columns() + .transform(NewColumnTransform::AllNulls(schema)) + .execute() + .await + .unwrap(); + + assert_eq!(result.version, 7); +} + +/// A declaration is sent as `{name, computed}` for the server to plan; the +/// client never types the expression itself. +#[tokio::test] +async fn test_add_computed_columns_sends_the_expression() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" => { + assert_eq!(request.method(), "POST"); + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + value["new_columns"], + serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + + let result = table + .add_columns() + .computed("doubled", "x * 2") + .execute() + .await + .unwrap(); + assert_eq!(result.version, 7); +} + +#[tokio::test] +async fn test_add_scalar_function_column_sends_atomic_null_declaration() { + let table = Table::new_with_handler("my_table", |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#, + ) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = serde_json::from_slice( + request.body().unwrap().as_bytes().unwrap(), + ) + .unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json" + )) + .unwrap(); + assert_eq!(actual, expected); + http::Response::builder() + .status(200) + .body(r#"{"version":8}"#) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + } + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"embed","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, + "inputs":[{"parameter":"text","kind":"column","value":{"path":"description"}}], + "output":{"kind":"scalar","arrow_type":"list","nullable":false} + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function_as("embedding", application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 8); +} + +/// The fixture binding's table: `title` and `body` bound as inputs, its +/// two outputs declared, plus an unbound `spare`. +fn fixture_bound_schema() -> Schema { + let binding = crate::function::FunctionBinding::from_json(include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + let binding_metadata = + crate::table::computed_columns::function_bindings_metadata(std::slice::from_ref(&binding)) + .unwrap(); + let mut fields = vec![ + Field::new("title", DataType::Utf8, true), + Field::new("body", DataType::Utf8, true), + ]; + fields.extend(binding.outputs().iter().map(|output| { + let data_type = match output.arrow_type.as_str() { + "utf8" => DataType::Utf8, + "int64" => DataType::Int64, + other => panic!("unexpected fixture output type {other}"), + }; + Field::new(&output.output_name, data_type, true).with_metadata( + crate::table::computed_columns::function_computed_column_metadata( + binding.binding_id(), + output.output_ordinal, + &["title".into(), "body".into()], + ), + ) + })); + fields.push(Field::new("spare", DataType::Int32, true)); + Schema::new_with_metadata( + fields, + HashMap::from([( + crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(), + binding_metadata, + )]), + ) +} + +/// Only a column the binding uses is refused, and it is refused before +/// any request goes out; the rest reach the server as usual. +#[tokio::test] +async fn test_add_columns_scopes_to_the_columns_a_binding_uses() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&fixture_bound_schema())) + .unwrap(), + "/v1/table/my_table/add_columns/" => http::Response::builder() + .status(200) + .body(r#"{"version":10}"#.to_string()) + .unwrap(), + path => panic!("Unexpected path: {path}"), + }); + table + .add_columns() + .computed("doubled", "spare * 2") + .execute() + .await + .unwrap(); + table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "eager".into(), + "spare + 1".into(), + )])) + .execute() + .await + .unwrap(); + + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&fixture_bound_schema())) + .unwrap(), + path => panic!("mutation request must not be sent: {path}"), + }); + for name in ["title", "search_text"] { + let err = table + .add_columns() + .computed(name, "1") + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); + let err = table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + name.into(), + "1".into(), + )])) + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); + } +} + +#[tokio::test] +async fn test_add_function_column_allows_an_existing_binding() { + let schema = fixture_bound_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_response(&schema)) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + actual["new_columns"], + serde_json::json!([ + {"name":"secondary_text","all_null":true}, + {"name":"secondary_token_count","all_null":true} + ]) + ); + http::Response::builder() + .status(200) + .body(r#"{"version":10}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"text_features","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, + "inputs":[ + {"parameter":"title","kind":"column","value":{"path":"title"}}, + {"parameter":"body","kind":"column","value":{"path":"body"}} + ], + "output":{"kind":"named_struct","fields":[ + {"name":"normalized_text","arrow_type":"utf8","nullable":false}, + {"name":"token_count","arrow_type":"int64","nullable":false} + ]}, + "columns":{ + "normalized_text":"secondary_text", + "token_count":"secondary_token_count" + } + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function(application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 10); +} + +#[tokio::test] +async fn test_add_fixed_size_list_function_column_declares_the_vector_type() { + let table = Table::new_with_handler("my_table", |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#, + ) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = serde_json::from_slice( + request.body().unwrap().as_bytes().unwrap(), + ) + .unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" + )) + .unwrap(); + assert_eq!(actual, expected); + http::Response::builder() + .status(200) + .body(r#"{"version":8}"#) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + } + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"embed","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, + "inputs":[{"parameter":"text","kind":"column","value":{"path":"description"}}], + "output":{"kind":"scalar","arrow_type":"fixed_size_list","nullable":false} + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function_as("embedding", application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 8); +} + +#[tokio::test] +async fn test_add_named_struct_function_expands_one_atomic_binding() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[ + {"name":"title","nullable":true,"type":{"type":"string"}}, + {"name":"body","nullable":true,"type":{"type":"string"}} + ]}}"#, + ) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_multi_output_declaration_request.json" + )) + .unwrap(); + assert_eq!(actual, expected); + http::Response::builder() + .status(200) + .body(r#"{"version":9}"#) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"text_features","version":"1","object_id":"fixture","location":"memory:///fixture","manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, + "inputs":[ + {"parameter":"title","kind":"column","value":{"path":"title"}}, + {"parameter":"body","kind":"column","value":{"path":"body"}} + ], + "output":{"kind":"named_struct","fields":[ + {"name":"normalized_text","arrow_type":"utf8","nullable":false}, + {"name":"token_count","arrow_type":"int64","nullable":false} + ]}, + "columns":{"normalized_text":"search_text"} + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function(application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 9); +} + +#[tokio::test] +async fn test_add_columns_fails_closed_on_newer_function_binding_metadata() { + let table = Table::new_with_handler("my_table", |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"x","nullable":true,"type":{"type":"int32"}}],"metadata":{"lancedb::function_bindings":"{\"version\":2,\"bindings\":[]}"}}}"#, + ) + .unwrap(), + path => panic!("mutation request must not be sent: {path}"), + } + }); + + let err = table + .add_columns() + .computed("doubled", "x * 2") + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); +} + +/// A remote refresh is a server job: the async form returns its handle, +/// and the blocking form refuses rather than invent a fill count. +#[tokio::test] +async fn test_refresh_column_async_submits_a_backfill_job() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column"); + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(value["column"], "doubled"); + http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-42"}"#) + .unwrap() + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + assert_eq!(job.id(), Some("j-42")); + + let err = table.refresh_column("doubled").await.unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } + if message.contains("refresh_column_async")), + "{err:?}" + ); +} + +#[rstest] +#[case(false, "Function is unavailable")] +#[case(true, "Function name was not found")] +#[tokio::test] +async fn test_function_column_operations_preserve_dependency_not_found( + #[case] declare_column: bool, + #[case] message: &'static str, +) { + let schema_requests = Arc::new(AtomicUsize::new(0)); + let requests = schema_requests.clone(); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + requests.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#.to_string(), + ) + .unwrap() + } + "/v1/table/my_table/backfill_column" | "/v1/table/my_table/add_columns/" => { + http::Response::builder() + .status(404) + .body(json!({"code": 4, "error": message}).to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + table.schema().await.unwrap(); + assert_eq!(schema_requests.load(Ordering::SeqCst), 1); + + let error = if declare_column { + let fixture: serde_json::Value = serde_json::from_str(include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" + )) + .unwrap(); + let application = crate::function::FunctionApplication::from_json( + &fixture["function"]["application"].to_string(), + ) + .unwrap(); + table + .add_columns() + .function_as("embedding", application) + .execute() + .await + .unwrap_err() + } else { + table.refresh_column_async("embedding").await.unwrap_err() + }; + let Error::Http { + source, + status_code, + request_id, + } = error + else { + panic!("dependency 404 was misclassified: {error:?}"); + }; + assert_eq!(status_code, Some(StatusCode::NOT_FOUND)); + assert!(source.to_string().contains(message)); + assert!(!request_id.is_empty()); + + table.schema().await.unwrap(); + assert_eq!(schema_requests.load(Ordering::SeqCst), 2); +} + +/// The error listing is table-addressed with optional job and column +/// filters, mirroring the server's SQL surface, and the two non-record +/// signals come back as their own fields rather than as rows. +#[tokio::test] +async fn test_function_errors_lists_the_rows_a_refresh_skipped() { + use crate::function::{FunctionErrorFragment, FunctionErrorRecord, FunctionErrorsRequest}; + + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/errors"); + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + value, + serde_json::json!({"job_id": "j-7", "column": "embedding", "limit": 2}) + ); + http::Response::builder() + .status(200) + .body( + r#"{"records": [{"job_id": "j-7", "fragment_id": 3, "row_offset": 9, + "column": "embedding", "function": "embed", "function_version": "2", + "table_version": 11, "error_type": "ValueError", + "error_message": "bad input 'x'", "created_at_millis": 1700000000000}], + "fragments": [{"job_id": "j-7", "fragment_id": 4, "rows_skipped": 500, + "rows_recorded": 100}], + "truncated": true}"#, + ) + .unwrap() + }); + + let errors = table + .function_errors( + FunctionErrorsRequest::new() + .job_id("j-7") + .column("embedding") + .limit(2), + ) + .await + .unwrap(); + assert_eq!( + errors.records, + [FunctionErrorRecord { + job_id: "j-7".into(), + fragment_id: 3, + row_offset: Some(9), + column: "embedding".into(), + function: "embed".into(), + function_version: "2".into(), + table_version: 11, + error_type: "ValueError".into(), + error_message: "bad input 'x'".into(), + created_at_millis: 1_700_000_000_000, + }] + ); + assert_eq!( + errors.fragments, + [FunctionErrorFragment { + job_id: "j-7".into(), + fragment_id: 4, + rows_skipped: 500, + rows_recorded: 100, + }] + ); + assert!(errors.truncated); + + // No filter sends no filter, and an empty listing reads as such. + let table = Table::new_with_handler("my_table", |request| { + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(value, serde_json::json!({})); + http::Response::builder() + .status(200) + .body(r#"{"records": []}"#) + .unwrap() + }); + let errors = table + .function_errors(FunctionErrorsRequest::new()) + .await + .unwrap(); + assert_eq!(errors, crate::function::FunctionErrors::default()); +} + +/// The refresh handle is wrapped for read-freshness tracking, so it has to +/// forward the detail APIs too -- this is the job an operator is holding +/// when a backfill goes quiet. +#[tokio::test] +async fn test_refresh_job_handle_reports_detail_and_events() { + 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 table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(br#"{"job_id": "j-42"}"#.to_vec()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body( + r#"{"job_id": "j-42", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 7, "spec": {"column": "doubled"}}"# + .as_bytes() + .to_vec(), + ) + .unwrap(), + "/v1/jobs/query_events" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["job_id"], "j-42"); + http::Response::builder() + .status(200) + .body(events.clone()) + .unwrap() + } + other => panic!("unexpected path {other}"), + } + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + job.refresh().await.unwrap(); + assert_eq!(job.state().as_deref(), Some("running")); + assert_eq!(job.job_type().as_deref(), Some("refresh_column")); + assert_eq!(job.creation_ms(), Some(7)); + assert_eq!(job.spec().unwrap()["column"], "doubled"); + + let batches = job + .events(crate::job::JobEventsRequest::default()) + .await + .unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); +} + +#[tokio::test] +async fn test_refresh_submission_uses_add_columns_version_fence() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" => http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap(), + "/v1/table/my_table/backfill_column" => { + let min_version = request + .headers() + .get("x-lancedb-min-version") + .and_then(|value| value.to_str().ok()); + if min_version != Some("7") { + return http::Response::builder() + .status(400) + .body(r#"{"error":"Column not found: doubled"}"#.to_string()) + .unwrap(); + } + http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-43"}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + + let result = table + .add_columns() + .computed("doubled", "a * 2") + .execute() + .await + .unwrap(); + assert_eq!(result.version, 7); + + let job = table.refresh_column_async("doubled").await.unwrap(); + assert_eq!(job.id(), Some("j-43")); +} + +/// The gate's reproducer: after a successful wait, a same-handle read +/// must carry the exact published version so a stale server cache cannot +/// serve the pre-backfill snapshot. +#[tokio::test] +async fn test_backfill_wait_establishes_read_freshness() { + let saw_published_version = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw = saw_published_version.clone(); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-7"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(refresh_done("j-7")) + .unwrap(), + "/v1/table/my_table/count_rows/" => { + saw.store( + request + .headers() + .get("x-lancedb-min-read-version") + .and_then(|value| value.to_str().ok()) + == Some("8"), + std::sync::atomic::Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + let result = job.wait().await.unwrap(); + assert_eq!(result.rows_assigned, 12); + assert_eq!(result.published_version, Some(8)); + table.count_rows(None).await.unwrap(); + assert!( + saw_published_version.load(std::sync::atomic::Ordering::SeqCst), + "read after wait did not carry the published version" + ); +} + +/// A checkout after submission wins over the completion fence: the +/// pinned view must not regain a timestamp floor from the job. +#[tokio::test] +async fn test_checkout_after_submit_beats_the_completion_fence() { + let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw = saw_min_timestamp.clone(); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-8"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(refresh_done("j-8")) + .unwrap(), + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/count_rows/" => { + saw.store( + request.headers().contains_key("x-lancedb-min-timestamp"), + std::sync::atomic::Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + table.checkout(3).await.unwrap(); + job.wait().await.unwrap(); + table.count_rows(None).await.unwrap(); + assert!( + !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), + "completion fence overrode an explicit checkout" + ); +} + +/// Tag checkout resets freshness state wholesale; the fence must not +/// survive it. +#[tokio::test] +async fn test_tag_checkout_after_submit_beats_the_completion_fence() { + let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw = saw_min_timestamp.clone(); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-9"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(refresh_done("j-9")) + .unwrap(), + "/v1/table/my_table/tags/version/" => http::Response::builder() + .status(200) + .body(r#"{"version": 5}"#.to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/count_rows/" => { + saw.store( + request.headers().contains_key("x-lancedb-min-timestamp"), + std::sync::atomic::Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + table.checkout_tag("v1").await.unwrap(); + job.wait().await.unwrap(); + table.count_rows(None).await.unwrap(); + assert!( + !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), + "completion fence overrode a tag checkout" + ); +} + +/// A checkout landing while the submission request is in flight advances +/// the epoch past the token captured at submit. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_checkout_during_submission_beats_the_completion_fence() { + let saw_min_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw = saw_min_timestamp.clone(); + 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 table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/backfill_column" => { + // Signal arrival, then hold the response until the + // test's checkout completes. + arrived_tx.lock().unwrap().send(()).unwrap(); + release_rx + .lock() + .unwrap() + .recv_timeout(std::time::Duration::from_secs(10)) + .unwrap(); + http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-10"}"#.to_string()) + .unwrap() + } + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(refresh_done("j-10")) + .unwrap(), + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/count_rows/" => { + saw.store( + request.headers().contains_key("x-lancedb-min-timestamp"), + std::sync::atomic::Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + } + }); + + let submit = tokio::spawn({ + let table = table.clone(); + async move { table.refresh_column_async("doubled").await } + }); + tokio::task::spawn_blocking(move || { + arrived_rx + .recv_timeout(std::time::Duration::from_secs(10)) + .unwrap() + }) + .await + .unwrap(); + table.checkout(7).await.unwrap(); + release_tx.send(()).unwrap(); + + let job = submit.await.unwrap().unwrap(); + job.wait().await.unwrap(); + table.count_rows(None).await.unwrap(); + assert!( + !saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst), + "completion fence overrode a checkout that landed mid-submission" + ); +} + +/// checkout_latest keeps the handle on latest, so a completed backfill +/// must retain the checkout timestamp and add its exact published version. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_checkout_latest_during_submission_keeps_the_fence() { + let seen_headers = Arc::new(std::sync::Mutex::new(None::)); + let saw = seen_headers.clone(); + 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 table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/backfill_column" => { + arrived_tx.lock().unwrap().send(()).unwrap(); + release_rx + .lock() + .unwrap() + .recv_timeout(std::time::Duration::from_secs(10)) + .unwrap(); + http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-11"}"#.to_string()) + .unwrap() + } + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(refresh_done("j-11")) + .unwrap(), + "/v1/table/my_table/count_rows/" => { + *saw.lock().unwrap() = Some(request.headers().clone()); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + + let submit = tokio::spawn({ + let table = table.clone(); + async move { table.refresh_column_async("doubled").await } + }); + tokio::task::spawn_blocking(move || { + arrived_rx + .recv_timeout(std::time::Duration::from_secs(10)) + .unwrap() + }) + .await + .unwrap(); + table.checkout_latest().await.unwrap(); + release_tx.send(()).unwrap(); + + let job = submit.await.unwrap().unwrap(); + job.wait().await.unwrap(); + table.count_rows(None).await.unwrap(); + let headers = seen_headers.lock().unwrap().clone().expect("no request"); + assert!(headers.contains_key("x-lancedb-min-timestamp")); + assert_eq!( + headers + .get("x-lancedb-min-read-version") + .and_then(|value| value.to_str().ok()), + Some("8") + ); +} diff --git a/rust/lancedb/src/remote/table/tests/freshness.rs b/rust/lancedb/src/remote/table/tests/freshness.rs new file mode 100644 index 000000000..1c0b34849 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/freshness.rs @@ -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( + body_for: F, +) -> ( + impl Fn(reqwest::Request) -> http::Response + Clone + Send + Sync + 'static, + Arc>>, +) +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 + Clone + Send + Sync + 'static, + Arc>>, +) { + 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( + body_for: F, +) -> ( + impl Fn(reqwest::Request) -> http::Response + Clone + Send + Sync + 'static, + Arc>>, +) +where + F: Fn(&str) -> String + Clone + Send + Sync + 'static, +{ + let captured: Arc>> = + 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")); +} diff --git a/rust/lancedb/src/remote/table/tests/hybrid.rs b/rust/lancedb/src/remote/table/tests/hybrid.rs new file mode 100644 index 000000000..02af81d91 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/hybrid.rs @@ -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>>) -> 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::>() + .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::>() + .await + .unwrap(); + + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 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::>() + .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::>() + .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" + ); +} diff --git a/rust/lancedb/src/remote/table/tests/index.rs b/rust/lancedb/src/remote/table/tests/index.rs new file mode 100644 index 000000000..3b40ab454 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/index.rs @@ -0,0 +1,1280 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[tokio::test] +async fn test_create_index() { + let cases = [ + ( + "IVF_FLAT", + json!({ + "metric_type": "hamming", + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfFlat(IvfFlatIndexBuilder::default().distance_type(DistanceType::Hamming)), + ), + ( + "IVF_FLAT", + json!({ + "metric_type": "hamming", + "num_partitions": 128, + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfFlat( + IvfFlatIndexBuilder::default() + .distance_type(DistanceType::Hamming) + .num_partitions(128), + ), + ), + ( + "IVF_PQ", + json!({ + "metric_type": "l2", + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfPq(Default::default()), + ), + ( + "IVF_PQ", + json!({ + "metric_type": "cosine", + "num_partitions": 128, + "num_bits": 4, + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfPq( + IvfPqIndexBuilder::default() + .distance_type(DistanceType::Cosine) + .num_partitions(128) + .num_bits(4), + ), + ), + ( + "IVF_PQ", + json!({ + "metric_type": "l2", + "num_sub_vectors": 16, + "sample_rate": 512, + "max_iterations": 100, + }), + Index::IvfPq( + IvfPqIndexBuilder::default() + .num_sub_vectors(16) + .sample_rate(512) + .max_iterations(100), + ), + ), + ( + "IVF_HNSW_SQ", + json!({ + "metric_type": "l2", + "sample_rate": 256, + "max_iterations": 50, + "m": 20, + "ef_construction": 300, + }), + Index::IvfHnswSq(Default::default()), + ), + ( + "IVF_HNSW_SQ", + json!({ + "metric_type": "l2", + "num_partitions": 128, + "sample_rate": 256, + "max_iterations": 50, + "m": 40, + "ef_construction": 500, + }), + Index::IvfHnswSq( + IvfHnswSqIndexBuilder::default() + .distance_type(DistanceType::L2) + .num_partitions(128) + .num_edges(40) + .ef_construction(500), + ), + ), + ( + "IVF_HNSW_FLAT", + json!({ + "metric_type": "l2", + "sample_rate": 256, + "max_iterations": 50, + "m": 20, + "ef_construction": 300, + }), + Index::IvfHnswFlat(Default::default()), + ), + ( + "IVF_HNSW_FLAT", + json!({ + "metric_type": "cosine", + "num_partitions": 64, + "sample_rate": 256, + "max_iterations": 50, + "m": 40, + "ef_construction": 500, + }), + Index::IvfHnswFlat( + IvfHnswFlatIndexBuilder::default() + .distance_type(DistanceType::Cosine) + .num_partitions(64) + .num_edges(40) + .ef_construction(500), + ), + ), + ( + "IVF_SQ", + json!({ + "metric_type": "l2", + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfSq(Default::default()), + ), + ( + "IVF_SQ", + json!({ + "metric_type": "cosine", + "num_partitions": 64, + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfSq( + IvfSqIndexBuilder::default() + .distance_type(DistanceType::Cosine) + .num_partitions(64), + ), + ), + ( + "IVF_RQ", + json!({ + "metric_type": "l2", + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfRq(Default::default()), + ), + ( + "IVF_RQ", + json!({ + "metric_type": "cosine", + "num_partitions": 64, + "num_bits": 8, + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfRq( + IvfRqIndexBuilder::default() + .distance_type(DistanceType::Cosine) + .num_partitions(64) + .num_bits(8), + ), + ), + // HNSW_PQ isn't yet supported on SaaS + ("BTREE", json!({}), Index::BTree(Default::default())), + ("BITMAP", json!({}), Index::Bitmap(Default::default())), + ("ZONEMAP", json!({}), Index::ZoneMap(Default::default())), + ("NGRAM", json!({}), Index::NGram(Default::default())), + ( + "BLOOM_FILTER", + json!({}), + Index::BloomFilter(Default::default()), + ), + ("RTREE", json!({}), Index::RTree(Default::default())), + ( + "BLOOM_FILTER", + json!({"number_of_items": 4096, "probability": 0.01}), + Index::BloomFilter( + crate::index::scalar::BloomFilterIndexBuilder::default() + .number_of_items(4096) + .unwrap() + .probability(0.01) + .unwrap(), + ), + ), + ( + "RTREE", + json!({"page_size": 1024}), + Index::RTree( + crate::index::scalar::RTreeIndexBuilder::default() + .page_size(1024) + .unwrap(), + ), + ), + ( + "LABEL_LIST", + json!({}), + Index::LabelList(Default::default()), + ), + ( + "FTS", + serde_json::to_value(InvertedIndexParams::default()).unwrap(), + Index::FTS(Default::default()), + ), + ( + "FTS", + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["block_size"] = 256.into(); + body + }, + Index::FTS(InvertedIndexParams::default().block_size(256).unwrap()), + ), + ( + "FTS", + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["custom_stop_words"] = json!(["cat", " cat ", "CAT"]); + body + }, + Index::FTS(InvertedIndexParams::default().custom_stop_words(Some(vec![ + "cat".to_string(), + " cat ".to_string(), + "CAT".to_string(), + ]))), + ), + ( + "FTS", + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["document_granularity"] = "list_element".into(); + body + }, + Index::FTS( + InvertedIndexParams::default() + .document_granularity(DocumentGranularity::ListElement), + ), + ), + ]; + + for (index_type, expected_body, index) in cases { + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 6, 0), + move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => { + 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 = expected_body.clone(); + expected_body["column"] = "a".into(); + expected_body[INDEX_TYPE_KEY] = index_type.into(); + + assert_eq!(body, expected_body); + + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }, + ); + + table.create_index(&["a"], index).execute().await.unwrap(); + } +} + +#[tokio::test] +async fn test_create_index_forwards_replace_false_on_existing_route() { + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["replace"], json!(false)); + + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(false) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_create_index_returns_job() { + let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let describe_calls_in_handler = describe_calls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-123"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["job_id"], "job-123"); + let state = if describe_calls_in_handler + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + "IN_PROGRESS" + } else { + "DONE" + }; + http::Response::builder() + .status(200) + .body(format!( + r#"{{"job_id": "job-123", "job_state": "{state}"}}"# + )) + .unwrap() + } + "/v1/jobs/cancel" => { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["job_id"], "job-123"); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + assert_eq!(job.id(), Some("job-123")); + job.wait().await.unwrap(); + assert_eq!(describe_calls.load(std::sync::atomic::Ordering::SeqCst), 2); + job.cancel().await.unwrap(); +} + +/// An unrecognized state is treated as still running, so the client keeps +/// polling rather than reporting a wrong terminal outcome. +#[tokio::test] +async fn test_job_wait_treats_unknown_state_as_running() { + let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let calls_in_handler = calls.clone(); + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-unknown"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let state = if calls_in_handler.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + "SOMETHING_NEW" + } else { + "DONE" + }; + http::Response::builder() + .status(200) + .body(format!( + r#"{{"job_id": "job-unknown", "job_state": "{state}"}}"# + )) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + job.wait().await.unwrap(); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn test_job_wait_surfaces_failure() { + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-err"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-err", "job_state": "FAILED"}"#.to_string()) + .unwrap(), + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + let err = job.wait().await.unwrap_err(); + let crate::Error::JobFailed { failure, .. } = &err else { + panic!("expected JobFailed, got {err:?}"); + }; + // The server said only that it failed, so nothing may be invented. + assert!(failure.message.is_none(), "{failure:?}"); + assert!(failure.phase.is_none(), "{failure:?}"); + assert!(failure.retryable.is_none(), "{failure:?}"); + assert_eq!(err.to_string(), "Job job-err failed"); +} + +/// A server that reports why the job failed has that reason surfaced +/// verbatim rather than replaced with a generic message. +#[tokio::test] +async fn test_job_wait_reports_the_server_failure_reason() { + let table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-err"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body( + r#"{"job_id": "job-err", "job_state": "FAILED", "failure": {"phase": "commit", "message": "preempted", "retryable": true}}"# + .to_string(), + ) + .unwrap(), + path => panic!("Unexpected path: {}", path), + } + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + let err = job.wait().await.unwrap_err(); + let crate::Error::JobFailed { failure, .. } = &err else { + panic!("expected JobFailed, got {err:?}"); + }; + assert_eq!(failure.message.as_deref(), Some("preempted")); + assert_eq!(failure.phase.as_deref(), Some("commit")); + assert_eq!(failure.retryable, Some(true)); + assert_eq!(err.to_string(), "Job job-err failed: preempted (in commit)"); +} + +/// Servers that return no job id (e.g. an empty create-index response) +/// yield an already-done job. +#[tokio::test] +async fn test_create_index_without_job_id_is_done() { + let table = Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + job.wait().await.unwrap(); + job.cancel().await.unwrap(); +} + +#[tokio::test] +async fn test_create_index_nested_field_paths() { + let schema = nested_index_schema(); + let expected_requests = Arc::new(vec![ + json!({ + "column": "rowId", + "index_type": "BTREE", + }), + json!({ + "column": "`row-id`", + "index_type": "BTREE", + }), + json!({ + "column": "userId", + "index_type": "BTREE", + }), + json!({ + "column": "MetaData.userId", + "index_type": "BTREE", + }), + json!({ + "column": "metadata.user_id", + "index_type": "BTREE", + }), + json!({ + "column": "image.embedding", + "index_type": "IVF_PQ", + "metric_type": "l2", + }), + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["column"] = "payload.text".into(); + body["index_type"] = "FTS".into(); + body + }, + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["column"] = "docs.content".into(); + body["index_type"] = "FTS".into(); + body + }, + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["column"] = "docs.content".into(); + body["index_type"] = "FTS".into(); + body["document_granularity"] = "list_element".into(); + body + }, + json!({ + "column": "`meta-data`.`user-id`", + "index_type": "BTREE", + }), + json!({ + "column": "literal.`a.b`", + "index_type": "BTREE", + }), + ]); + let request_idx = Arc::new(AtomicUsize::new(0)); + let table = Table::new_with_handler_version("my_table", semver::Version::new(0, 6, 0), { + let schema = schema.clone(); + let expected_requests = expected_requests.clone(); + let request_idx = request_idx.clone(); + move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/create_index/" => { + assert_eq!( + request.headers().get("Content-Type").unwrap(), + JSON_CONTENT_TYPE + ); + let idx = request_idx.fetch_add(1, Ordering::SeqCst); + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body, expected_requests[idx]); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + } + }); + + table + .create_index(&["rowId"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["`ROW-ID`"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["userId"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["MetaData.userId"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["Metadata.USER_ID"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["Image.Embedding"], Index::Auto) + .execute() + .await + .unwrap(); + table + .create_index(&["Payload.Text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["Docs.Content"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index( + &["Docs.Content"], + Index::FTS( + InvertedIndexParams::default() + .document_granularity(DocumentGranularity::ListElement), + ), + ) + .execute() + .await + .unwrap(); + table + .create_index(&["`META-DATA`.`USER-ID`"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["literal.`A.B`"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + + assert_eq!(request_idx.load(Ordering::SeqCst), expected_requests.len()); +} + +#[tokio::test] +async fn test_create_list_element_fts_requires_server_support() { + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + |_| -> http::Response { + panic!("unsupported index creation must fail before sending a request") + }, + ); + + let result = table + .create_index( + &["docs.content"], + Index::FTS( + InvertedIndexParams::default() + .document_granularity(DocumentGranularity::ListElement), + ), + ) + .execute() + .await; + + assert!( + result + .unwrap_err() + .to_string() + .contains("document granularity requires remote server version 0.6.0 or later") + ); +} + +#[tokio::test] +async fn test_list_indices() { + let schema = Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8), + false, + ), + Field::new( + "metadata", + DataType::Struct(vec![Field::new("my.column", DataType::Utf8, true)].into()), + false, + ), + ]); + let 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/" => { + return http::Response::builder() + .status(200) + .body(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": ["metadata.`my.column`"], + "index_status": "done", + }, + ] + }) + } + "/v1/table/my_table/index/vector_idx/stats/" => { + serde_json::json!({ + "num_indexed_rows": 100000, + "num_unindexed_rows": 0, + "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": 0, + "index_type": "LABEL_LIST" + }) + } + path => panic!("Unexpected path: {}", path), + }; + http::Response::builder() + .status(200) + .body(serde_json::to_string(&response_body).unwrap()) + .unwrap() + }); + + let indices = table.list_indices().await.unwrap(); + let expected = vec![ + IndexConfig { + name: "vector_idx".into(), + index_type: IndexType::IvfPq, + columns: vec!["vector".into()], + index_uuid: None, + type_url: None, + created_at: None, + num_indexed_rows: None, + num_unindexed_rows: None, + size_bytes: None, + num_segments: None, + index_version: None, + index_details: None, + }, + IndexConfig { + name: "my_idx".into(), + index_type: IndexType::LabelList, + columns: vec!["metadata.`my.column`".into()], + index_uuid: None, + type_url: None, + created_at: None, + num_indexed_rows: None, + num_unindexed_rows: None, + size_bytes: None, + num_segments: None, + index_version: None, + index_details: None, + }, + ]; + assert_eq!(indices, expected); +} + +#[rstest] +#[case::legacy(false)] +#[case::enriched(true)] +#[tokio::test] +async fn test_list_indices_fts_public_list_path(#[case] enriched: bool) { + let schema = nested_index_schema(); + let table = Table::new_with_handler("my_table", move |request| { + let body = match request.url().path() { + "/v1/table/my_table/describe/" => describe_response(&schema), + "/v1/table/my_table/index/list/" => serde_json::json!({ + "indexes": [{ + "index_name": "docs_idx", + "columns": ["docs.item.content"], + "index_type": enriched.then_some("FTS"), + }], + }) + .to_string(), + "/v1/table/my_table/index/docs_idx/stats/" => { + assert!(!enriched, "enriched responses must not fetch index stats"); + serde_json::json!({ + "num_indexed_rows": 1, + "num_unindexed_rows": 0, + "index_type": "FTS", + }) + .to_string() + } + path => panic!("Unexpected path: {path}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + + let indices = table.list_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].index_type, IndexType::FTS); + assert_eq!(indices[0].columns, vec!["docs.content"]); +} + +#[tokio::test] +async fn test_list_indices_nested_field_paths() { + let schema = nested_index_schema(); + let 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/" => { + return http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(); + } + "/v1/table/my_table/index/list/" => { + serde_json::json!({ + "indexes": [ + { + "index_name": "row_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000001", + "columns": ["rowId"], + "index_status": "done", + }, + { + "index_name": "row_dash_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000002", + "columns": ["`ROW-ID`"], + "index_status": "done", + }, + { + "index_name": "user_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000003", + "columns": ["userId"], + "index_status": "done", + }, + { + "index_name": "mixed_case_metadata_user_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000004", + "columns": ["MetaData.userId"], + "index_status": "done", + }, + { + "index_name": "metadata_user_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000005", + "columns": ["Metadata.USER_ID"], + "index_status": "done", + }, + { + "index_name": "image_embedding_idx", + "index_uuid": "00000000-0000-0000-0000-000000000006", + "columns": ["Image.Embedding"], + "index_status": "done", + }, + { + "index_name": "payload_text_idx", + "index_uuid": "00000000-0000-0000-0000-000000000007", + "columns": ["Payload.Text"], + "index_status": "done", + }, + { + "index_name": "meta_data_user_id_idx", + "index_uuid": "00000000-0000-0000-0000-000000000008", + "columns": ["`META-DATA`.`USER-ID`"], + "index_status": "done", + }, + { + "index_name": "literal_dot_idx", + "index_uuid": "00000000-0000-0000-0000-000000000009", + "columns": ["literal.`A.B`"], + "index_status": "done", + }, + ] + }) + } + "/v1/table/my_table/index/row_id_idx/stats/" + | "/v1/table/my_table/index/row_dash_id_idx/stats/" + | "/v1/table/my_table/index/user_id_idx/stats/" + | "/v1/table/my_table/index/mixed_case_metadata_user_id_idx/stats/" + | "/v1/table/my_table/index/metadata_user_id_idx/stats/" + | "/v1/table/my_table/index/meta_data_user_id_idx/stats/" + | "/v1/table/my_table/index/literal_dot_idx/stats/" => { + serde_json::json!({ + "num_indexed_rows": 100000, + "num_unindexed_rows": 0, + "index_type": "BTREE" + }) + } + "/v1/table/my_table/index/image_embedding_idx/stats/" => { + serde_json::json!({ + "num_indexed_rows": 100000, + "num_unindexed_rows": 0, + "index_type": "IVF_PQ", + "distance_type": "l2" + }) + } + "/v1/table/my_table/index/payload_text_idx/stats/" => { + serde_json::json!({ + "num_indexed_rows": 100000, + "num_unindexed_rows": 0, + "index_type": "FTS" + }) + } + path => panic!("Unexpected path: {}", path), + }; + http::Response::builder() + .status(200) + .body(serde_json::to_string(&response_body).unwrap()) + .unwrap() + }); + + let indices = table.list_indices().await.unwrap(); + // The remote path leaves the rich metadata fields None until the server + // wires them through. See https://github.com/lancedb/lancedb/issues/3494 + let expected: Vec = [ + ("row_id_idx", IndexType::BTree, "rowId"), + ("row_dash_id_idx", IndexType::BTree, "`row-id`"), + ("user_id_idx", IndexType::BTree, "userId"), + ( + "mixed_case_metadata_user_id_idx", + IndexType::BTree, + "MetaData.userId", + ), + ("metadata_user_id_idx", IndexType::BTree, "metadata.user_id"), + ("image_embedding_idx", IndexType::IvfPq, "image.embedding"), + ("payload_text_idx", IndexType::FTS, "payload.text"), + ( + "meta_data_user_id_idx", + IndexType::BTree, + "`meta-data`.`user-id`", + ), + ("literal_dot_idx", IndexType::BTree, "literal.`a.b`"), + ] + .into_iter() + .map(|(name, index_type, column)| IndexConfig { + name: name.into(), + index_type, + columns: vec![column.into()], + index_uuid: None, + type_url: None, + created_at: None, + num_indexed_rows: None, + num_unindexed_rows: None, + size_bytes: None, + num_segments: None, + index_version: None, + index_details: None, + }) + .collect(); + assert_eq!(indices, expected); +} + +/// Verifies that when the server returns `index_type` in the list response, +/// `list_indices` uses all enriched fields directly and does **not** make a +/// per-index `/index/{name}/stats/` call. +#[tokio::test] +async fn test_list_indices_enriched() { + let schema = Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8), + false, + ), + Field::new("text", DataType::Utf8, false), + ]); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "vector_idx", + "index_uuid": "3fa85f64-5717-4562-b3fc-2c963f66afa6", + "columns": ["vector"], + "index_type": "IVF_PQ", + "index_status": "done", + "num_indexed_rows": 1000, + "num_unindexed_rows": 50, + "size_bytes": 204800, + "num_segments": 2, + "index_version": 1, + "index_details": "{\"num_partitions\":16}", + "created_at": "2026-06-18T21:37:36.637Z", + "type_url": "type.googleapis.com/lance.index.vector.IvfPq", + }, + { + "index_name": "text_idx", + "index_uuid": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", + "columns": ["text"], + "index_type": "FTS", + "index_status": "done", + "num_indexed_rows": 1000, + "num_unindexed_rows": 0, + "size_bytes": 8192, + "num_segments": 1, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + // stats endpoint must NOT be called for enriched responses + path => panic!("Unexpected path (stats should not be called): {}", path), + } + }); + + let indices = table.list_indices().await.unwrap(); + assert_eq!(indices.len(), 2); + + let vec_idx = &indices[0]; + assert_eq!(vec_idx.name, "vector_idx"); + assert_eq!(vec_idx.index_type, IndexType::IvfPq); + assert_eq!(vec_idx.columns, vec!["vector".to_string()]); + assert_eq!( + vec_idx.index_uuid, + Some("3fa85f64-5717-4562-b3fc-2c963f66afa6".to_string()) + ); + assert_eq!(vec_idx.num_indexed_rows, Some(1000)); + assert_eq!(vec_idx.num_unindexed_rows, Some(50)); + assert_eq!(vec_idx.size_bytes, Some(204800)); + assert_eq!(vec_idx.num_segments, Some(2)); + assert_eq!(vec_idx.index_version, Some(1)); + assert_eq!( + vec_idx.index_details, + Some("{\"num_partitions\":16}".to_string()) + ); + assert_eq!( + vec_idx.type_url, + Some("type.googleapis.com/lance.index.vector.IvfPq".to_string()) + ); + assert_eq!( + vec_idx.created_at, + Some("2026-06-18T21:37:36.637Z".parse::>().unwrap()) + ); + + let text_idx = &indices[1]; + assert_eq!(text_idx.name, "text_idx"); + assert_eq!(text_idx.index_type, IndexType::FTS); + assert_eq!(text_idx.columns, vec!["text".to_string()]); + assert_eq!( + text_idx.index_uuid, + Some("aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee".to_string()) + ); + assert_eq!(text_idx.num_indexed_rows, Some(1000)); + assert_eq!(text_idx.num_unindexed_rows, Some(0)); + assert_eq!(text_idx.size_bytes, Some(8192)); + assert_eq!(text_idx.num_segments, Some(1)); + // optional fields not present in the response are None + assert_eq!(text_idx.index_version, None); + assert_eq!(text_idx.index_details, None); + assert_eq!(text_idx.type_url, None); + assert_eq!(text_idx.created_at, None); +} + +#[tokio::test] +async fn test_tokenize_uses_remote_index_details() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "icu", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": true, + "ascii_folding": true, + "custom_stop_words": ["hello"], + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let tokens = table + .tokenize("Hello, こんにちは世界!", "text_idx") + .await + .unwrap(); + + // Positions are relative to the first retained token, so dropping the + // leading "hello" stop word does not shift the remaining tokens. + assert_eq!( + tokens, + vec![ + FtsToken { + text: "こんにちは".to_string(), + position: 0, + }, + FtsToken { + text: "世界".to_string(), + position: 1, + }, + ] + ); +} + +#[tokio::test] +async fn test_tokenize_requires_existing_index_name() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let table = Table::new_with_handler("my_table", move |request| -> http::Response { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ "indexes": [] }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table.tokenize("hello", "text_idx").await.unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("No index named 'text_idx'") + )); +} + +#[tokio::test] +async fn test_tokenize_with_column_remote_requires_index_details() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table + .tokenize_with_column("hello", "text") + .await + .unwrap_err(); + + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("does not include tokenizer details") + )); +} + +#[test] +fn test_deserialize_created_at() { + #[derive(Deserialize)] + struct Wrapper { + #[serde(default, deserialize_with = "deserialize_created_at")] + created_at: Option>, + } + + // RFC 3339 string (current server format). + let w: Wrapper = serde_json::from_str(r#"{"created_at": "2026-06-18T21:37:36.637Z"}"#).unwrap(); + assert_eq!( + w.created_at, + Some("2026-06-18T21:37:36.637Z".parse::>().unwrap()) + ); + + // Unix milliseconds (legacy server format). + let w: Wrapper = serde_json::from_str(r#"{"created_at": 1700000000000}"#).unwrap(); + assert_eq!(w.created_at, DateTime::from_timestamp_millis(1700000000000)); + + // Null and missing both yield None. + let w: Wrapper = serde_json::from_str(r#"{"created_at": null}"#).unwrap(); + assert_eq!(w.created_at, None); + let w: Wrapper = serde_json::from_str(r#"{}"#).unwrap(); + assert_eq!(w.created_at, None); + + // A malformed string is rejected rather than silently dropped to None. + assert!(serde_json::from_str::(r#"{"created_at": "not-a-date"}"#).is_err()); +} diff --git a/rust/lancedb/src/remote/table/tests/lsm.rs b/rust/lancedb/src/remote/table/tests/lsm.rs new file mode 100644 index 000000000..17c2e7b94 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/lsm.rs @@ -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>, +) { + 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 { + http::Response::builder() + .status(202) + .body(String::new()) + .unwrap() +} + +fn ok_json(body: String) -> http::Response { + 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 = [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 = 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::(&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::) + } + }; + 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() + }) +} diff --git a/rust/lancedb/src/remote/table/tests/mod.rs b/rust/lancedb/src/remote/table/tests/mod.rs new file mode 100644 index 000000000..1f229dfdc --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/mod.rs @@ -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 { + 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 { + 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 { + 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 { + 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::>(), + "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 { + 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 = 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; diff --git a/rust/lancedb/src/remote/table/tests/multipart.rs b/rust/lancedb/src/remote/table/tests/multipart.rs new file mode 100644 index 000000000..b2b3f8958 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/multipart.rs @@ -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::>() + .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::::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> = Vec::new(); + let data: Box = + 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); +} diff --git a/rust/lancedb/src/remote/table/tests/query.rs b/rust/lancedb/src/remote/table/tests/query.rs new file mode 100644 index 000000000..a458eb7c4 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/query.rs @@ -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 { + 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::::None, + "upper_bound": Option::::None, + "k": 10, + "ef": Option::::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::>().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::::None, + "upper_bound": Option::::None, + "k": 10, + "ef": Option::::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::>().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::>().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::::None, + "upper_bound": Option::::None, + "ef": Option::::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::::None, + "upper_bound": Option::::None, + "ef": Option::::None, + "refine_factor": Option::::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::>() + .await + .unwrap(); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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 { + 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::>() + .await + .unwrap(); + let results = concat_batches(&results[0].schema(), &results).unwrap(); + + let query_index = results["query_index"].as_primitive::(); + // We don't guarantee order. + assert!(query_index.values().contains(&0)); + assert!(query_index.values().contains(&1)); +} diff --git a/rust/lancedb/src/remote/table/tests/schema_cache.rs b/rust/lancedb/src/remote/table/tests/schema_cache.rs new file mode 100644 index 000000000..75291a481 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/schema_cache.rs @@ -0,0 +1,1056 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[tokio::test] +async fn test_table_with_namespace_identifier() { + // Test that a table created with namespace uses the correct identifier in API calls + let table = Table::new_with_handler("ns1$ns2$table1", |request| { + assert_eq!(request.method(), "POST"); + // All API calls should use the full identifier in the path + assert_eq!(request.url().path(), "/v1/table/ns1$ns2$table1/describe/"); + + http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": { "fields": [] }}"#) + .unwrap() + }); + + // The name() method should return just the base name, not the full identifier + assert_eq!(table.name(), "ns1$ns2$table1"); + + // API operations should work correctly + let version = table.version().await.unwrap(); + assert_eq!(version, 1); +} + +#[tokio::test] +async fn test_query_with_namespace() { + let table = Table::new_with_handler("analytics$events", |request| { + match request.url().path() { + "/v1/table/analytics$events/query/" => { + assert_eq!(request.method(), "POST"); + + // Return empty arrow stream + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let body = write_ipc_file(&data); + + http::Response::builder() + .status(200) + .header("Content-Type", ARROW_FILE_CONTENT_TYPE) + .body(body) + .unwrap() + } + _ => { + panic!("Unexpected path: {}", request.url().path()); + } + } + }); + + let results = table.query().execute().await.unwrap(); + let batches = results.try_collect::>().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); +} + +#[tokio::test] +async fn test_add_data_with_namespace() { + 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("prod$metrics", move |mut request| { + match request.url().path() { + "/v1/table/prod$metrics/describe/" => http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap(), + "/v1/table/prod$metrics/insert/" => { + assert_eq!(request.method(), "POST"); + 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(r#"{"version": 2}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected request path: {}", path), + } + }); + + let result = table.add(data.clone()).execute().await.unwrap(); + + assert_eq!(result.version, 2); + + 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_create_index_with_namespace() { + let table = Table::new_with_handler("dev$users", |request| { + match request.url().path() { + "/v1/table/dev$users/create_index/" => { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.headers().get("Content-Type").unwrap(), + JSON_CONTENT_TYPE + ); + + // Verify the request body contains the column name + 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(); + assert_eq!(value["column"], "embedding"); + assert_eq!(value["index_type"], "IVF_PQ"); + } + + http::Response::builder() + .status(200) + .body("".to_string()) + .unwrap() + } + "/v1/table/dev$users/describe/" => { + let schema = Schema::new(vec![Field::new( + "embedding", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + 8, + ), + false, + )]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + _ => { + panic!("Unexpected path: {}", request.url().path()); + } + } + }); + + table + .create_index(&["embedding"], Index::IvfPq(IvfPqIndexBuilder::default())) + .execute() + .await + .unwrap(); +} + +#[tokio::test] +async fn test_drop_columns_with_namespace() { + let table = Table::new_with_handler("test$schema_ops", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.url().path(), + "/v1/table/test$schema_ops/drop_columns/" + ); + 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 columns = value["columns"].as_array().unwrap(); + assert_eq!(columns.len(), 2); + assert_eq!(columns[0], "old_col1"); + assert_eq!(columns[1], "old_col2"); + } + + http::Response::builder() + .status(200) + .body(r#"{"version": 5}"#) + .unwrap() + }); + + let result = table.drop_columns(&["old_col1", "old_col2"]).await.unwrap(); + assert_eq!(result.version, 5); +} + +#[tokio::test] +async fn test_uri() { + 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": 1, "schema": {"fields": []}, "location": "s3://bucket/path/to/table"}"#) + .unwrap() + }); + + let uri = table.uri().await.unwrap(); + assert_eq!(uri, "s3://bucket/path/to/table"); +} + +#[tokio::test] +async fn test_uri_missing_location() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + + // Server returns response without location field + http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}}"#) + .unwrap() + }); + + let result = table.uri().await; + assert!(result.is_err()); + assert!(matches!(&result, Err(Error::NotSupported { .. }))); +} + +#[tokio::test] +async fn test_uri_caching() { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + call_count_clone.fetch_add(1, Ordering::SeqCst); + + http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}, "location": "gs://bucket/table"}"#) + .unwrap() + }); + + // First call should fetch from server + let uri1 = table.uri().await.unwrap(); + assert_eq!(uri1, "gs://bucket/table"); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second call should use cached value + let uri2 = table.uri().await.unwrap(); + assert_eq!(uri2, "gs://bucket/table"); + assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call +} + +/// Test that schema is fetched once and cached for subsequent calls +#[tokio::test] +async fn test_schema_caching() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + call_count_clone.fetch_add(1, Ordering::SeqCst); + + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + }); + + // First call should fetch from server + let schema1 = table.schema().await.unwrap(); + assert_eq!(schema1.fields().len(), 1); + assert_eq!(schema1.field(0).name(), "a"); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second call should use cached value + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call + + // Third call should still use cached value + let schema3 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call +} + +/// Test that schema cache expires after 30 seconds TTL +#[tokio::test] +async fn test_schema_cache_invalidation_after_ttl() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + call_count_clone.fetch_add(1, Ordering::SeqCst); + + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + }); + + // First call should fetch from server + let _schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second call should use cached value (within TTL) + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&_schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Advance mock time past TTL (no real wait) + clock::advance_by(Duration::from_secs(31)); + + // Third call should re-fetch from server (TTL expired) + let schema3 = table.schema().await.unwrap(); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&_schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 2); +} + +/// Test that schema cache is invalidated after schema-changing operations +#[rstest] +#[case("overwrite")] +#[case("add_columns")] +#[case("drop_columns")] +#[case("alter_columns")] +#[tokio::test] +async fn test_schema_cache_invalidation_after_operation(#[case] operation: &str) { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false}, + {"name": "b", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else if path == "/v1/table/my_table/insert/" + || path == "/v1/table/my_table/add_columns/" + || path == "/v1/table/my_table/drop_columns/" + || path == "/v1/table/my_table/alter_columns/" + { + http::Response::builder() + .status(200) + .body(r#"{"version": 2}"#) + .unwrap() + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call should fetch from server + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second schema call should use cached value + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Perform the schema-changing operation + match operation { + "overwrite" => { + let data = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); + let _ = table.add(data).mode(AddDataMode::Overwrite).execute().await; + } + "add_columns" => { + let _ = table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "c".into(), + "a + 1".into(), + )])) + .execute() + .await; + } + "drop_columns" => { + let _ = table.drop_columns(&["b"]).await; + } + "alter_columns" => { + let alterations = vec![ColumnAlteration::new("a".into()).rename("new_a".into())]; + let _ = table.alter_columns(&alterations).await; + } + _ => panic!("Unknown operation: {}", operation), + } + + // Schema call after operation should re-fetch from server (cache invalidated) + let schema3 = table.schema().await.unwrap(); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 2); +} + +/// Test that schema cache is invalidated when server returns certain error codes +#[rstest] +#[case(400, true)] // 400 Bad Request should invalidate cache +#[case(401, false)] // 401 Unauthorized should NOT invalidate cache +#[case(403, false)] // 403 Forbidden should NOT invalidate cache +#[case(404, true)] // 404 Not Found should invalidate (table might be recreated) +#[case(500, true)] // 500 Internal Server Error should invalidate cache +#[case(503, false)] // 503 Service Unavailable should NOT invalidate cache +#[tokio::test] +async fn test_schema_cache_invalidation_on_errors( + #[case] error_status: u16, + #[case] should_invalidate: bool, +) { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + let current_count = call_count_clone.load(Ordering::SeqCst); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else if path == "/v1/table/my_table/count_rows/" { + // Return error on first count_rows call + if current_count == 1 { + http::Response::builder() + .status(error_status) + .body("error") + .unwrap() + } else { + http::Response::builder().status(200).body("10").unwrap() + } + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call should fetch from server + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second schema call should use cached value + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Perform operation that returns error + let result = table.count_rows(None).await; + assert!(result.is_err()); + + // Schema call after error - check if cache was invalidated + let schema3 = table.schema().await.unwrap(); + if should_invalidate { + assert_eq!( + call_count.load(Ordering::SeqCst), + 2, + "Cache should be invalidated for {} error", + error_status + ); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + } else { + assert_eq!( + call_count.load(Ordering::SeqCst), + 1, + "Cache should NOT be invalidated for {} error", + error_status + ); + assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + } +} + +/// A pinned snapshot should reuse the version and schema returned by its +/// initial describe instead of issuing two more describe requests. +#[tokio::test] +async fn test_checkout_current_seeds_schema_from_single_describe() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let calls = describe_calls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + calls.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version":42,"schema":{"fields":[{"name":"a","type":{"type":"int32"},"nullable":false}]}}"#, + ) + .unwrap() + }); + + let snapshot = table.checkout_current().await.unwrap(); + assert_eq!(snapshot.schema().await.unwrap().fields().len(), 1); + assert_eq!(describe_calls.load(Ordering::SeqCst), 1); +} + +/// Test that schema cache is invalidated after checkout +#[tokio::test] +async fn test_schema_cache_invalidation_on_checkout() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + call_count_clone.fetch_add(1, Ordering::SeqCst); + let count = call_count_clone.load(Ordering::SeqCst); + + if path == "/v1/table/my_table/describe/" { + // Return different schemas for different calls + if count <= 2 { + // First schema call and checkout validation + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + // After checkout + http::Response::builder() + .status(200) + .body( + r#"{"version": 2, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false}, + {"name": "b", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call + let schema1 = table.schema().await.unwrap(); + assert_eq!(schema1.fields().len(), 1); + + // Second schema call should use cached value (no new call) + let call_count_before = call_count.load(Ordering::SeqCst); + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), call_count_before); + + // Checkout to version 2 (makes a describe call to validate) + let _ = table.checkout(2).await; + + // Schema call after checkout should re-fetch (cache was invalidated) + let schema3 = table.schema().await.unwrap(); + assert_eq!(schema3.fields().len(), 2); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_schema_fetch_does_not_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 table = Table::new_with_handler("my_table", move |request| { + let body = request_body_json(&request); + let pinned = body["version"].as_u64() == Some(5); + if !pinned { + arrived_tx.lock().unwrap().send(()).unwrap(); + release_rx + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(10)) + .unwrap(); + } + let field = if pinned { "pinned" } else { "latest" }; + http::Response::builder() + .status(200) + .body(format!( + r#"{{"version":5,"schema":{{"fields":[{{"name":"{field}","type":{{"type":"int32"}},"nullable":false}}]}}}}"# + )) + .unwrap() + }); + + let schema_fetch = 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(); + + let schema = schema_fetch.await.unwrap().unwrap(); + assert!(schema.field_with_name("pinned").is_ok()); + let cached = table.schema().await.unwrap(); + assert!(cached.field_with_name("pinned").is_ok()); +} + +/// Test that schema cache is invalidated after checkout_latest +#[tokio::test] +async fn test_schema_cache_invalidation_on_checkout_latest() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second schema call should use cached value + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Checkout latest + let _ = table.checkout_latest().await; + + // Schema call after checkout_latest should re-fetch (cache was invalidated) + let schema3 = table.schema().await.unwrap(); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 2); +} + +/// Test that schema cache is invalidated after checkout_tag +#[tokio::test] +async fn test_schema_cache_invalidation_on_checkout_tag() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else if path == "/v1/table/my_table/tags/list/" { + http::Response::builder() + .status(200) + .body(r#"{"tags": {"v2": {"version": 2}}}"#) + .unwrap() + } else if path == "/v1/table/my_table/tags/version/" { + http::Response::builder() + .status(200) + .body(r#"{"version": 2}"#) + .unwrap() + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second schema call should use cached value + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Checkout tag + table + .checkout_tag("v2") + .await + .expect("checkout_tag should succeed"); + + // Schema call after checkout_tag should re-fetch (cache was invalidated) + let schema3 = table.schema().await.unwrap(); + assert_eq!( + call_count.load(Ordering::SeqCst), + 2, + "Cache should have been invalidated and re-fetched" + ); + assert_ne!( + Arc::as_ptr(&schema3), + Arc::as_ptr(&schema1), + "Should be different Arc instances" + ); +} + +/// Test that restore invalidates cache (via checkout_latest) +#[tokio::test] +async fn test_schema_cache_invalidation_on_restore() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else if path == "/v1/table/my_table/restore/" { + http::Response::builder() + .status(200) + .body(r#"{"version": 1}"#) + .unwrap() + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + table.checkout(1).await.unwrap(); + + // First schema call + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 2); + + // Second schema call uses cache + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 2); + + // Restore operation + table.restore().await.unwrap(); + + // Schema call after restore should re-fetch (cache invalidated) + let schema3 = table.schema().await.unwrap(); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 3); +} + +/// Test that centralized error handling invalidates cache on query errors +#[tokio::test] +async fn test_centralized_error_invalidation_on_query() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + let current_count = call_count_clone.load(Ordering::SeqCst); + + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else if path == "/v1/table/my_table/query/" { + // Return 400 error on first query (could be schema mismatch) + if current_count == 1 { + http::Response::builder() + .status(400) + .body("Bad request") + .unwrap() + } else { + // Return empty result for successful query + http::Response::builder() + .status(200) + .header("content-type", "application/vnd.apache.arrow.stream") + .body("") + .unwrap() + } + } else { + http::Response::builder() + .status(404) + .body("not found") + .unwrap() + } + }); + + // First schema call + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second schema call uses cache + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Query that returns 400 error + let result = table.query().execute().await; + assert!(result.is_err()); + + // Schema call after error should re-fetch (cache invalidated by centralized handler) + let schema3 = table.schema().await.unwrap(); + assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + assert_eq!(call_count.load(Ordering::SeqCst), 2); +} + +/// Test that concurrent schema() calls with an empty cache only trigger one fetch. +#[tokio::test] +async fn test_concurrent_schema_calls_single_fetch() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Arc::new(Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + panic!("Unexpected request: {}", path); + } + })); + + let mut handles = Vec::new(); + for _ in 0..10 { + let table = table.clone(); + handles.push(tokio::spawn(async move { table.schema().await.unwrap() })); + } + + let schemas: Vec = futures::future::try_join_all(handles).await.unwrap(); + + // All callers should get the same Arc + for schema in &schemas { + assert_eq!(Arc::as_ptr(schema), Arc::as_ptr(&schemas[0])); + } + // Only one describe call should have been made + assert_eq!(call_count.load(Ordering::SeqCst), 1); +} + +/// Test that a background refresh is triggered in the refresh window and +/// returns the cached value immediately. +#[tokio::test] +async fn test_background_refresh_triggers_in_window() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + let count = call_count_clone.fetch_add(1, Ordering::SeqCst); + if count == 0 { + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + http::Response::builder() + .status(200) + .body( + r#"{"version": 2, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false}, + {"name": "b", "type": { "type": "string" }, "nullable": true} + ]}}"#, + ) + .unwrap() + } + } else { + panic!("Unexpected request: {}", path); + } + }); + + // Populate cache and trigger peek transition to Current state + let schema1 = table.schema().await.unwrap(); + assert_eq!(schema1.fields().len(), 1); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + // Second call transitions cache from Refreshing to Current via peek() + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + + // Advance into refresh window (TTL=30s, window=5s, so 26s is in window) + clock::advance_by(Duration::from_secs(26)); + + // This call enters the refresh window: returns cached value and creates + // a background shared future (Refreshing state with previous). + let schema3 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + // Only the initial fetch so far + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Advance past TTL so the previous value expires. This forces the next + // schema() to Wait on the in-flight shared future, driving it to completion. + clock::advance_by(Duration::from_secs(30)); + + let schema4 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 2); + assert_eq!(schema4.fields().len(), 2); + assert_ne!(Arc::as_ptr(&schema4), Arc::as_ptr(&schema1)); +} + +/// Test that multiple calls during the refresh window don't trigger +/// duplicate background refreshes. +#[tokio::test] +async fn test_no_duplicate_background_refreshes() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + panic!("Unexpected request: {}", path); + } + }); + + // Populate cache and transition to Current state + let schema1 = table.schema().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + let _ = table.schema().await.unwrap(); // peek transition + + // Advance into refresh window + clock::advance_by(Duration::from_secs(26)); + + // Multiple rapid calls should all return cached. The first one enters + // the refresh window and starts a background fetch (Refreshing state). + // Subsequent calls see Refreshing with a valid previous and return it. + for _ in 0..5 { + let schema = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema), Arc::as_ptr(&schema1)); + } + + // Advance past TTL and drive the shared future to completion + clock::advance_by(Duration::from_secs(30)); + let _ = table.schema().await.unwrap(); + + // Only one additional describe call (the background refresh), + // not five separate ones + assert_eq!(call_count.load(Ordering::SeqCst), 2); +} + +/// Test that if a background refresh fails, the previously cached value +/// is preserved and still returned. +#[tokio::test] +async fn test_background_refresh_error_preserves_cache() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + let count = call_count_clone.fetch_add(1, Ordering::SeqCst); + if count == 0 { + // First call succeeds + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "a", "type": { "type": "int32" }, "nullable": false} + ]}}"#, + ) + .unwrap() + } else { + // Subsequent calls fail (422 is not retried) + http::Response::builder() + .status(422) + .body("Unprocessable Entity") + .unwrap() + } + } else { + panic!("Unexpected request: {}", path); + } + }); + + // Populate cache and transition to Current state + let schema1 = table.schema().await.unwrap(); + assert_eq!(schema1.fields().len(), 1); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + let _ = table.schema().await.unwrap(); // peek transition + + // Advance into refresh window + clock::advance_by(Duration::from_secs(26)); + + // Trigger background refresh (returns cached value). The background + // fetch will fail but the previous value should be preserved. + let schema2 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1)); + + // Still in the refresh window: the previous value is valid, + // so calling schema() should still return it. + let schema3 = table.schema().await.unwrap(); + assert_eq!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1)); + + // Advance past TTL. The shared future will be driven and fail. + // The peek() error path should revert to the previous cached value. + clock::advance_by(Duration::from_secs(30)); + + // After the error, the previous is restored but its timestamp is old, + // so the next call triggers a new fetch which also fails. + let result = table.schema().await; + assert_eq!(call_count.load(Ordering::SeqCst), 2); + // The error from the failed fetch should be propagated + assert!(result.is_err()); +} diff --git a/rust/lancedb/src/remote/table/tests/versions.rs b/rust/lancedb/src/remote/table/tests/versions.rs new file mode 100644 index 000000000..ef69c7cab --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/versions.rs @@ -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::>().unwrap() + ); + assert_eq!(versions[1].version, 2); + assert_eq!( + versions[1].timestamp, + "2024-02-01T00:00:00Z".parse::>().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::>().unwrap() + ); + assert!(versions[0].metadata.is_empty()); + assert_eq!(versions[1].version, 2); + assert_eq!( + versions[1].timestamp, + "2024-02-01T00:00:00Z".parse::>().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::("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 = 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 { .. }))); +} diff --git a/rust/lancedb/src/remote/table/tests/write.rs b/rust/lancedb/src/remote/table/tests/write.rs new file mode 100644 index 000000000..4a058b089 --- /dev/null +++ b/rust/lancedb/src/remote/table/tests/write.rs @@ -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 { + Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )) + }; + + // These table operations should translate 404 to TableNotFound. + let results: Vec>> = 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::("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::() + .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[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[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::("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 = 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::>(); + 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 = + 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::>(); + 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 = 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::>(); + assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]); + + let params = request.url().query_pairs().collect::>(); + 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 = 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 = 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::>(); + 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 }); +} diff --git a/rust/lancedb/src/remote/token_cache.rs b/rust/lancedb/src/remote/token_cache.rs index 6f34ab73c..35c87b92e 100644 --- a/rust/lancedb/src/remote/token_cache.rs +++ b/rust/lancedb/src/remote/token_cache.rs @@ -979,918 +979,4 @@ pub fn token_cache_for_config(config: &OAuthConfig) -> Result 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>>, - device_authorizations: Arc, - refresh_attempts: Arc, - invalid_grant_rejections: Arc, - access_tokens_issued: Arc, - current_refresh: Arc>>, - fail_refreshes: Arc, - issue_refresh_tokens: Arc, - } - - 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(¤t_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, - ¤t_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, - ¤t_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>, - 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::().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 { - 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("")); - // 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; diff --git a/rust/lancedb/src/remote/token_cache/tests.rs b/rust/lancedb/src/remote/token_cache/tests.rs new file mode 100644 index 000000000..a1a5952d8 --- /dev/null +++ b/rust/lancedb/src/remote/token_cache/tests.rs @@ -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>>, + device_authorizations: Arc, + refresh_attempts: Arc, + invalid_grant_rejections: Arc, + access_tokens_issued: Arc, + current_refresh: Arc>>, + fail_refreshes: Arc, + issue_refresh_tokens: Arc, +} + +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(¤t_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, + ¤t_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, + ¤t_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>, + 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::().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 { + 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("")); + // 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")); +}