fix(sqlserver): paginate joins with duplicate result columns

Closes #5284

crates/dbx-core/src/db/sqlserver.rs

@@ -471,9 +471,11 @@ async fn collect_first_result_limited(
     mut stream: QueryStream<'_>,
     start: Instant,
     max_rows: Option<usize>,
+    result_offset: usize,
     spatial_columns: &[SqlServerSpatialColumn],
 ) -> Result<QueryResult, String> {
     let row_limit = query_result_row_limit(max_rows);
+    let mut remaining_offset = result_offset;
     let mut columns: Vec<String> = vec![];
     let mut column_types: Vec<String> = vec![];
     let mut rows: Vec<Vec<serde_json::Value>> = Vec::new();
@@ -491,6 +493,10 @@ async fn collect_first_result_limited(
             }
             QueryItem::Metadata(_) => {}
             QueryItem::Row(row) if row.result_index() == 0 => {
+                if remaining_offset > 0 {
+                    remaining_offset -= 1;
+                    continue;
+                }
                 if rows.len() < row_limit {
                     let (values, srids) =
                         row_to_json_with_spatial_metadata(&row, spatial_columns, |column_index, srid| {
@@ -529,6 +535,7 @@ struct SqlServerResultSet {
     column_types: Vec<String>,
     rows: Vec<Vec<serde_json::Value>>,
     truncated: bool,
+    remaining_offset: usize,
 }
 
 pub struct SqlServerStreamExportSummary {
@@ -1226,20 +1233,25 @@ async fn collect_result_sets_limited(
     mut stream: QueryStream<'_>,
     start: Instant,
     max_rows: Option<usize>,
+    result_offset: usize,
 ) -> Result<Vec<QueryResult>, String> {
     let row_limit = query_result_row_limit(max_rows);
     let mut results = Vec::new();
     let mut current: Option<SqlServerResultSet> = None;
+    let mut saw_result_set = false;
 
     while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? {
         match item {
             QueryItem::Metadata(metadata) => {
                 push_sqlserver_result_set(&mut results, current.take(), start);
+                let remaining_offset = if saw_result_set { 0 } else { result_offset };
+                saw_result_set = true;
                 current = Some(SqlServerResultSet {
                     columns: columns_from_metadata(&metadata),
                     column_types: column_types_from_metadata(&metadata),
                     rows: Vec::new(),
                     truncated: false,
+                    remaining_offset,
                 });
             }
             QueryItem::Row(row) => {
@@ -1248,7 +1260,13 @@ async fn collect_result_sets_limited(
                     column_types: row.columns().iter().map(sqlserver_column_type_name).collect(),
                     rows: Vec::new(),
                     truncated: false,
+                    remaining_offset: if saw_result_set { 0 } else { result_offset },
                 });
+                saw_result_set = true;
+                if result.remaining_offset > 0 {
+                    result.remaining_offset -= 1;
+                    continue;
+                }
                 if result.rows.len() < row_limit {
                     result.rows.push(row_to_json(&row));
                 } else {
@@ -2478,6 +2496,7 @@ pub async fn execute_query_with_max_rows(
     max_rows: Option<usize>,
 ) -> Result<QueryResult, String> {
     let start = Instant::now();
+    let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
 
     if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"])
         || sqlserver_dml_output_returns_rows(sql)
@@ -2490,7 +2509,14 @@ pub async fn execute_query_with_max_rows(
         };
         let (result, messages) = capture_sqlserver_messages(async {
             let stream = sqlserver_driver_result(client.query(query.sql.as_str(), &[])).await?;
-            sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, &query.spatial_columns)).await
+            sqlserver_driver_result(collect_first_result_limited(
+                stream,
+                start,
+                max_rows,
+                result_offset,
+                &query.spatial_columns,
+            ))
+            .await
         })
         .await;
         let mut result = query_result_with_server_messages(result?, messages);
@@ -2531,6 +2557,7 @@ pub async fn execute_batch_with_max_rows(
     max_rows: Option<usize>,
 ) -> Result<Vec<QueryResult>, String> {
     let start = Instant::now();
+    let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
     if sqlserver_batch_can_use_execute(sql) {
         let (result, messages) = capture_sqlserver_messages(sqlserver_driver_result(client.execute(sql, &[]))).await;
         let result = result?;
@@ -2562,6 +2589,7 @@ pub async fn execute_batch_with_max_rows(
                         stream,
                         start,
                         max_rows,
+                        result_offset,
                         &query.spatial_columns,
                     ))
                     .await
@@ -2591,9 +2619,10 @@ pub async fn execute_simple_batch_with_max_rows(
     max_rows: Option<usize>,
 ) -> Result<Vec<QueryResult>, String> {
     let start = Instant::now();
+    let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
     let (results, messages) = capture_sqlserver_messages(async {
         let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
-        sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await
+        sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows, result_offset)).await
     })
     .await;
     let mut results = results?;
@@ -2629,9 +2658,10 @@ async fn execute_simple_batch_first_result_with_max_rows(
     max_rows: Option<usize>,
 ) -> Result<QueryResult, String> {
     let start = Instant::now();
+    let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
     let (result, messages) = capture_sqlserver_messages(async {
         let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
-        sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, &[])).await
+        sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, &[])).await
     })
     .await;
     let mut result = query_result_with_server_messages(result?, messages);
@@ -3214,7 +3244,7 @@ mod tests {
 
         let first_result = source.split("async fn execute_simple_batch_first_result_with_max_rows").nth(1).unwrap();
         let first_result = first_result.split("fn strip_dbx_sqlserver_row_number_column").next().unwrap();
-        assert!(first_result.contains("collect_first_result_limited(stream, start, max_rows, &[])"));
+        assert!(first_result.contains("collect_first_result_limited(stream, start, max_rows, result_offset, &[])"));
         assert!(!first_result.contains("collect_result_sets_limited"));
     }
 
@@ -4106,6 +4136,7 @@ mod tests {
                 column_types: vec![],
                 rows: vec![],
                 truncated: false,
+                remaining_offset: 0,
             }),
             Instant::now(),
         );
@@ -4120,7 +4151,13 @@ mod tests {
         let mut results = Vec::new();
         super::push_sqlserver_result_set(
             &mut results,
-            Some(SqlServerResultSet { columns: vec![], column_types: vec![], rows: vec![], truncated: false }),
+            Some(SqlServerResultSet {
+                columns: vec![],
+                column_types: vec![],
+                rows: vec![],
+                truncated: false,
+                remaining_offset: 0,
+            }),
             Instant::now(),
         );