Skip to content

Commit c50da41

Browse files
olavloiteolavloite
authored
feat: support partitioned queries (#1300)
Adds support for partitioned queries to the JDBC driver. Partitioned queries can be executed either using SQL statements or specific Cloud Spanner methods that are exposed by the specific Cloud Spanner JDBC interfaces.
1 parent c3be4d4 commit c50da41

12 files changed

+964
-7
lines changed

‎clirr-ignored-differences.xml

+42
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,46 @@
1616
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
1717
<method>void setSavepointSupport(com.google.cloud.spanner.connection.SavepointSupport)</method>
1818
</difference>
19+
20+
<difference>
21+
<differenceType>7012</differenceType>
22+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
23+
<method>int getMaxPartitionedParallelism()</method>
24+
</difference>
25+
<difference>
26+
<differenceType>7012</differenceType>
27+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
28+
<method>int getMaxPartitions()</method>
29+
</difference>
30+
<difference>
31+
<differenceType>7012</differenceType>
32+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
33+
<method>boolean isAutoPartitionMode()</method>
34+
</difference>
35+
<difference>
36+
<differenceType>7012</differenceType>
37+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
38+
<method>boolean isDataBoostEnabled()</method>
39+
</difference>
40+
<difference>
41+
<differenceType>7012</differenceType>
42+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
43+
<method>void setAutoPartitionMode(boolean)</method>
44+
</difference>
45+
<difference>
46+
<differenceType>7012</differenceType>
47+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
48+
<method>void setDataBoostEnabled(boolean)</method>
49+
</difference>
50+
<difference>
51+
<differenceType>7012</differenceType>
52+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
53+
<method>void setMaxPartitionedParallelism(int)</method>
54+
</difference>
55+
<difference>
56+
<differenceType>7012</differenceType>
57+
<className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
58+
<method>void setMaxPartitions(int)</method>
59+
</difference>
60+
1961
</differences>

‎src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java

+19
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,25 @@ private StatementTimeout(long timeout, TimeUnit unit) {
120120
}
121121
}
122122

123+
/** Functional interface that throws {@link SQLException}. */
124+
interface JdbcFunction<T, R> {
125+
R apply(T t) throws SQLException;
126+
}
127+
128+
/** Runs the given function with the timeout that has been set on this statement. */
129+
protected <T> T runWithStatementTimeout(JdbcFunction<Connection, T> function)
130+
throws SQLException {
131+
checkClosed();
132+
StatementTimeout originalTimeout = setTemporaryStatementTimeout();
133+
try {
134+
return function.apply(getConnection().getSpannerConnection());
135+
} catch (SpannerException spannerException) {
136+
throw JdbcSqlExceptionFactory.of(spannerException);
137+
} finally {
138+
resetStatementTimeout(originalTimeout);
139+
}
140+
}
141+
123142
/**
124143
* Sets the statement timeout of the Spanner {@link Connection} to the query timeout of this JDBC
125144
* {@link Statement} and returns the original timeout of the Spanner {@link Connection} so it can

‎src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java

+70-1
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,14 @@
2222
import com.google.cloud.spanner.CommitStats;
2323
import com.google.cloud.spanner.Dialect;
2424
import com.google.cloud.spanner.Mutation;
25-
import com.google.cloud.spanner.ResultSet;
25+
import com.google.cloud.spanner.Options.QueryOption;
26+
import com.google.cloud.spanner.PartitionOptions;
2627
import com.google.cloud.spanner.TimestampBound;
2728
import com.google.cloud.spanner.connection.AutocommitDmlMode;
2829
import com.google.cloud.spanner.connection.SavepointSupport;
2930
import com.google.cloud.spanner.connection.TransactionMode;
3031
import java.sql.Connection;
32+
import java.sql.ResultSet;
3133
import java.sql.SQLException;
3234
import java.sql.Timestamp;
3335
import java.util.Iterator;
@@ -334,6 +336,73 @@ default Dialect getDialect() {
334336
return Dialect.GOOGLE_STANDARD_SQL;
335337
}
336338

339+
/**
340+
* Enable data boost for partitioned queries. See also {@link
341+
* CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
342+
* CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
343+
*/
344+
default void setDataBoostEnabled(boolean dataBoostEnabled) throws SQLException {
345+
throw new UnsupportedOperationException();
346+
}
347+
348+
/**
349+
* Returns whether data boost is enabled for partitioned queries. See also {@link
350+
* CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
351+
* CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
352+
*/
353+
default boolean isDataBoostEnabled() throws SQLException {
354+
throw new UnsupportedOperationException();
355+
}
356+
357+
/**
358+
* Sets whether this connection should always use partitioned queries when a query is executed on
359+
* this connection. Setting this flag to <code>true</code> and then executing a query that cannot
360+
* be partitioned, or executing a query in a read/write transaction, will cause an error. Use this
361+
* flag in combination with {@link #setDataBoostEnabled(boolean)} to force all queries on this
362+
* connection to use data boost.
363+
*/
364+
default void setAutoPartitionMode(boolean alwaysUsePartitionedQueries) throws SQLException {
365+
throw new UnsupportedOperationException();
366+
}
367+
368+
/** Returns whether this connection will execute all queries as partitioned queries. */
369+
default boolean isAutoPartitionMode() throws SQLException {
370+
throw new UnsupportedOperationException();
371+
}
372+
373+
/**
374+
* Sets the maximum number of partitions that should be included as a hint to Cloud Spanner when
375+
* partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
376+
* choose to ignore the hint.
377+
*/
378+
default void setMaxPartitions(int maxPartitions) throws SQLException {
379+
throw new UnsupportedOperationException();
380+
}
381+
382+
/**
383+
* Gets the maximum number of partitions that should be included as a hint to Cloud Spanner when
384+
* partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
385+
* choose to ignore the hint.
386+
*/
387+
default int getMaxPartitions() throws SQLException {
388+
throw new UnsupportedOperationException();
389+
}
390+
391+
/**
392+
* Sets the maximum degree of parallelism that is used when executing a partitioned query. A
393+
* partitioned query will use up to <code>maxThreads</code> to execute and retrieve the results
394+
* from Cloud Spanner. Set this value to <code>0</code>> to use the number of available processors
395+
* as returned by {@link Runtime#availableProcessors()}.
396+
*/
397+
default void setMaxPartitionedParallelism(int maxThreads) throws SQLException {
398+
throw new UnsupportedOperationException();
399+
}
400+
401+
/** Returns the maximum degree of parallelism that is used for partitioned queries. */
402+
default int getMaxPartitionedParallelism() throws SQLException {
403+
throw new UnsupportedOperationException();
404+
}
405+
337406
/**
338407
* @see
339408
* com.google.cloud.spanner.connection.Connection#addTransactionRetryListener(com.google.cloud.spanner.connection.TransactionRetryListener)

‎src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.spanner.jdbc;
17+
18+
import com.google.cloud.spanner.Options.QueryOption;
19+
import com.google.cloud.spanner.PartitionOptions;
20+
import java.sql.ResultSet;
21+
22+
/**
23+
* Result set that is returned for partitioned queries, e.g. for 'run partitioned query select ...'
24+
* or for {@link CloudSpannerJdbcPreparedStatement#runPartitionedQuery(PartitionOptions,
25+
* QueryOption...)}.
26+
*/
27+
public interface CloudSpannerJdbcPartitionedQueryResultSet extends ResultSet {
28+
/** Returns the number of partitions that this result set contains. */
29+
int getNumPartitions();
30+
31+
/** Returns the degree of parallelism that this result set uses. */
32+
int getParallelism();
33+
}

