From efef4cc8d905fa43316092c9595510a01a87faf5 Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Wed, 13 Dec 2023 18:52:04 +0000 Subject: [PATCH 1/6] Support stateless queries. From https://cloud.google.com/java/docs/reference/google-cloud-bigquery/2.34.1/com.google.cloud.bigquery.BigQuery > Stateless queries: query execution without corresponding job metadata Stateless queries are currently in preview. This PR introduces support for stateless queries via a new JDBC query parameter. Set `jobcreationmode` to a string that matches one of the standard `JobCreationMode` enum values: - `JOB_CREATION_MODE_UNSPECIFIED`: Unspecified JobCreationMode, defaults to `JOB_CREATION_REQUIRED`. - `JOB_CREATION_REQUIRED`: Default. Job creation is always required. - `JOB_CREATION_OPTIONAL`: Job creation is optional. Returning immediate results is prioritized. BigQuery will automatically determine if a Job needs to be created. The conditions under which BigQuery can decide to not create a Job are subject to change. If Job creation is required, JOB_CREATION_REQUIRED mode should be used, which is the default. See https://github.com/googleapis/java-bigquery/blob/v2.34.0/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java#L98-L111 For example, the following create a connection with optional job creation: DriverManager.getConnection( "jdbc:BQDriver:my-project?jobcreationmode=JOB_CREATION_OPTIONAL", driverProperties); Note that this will cause queries to fail if the provided project does not have stateless queries preview enabled. --- .../starschema/clouddb/jdbc/BQConnection.java | 48 +++++++++++++ .../clouddb/jdbc/BQForwardOnlyResultSet.java | 10 +++ .../clouddb/jdbc/BQPreparedStatement.java | 2 +- .../clouddb/jdbc/BQScrollableResultSet.java | 16 ++++- .../starschema/clouddb/jdbc/BQStatement.java | 9 ++- .../clouddb/jdbc/BQStatementRoot.java | 11 +-- .../clouddb/jdbc/BQSupportFuncts.java | 7 +- .../BQForwardOnlyResultSetFunctionTest.java | 68 +++++++++++++++---- .../BQScrollableResultSetFunctionTest.java | 33 ++++++++- .../starschema/clouddb/jdbc/JdbcUrlTest.java | 24 +++++++ .../clouddb/jdbc/StatelessQuery.java | 42 ++++++++++++ 11 files changed, 242 insertions(+), 28 deletions(-) create mode 100644 src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java index 48d0b8eb..363b7014 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java @@ -72,6 +72,32 @@ public class BQConnection implements Connection { /** Boolean to determine whether or not to use legacy sql (default: false) * */ private final boolean useLegacySql; + /** + * Enum that describes whether to create a job in projects that support stateless queries. Copied + * from BigQueryImpl + */ + public static enum JobCreationMode { + /** If unspecified JOB_CREATION_REQUIRED is the default. */ + JOB_CREATION_MODE_UNSPECIFIED, + /** Default. Job creation is always required. */ + JOB_CREATION_REQUIRED, + + /** + * Job creation is optional. Returning immediate results is prioritized. BigQuery will + * automatically determine if a Job needs to be created. The conditions under which BigQuery can + * decide to not create a Job are subject to change. If Job creation is required, + * JOB_CREATION_REQUIRED mode should be used, which is the default. + * + *

