Skip to content

Commit db22477

Browse files
Neenu1995Neenu1995
authored
feat(bigquery-jdbc): harden JDBC ResultSet compliance and optimize warning/statistics caching (#13348)
b/516471577 This PR improves the standard JDBC compliance of the BigQuery JDBC driver's `ResultSet` implementations and reduces redundant network traffic by introducing a lazy-loaded `Job` caching mechanism. These changes prevent driver crashes when running in database clients or ORM frameworks that automatically call standard pagination/warning getters, and cut down network latency by 50% when fetching query warnings and statistics. ## 🛠️ Key Changes ### 1. ResultSet JDBC Compatibility Stubs - Implemented standard, non-throwing method stubs in `BigQueryBaseResultSet`: - `getFetchSize()`: Returns the locally set fetch size, defaulting to the statement's fetch size or the default driver buffer size (20,000). - `setFetchSize(int rows)`: Stores the value in a local field (as a standard no-op placeholder). - `getFetchDirection()`: Returns `ResultSet.FETCH_FORWARD` (forward-only is the only supported direction). - `setFetchDirection(int direction)`: Restricts parameter to `FETCH_FORWARD`, throwing a `SQLException` otherwise. - `getWarnings()`: Lazily fetches warnings from the query's associated job. - `clearWarnings()`: Stubs out as a no-op. - Removed throwing implementations of the above methods from `BigQueryNoOpsResultSet`. ### 2. Network Optimization via Job Caching - Introduced a cached `job` field and `setJob(Job)` setter on `BigQueryBaseResultSet`. - Updated `getQueryStatistics()` and `getWarnings()` to check and reuse the pre-cached `Job` object in memory before performing a remote lookup via `bigQuery.getJob(jobId)`. - Updated `BigQueryStatement` to automatically assign the query execution `Job` to the statement's `currentResultSet` during execution. - Left the signatures of all result set factory methods (`processJsonResultSet`, `processArrowResultSet`, `processQueryResponse`) unchanged to preserve compatibility with existing Mockito spy setups.
1 parent 4236143 commit db22477

11 files changed

Lines changed: 482 additions & 102 deletions

File tree

‎java-bigquery-jdbc/.gitignore‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,7 @@ target-it/**
77
tools/**/*.class
88
tools/**/drivers/**
99
tools/**/logs/**
10-
tools/**/*.jfr
10+
tools/**/*.jfr
11+
12+
# Gemini/Jetski agent custom skills
13+
.agents/

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import com.google.cloud.bigquery.BigQuery;
2323
import com.google.cloud.bigquery.Field;
24+
import com.google.cloud.bigquery.Job;
2425
import com.google.cloud.bigquery.Schema;
2526
import com.google.cloud.bigquery.StandardSQLTypeName;
2627
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
@@ -93,9 +94,10 @@ private BigQueryArrowResultSet(
9394
int fromIndex,
9495
int toIndexExclusive,
9596
Thread ownedThread,
96-
BigQuery bigQuery)
97+
BigQuery bigQuery,
98+
Job job)
9799
throws SQLException {
98-
super(bigQuery, statement, schema, isNested);
100+
super(bigQuery, statement, schema, isNested, job);
99101
LOG.finestTrace("<init>");
100102
this.totalRows = totalRows;
101103
this.buffer = buffer;
@@ -128,6 +130,19 @@ static BigQueryArrowResultSet of(
128130
Thread ownedThread,
129131
BigQuery bigQuery)
130132
throws SQLException {
133+
return of(schema, arrowSchema, totalRows, statement, buffer, ownedThread, bigQuery, null);
134+
}
135+
136+
static BigQueryArrowResultSet of(
137+
Schema schema,
138+
ArrowSchema arrowSchema,
139+
long totalRows,
140+
BigQueryStatement statement,
141+
BlockingQueue<BigQueryArrowBatchWrapper> buffer,
142+
Thread ownedThread,
143+
BigQuery bigQuery,
144+
Job job)
145+
throws SQLException {
131146
return new BigQueryArrowResultSet(
132147
schema,
133148
arrowSchema,
@@ -139,7 +154,8 @@ static BigQueryArrowResultSet of(
139154
-1,
140155
-1,
141156
ownedThread,
142-
bigQuery);
157+
bigQuery,
158+
job);
143159
}
144160

145161
BigQueryArrowResultSet() throws SQLException {
@@ -159,7 +175,18 @@ static BigQueryArrowResultSet getNestedResultSet(
159175
Schema schema, BigQueryArrowBatchWrapper nestedBatch, int fromIndex, int toIndexExclusive)
160176
throws SQLException {
161177
return new BigQueryArrowResultSet(
162-
schema, null, -1, null, null, nestedBatch, true, fromIndex, toIndexExclusive, null, null);
178+
schema,
179+
null,
180+
-1,
181+
null,
182+
null,
183+
nestedBatch,
184+
true,
185+
fromIndex,
186+
toIndexExclusive,
187+
null,
188+
null,
189+
null);
163190
}
164191

165192
private class ArrowDeserializer implements AutoCloseable {

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java‎

Lines changed: 129 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
package com.google.cloud.bigquery.jdbc;
1818

1919
import com.google.cloud.bigquery.BigQuery;
20+
import com.google.cloud.bigquery.BigQueryError;
21+
import com.google.cloud.bigquery.BigQueryException;
2022
import com.google.cloud.bigquery.Field;
2123
import com.google.cloud.bigquery.FieldList;
2224
import com.google.cloud.bigquery.Job;
@@ -40,10 +42,12 @@
4042
import java.sql.ResultSet;
4143
import java.sql.ResultSetMetaData;
4244
import java.sql.SQLException;
45+
import java.sql.SQLWarning;
4346
import java.sql.Statement;
4447
import java.sql.Time;
4548
import java.sql.Timestamp;
4649
import java.util.Calendar;
50+
import java.util.List;
4751

4852
public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
4953
implements BigQueryResultSet {
@@ -58,15 +62,28 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
5862
protected final boolean isNested;
5963
protected boolean isClosed = false;
6064
protected boolean wasNull = false;
65+
private int fetchSize = -1;
66+
private Job job;
67+
private SQLWarning warnings;
68+
private boolean warningsLoaded = false;
6169
protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE;
6270

6371
protected BigQueryBaseResultSet(
6472
BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) {
73+
this(bigQuery, statement, schema, isNested, null);
74+
}
75+
76+
protected BigQueryBaseResultSet(
77+
BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested, Job job) {
6578
this.bigQuery = bigQuery;
6679
this.statement = statement;
6780
this.schema = schema;
6881
this.schemaFieldList = schema != null ? schema.getFields() : null;
6982
this.isNested = isNested;
83+
this.job = job;
84+
if (job != null) {
85+
this.jobId = job.getJobId();
86+
}
7087
this.LOG =
7188
BigQueryJdbcResultSetLogger.getLogger(
7289
this.getClass(), statement != null ? statement.connectionId : null);
@@ -76,11 +93,12 @@ public QueryStatistics getQueryStatistics() {
7693
if (queryStatistics != null) {
7794
return queryStatistics;
7895
}
79-
if (jobId == null || bigQuery == null) {
80-
return null;
96+
Job activeJob = this.job;
97+
if (activeJob == null && jobId != null && bigQuery != null) {
98+
this.job = bigQuery.getJob(jobId);
99+
activeJob = this.job;
81100
}
82-
Job job = bigQuery.getJob(jobId);
83-
queryStatistics = job != null ? job.getStatistics() : null;
101+
queryStatistics = activeJob != null ? activeJob.getStatistics() : null;
84102
return queryStatistics;
85103
}
86104

@@ -92,6 +110,20 @@ public JobId getJobId() {
92110
return jobId;
93111
}
94112

113+
public void setJob(Job job) {
114+
this.job = job;
115+
if (job != null) {
116+
this.jobId = job.getJobId();
117+
}
118+
this.queryStatistics = null;
119+
this.warnings = null;
120+
this.warningsLoaded = false;
121+
}
122+
123+
public Job getJob() {
124+
return job;
125+
}
126+
95127
public void setQueryId(String queryId) {
96128
this.queryId = queryId;
97129
}
@@ -688,4 +720,97 @@ public <T> T unwrap(Class<T> iface) throws SQLException {
688720
public boolean isWrapperFor(Class<?> iface) throws SQLException {
689721
return iface != null && iface.isInstance(this);
690722
}
723+
724+
@Override
725+
public int getFetchDirection() throws SQLException {
726+
checkClosed();
727+
// Fetch direction is restricted to forward-only.
728+
return ResultSet.FETCH_FORWARD;
729+
}
730+
731+
@Override
732+
public void setFetchDirection(int direction) throws SQLException {
733+
checkClosed();
734+
// Restricts the fetch direction to FETCH_FORWARD. Other directions are not supported.
735+
if (direction != ResultSet.FETCH_FORWARD) {
736+
throw new SQLException("Only FETCH_FORWARD is supported");
737+
}
738+
}
739+
740+
@Override
741+
public void setFetchSize(int rows) throws SQLException {
742+
checkClosed();
743+
if (rows < 0) {
744+
throw new SQLException("Fetch size must be >= 0");
745+
}
746+
// This is a no-op placeholder for JDBC API compliance to prevent crashes in
747+
// third-party client tools that call this automatically.
748+
// The driver manages pagination internally under the hood.
749+
this.fetchSize = rows;
750+
}
751+
752+
@Override
753+
public int getFetchSize() throws SQLException {
754+
checkClosed();
755+
// Returns the fetch size set on this ResultSet, or falls back to the statement's
756+
// fetch size, defaulting to the internal row buffer size of
757+
// BigQueryStatement.DEFAULT_BUFFER_SIZE.
758+
if (this.fetchSize > 0) {
759+
return this.fetchSize;
760+
}
761+
if (statement != null) {
762+
int statementFetchSize = statement.getFetchSize();
763+
if (statementFetchSize > 0) {
764+
return statementFetchSize;
765+
}
766+
}
767+
return BigQueryStatement.DEFAULT_BUFFER_SIZE;
768+
}
769+
770+
@Override
771+
public SQLWarning getWarnings() throws SQLException {
772+
checkClosed();
773+
// Dynamically fetches and chains non-fatal execution errors from the BigQuery Job
774+
// as SQLWarning objects, using lazy-loading and local caching for performance.
775+
if (warningsLoaded) {
776+
return warnings;
777+
}
778+
Job activeJob = this.job;
779+
if (activeJob == null && jobId != null && bigQuery != null) {
780+
try {
781+
this.job = bigQuery.getJob(jobId);
782+
activeJob = this.job;
783+
} catch (BigQueryException e) {
784+
throw new BigQueryJdbcException("Failed to retrieve job for warnings", e);
785+
}
786+
}
787+
if (activeJob != null
788+
&& activeJob.getStatus() != null
789+
&& activeJob.getStatus().getExecutionErrors() != null) {
790+
List<BigQueryError> errors = activeJob.getStatus().getExecutionErrors();
791+
SQLWarning head = null;
792+
SQLWarning tail = null;
793+
for (BigQueryError error : errors) {
794+
SQLWarning warning = new SQLWarning(error.getMessage(), error.getReason());
795+
if (head == null) {
796+
head = warning;
797+
tail = warning;
798+
} else {
799+
tail.setNextWarning(warning);
800+
tail = warning;
801+
}
802+
}
803+
this.warnings = head;
804+
}
805+
this.warningsLoaded = true;
806+
return warnings;
807+
}
808+
809+
@Override
810+
public void clearWarnings() throws SQLException {
811+
checkClosed();
812+
// Clears the cached warnings chain.
813+
this.warnings = null;
814+
this.warningsLoaded = true;
815+
}
691816
}

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.cloud.bigquery.Field;
2424
import com.google.cloud.bigquery.FieldValue;
2525
import com.google.cloud.bigquery.FieldValue.Attribute;
26+
import com.google.cloud.bigquery.Job;
2627
import com.google.cloud.bigquery.Schema;
2728
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
2829
import java.sql.ResultSet;
@@ -54,8 +55,9 @@ private BigQueryJsonResultSet(
5455
int fromIndex,
5556
int toIndexExclusive,
5657
Thread[] ownedThreads,
57-
BigQuery bigQuery) {
58-
super(bigQuery, statement, schema, isNested);
58+
BigQuery bigQuery,
59+
Job job) {
60+
super(bigQuery, statement, schema, isNested, job);
5961
this.totalRows = totalRows;
6062
this.buffer = buffer;
6163
this.cursor = cursor;
@@ -79,8 +81,20 @@ static BigQueryJsonResultSet of(
7981
Thread[] ownedThreads,
8082
BigQuery bigQuery) {
8183

84+
return of(schema, totalRows, buffer, statement, ownedThreads, bigQuery, null);
85+
}
86+
87+
static BigQueryJsonResultSet of(
88+
Schema schema,
89+
long totalRows,
90+
BlockingQueue<BigQueryFieldValueListWrapper> buffer,
91+
BigQueryStatement statement,
92+
Thread[] ownedThreads,
93+
BigQuery bigQuery,
94+
Job job) {
95+
8296
return new BigQueryJsonResultSet(
83-
schema, totalRows, buffer, statement, false, null, -1, -1, ownedThreads, bigQuery);
97+
schema, totalRows, buffer, statement, false, null, -1, -1, ownedThreads, bigQuery, job);
8498
}
8599

86100
static BigQueryJsonResultSet of(
@@ -91,7 +105,7 @@ static BigQueryJsonResultSet of(
91105
Thread[] ownedThreads) {
92106

93107
return new BigQueryJsonResultSet(
94-
schema, totalRows, buffer, statement, false, null, -1, -1, ownedThreads, null);
108+
schema, totalRows, buffer, statement, false, null, -1, -1, ownedThreads, null, null);
95109
}
96110

97111
BigQueryJsonResultSet() {
@@ -127,6 +141,7 @@ static BigQueryJsonResultSet getNestedResultSet(
127141
fromIndex,
128142
toIndexExclusive,
129143
null,
144+
null,
130145
null);
131146
}
132147

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryNoOpsResultSet.java‎

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import java.sql.ResultSet;
3434
import java.sql.RowId;
3535
import java.sql.SQLException;
36-
import java.sql.SQLWarning;
3736
import java.sql.SQLXML;
3837
import java.sql.Time;
3938
import java.sql.Timestamp;
@@ -42,21 +41,6 @@
4241
/** NoOps Abstract base class for BigQuery JDBC ResultSet(s). */
4342
abstract class BigQueryNoOpsResultSet implements ResultSet {
4443

45-
@Override
46-
public int getFetchDirection() throws SQLException {
47-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
48-
}
49-
50-
@Override
51-
public void setFetchSize(int rows) throws SQLException {
52-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
53-
}
54-
55-
@Override
56-
public int getFetchSize() throws SQLException {
57-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
58-
}
59-
6044
@Override
6145
public String getCursorName() throws SQLException {
6246
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
@@ -102,11 +86,6 @@ public boolean previous() throws SQLException {
10286
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
10387
}
10488

105-
@Override
106-
public void setFetchDirection(int direction) throws SQLException {
107-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
108-
}
109-
11089
@Override
11190
public boolean rowUpdated() throws SQLException {
11291
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
@@ -655,16 +634,6 @@ public void updateNClob(String columnLabel, Reader reader) throws SQLException {
655634
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
656635
}
657636

658-
@Override
659-
public SQLWarning getWarnings() throws SQLException {
660-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
661-
}
662-
663-
@Override
664-
public void clearWarnings() throws SQLException {
665-
throw new BigQueryJdbcSqlFeatureNotSupportedException(METHOD_NOT_IMPLEMENTED);
666-
}
667-
668637
void checkClosed() throws SQLException {
669638
if (isClosed()) {
670639
throw new BigQueryJdbcException("This " + getClass().getName() + " has been closed");

0 commit comments

Comments
 (0)