‎src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java

+61
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.jdbc;
18+
19+
import com.google.cloud.spanner.Options.QueryOption;
20+
import com.google.cloud.spanner.PartitionOptions;
21+
import java.sql.PreparedStatement;
22+
import java.sql.ResultSet;
23+
import java.sql.SQLException;
24+
25+
/**
26+
* This interface is implemented by {@link PreparedStatement}s that are created on Cloud Spanner
27+
* JDBC connections.
28+
*/
29+
public interface CloudSpannerJdbcPreparedStatement extends PreparedStatement {
30+
31+
/**
32+
* Partitions this query, so it can be executed in parallel. This method returns a {@link
33+
* ResultSet} with a string-representation of the partitions that were created. These strings can
34+
* be used to execute a partition either on this connection or an any other connection (on this
35+
* host or an any other host) by calling the method {@link #runPartition()}. This method will
36+
* automatically enable data boost for the query if {@link
37+
* CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
38+
*/
39+
ResultSet partitionQuery(PartitionOptions partitionOptions, QueryOption... options)
40+
throws SQLException;
41+
42+
/**
43+
* Executes the given partition of a query. The partition that should be executed must be set as a
44+
* string parameter on this {@link PreparedStatement} using {@link #setString(int, String)}. The
45+
* value should be a string that was returned by {@link #partitionQuery(PartitionOptions,
46+
* QueryOption...)}.
47+
*/
48+
ResultSet runPartition() throws SQLException;
49+
50+
/**
51+
* Executes the given query as a partitioned query. The query will first be partitioned using the
52+
* {@link #partitionQuery(PartitionOptions, QueryOption...)} method. Each of the partitions will
53+
* then be executed in the background, and the results will be merged into a single result set.
54+
*
55+
* <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
56+
* threads to execute the partitioned query. Set this variable to a higher/lower value to
57+
* increase/decrease the degree of parallelism used for execution.
58+
*/
59+
CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
60+
PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
61+
}

‎src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.jdbc;
18+
19+
import com.google.cloud.spanner.Options.QueryOption;
20+
import com.google.cloud.spanner.PartitionOptions;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import java.sql.Statement;
24+
25+
/**
26+
* This interface is implemented by {@link Statement}s that are created on Cloud Spanner JDBC
27+
* connections.
28+
*/
29+
public interface CloudSpannerJdbcStatement extends Statement {
30+
31+
/**
32+
* Partitions the given query, so it can be executed in parallel. This method returns a {@link
33+
* ResultSet} with a string-representation of the partitions that were created. These strings can
34+
* be used to execute a partition either on this connection or an any other connection (on this
35+
* host or an any other host) by calling the method {@link #runPartition(String)}. This method
36+
* will automatically enable data boost for the query if {@link
37+
* CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
38+
*/
39+
ResultSet partitionQuery(String query, PartitionOptions partitionOptions, QueryOption... options)
40+
throws SQLException;
41+
42+
/**
43+
* Executes the given partition of a query. The encodedPartitionId should be a string that was
44+
* returned by {@link #partitionQuery(String, PartitionOptions, QueryOption...)}.
45+
*/
46+
ResultSet runPartition(String encodedPartitionId) throws SQLException;
47+
48+
/**
49+
* Executes the given query as a partitioned query. The query will first be partitioned using the
50+
* {@link #partitionQuery(String, PartitionOptions, QueryOption...)} method. Each of the
51+
* partitions will then be executed in the background, and the results will be merged into a
52+
* single result set.
53+
*
54+
* <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
55+
* threads to execute the partitioned query. Set this variable to a higher/lower value to
56+
* increase/decrease the degree of parallelism used for execution.
57+
*/
58+
CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
59+
String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
60+
}

0 commit comments

Comments
 (0)