fix(highgo): paginate queries via limit offset for performance

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<SqlInsertWriter> =
@@ -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);

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 {

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));