diff --git a/crates/dbx-sql/src/query_result_sql.rs b/crates/dbx-sql/src/query_result_sql.rs index f3698a44fe..4a72827ba9 100644 --- a/crates/dbx-sql/src/query_result_sql.rs +++ b/crates/dbx-sql/src/query_result_sql.rs @@ -192,14 +192,14 @@ pub fn build_query_pagination_execution_plan( } let can_use_first_page_cursor = options.use_agent_cursor && options.pagination.offset == 0; - // HighGo and OceanBase Oracle can spend substantially more time executing - // an unbounded query before the Agent exposes its first cursor page. Prefer - // a bounded SQL query whenever it can be rewritten safely. Independent - // pages of an unordered query do not have a stable row order; callers - // should add ORDER BY when that matters. + // HighGo, OceanBase Oracle, and Xugu can spend substantially more time + // executing an unbounded query before the Agent exposes its first cursor + // page. Prefer a bounded SQL query whenever it can be rewritten safely. + // Independent pages of an unordered query do not have a stable row order; + // callers should add ORDER BY when that matters. // Kingbase keeps the cursor for unordered queries to preserve its behavior. let prefer_server_pagination = match options.database_type { - Some(DatabaseType::Highgo | DatabaseType::OceanbaseOracle) => true, + Some(DatabaseType::Highgo | DatabaseType::OceanbaseOracle | DatabaseType::Xugu) => true, Some(DatabaseType::Kingbase) => kingbase_server_pagination_is_stable(&options.query_base_sql), _ => false, }; @@ -261,11 +261,11 @@ pub fn build_query_pagination_execution_plan( plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; - } else if options.database_type == Some(DatabaseType::Kingbase) + } else if matches!(options.database_type, Some(DatabaseType::Kingbase | DatabaseType::Xugu)) && single_selectable_statement(&options.sql, options.database_type).is_ok() && has_top_level_top(&options.sql) { - // Kingbase SQL Server compatibility mode rejects a statement that mixes a + // Kingbase SQL Server compatibility mode and Xugu reject statements mixing a // top-level TOP with a sibling LIMIT/OFFSET. Without an Agent cursor the // query-result export executes the statement once and streams the whole // result; the TOP clause already bounds the row count. @@ -313,11 +313,13 @@ pub fn build_paginated_query_sql(options: PaginatedQuerySqlOptions) -> QuerySqlB TablePaginationStrategy::AgentMaxRows | TablePaginationStrategy::Unbounded => ok(format!("{statement};")), TablePaginationStrategy::IrisTop => ok(add_iris_top_limit(&statement, safe_limit)), TablePaginationStrategy::LimitOffset => { - // Kingbase SQL Server compatibility mode accepts TOP as a real clause. + // Kingbase SQL Server compatibility mode and Xugu accept TOP as a real clause. // Appending LIMIT/OFFSET alongside a top-level TOP would be rejected by // the server ("multiple TOP/LIMIT clauses not allowed"), so fall back to // the Agent cursor / client-side row cap for such statements. - if options.database_type == Some(DatabaseType::Kingbase) && has_top_level_top(&statement) { + if matches!(options.database_type, Some(DatabaseType::Kingbase | DatabaseType::Xugu)) + && has_top_level_top(&statement) + { return err("unsupported"); } let dedup_order_by = dedup_projection_count_without_order_by(&options.original_sql); @@ -4822,6 +4824,117 @@ WHERE u.id = picked.id; assert!(!second_page.use_agent_result_session); } + #[test] + fn xugu_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::Xugu), + pagination: QueryPagination { limit: 100, 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 100;"); + assert_eq!(first_page.page_sql, Some(first_page.sql_to_execute.clone())); + assert_eq!(first_page.page_limit, Some(100)); + 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::Xugu), + pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(second_page.sql_to_execute, "SELECT * FROM events LIMIT 100 OFFSET 100;"); + assert_eq!(second_page.page_sql, Some(second_page.sql_to_execute.clone())); + assert_eq!(second_page.page_limit, Some(100)); + assert_eq!(second_page.page_offset, Some(100)); + assert!(!second_page.use_agent_result_session); + } + + #[test] + fn xugu_keeps_agent_fallback_for_unrewritable_query() { + 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::Xugu), + pagination: QueryPagination { limit: 100, 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.use_agent_result_session); + } + + #[test] + fn xugu_top_clause_keeps_first_page_agent_cursor() { + for sql in [ + "SELECT TOP 10 * FROM events", + "select top 10 * from events", + "/* first page */ SELECT /* bounded */ TOP 10 * FROM events", + ] { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: Some(DatabaseType::Xugu), + pagination: QueryPagination { limit: 100, 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_eq!(plan.page_limit, Some(100)); + assert_eq!(plan.page_offset, Some(0)); + assert!(plan.use_agent_result_session); + assert!(!plan.single_execution); + } + } + + #[test] + fn xugu_top_clause_without_agent_uses_single_execution() { + let sql = "SELECT TOP 10 * FROM events"; + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: Some(DatabaseType::Xugu), + pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: true, + }); + + assert_eq!(plan.sql_to_execute, sql); + assert!(plan.page_sql.is_none()); + assert_eq!(plan.page_limit, Some(100)); + assert_eq!(plan.page_offset, Some(0)); + assert!(!plan.use_agent_result_session); + assert!(plan.single_execution); + } + + #[test] + fn xugu_top_clause_rejects_sibling_limit_on_every_page() { + for offset in [0, 100] { + let result = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: "SELECT TOP 10 * FROM events".to_string(), + database_type: Some(DatabaseType::Xugu), + limit: 100, + offset, + }); + + assert!(!result.ok); + assert!(result.sql.is_none()); + assert_eq!(result.reason.as_deref(), Some("unsupported")); + } + } + #[test] fn oceanbase_oracle_prefers_bounded_first_page_and_keeps_cursor_fallback() { let sql = "SELECT * FROM events";