diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 1a97121993..6961d2405c 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -39,6 +39,8 @@ use tokio_util::sync::CancellationToken; const AGENT_UNBOUNDED_ROW_LIMIT: usize = i32::MAX as usize; const STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str = "Streaming export is unsupported for this query. Simplify it or use a supported driver."; +const HIGHGO_STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str = + "HighGo cannot safely paginate this query for export. Run and export a single SELECT statement that supports LIMIT/OFFSET."; const AGENT_SESSION_MISSING_ERROR: &str = "Streaming export needs a result-set session, but this driver returned no session_id."; const STREAM_PROGRESS_TIME_INTERVAL: Duration = Duration::from_secs(1); @@ -506,6 +508,21 @@ fn supports_streaming_offset_pagination(request: &QueryResultExportRequest, page && !first_sql.trim().eq_ignore_ascii_case(second_sql.trim()) } +fn streaming_pagination_preflight_error( + request: &QueryResultExportRequest, + page_size: usize, + has_keyset_plan: bool, + has_single_execution_bound: bool, +) -> Option<&'static str> { + if has_keyset_plan || supports_streaming_offset_pagination(request, page_size) || has_single_execution_bound { + return None; + } + if request.database_type == DatabaseType::Highgo { + return Some(HIGHGO_STREAMING_PAGINATION_UNSUPPORTED_ERROR); + } + (!request.use_agent_cursor).then_some(STREAMING_PAGINATION_UNSUPPORTED_ERROR) +} + /// Enforceable in-memory row bound for a single-execution export, or `None` /// when the query cannot be safely streamed in one shot without an Agent /// cursor. Kingbase SQL Server compatibility mode TOP queries cannot be @@ -747,12 +764,13 @@ async fn export_query_result_core_inner( // bound (concrete TOP count and/or the configured export row limit), then it // stops after the first response. let single_execution_bound = single_execution_page_limit(request, page_size); - if keyset_plan.is_none() - && !request.use_agent_cursor - && !supports_streaming_offset_pagination(request, page_size) - && single_execution_bound.is_none() - { - return Err(STREAMING_PAGINATION_UNSUPPORTED_ERROR.to_string()); + if let Some(error) = streaming_pagination_preflight_error( + request, + page_size, + keyset_plan.is_some(), + single_execution_bound.is_some(), + ) { + return Err(error.to_string()); } let mut sql_writer: Option = @@ -1917,6 +1935,32 @@ mod tests { } } + #[test] + fn highgo_unrewritable_query_export_reports_actionable_pagination_error() { + let mut req = request("csv", None, None); + req.database_type = DatabaseType::Highgo; + req.use_agent_cursor = true; + req.sql = "SELECT * FROM users; SELECT 1".to_string(); + req.query_base_sql = req.sql.clone(); + + assert!(!supports_streaming_offset_pagination(&req, 100)); + assert_eq!( + streaming_pagination_preflight_error(&req, 100, false, false), + Some(HIGHGO_STREAMING_PAGINATION_UNSUPPORTED_ERROR) + ); + } + + #[test] + fn agent_cursor_export_keeps_unrewritable_fallback_for_other_databases() { + let mut req = request("csv", None, None); + req.database_type = DatabaseType::Oracle; + req.use_agent_cursor = true; + req.sql = "SELECT * FROM users; SELECT 1".to_string(); + req.query_base_sql = req.sql.clone(); + + assert_eq!(streaming_pagination_preflight_error(&req, 100, false, false), None); + } + #[test] fn csv_unlimited_export_has_no_effective_row_limit() { assert_eq!(effective_row_limit(&request("csv", None, None)), None); diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index 6010426ed5..4eb0b06c5e 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -167,8 +167,18 @@ pub fn build_query_pagination_execution_plan( } let can_use_first_page_cursor = options.use_agent_cursor && options.pagination.offset == 0; - let prefer_server_pagination = options.database_type == Some(DatabaseType::Kingbase) - && kingbase_server_pagination_is_stable(&options.query_base_sql); + // HighGo's PostgreSQL-compatible JDBC driver can buffer an unbounded result + // before the Agent has a chance to expose its cursor page. Prefer an actual + // LIMIT/OFFSET query whenever it can be rewritten safely. For an unordered + // query, independent pages are not guaranteed to preserve row order; this is + // an intentional tradeoff to keep HighGo execution bounded. Kingbase keeps + // the cursor for unordered queries because separate executions may not + // preserve row order there. + let prefer_server_pagination = match options.database_type { + Some(DatabaseType::Highgo) => true, + Some(DatabaseType::Kingbase) => kingbase_server_pagination_is_stable(&options.query_base_sql), + _ => false, + }; if can_use_first_page_cursor && !prefer_server_pagination { if !options.first_page_uses_actual_sql && options.sql == options.query_base_sql { plan.sql_to_execute = options.query_base_sql; @@ -190,11 +200,15 @@ pub fn build_query_pagination_execution_plan( plan.page_sql = paginated.sql; plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); - } else if can_use_first_page_cursor { + } else if can_use_first_page_cursor && options.database_type != Some(DatabaseType::Highgo) { // Kingbase JDBC may buffer an entire result in auto-commit mode, so use // LIMIT/OFFSET whenever the statement can be rewritten safely. Keep the // Agent cursor as a bounded fallback for multi-statement or dialect- - // specific SQL that the pagination parser cannot transform. + // specific SQL that the pagination parser cannot transform. HighGo does + // not take this fallback: its JDBC driver may materialize the unbounded + // result before cursor paging starts. Leaving the page metadata unset + // routes it through regular execution and the configured JDBC maxRows + // safeguard instead. if !options.first_page_uses_actual_sql && options.sql == options.query_base_sql { plan.sql_to_execute = options.query_base_sql; } @@ -4124,6 +4138,58 @@ WHERE u.id = picked.id; assert!(plan.use_agent_result_session); } + #[test] + fn highgo_prefers_server_pagination_over_agent_cursor() { + let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Highgo), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(first_page.sql_to_execute, "SELECT * FROM events LIMIT 500;"); + assert_eq!(first_page.page_sql, Some(first_page.sql_to_execute.clone())); + assert_eq!(first_page.page_limit, Some(500)); + assert_eq!(first_page.page_offset, Some(0)); + assert!(!first_page.use_agent_result_session); + + let second_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Highgo), + pagination: QueryPagination { limit: 500, offset: 500, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(second_page.sql_to_execute, "SELECT * FROM events LIMIT 500 OFFSET 500;"); + assert_eq!(second_page.page_sql, Some(second_page.sql_to_execute.clone())); + assert_eq!(second_page.page_limit, Some(500)); + assert_eq!(second_page.page_offset, Some(500)); + assert!(!second_page.use_agent_result_session); + } + + #[test] + fn highgo_avoids_agent_cursor_for_unrewritable_queries() { + let sql = "SELECT * FROM events; SELECT 1"; + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: Some(DatabaseType::Highgo), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, sql); + assert!(plan.page_sql.is_none()); + assert!(plan.page_limit.is_none()); + assert!(plan.page_offset.is_none()); + assert!(!plan.use_agent_result_session); + } + #[test] fn kingbase_prefers_server_pagination_for_ordered_query() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index 31c8485ded..b54244251c 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -505,6 +505,13 @@ enum TableExportCursorKind { ExternalDriver, } +fn table_export_cursor_allowed(database_type: DatabaseType, cursor_kind: TableExportCursorKind) -> bool { + // HighGo's JDBC driver may materialize the complete unbounded result before + // startTableRead can return its first cursor page. Use the existing bounded + // keyset/LIMIT-OFFSET table export path instead. + database_type != DatabaseType::Highgo || cursor_kind != TableExportCursorKind::Agent +} + #[derive(Clone, Debug, PartialEq, Eq)] enum TableExportCursorSession { Agent(String), @@ -652,7 +659,10 @@ async fn fetch_table_export_batch( } if !*table_read_attempted { - match table_export_cursor_kind(state, pool_key).await { + let cursor_kind = table_export_cursor_kind(state, pool_key) + .await + .filter(|cursor_kind| table_export_cursor_allowed(*db_type, *cursor_kind)); + match cursor_kind { Some(TableExportCursorKind::Agent) => { *table_read_attempted = true; let sql = table_cursor_sql(request, sql_context, query_col_names, column_types, primary_keys); @@ -2348,6 +2358,13 @@ mod tests { assert_eq!(format_tsv_rows(&rows), ""); } + #[test] + fn highgo_table_export_avoids_unbounded_agent_cursor() { + assert!(!table_export_cursor_allowed(DatabaseType::Highgo, TableExportCursorKind::Agent)); + assert!(table_export_cursor_allowed(DatabaseType::Highgo, TableExportCursorKind::ExternalDriver)); + assert!(table_export_cursor_allowed(DatabaseType::Oracle, TableExportCursorKind::Agent)); + } + #[test] fn export_batch_size_respects_row_limit_remaining_rows() { assert_eq!(next_export_batch_size(None, 12_000, 10_000), Some(10_000));