fix(bqjdbc): pass rowsInPage with TableResult

.gitignore

@@ -86,3 +86,7 @@ monorepo
 *.tfstate.lock.info
 
 .jqwik-database
+
+# Agent configuration and local skills
+.agent/
+.agents/

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -981,7 +981,9 @@ private boolean meetsReadRatio(TableResult results) {
       return false;
     }
 
-    long pageSize = querySettings.getMaxResultPerPage();
+    Long rowsInPage = results.getRowsInPage();
+    long pageSize =
+        (rowsInPage != null && rowsInPage > 0) ? rowsInPage : querySettings.getMaxResultPerPage();
 
     // Prevent division by zero due to potential overflows/empty sets:
     if (pageSize <= 0) {

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -1685,6 +1685,7 @@ && getOptions().getOpenTelemetryTracer() != null) {
           .setSchema(schema)
           .setTotalRows(data.y())
           .setPageNoSchema(data.x())
+          .setRowsInPage(data.x() != null ? (long) Iterables.size(data.x().getValues()) : 0L)
           .build();
     } finally {
       if (tableDataList != null) {
@@ -2016,6 +2017,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
           .setJobId(jobId)
           .setQueryId(results.getQueryId())
           .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
+          .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
           .build();
     }
     // only 1 page of result
@@ -2035,6 +2037,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
             results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
         .setQueryId(results.getQueryId())
         .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
+        .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
         .build();
   }
 

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java

@@ -424,6 +424,7 @@ public TableResult getQueryResults(QueryResultsOption... options)
                 .setJobId(job.getJobId())
                 .setTotalRows(0L)
                 .setPageNoSchema(new PageImpl<FieldValueList>(null, "", null))
+                .setRowsInPage(0L)
                 .build();
         return emptyTableResult;
       }

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java

@@ -49,6 +49,8 @@ public abstract static class Builder {
 
     public abstract TableResult.Builder setJobCreationReason(JobCreationReason jobCreationReason);
 
+    public abstract TableResult.Builder setRowsInPage(Long rowsInPage);
+
     /** Creates a @code TableResult} object. */
     public abstract TableResult build();
   }
@@ -81,6 +83,10 @@ public static Builder newBuilder() {
   @Nullable
   public abstract JobCreationReason getJobCreationReason();
 
+  /** Returns the number of rows in the current page of results. */
+  @Nullable
+  public abstract Long getRowsInPage();
+
   @Override
   public boolean hasNextPage() {
     return getPageNoSchema().hasNextPage();
@@ -94,12 +100,18 @@ public String getNextPageToken() {
   @Override
   public TableResult getNextPage() {
     if (getPageNoSchema().hasNextPage()) {
+      Page<FieldValueList> nextPageNoSchema = getPageNoSchema().getNextPage();
+      long nextRows =
+          nextPageNoSchema.getValues() != null
+              ? (long) Iterables.size(nextPageNoSchema.getValues())
+              : 0L;
       return TableResult.newBuilder()
           .setSchema(getSchema())
           .setTotalRows(getTotalRows())
-          .setPageNoSchema(getPageNoSchema().getNextPage())
+          .setPageNoSchema(nextPageNoSchema)
           .setQueryId(getQueryId())
           .setJobCreationReason(getJobCreationReason())
+          .setRowsInPage(nextRows)
           .build();
     }
     return null;
@@ -137,12 +149,14 @@ public String toString() {
         .add("totalRows", getTotalRows())
         .add("cursor", getNextPageToken())
         .add("queryId", getQueryId())
+        .add("rowsInPage", getRowsInPage())
         .toString();
   }
 
   @Override
   public final int hashCode() {
-    return Objects.hash(getPageNoSchema(), getSchema(), getTotalRows(), getQueryId());
+    return Objects.hash(
+        getPageNoSchema(), getSchema(), getTotalRows(), getQueryId(), getRowsInPage());
   }
 
   @Override
@@ -158,6 +172,7 @@ public final boolean equals(Object obj) {
         && Iterators.elementsEqual(getValues().iterator(), response.getValues().iterator())
         && Objects.equals(getSchema(), response.getSchema())
         && getTotalRows() == response.getTotalRows()
-        && getQueryId() == response.getQueryId();
+        && Objects.equals(getQueryId(), response.getQueryId())
+        && Objects.equals(getRowsInPage(), response.getRowsInPage());
   }
 }

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java

@@ -210,6 +210,7 @@ public class SerializationTest extends BaseSerializationTest {
           .setSchema(Schema.of())
           .setTotalRows(0L)
           .setPageNoSchema(new PageImpl(null, "", ImmutableList.of()))
+          .setRowsInPage(0L)
           .build();
   private static final BigQuery BIGQUERY =
       BigQueryOptions.newBuilder().setProjectId("p1").build().getService();

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java

@@ -54,10 +54,15 @@ private static FieldValueList newFieldValueList(String s) {
   @Test
   void testNullSchema() {
     TableResult result =
-        TableResult.newBuilder().setTotalRows(3L).setPageNoSchema(INNER_PAGE_0).build();
+        TableResult.newBuilder()
+            .setTotalRows(3L)
+            .setPageNoSchema(INNER_PAGE_0)
+            .setRowsInPage(2L)
+            .build();
     assertThat(result.getSchema()).isNull();
     assertThat(result.hasNextPage()).isTrue();
     assertThat(result.getNextPageToken()).isNotNull();
+    assertThat(result.getRowsInPage()).isEqualTo(2L);
     assertThat(result.getValues())
         .containsExactly(newFieldValueList("0"), newFieldValueList("1"))
         .inOrder();
@@ -66,6 +71,7 @@ void testNullSchema() {
     assertThat(next.getSchema()).isNull();
     assertThat(next.hasNextPage()).isFalse();
     assertThat(next.getNextPageToken()).isNull();
+    assertThat(next.getRowsInPage()).isEqualTo(1L);
     assertThat(next.getValues()).containsExactly(newFieldValueList("2"));
     assertThat(next.getNextPage()).isNull();
 
@@ -81,10 +87,12 @@ void testSchema() {
             .setSchema(SCHEMA)
             .setTotalRows(3L)
             .setPageNoSchema(INNER_PAGE_0)
+            .setRowsInPage(2L)
             .build();
     assertThat(result.getSchema()).isEqualTo(SCHEMA);
     assertThat(result.hasNextPage()).isTrue();
     assertThat(result.getNextPageToken()).isNotNull();
+    assertThat(result.getRowsInPage()).isEqualTo(2L);
     assertThat(result.getValues())
         .containsExactly(
             newFieldValueList("0").withSchema(SCHEMA.getFields()),
@@ -95,6 +103,7 @@ void testSchema() {
     assertThat(next.getSchema()).isEqualTo(SCHEMA);
     assertThat(next.hasNextPage()).isFalse();
     assertThat(next.getNextPageToken()).isNull();
+    assertThat(next.getRowsInPage()).isEqualTo(1L);
     assertThat(next.getValues())
         .containsExactly(newFieldValueList("2").withSchema(SCHEMA.getFields()));
     assertThat(next.getNextPage()).isNull();