diff --git a/crates/sqlbuild-cli/src/native_commands/_helpers/execution.rs b/crates/sqlbuild-cli/src/native_commands/_helpers/execution.rs index a7a8b314c..9d5ce63b7 100644 --- a/crates/sqlbuild-cli/src/native_commands/_helpers/execution.rs +++ b/crates/sqlbuild-cli/src/native_commands/_helpers/execution.rs @@ -346,6 +346,7 @@ pub(crate) fn metadata_sync( reconcile_at: Instant::now(), reconcile_seconds, limit: MAX_HISTORY_RECORDS, + statistics_limit: MAX_HISTORY_RECORDS, pending: None, next_at: Instant::now(), } @@ -389,7 +390,7 @@ pub(crate) async fn statistics_page( generation: state.generation, location: state.location().clone(), after, - limit: state.limit, + limit: state.statistics_limit, }, ) .await @@ -423,8 +424,13 @@ pub(crate) fn freeze_history_page( > BODY_BYTES { state.statistics = None; - if state.limit > 1 { - state.limit = (state.limit / 2).max(1); + let limit = if state.kind() == WarehouseHistoryKind::Statistics { + &mut state.statistics_limit + } else { + &mut state.limit + }; + if *limit > 1 { + *limit /= 2; } else { state.next_at = Instant::now() + Duration::from_secs(METADATA_POLL_SECONDS); drop(write_output(&[( @@ -450,6 +456,9 @@ pub(crate) fn acknowledge_metadata_page(mut state: LocationMetadataSync) -> Loca state.statistics = None; } if state.source_after.is_none() { + if state.kind() == WarehouseHistoryKind::Statistics { + state.statistics_limit = MAX_HISTORY_RECORDS; + } state.index += 1; state.scan_id = Uuid::now_v7(); state.observed_since = None; diff --git a/crates/sqlbuild-cli/src/native_commands/models.rs b/crates/sqlbuild-cli/src/native_commands/models.rs index 61f4d42fd..a1bbf0838 100644 --- a/crates/sqlbuild-cli/src/native_commands/models.rs +++ b/crates/sqlbuild-cli/src/native_commands/models.rs @@ -200,6 +200,8 @@ pub(crate) struct LocationMetadataSync { pub(crate) reconcile_at: tokio::time::Instant, pub(crate) reconcile_seconds: u32, pub(crate) limit: usize, + /// Survives a failed statistics read's rewind through catalog, which resets `limit`. + pub(crate) statistics_limit: usize, pub(crate) pending: Option, pub(crate) next_at: tokio::time::Instant, } @@ -207,14 +209,16 @@ pub(crate) struct LocationMetadataSync { impl LocationMetadataSync { /// Failed samples revisit catalog before retrying; no frozen delivery is discarded. pub(crate) fn retry_read(&mut self) { - self.limit = (self.limit / 2).max(1); self.statistics = None; if self.kind() == WarehouseHistoryKind::Statistics { + self.statistics_limit = (self.statistics_limit / 2).max(1); self.index -= 1; self.source_after = None; self.scan_id = Uuid::now_v7(); self.observed_since = None; self.watermark_loaded = false; + } else { + self.limit = (self.limit / 2).max(1); } self.next_at = tokio::time::Instant::now() + std::time::Duration::from_secs(super::constants::METADATA_POLL_SECONDS); diff --git a/crates/sqlbuild-cli/src/native_commands/tests.rs b/crates/sqlbuild-cli/src/native_commands/tests.rs index 36ed1676b..6aecb7754 100644 --- a/crates/sqlbuild-cli/src/native_commands/tests.rs +++ b/crates/sqlbuild-cli/src/native_commands/tests.rs @@ -6,6 +6,8 @@ mod decision; mod helpers; #[path = "tests/test_init.rs"] mod init; +#[path = "tests/test_location_metadata.rs"] +mod location_metadata; #[path = "tests/test_location_plan.rs"] mod location_plan; #[path = "tests/test_paths.rs"] diff --git a/crates/sqlbuild-cli/src/native_commands/tests/helpers.rs b/crates/sqlbuild-cli/src/native_commands/tests/helpers.rs index ec817e0ae..ed08157d2 100644 --- a/crates/sqlbuild-cli/src/native_commands/tests/helpers.rs +++ b/crates/sqlbuild-cli/src/native_commands/tests/helpers.rs @@ -1,8 +1,13 @@ use std::path::Path; +use sqlbuild_orchestration::warehouse_metadata::models::MetadataBatch; +use uuid::Uuid; + use crate::argv::main::parse_cli_arguments::parse_cli_arguments; use crate::argv::models::{ArgvOutcome, ArgvSettings}; use crate::native_commands::_helpers::argv_host::{NativeArgvHost, namespace_defaults}; +use crate::native_commands::_helpers::execution::acknowledge_metadata_page; +use crate::native_commands::models::{LocationMetadataPage, LocationMetadataSync}; /// Parse `argv` onto the `CliNamespace` defaults with the native host. pub(super) fn parse(argv: &[&str]) -> (ArgvOutcome, NativeArgvHost) { @@ -35,3 +40,23 @@ pub(super) fn plan_output(diff_line: &str, count: usize) -> Vec { serde_json::to_vec(&serde_json::json!({"selected_count": count, "models": models})) .expect("plan output") } + +/// Acknowledges `pages` complete, empty pages, each finishing the kind it was read for. +pub(super) fn acknowledge_complete_pages( + state: LocationMetadataSync, + pages: usize, +) -> LocationMetadataSync { + (0..pages).fold(state, |mut state, _| { + state.pending = Some(LocationMetadataPage { + batch: MetadataBatch { + kind: state.kind(), + session_id: Uuid::nil(), + generation: 1, + records: Vec::new(), + scan: None, + }, + next_source_after: None, + }); + acknowledge_metadata_page(state) + }) +} diff --git a/crates/sqlbuild-cli/src/native_commands/tests/test_location_metadata.rs b/crates/sqlbuild-cli/src/native_commands/tests/test_location_metadata.rs new file mode 100644 index 000000000..a888efba5 --- /dev/null +++ b/crates/sqlbuild-cli/src/native_commands/tests/test_location_metadata.rs @@ -0,0 +1,55 @@ +use sqlbuild_orchestration::warehouse_metadata::constants::MAX_HISTORY_RECORDS; +use sqlbuild_orchestration::warehouse_metadata::models::{ + WarehouseHistoryKind, WarehouseStateLocation, +}; + +use crate::native_commands::_helpers::execution::metadata_sync; +use crate::native_commands::tests::helpers::acknowledge_complete_pages; +use crate::native_commands::tests::test_types::MetadataRetryTestCase; + +#[test] +fn given_failed_metadata_read_when_back_at_statistics_then_page_shrinks_until_scan_completes() { + let test_cases = [ + MetadataRetryTestCase { + description: "a timed-out statistics page is retried smaller after catalog is revisited", + failed_kind: WarehouseHistoryKind::Statistics, + acknowledged_pages: 1, + expected_statistics_limit: MAX_HISTORY_RECORDS / 2, + }, + MetadataRetryTestCase { + description: "a failed catalog page leaves the statistics page size alone", + failed_kind: WarehouseHistoryKind::Catalog, + acknowledged_pages: 1, + expected_statistics_limit: MAX_HISTORY_RECORDS, + }, + MetadataRetryTestCase { + description: "a completed statistics scan restores the full statistics page", + failed_kind: WarehouseHistoryKind::Statistics, + acknowledged_pages: 1 + WarehouseHistoryKind::ALL.len(), + expected_statistics_limit: MAX_HISTORY_RECORDS, + }, + ]; + for test_case in test_cases { + let location = WarehouseStateLocation { + database: Some("analytics".into()), + schema: "state".into(), + }; + let mut state = metadata_sync((vec![location], 1), 3600); + state.index = WarehouseHistoryKind::ALL + .iter() + .position(|kind| *kind == test_case.failed_kind) + .expect("kind is listed"); + state.retry_read(); + let state = acknowledge_complete_pages(state, test_case.acknowledged_pages); + assert_eq!( + (state.kind(), state.statistics_limit, state.limit), + ( + WarehouseHistoryKind::Statistics, + test_case.expected_statistics_limit, + MAX_HISTORY_RECORDS, + ), + "{description}", + description = test_case.description + ); + } +} diff --git a/crates/sqlbuild-cli/src/native_commands/tests/test_types.rs b/crates/sqlbuild-cli/src/native_commands/tests/test_types.rs index 93270cf91..580e0ee5e 100644 --- a/crates/sqlbuild-cli/src/native_commands/tests/test_types.rs +++ b/crates/sqlbuild-cli/src/native_commands/tests/test_types.rs @@ -1,3 +1,5 @@ +use sqlbuild_orchestration::warehouse_metadata::models::WarehouseHistoryKind; + use crate::argv::models::{TypeFailure, Value}; use crate::native_commands::errors::PythonNeeded; use crate::native_commands::models::{CommandOutput, NativeCommand, RulesAction}; @@ -89,3 +91,11 @@ pub(super) struct PlanDocumentTestCase { pub(super) planned: bool, pub(super) expected_first_model: serde_json::Value, } + +/// One failed metadata read and the statistics page size the sync uses afterwards. +pub(super) struct MetadataRetryTestCase { + pub(super) description: &'static str, + pub(super) failed_kind: WarehouseHistoryKind, + pub(super) acknowledged_pages: usize, + pub(super) expected_statistics_limit: usize, +}