Note that no job ID will be created if the results were returned immediately. + */ + JOB_CREATION_OPTIONAL; + + private JobCreationMode() {} + } + + /** The job creation mode - */ + private JobCreationMode jobCreationMode = JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED; + /** getter for useLegacySql */ public boolean getUseLegacySql() { return useLegacySql; @@ -210,6 +236,9 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor this.useQueryCache = parseBooleanQueryParam(caseInsensitiveProps.getProperty("querycache"), true); + this.jobCreationMode = + parseJobCreationMode(caseInsensitiveProps.getProperty("jobcreationmode")); + // Create Connection to BigQuery if (serviceAccount) { try { @@ -322,6 +351,21 @@ private static List parseArrayQueryParam(@Nullable String string, Charac : Arrays.asList(string.split(delimiter + "\\s*")); } + /** + * Return a {@link JobCreationMode} or raise an exception if the string does not match a variant. + */ + private static JobCreationMode parseJobCreationMode(@Nullable String string) + throws BQSQLException { + if (string == null) { + return null; + } + try { + return JobCreationMode.valueOf(string); + } catch (IllegalArgumentException e) { + throw new BQSQLException("could not parse " + string + " as job creation mode", e); + } + } + /** * * @@ -1214,4 +1258,8 @@ public Long getMaxBillingBytes() { public Integer getTimeoutMs() { return timeoutMs; } + + public JobCreationMode getJobCreationMode() { + return jobCreationMode; + } } diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java b/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java index 9000b056..5b0bbf17 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java @@ -105,6 +105,8 @@ public class BQForwardOnlyResultSet implements java.sql.ResultSet { private String projectId; /** Reference for the Job */ private @Nullable Job completedJob; + /** The BigQuery query ID; set if the query completed without a Job */ + private final @Nullable String queryId; /** The total number of bytes processed while creating this ResultSet */ private final @Nullable Long totalBytesProcessed; /** Whether the ResultSet came from BigQuery's cache */ @@ -127,12 +129,14 @@ public BQForwardOnlyResultSet( Bigquery bigquery, String projectId, @Nullable Job completedJob, + String queryId, BQStatementRoot bqStatementRoot) throws SQLException { this( bigquery, projectId, completedJob, + queryId, bqStatementRoot, null, false, @@ -160,6 +164,7 @@ public BQForwardOnlyResultSet( Bigquery bigquery, String projectId, @Nullable Job completedJob, + @Nullable String queryId, BQStatementRoot bqStatementRoot, List prefetchedRows, boolean prefetchedAllRows, @@ -172,6 +177,7 @@ public BQForwardOnlyResultSet( logger.debug("Created forward only resultset TYPE_FORWARD_ONLY"); this.Statementreference = (Statement) bqStatementRoot; this.completedJob = completedJob; + this.queryId = queryId; this.projectId = projectId; if (bigquery == null) { throw new BQSQLException("Failed to fetch results. Connection is closed."); @@ -2992,4 +2998,8 @@ public boolean wasNull() throws SQLException { return null; } } + + public @Nullable String getQueryId() { + return queryId; + } } diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQPreparedStatement.java b/src/main/java/net/starschema/clouddb/jdbc/BQPreparedStatement.java index 476e2da7..a191caa4 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQPreparedStatement.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQPreparedStatement.java @@ -254,7 +254,7 @@ public ResultSet executeQuery() throws SQLException { this); } else { return new BQForwardOnlyResultSet( - this.connection.getBigquery(), this.projectId, referencedJob, this); + this.connection.getBigquery(), this.projectId, referencedJob, null, this); } } // Pause execution for half second before polling job status diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQScrollableResultSet.java b/src/main/java/net/starschema/clouddb/jdbc/BQScrollableResultSet.java index b938852c..d51fef05 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQScrollableResultSet.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQScrollableResultSet.java @@ -66,7 +66,10 @@ public class BQScrollableResultSet extends ScrollableResultset */ private final @Nullable List biEngineReasons; - private final JobReference jobReference; + private final @Nullable JobReference jobReference; + + /** The BigQuery query ID; set if the query completed without a Job */ + private final @Nullable String queryId; private TableSchema schema; @@ -86,7 +89,8 @@ public BQScrollableResultSet( bigQueryGetQueryResultResponse.getCacheHit(), null, null, - bigQueryGetQueryResultResponse.getJobReference()); + bigQueryGetQueryResultResponse.getJobReference(), + null); BigInteger maxrow; try { @@ -104,7 +108,8 @@ public BQScrollableResultSet( @Nullable Boolean cacheHit, @Nullable String biEngineMode, @Nullable List biEngineReasons, - JobReference jobReference) { + @Nullable JobReference jobReference, + @Nullable String queryId) { logger.debug("Created Scrollable resultset TYPE_SCROLL_INSENSITIVE"); try { maxFieldSize = bqStatementRoot.getMaxFieldSize(); @@ -126,6 +131,7 @@ public BQScrollableResultSet( this.biEngineMode = biEngineMode; this.biEngineReasons = biEngineReasons; this.jobReference = jobReference; + this.queryId = queryId; } /** {@inheritDoc} */ @@ -302,4 +308,8 @@ public String getString(int columnIndex) throws SQLException { return null; } } + + public @Nullable String getQueryId() { + return queryId; + } } diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java b/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java index 2350a896..6310a993 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java @@ -214,6 +214,7 @@ private ResultSet executeQueryHelper(String querySql, boolean unlimitedBillingBy this.connection.getBigquery(), projectId, referencedJob, + qr.getQueryId(), this, rows, fetchedAll, @@ -234,7 +235,8 @@ private ResultSet executeQueryHelper(String querySql, boolean unlimitedBillingBy qr.getCacheHit(), biEngineMode, biEngineReasons, - qr.getJobReference()); + qr.getJobReference(), + qr.getQueryId()); } jobAlreadyCompleted = true; } @@ -285,7 +287,7 @@ private ResultSet executeQueryHelper(String querySql, boolean unlimitedBillingBy this); } else { return new BQForwardOnlyResultSet( - this.connection.getBigquery(), projectId, referencedJob, this); + this.connection.getBigquery(), projectId, referencedJob, null, this); } } // Pause execution for half second before polling job status @@ -345,7 +347,8 @@ protected QueryResponse runSyncQuery(String querySql, boolean unlimitedBillingBy // socket timeouts (long) getMaxRows(), this.getAllLabels(), - this.connection.getUseQueryCache()); + this.connection.getUseQueryCache(), + this.connection.getJobCreationMode()); syncResponseFromCurrentQuery.set(resp); this.mostRecentJobReference.set(resp.getJobReference()); } catch (Exception e) { diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java b/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java index 627dbf17..bf8d08d1 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java @@ -265,7 +265,8 @@ private int executeDML(String sql) throws SQLException { (long) querytimeout * 1000, (long) getMaxRows(), this.getAllLabels(), - this.connection.getUseQueryCache()); + this.connection.getUseQueryCache(), + this.connection.getJobCreationMode()); this.mostRecentJobReference.set(qr.getJobReference()); if (defaultValueIfNull(qr.getJobComplete(), false)) { @@ -327,7 +328,8 @@ public ResultSet executeQuery(String querySql, boolean unlimitedBillingBytes) (long) querytimeout * 1000, (long) getMaxRows(), this.getAllLabels(), - this.connection.getUseQueryCache()); + this.connection.getUseQueryCache(), + this.connection.getJobCreationMode()); this.mostRecentJobReference.set(qr.getJobReference()); referencedJob = @@ -362,7 +364,8 @@ public ResultSet executeQuery(String querySql, boolean unlimitedBillingBytes) qr.getCacheHit(), biEngineMode, biEngineReasons, - referencedJob.getJobReference()); + referencedJob.getJobReference(), + qr.getQueryId()); } jobAlreadyCompleted = true; } @@ -384,7 +387,7 @@ public ResultSet executeQuery(String querySql, boolean unlimitedBillingBytes) this); } else { return new BQForwardOnlyResultSet( - this.connection.getBigquery(), projectId, referencedJob, this); + this.connection.getBigquery(), projectId, referencedJob, null, this); } } // Pause execution for half second before polling job status diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java b/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java index 4f39d0a2..bc869990 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java @@ -44,6 +44,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Stream; +import net.starschema.clouddb.jdbc.BQConnection.JobCreationMode; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -644,7 +645,8 @@ static QueryResponse runSyncQuery( Long queryTimeoutMs, Long maxResults, Map labels, - boolean useQueryCache) + boolean useQueryCache, + JobCreationMode jobCreationMode) throws IOException { QueryRequest qr = new QueryRequest() @@ -654,6 +656,9 @@ static QueryResponse runSyncQuery( .setQuery(querySql) .setUseLegacySql(useLegacySql) .setMaximumBytesBilled(maxBillingBytes); + if (jobCreationMode != null) { + qr = qr.setJobCreationMode(jobCreationMode.name()); + } if (dataSet != null) { qr.setDefaultDataset(new DatasetReference().setDatasetId(dataSet).setProjectId(projectId)); } diff --git a/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java b/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java index 41195c81..71626bee 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java +++ b/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java @@ -28,7 +28,12 @@ import com.google.gson.Gson; import java.io.IOException; import java.math.BigDecimal; -import java.sql.*; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Time; +import java.sql.Timestamp; import java.text.DateFormat; import java.text.ParseException; import java.text.SimpleDateFormat; @@ -40,6 +45,8 @@ import java.util.Properties; import java.util.TimeZone; import junit.framework.Assert; +import org.assertj.core.api.Assertions; +import org.junit.After; import org.junit.Before; import org.junit.Test; import org.slf4j.Logger; @@ -72,6 +79,14 @@ public void setup() throws SQLException, IOException { this.defaultConn = new BQConnection(url, new Properties()); } + @After + public void teardown() throws SQLException { + if (defaultConn != null) { + defaultConn.close(); + defaultConn = null; + } + } + private BQConnection conn() throws SQLException, IOException { return this.defaultConn; } @@ -202,21 +217,26 @@ public void isClosedValidtest() { */ @Before public void NewConnection() { - NewConnection(true); + NewConnection("&useLegacySql=true"); } - void NewConnection(boolean useLegacySql) { - + void NewConnection(String extraUrl) { this.logger.info("Testing the JDBC driver"); try { Class.forName("net.starschema.clouddb.jdbc.BQDriver"); Properties props = BQSupportFuncts.readFromPropFile( getClass().getResource("/installedaccount1.properties").getFile()); - props.setProperty("useLegacySql", String.valueOf(useLegacySql)); + String jdcbUrl = BQSupportFuncts.constructUrlFromPropertiesFile(props); + if (extraUrl != null) { + jdcbUrl += extraUrl; + } + if (BQForwardOnlyResultSetFunctionTest.con != null) { + BQForwardOnlyResultSetFunctionTest.con.close(); + } BQForwardOnlyResultSetFunctionTest.con = DriverManager.getConnection( - BQSupportFuncts.constructUrlFromPropertiesFile(props), + jdcbUrl, BQSupportFuncts.readFromPropFile( getClass().getResource("/installedaccount1.properties").getFile())); } catch (Exception e) { @@ -438,7 +458,7 @@ public void testResultSetTypesInGetString() throws SQLException { + "STRUCT(1 as a, ['an', 'array'] as b)," + "TIMESTAMP('2012-01-01 00:00:03.0000') as t"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); java.sql.ResultSet result = null; try { Statement stmt = @@ -500,7 +520,7 @@ public void testResultSetDateTimeType() throws SQLException, ParseException { final DateFormat utcDateFormatter = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS"); utcDateFormatter.setTimeZone(TimeZone.getTimeZone("UTC")); - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); Statement stmt = BQForwardOnlyResultSetFunctionTest.con.createStatement( ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); @@ -534,7 +554,7 @@ public void testResultSetDateTimeType() throws SQLException, ParseException { public void testResultSetTimestampType() throws SQLException, ParseException { final String sql = "SELECT TIMESTAMP('2012-01-01 01:02:03.04567')"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); Statement stmt = BQForwardOnlyResultSetFunctionTest.con.createStatement( ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); @@ -567,7 +587,7 @@ public void testResultSetTypesInGetObject() throws SQLException, ParseException + "CAST('2011-04-03' AS DATE), " + "CAST('nan' AS FLOAT)"; - this.NewConnection(true); + this.NewConnection("&useLegacySql=true"); java.sql.ResultSet result = null; try { Statement stmt = @@ -595,7 +615,7 @@ public void testResultSetTypesInGetObject() throws SQLException, ParseException public void testResultSetArraysInGetObject() throws SQLException, ParseException { final String sql = "SELECT [1, 2, 3], [TIMESTAMP(\"2010-09-07 15:30:00 America/Los_Angeles\")]"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); java.sql.ResultSet result = null; try { Statement stmt = @@ -629,7 +649,8 @@ public void testResultSetArraysInGetObject() throws SQLException, ParseException @Test public void testResultSetTimeType() throws SQLException, ParseException { final String sql = "select current_time(), CAST('00:00:02.12345' AS TIME)"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); + java.sql.ResultSet result = null; try { Statement stmt = @@ -686,7 +707,7 @@ public void testResultSetProcedures() throws SQLException, ParseException { final String sql = "CREATE PROCEDURE looker_test.procedure_test(target_id INT64)\n" + "BEGIN\n" + "END;"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); java.sql.ResultSet result = null; try { Statement stmt = @@ -713,7 +734,7 @@ public void testResultSetProcedures() throws SQLException, ParseException { public void testResultSetProceduresAsync() throws SQLException { final String sql = "CREATE PROCEDURE looker_test.long_procedure(target_id INT64)\n" + "BEGIN\n" + "END;"; - this.NewConnection(false); + this.NewConnection("&useLegacySql=false"); try { BQConnection bq = conn(); @@ -756,7 +777,7 @@ public void testBQForwardOnlyResultSetDoesntThrowNPE() throws Exception { // before the results have been fetched. This was throwing a NPE. bq.close(); try { - new BQForwardOnlyResultSet(bq.getBigquery(), defaultProjectId, ref, stmt); + new BQForwardOnlyResultSet(bq.getBigquery(), defaultProjectId, ref, null, stmt); Assert.fail("Initalizing BQForwardOnlyResultSet should throw something other than a NPE."); } catch (SQLException e) { Assert.assertEquals(e.getMessage(), "Failed to fetch results. Connection is closed."); @@ -807,4 +828,21 @@ private void mockResponse(String jsonResponse) throws Exception { results.getBiEngineMode(); results.getBiEngineReasons(); } + + @Test + public void testStatelessQuery() throws SQLException { + NewConnection("&useLegacySql=false&jobcreationmode=JOB_CREATION_OPTIONAL"); + StatelessQuery.assumeStatelessQueriesEnabled( + BQForwardOnlyResultSetFunctionTest.con.getCatalog()); + final Statement stmt = + BQForwardOnlyResultSetFunctionTest.con.createStatement( + ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); + final ResultSet result = stmt.executeQuery(StatelessQuery.exampleQuery()); + final String[][] rows = BQSupportMethods.GetQueryResult(result); + Assertions.assertThat(rows).isEqualTo(StatelessQuery.exampleValues()); + + final BQForwardOnlyResultSet bqResultSet = (BQForwardOnlyResultSet) result; + Assertions.assertThat(bqResultSet.getJobId()).isNull(); + Assertions.assertThat(bqResultSet.getQueryId()).contains("!"); + } } diff --git a/src/test/java/net/starschema/clouddb/jdbc/BQScrollableResultSetFunctionTest.java b/src/test/java/net/starschema/clouddb/jdbc/BQScrollableResultSetFunctionTest.java index 2f40868a..c7e81907 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/BQScrollableResultSetFunctionTest.java +++ b/src/test/java/net/starschema/clouddb/jdbc/BQScrollableResultSetFunctionTest.java @@ -28,6 +28,8 @@ import java.sql.Statement; import java.util.Properties; import junit.framework.Assert; +import org.assertj.core.api.Assertions; +import org.junit.After; import org.junit.Before; import org.junit.Test; import org.slf4j.Logger; @@ -242,7 +244,16 @@ public void isClosedValidtest() { */ @Before public void NewConnection() { + NewConnection("&useLegacySql=true"); + } + + @After + public void closeConnection() throws SQLException { + BQScrollableResultSetFunctionTest.con.close(); + BQScrollableResultSetFunctionTest.con = null; + } + public void NewConnection(String extraUrl) { try { if (BQScrollableResultSetFunctionTest.con == null || !BQScrollableResultSetFunctionTest.con.isValid(0)) { @@ -253,7 +264,9 @@ public void NewConnection() { BQSupportFuncts.constructUrlFromPropertiesFile( BQSupportFuncts.readFromPropFile( getClass().getResource("/installedaccount1.properties").getFile())); - jdbcUrl += "&useLegacySql=true"; + if (jdbcUrl != null) { + jdbcUrl += extraUrl; + } BQScrollableResultSetFunctionTest.con = DriverManager.getConnection( jdbcUrl, @@ -714,4 +727,22 @@ private void mockResponse(String jsonResponse) throws Exception { results.getBiEngineMode(); results.getBiEngineReasons(); } + + @Test + public void testStatelessQuery() throws SQLException { + closeConnection(); + NewConnection("&useLegacySql=true&jobcreationmode=JOB_CREATION_OPTIONAL"); + StatelessQuery.assumeStatelessQueriesEnabled( + BQScrollableResultSetFunctionTest.con.getCatalog()); + final Statement stmt = + BQScrollableResultSetFunctionTest.con.createStatement( + ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY); + final ResultSet result = stmt.executeQuery(StatelessQuery.exampleQuery()); + final String[][] rows = BQSupportMethods.GetQueryResult(result); + Assertions.assertThat(rows).isEqualTo(StatelessQuery.exampleValues()); + + final BQScrollableResultSet bqResultSet = (BQScrollableResultSet) result; + Assertions.assertThat(bqResultSet.getJobId()).isNull(); + Assertions.assertThat(bqResultSet.getQueryId()).contains("!"); + } } diff --git a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java b/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java index 3732de3a..6b32a9d3 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java +++ b/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java @@ -13,6 +13,7 @@ import java.util.Map; import java.util.Properties; import junit.framework.Assert; +import net.starschema.clouddb.jdbc.BQConnection.JobCreationMode; import org.assertj.core.api.Assertions; import org.junit.Before; import org.junit.Rule; @@ -508,4 +509,27 @@ private Properties getProperties(String pathToProp) throws IOException { private String getUrl(String pathToProp, String dataset) throws IOException { return BQSupportFuncts.constructUrlFromPropertiesFile(getProperties(pathToProp), true, dataset); } + + @Test + public void missingJobCreationModeDefaultsToNull() throws Exception { + final String url = getUrl("/protectedaccount.properties", null); + Assertions.assertThat(url).doesNotContain("jobcreationmode"); + bq = new BQConnection(url, new Properties()); + final JobCreationMode mode = bq.getJobCreationMode(); + Assertions.assertThat(mode).isNull(); + } + + @Test + public void jobCreationModeTest() throws Exception { + final String url = getUrl("/protectedaccount.properties", null); + Assertions.assertThat(url).doesNotContain("jobcreationmode"); + final JobCreationMode[] modes = JobCreationMode.values(); + for (JobCreationMode mode : modes) { + final String fullURL = String.format("%s&jobcreationmode=%s", url, mode.name()); + try (BQConnection bq = new BQConnection(fullURL, new Properties())) { + final JobCreationMode parsedMode = bq.getJobCreationMode(); + Assertions.assertThat(parsedMode).isEqualTo(mode); + } + } + } } diff --git a/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java b/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java new file mode 100644 index 00000000..f86f2a6e --- /dev/null +++ b/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java @@ -0,0 +1,42 @@ +package net.starschema.clouddb.jdbc; + +import com.google.common.collect.ImmutableSet; +import java.util.Set; +import org.junit.Assume; + +/** Helpers for tests that require projects with stateless queries enabled */ +public final class StatelessQuery { + + private StatelessQuery() {} + + private static final Set ENABLED_PROJECTS = + ImmutableSet.of("disco-parsec-659", "looker-db-test"); + + /** + * Raise an {@link org.junit.AssumptionViolatedException} if the provided project isn't one that's + * known to have stateless queries enabled + * + * @param project the project to check - get it from {@link BQConnection.getCatalog() } + */ + public static void assumeStatelessQueriesEnabled(String project) { + Assume.assumeTrue(ENABLED_PROJECTS.contains(project)); + } + + /** + * A small query that should run statelessly (that is, without a job). + * + * @return the query + */ + public static String exampleQuery() { + return "SELECT 9876"; + } + + /** + * The values returned by {@link StatelessQuery} + * + * @return An array of strings representing the returned values + */ + public static String[][] exampleValues() { + return new String[][] {new String[] {"9876"}}; + } +} From 71bc3a9d978b80d3a3de2f8c54a19ff96c86b295 Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Wed, 13 Dec 2023 23:55:18 +0000 Subject: [PATCH 2/6] Add microbenchmark for stateless queries. This adds a Java Microbenchmark Harness (JMH) benchmark that compares a small query with and without job creation. Documentation on running the microbenchmark can be found in its class' Javadoc. A small query with job creation takes almost twice as long as one with optional job creation. --- pom.xml | 37 +++++++- .../clouddb/jdbc/ConnectionFromResources.java | 44 ++++++++++ .../jdbc/StatelessSmallQueryBenchmark.java | 84 +++++++++++++++++++ 3 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 src/test/java/net/starschema/clouddb/jdbc/ConnectionFromResources.java create mode 100644 src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java diff --git a/pom.xml b/pom.xml index 5e248e87..dc202f93 100644 --- a/pom.xml +++ b/pom.xml @@ -142,6 +142,18 @@ 1.19.0 test + + org.openjdk.jmh + jmh-core + 1.37 + test + + + org.openjdk.jmh + jmh-generator-annprocess + 1.37 + test + @@ -221,8 +233,29 @@ + + org.codehaus.mojo + exec-maven-plugin + 3.1.0 + + + benchmark-stateless-small + exec + + test + java + + -classpath + + org.openjdk.jmh.Main + + net.starschema.clouddb.jdbc.StatelessSmallQueryBenchmark + + + + + + - - diff --git a/src/test/java/net/starschema/clouddb/jdbc/ConnectionFromResources.java b/src/test/java/net/starschema/clouddb/jdbc/ConnectionFromResources.java new file mode 100644 index 00000000..ad6120d3 --- /dev/null +++ b/src/test/java/net/starschema/clouddb/jdbc/ConnectionFromResources.java @@ -0,0 +1,44 @@ +package net.starschema.clouddb.jdbc; + +import java.io.IOException; +import java.io.InputStream; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.util.Properties; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Utility class to enable BigQuery connections from properties resources. */ +public class ConnectionFromResources { + private static final Logger logger = LoggerFactory.getLogger(ConnectionFromResources.class); + + /** + * Connect to a BigQuery project as described in a properties file + * + * @param propertiesFilePath the path to the properties in file in src/test/resources + * @param extraUrl extra URL arguments to add; must start with & + * @return A {@link BQConnection} connected to the given database + * @throws IOException if the properties file cannot be read + * @throws SQLException if the configuration can't connect to BigQuery + */ + public static BQConnection connect(String propertiesFilePath, String extraUrl) + throws IOException, SQLException { + final Properties properties = new Properties(); + final ClassLoader loader = ConnectionFromResources.class.getClassLoader(); + try (InputStream stream = loader.getResourceAsStream(propertiesFilePath)) { + properties.load(stream); + } + final StringBuilder jdcbUrlBuilder = + new StringBuilder(BQSupportFuncts.constructUrlFromPropertiesFile(properties)); + if (extraUrl != null) { + jdcbUrlBuilder.append(extraUrl); + } + final String jdbcUrl = jdcbUrlBuilder.toString(); + + final Connection connection = DriverManager.getConnection(jdbcUrl, properties); + final BQConnection bqConnection = (BQConnection) connection; + logger.info("Created connection from {} to {}", propertiesFilePath, bqConnection.getURLPART()); + return bqConnection; + } +} diff --git a/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java b/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java new file mode 100644 index 00000000..d07c07f3 --- /dev/null +++ b/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java @@ -0,0 +1,84 @@ +package net.starschema.clouddb.jdbc; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import org.assertj.core.api.Assertions; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Threads; + +/** + * Performance microbenchmark for stateless queries. This uses the example query from {@link + * StatelessQuery}, same as the tests. + * + *

Run with mvn exec:exec@benchmark-stateless - note that mvn install must have been run at least + * once before. + */ +@Fork(1) +@Threads(10) +@BenchmarkMode(Mode.Throughput) +public class StatelessSmallQueryBenchmark { + private static final String CONNECTION_PROPERTIES = "installedaccount1.properties"; + + @State(Scope.Thread) + public static class RequiredJob { + private BQConnection connection; + + @Setup(Level.Trial) + public void connect() throws SQLException, IOException { + connection = ConnectionFromResources.connect(CONNECTION_PROPERTIES, null); + } + + @TearDown + public void disconnect() throws SQLException { + connection.close(); + } + } + + @State(Scope.Thread) + public static class OptionalJob { + private BQConnection connection; + + @Setup(Level.Trial) + public void connect() throws SQLException, IOException { + connection = + ConnectionFromResources.connect( + CONNECTION_PROPERTIES, "&jobcreationmode=JOB_CREATION_OPTIONAL"); + } + + @TearDown + public void disconnect() throws SQLException { + connection.close(); + } + } + + private String[][] benchmarkSmallQuery(final Connection connection) throws SQLException { + final Statement statement = connection.createStatement(); + final ResultSet results = statement.executeQuery(StatelessQuery.exampleQuery()); + final String[][] rows = BQSupportMethods.GetQueryResult(results); + Assertions.assertThat(rows).isEqualTo(StatelessQuery.exampleValues()); + return rows; + } + + @Benchmark + public String[][] benchmarkSmallQueryRequiredJob(final RequiredJob requiredJob) + throws SQLException { + return benchmarkSmallQuery(requiredJob.connection); + } + + @Benchmark + public String[][] benchmarkSmallQueryOptionalJob(final OptionalJob optionalJob) + throws SQLException { + return benchmarkSmallQuery(optionalJob.connection); + } +} From 40e7b132af617a9b1adc56aa26d8c5941eeb3403 Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Thu, 14 Dec 2023 00:12:18 +0000 Subject: [PATCH 3/6] Address review comment about docs; fix Javadoc. Method references need # not . --- src/main/java/net/starschema/clouddb/jdbc/BQConnection.java | 4 +++- src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java index 363b7014..db624fda 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java @@ -74,7 +74,9 @@ public class BQConnection implements Connection { /** * Enum that describes whether to create a job in projects that support stateless queries. Copied - * from BigQueryImpl + * from google-cloud-bigquery + * 2.34.0 */ public static enum JobCreationMode { /** If unspecified JOB_CREATION_REQUIRED is the default. */ diff --git a/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java b/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java index f86f2a6e..f36dfb30 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java +++ b/src/test/java/net/starschema/clouddb/jdbc/StatelessQuery.java @@ -16,7 +16,7 @@ private StatelessQuery() {} * Raise an {@link org.junit.AssumptionViolatedException} if the provided project isn't one that's * known to have stateless queries enabled * - * @param project the project to check - get it from {@link BQConnection.getCatalog() } + * @param project the project to check - get it from {@link BQConnection#getCatalog() } */ public static void assumeStatelessQueriesEnabled(String project) { Assume.assumeTrue(ENABLED_PROJECTS.contains(project)); @@ -32,7 +32,7 @@ public static String exampleQuery() { } /** - * The values returned by {@link StatelessQuery} + * The values returned by {@link StatelessQuery#exampleQuery()} * * @return An array of strings representing the returned values */ From 3636f183fd65aaac202f0195ad8f8876677ef34a Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Thu, 14 Dec 2023 00:37:23 +0000 Subject: [PATCH 4/6] Include benchmark in Javadoc --- .../clouddb/jdbc/StatelessSmallQueryBenchmark.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java b/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java index d07c07f3..4ba0ac0a 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java +++ b/src/test/java/net/starschema/clouddb/jdbc/StatelessSmallQueryBenchmark.java @@ -23,6 +23,12 @@ * *

Run with mvn exec:exec@benchmark-stateless - note that mvn install must have been run at least * once before. + * + *

Representative summary as of 2023-12-14: + * Benchmark Mode Cnt Score Error Units + * StatelessSmallQueryBenchmark.benchmarkSmallQueryOptionalJob thrpt 5 67.994 ± 10.326 ops/s + * StatelessSmallQueryBenchmark.benchmarkSmallQueryRequiredJob thrpt 5 37.171 ± 3.041 ops/s + * */ @Fork(1) @Threads(10) From b861a6c39828603ef402e45b6c9ea658e6a2bd58 Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Fri, 15 Dec 2023 00:12:53 +0000 Subject: [PATCH 5/6] Add query-ID-less constructor for backwards compat. --- .../clouddb/jdbc/BQForwardOnlyResultSet.java | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java b/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java index 5b0bbf17..fe697246 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java @@ -125,11 +125,41 @@ public class BQForwardOnlyResultSet implements java.sql.ResultSet { */ private int Cursor = -1; + /** + * Constructor without query ID for backwards compatibility. + * + * @param bigquery Bigquery driver instance for which this is a result + * @param projectId the project from which these results were queried + * @param completedJob the query's job, if any + * @param bqStatementRoot the statement for which this is a result + * @throws SQLException thrown if the results can't be retrieved + */ public BQForwardOnlyResultSet( Bigquery bigquery, String projectId, @Nullable Job completedJob, - String queryId, + BQStatementRoot bqStatementRoot) + throws SQLException { + this( + bigquery, + projectId, + completedJob, + null, + bqStatementRoot, + null, + false, + null, + 0L, + false, + null, + null); + } + + public BQForwardOnlyResultSet( + Bigquery bigquery, + String projectId, + @Nullable Job completedJob, + @Nullable String queryId, BQStatementRoot bqStatementRoot) throws SQLException { this( From 091b88f98b696963fa36793e3a1f2fc232d5d441 Mon Sep 17 00:00:00 2001 From: Mark Williams Date: Mon, 18 Dec 2023 19:04:22 +0000 Subject: [PATCH 6/6] Remove job creation mode parsing method. Addresses review comment --- .../starschema/clouddb/jdbc/BQConnection.java | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java index db624fda..783609c0 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java +++ b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java @@ -238,8 +238,17 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor this.useQueryCache = parseBooleanQueryParam(caseInsensitiveProps.getProperty("querycache"), true); - this.jobCreationMode = - parseJobCreationMode(caseInsensitiveProps.getProperty("jobcreationmode")); + final String jobCreationModeString = caseInsensitiveProps.getProperty("jobcreationmode"); + if (jobCreationModeString == null) { + jobCreationMode = null; + } else { + try { + jobCreationMode = JobCreationMode.valueOf(jobCreationModeString); + } catch (IllegalArgumentException e) { + throw new BQSQLException( + "could not parse " + jobCreationModeString + " as job creation mode", e); + } + } // Create Connection to BigQuery if (serviceAccount) { @@ -353,21 +362,6 @@ private static List parseArrayQueryParam(@Nullable String string, Charac : Arrays.asList(string.split(delimiter + "\\s*")); } - /** - * Return a {@link JobCreationMode} or raise an exception if the string does not match a variant. - */ - private static JobCreationMode parseJobCreationMode(@Nullable String string) - throws BQSQLException { - if (string == null) { - return null; - } - try { - return JobCreationMode.valueOf(string); - } catch (IllegalArgumentException e) { - throw new BQSQLException("could not parse " + string + " as job creation mode", e); - } - } - /** * *