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..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 @@ -72,6 +72,7 @@ static SqlType getSqlTypeFromStatementType(StatementType statementType) { case "ROLLBACK_TRANSACTION": return SqlType.TCL; case "EXPORT_DATA": + return SqlType.EXPORT; case "EXPORT_MODEL": case "LOAD_DATA": default: 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..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 @@ -685,15 +685,12 @@ 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) + ? dmlStats.getNumDmlAffectedRows() + : 0L; + updateAffectedRowCount(dmlRowCount); break; case TCL: case DDL: @@ -725,8 +722,43 @@ void handleQueryResult(String query, TableResult results, SqlType queryType) throw new BigQueryJdbcException(ex); } break; + case EXPORT: + 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(exportRowCount); + break; case OTHER: - throw new BigQueryJdbcException(String.format("Unexpected value: " + queryType)); + String truncatedQuery = truncateQuery(query); + String id = + (results.getJobId() != null) ? results.getJobId().getJob() : results.getQueryId(); + LOG.warning( + "Encountered unmapped SQL statement type [Job/Query ID: %s]. Treating as update statement: %s", + id, 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); } } @@ -1589,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()) { @@ -1610,6 +1648,7 @@ enum SqlType { DDL, SCRIPT, TCL, + EXPORT, OTHER }