From c091d9dad0db102ec016c89cc4121a9dd9defd64 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 26 May 2026 00:00:26 +0000 Subject: [PATCH 1/4] fix: handle EXPORT DATA, EXPORT MODEL, and LOAD DATA statements --- .../bigquery/jdbc/BigQuerySqlTypeConverter.java | 2 ++ .../cloud/bigquery/jdbc/BigQueryStatement.java | 15 +++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java index cfdc64a14e7f..90e98a60a768 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java @@ -72,8 +72,10 @@ static SqlType getSqlTypeFromStatementType(StatementType statementType) { case "ROLLBACK_TRANSACTION": return SqlType.TCL; case "EXPORT_DATA": + return SqlType.EXPORT; case "EXPORT_MODEL": case "LOAD_DATA": + return SqlType.DDL; default: return SqlType.OTHER; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 41ec945bd68a..4ff21e84032b 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -725,6 +725,20 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) throw new BigQueryJdbcException(ex); } break; + case EXPORT: + try { + Job completedJob = this.bigQuery.getJob(results.getJobId()).waitFor(); + JobStatistics.QueryStatistics statistics = + (JobStatistics.QueryStatistics) completedJob.getStatistics(); + if (statistics.getExportDataStats() != null) { + updateAffectedRowCount(statistics.getExportDataStats().getRowCount()); + } else { + updateAffectedRowCount(0L); + } + } catch (InterruptedException ex) { + throw new BigQueryJdbcRuntimeException(ex); + } + break; case OTHER: throw new BigQueryJdbcException(String.format("Unexpected value: " + queryType)); } @@ -1610,6 +1624,7 @@ enum SqlType { DDL, SCRIPT, TCL, + EXPORT, OTHER } From 2f0c7c6a4df9f17ca816336360b9825515aba03b Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 26 May 2026 00:11:25 +0000 Subject: [PATCH 2/4] chore: address pr feedback --- .../bigquery/jdbc/BigQueryStatement.java | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 4ff21e84032b..60992b3c6091 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -727,15 +727,22 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) break; case EXPORT: try { - Job completedJob = this.bigQuery.getJob(results.getJobId()).waitFor(); - JobStatistics.QueryStatistics statistics = - (JobStatistics.QueryStatistics) completedJob.getStatistics(); - if (statistics.getExportDataStats() != null) { - updateAffectedRowCount(statistics.getExportDataStats().getRowCount()); - } else { - updateAffectedRowCount(0L); + Job job = this.bigQuery.getJob(results.getJobId()); + Job completedJob = (job != null) ? job.waitFor() : null; + JobStatistics stats = (completedJob != null) ? completedJob.getStatistics() : null; + + Long rowCount = 0L; + if (stats instanceof JobStatistics.QueryStatistics) { + JobStatistics.QueryStatistics queryStats = (JobStatistics.QueryStatistics) stats; + JobStatistics.QueryStatistics.ExportDataStats exportStats = + queryStats.getExportDataStats(); + if (exportStats != null && exportStats.getRowCount() != null) { + rowCount = exportStats.getRowCount(); + } } + updateAffectedRowCount(rowCount); } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); throw new BigQueryJdbcRuntimeException(ex); } break; From 25f61e6423e90777edd6c5a5e95627c510b757fc Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 26 May 2026 16:07:31 +0000 Subject: [PATCH 3/4] chore: remove throwing exception from OTHER --- .../jdbc/BigQuerySqlTypeConverter.java | 1 - .../bigquery/jdbc/BigQueryStatement.java | 62 +++++++++++-------- 2 files changed, 35 insertions(+), 28 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java index 90e98a60a768..a8e0ba3246df 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQuerySqlTypeConverter.java @@ -75,7 +75,6 @@ static SqlType getSqlTypeFromStatementType(StatementType statementType) { return SqlType.EXPORT; case "EXPORT_MODEL": case "LOAD_DATA": - return SqlType.DDL; default: return SqlType.OTHER; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 60992b3c6091..5ce07e15ebac 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -685,15 +685,9 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) break; case DML: case DML_EXTRA: - try { - Job completedJob = this.bigQuery.getJob(results.getJobId()).waitFor(); - JobStatistics.QueryStatistics statistics = completedJob.getStatistics(); - updateAffectedRowCount(statistics.getNumDmlAffectedRows()); - } catch (InterruptedException ex) { - throw new BigQueryJdbcRuntimeException(ex); - } catch (NullPointerException ex) { - throw new BigQueryJdbcException(ex); - } + QueryStatistics dmlStats = getQueryStatisticsFromJob(results); + Long dmlRowCount = (dmlStats != null) ? dmlStats.getNumDmlAffectedRows() : null; + updateAffectedRowCount(dmlRowCount); break; case TCL: case DDL: @@ -726,28 +720,42 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) } break; case EXPORT: - try { - Job job = this.bigQuery.getJob(results.getJobId()); - Job completedJob = (job != null) ? job.waitFor() : null; - JobStatistics stats = (completedJob != null) ? completedJob.getStatistics() : null; - - Long rowCount = 0L; - if (stats instanceof JobStatistics.QueryStatistics) { - JobStatistics.QueryStatistics queryStats = (JobStatistics.QueryStatistics) stats; - JobStatistics.QueryStatistics.ExportDataStats exportStats = - queryStats.getExportDataStats(); - if (exportStats != null && exportStats.getRowCount() != null) { - rowCount = exportStats.getRowCount(); - } + QueryStatistics exportStats = getQueryStatisticsFromJob(results); + Long exportRowCount = 0L; + if (exportStats != null) { + QueryStatistics.ExportDataStats dataStats = exportStats.getExportDataStats(); + if (dataStats != null && dataStats.getRowCount() != null) { + exportRowCount = dataStats.getRowCount(); } - updateAffectedRowCount(rowCount); - } catch (InterruptedException ex) { - Thread.currentThread().interrupt(); - throw new BigQueryJdbcRuntimeException(ex); } + updateAffectedRowCount(exportRowCount); break; case OTHER: - throw new BigQueryJdbcException(String.format("Unexpected value: " + queryType)); + String truncatedQuery = + (query != null && query.length() > 60) ? query.substring(0, 60) + "..." : query; + String jobId = (results.getJobId() != null) ? results.getJobId().getJob() : "unknown"; + LOG.warning( + "Encountered unmapped SQL statement type [Job ID: %s]. Treating as update statement: %s", + jobId, truncatedQuery); + updateAffectedRowCount(results.getTotalRows()); + break; + } + } + + private QueryStatistics getQueryStatisticsFromJob(TableResult results) throws SQLException { + try { + Job job = this.bigQuery.getJob(results.getJobId()); + Job completedJob = (job != null) ? job.waitFor() : null; + JobStatistics stats = (completedJob != null) ? completedJob.getStatistics() : null; + if (stats instanceof QueryStatistics) { + return (QueryStatistics) stats; + } + return null; + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + throw new BigQueryJdbcRuntimeException("Interrupted while waiting for job completion", ex); + } catch (BigQueryException ex) { + throw new BigQueryJdbcException("BigQueryException while waiting for job completion", ex); } } From c32819dc5033520110e62518351c1ec16a74417f Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Tue, 26 May 2026 21:43:16 +0000 Subject: [PATCH 4/4] chore: address pr feedback --- .../bigquery/jdbc/BigQueryStatement.java | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 5ce07e15ebac..52b856f1239e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -686,7 +686,10 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) case DML: case DML_EXTRA: QueryStatistics dmlStats = getQueryStatisticsFromJob(results); - Long dmlRowCount = (dmlStats != null) ? dmlStats.getNumDmlAffectedRows() : null; + Long dmlRowCount = + (dmlStats != null && dmlStats.getNumDmlAffectedRows() != null) + ? dmlStats.getNumDmlAffectedRows() + : 0L; updateAffectedRowCount(dmlRowCount); break; case TCL: @@ -731,12 +734,12 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) updateAffectedRowCount(exportRowCount); break; case OTHER: - String truncatedQuery = - (query != null && query.length() > 60) ? query.substring(0, 60) + "..." : query; - String jobId = (results.getJobId() != null) ? results.getJobId().getJob() : "unknown"; + String truncatedQuery = truncateQuery(query); + String id = + (results.getJobId() != null) ? results.getJobId().getJob() : results.getQueryId(); LOG.warning( - "Encountered unmapped SQL statement type [Job ID: %s]. Treating as update statement: %s", - jobId, truncatedQuery); + "Encountered unmapped SQL statement type [Job/Query ID: %s]. Treating as update statement: %s", + id, truncatedQuery); updateAffectedRowCount(results.getTotalRows()); break; } @@ -1618,13 +1621,19 @@ protected void logQueryExecutionStart(String sql) { if (sql == null) { return; } - String sanitizedSql = sql.trim().replaceAll("\\s+", " "); - String truncatedSql = - sanitizedSql.length() > 256 ? sanitizedSql.substring(0, 256) + "..." : sanitizedSql; + String truncatedSql = truncateQuery(sql); LOG.info("Executing query: " + truncatedSql); LOG.info("Using query settings: " + this.querySettings.toString()); } + private String truncateQuery(String sql) { + if (sql == null) { + return null; + } + String sanitizedSql = sql.trim().replaceAll("\\s+", " "); + return sanitizedSql.length() > 256 ? sanitizedSql.substring(0, 256) + "..." : sanitizedSql; + } + /** Throws a {@link BigQueryJdbcException} if this object is closed */ void checkClosed() throws SQLException { if (isClosed()) {