Skip to content

Commit 1c70654

Browse files
olavloiteolavloite
authored
feat: samples for partitioned queries on Spanner with PostgreSQL (GoogleCloudPlatform#8937)
* feat: samples for partitioned queries on Spanner with PostgreSQL Adds samples for running partitioned queries with the Spanner JDBC driver on PostgreSQL-dialect databases. * chore: fix license header
1 parent 09a0753 commit 1c70654

6 files changed

Lines changed: 346 additions & 1 deletion

File tree

‎spanner/jdbc/src/main/java/com/example/spanner/jdbc/PgAutoPartitionModeExample.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
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.example.spanner.jdbc;
18+
19+
import java.sql.Connection;
20+
import java.sql.DriverManager;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import java.sql.Statement;
24+
25+
public class PgAutoPartitionModeExample {
26+
27+
public static void main(String[] args) throws SQLException {
28+
autoPartitionMode();
29+
}
30+
31+
static void autoPartitionMode() throws SQLException {
32+
// TODO(developer): Replace these variables before running the sample.
33+
String projectId = "my-project";
34+
String instanceId = "my-instance";
35+
String databaseId = "my-database";
36+
autoPartitionMode(projectId, instanceId, databaseId);
37+
}
38+
39+
// This example shows how to use 'spanner.auto_partition_mode=true' to execute partitioned queries
40+
// with the JDBC driver for a PostgreSQL-dialect database.
41+
static void autoPartitionMode(String projectId, String instanceId, String databaseId)
42+
throws SQLException {
43+
String connectionUrl = String.format("jdbc:cloudspanner:/projects/%s/instances/%s/databases/%s",
44+
projectId, instanceId, databaseId);
45+
try (Connection connection = DriverManager.getConnection(
46+
connectionUrl); Statement statement = connection.createStatement()) {
47+
// A connection can also be set to 'spanner.auto_partition_mode', which will instruct it to
48+
// execute all queries as a partitioned query. This is essentially the same as automatically
49+
// prefixing all queries with 'RUN PARTITIONED QUERY ...'.
50+
statement.execute("set spanner.auto_partition_mode=true");
51+
// This will execute at most max_partitioned_parallelism partitions in parallel.
52+
statement.execute("set spanner.max_partitioned_parallelism=8");
53+
try (ResultSet resultSet = statement.executeQuery(
54+
"SELECT SingerId, FirstName, LastName FROM singers")) {
55+
while (resultSet.next()) {
56+
System.out.printf("%s %s %s%n", resultSet.getString(1), resultSet.getString(2),
57+
resultSet.getString(3));
58+
}
59+
}
60+
}
61+
}
62+
}

‎spanner/jdbc/src/main/java/com/example/spanner/jdbc/PgDataBoostExample.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
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.example.spanner.jdbc;
18+
19+
import java.sql.Connection;
20+
import java.sql.DriverManager;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import java.sql.Statement;
24+
25+
public class PgDataBoostExample {
26+
27+
public static void main(String[] args) throws SQLException {
28+
dataBoost();
29+
}
30+
31+
static void dataBoost() throws SQLException {
32+
// TODO(developer): Replace these variables before running the sample.
33+
String projectId = "my-project";
34+
String instanceId = "my-instance";
35+
String databaseId = "my-database";
36+
dataBoost(projectId, instanceId, databaseId);
37+
}
38+
39+
// This example shows how to execute queries with data boost on PostgreSQL-dialect databases using
40+
// the Google Cloud Spanner JDBC driver.
41+
static void dataBoost(String projectId, String instanceId, String databaseId)
42+
throws SQLException {
43+
String connectionUrl = String.format("jdbc:cloudspanner:/projects/%s/instances/%s/databases/%s",
44+
projectId, instanceId, databaseId);
45+
try (Connection connection = DriverManager.getConnection(
46+
connectionUrl); Statement statement = connection.createStatement()) {
47+
48+
// A connection can also be set to 'spanner.auto_partition_mode', which will instruct it to
49+
// execute all queries as a partitioned query. This is essentially the same as automatically
50+
// prefixing all queries with 'RUN PARTITIONED QUERY ...'.
51+
statement.execute("set spanner.auto_partition_mode=true");
52+
53+
// This will execute at most max_partitioned_parallelism partitions in parallel.
54+
statement.execute("set spanner.max_partitioned_parallelism=8");
55+
56+
// Setting 'spanner.data_boost_enabled' to true will instruct the JDBC connection to execute
57+
// all partitioned queries using Data Boost. This setting applies to all the above methods
58+
// that can be used to run a partitioned query:
59+
// 1. RUN PARTITION '...'
60+
// 2. RUN PARTITIONED QUERY ...
61+
// 3. SET AUTO_PARTITION_MODE=TRUE; SELECT ...
62+
statement.execute("set spanner.data_boost_enabled=true");
63+
64+
// This query will be executed as a partitioned query using data boost.
65+
try (ResultSet resultSet = statement.executeQuery(
66+
"SELECT SingerId, FirstName, LastName FROM singers")) {
67+
while (resultSet.next()) {
68+
System.out.printf("%s %s %s%n", resultSet.getString(1), resultSet.getString(2),
69+
resultSet.getString(3));
70+
}
71+
}
72+
}
73+
}
74+
}

‎spanner/jdbc/src/main/java/com/example/spanner/jdbc/PgPartitionQueryExample.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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.example.spanner.jdbc;
18+
19+
import java.sql.Connection;
20+
import java.sql.DriverManager;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import java.sql.Statement;
24+
import java.util.ArrayList;
25+
import java.util.List;
26+
27+
public class PgPartitionQueryExample {
28+
29+
public static void main(String[] args) throws SQLException {
30+
partitionQuery();
31+
}
32+
33+
static void partitionQuery() throws SQLException {
34+
// TODO(developer): Replace these variables before running the sample.
35+
String projectId = "my-project";
36+
String instanceId = "my-instance";
37+
String databaseId = "my-database";
38+
partitionQuery(projectId, instanceId, databaseId);
39+
}
40+
41+
// This example shows how to partition a query and execute each returned partition on a
42+
// PostgreSQL-dialect database with the JDBC driver.
43+
static void partitionQuery(String projectId, String instanceId, String databaseId)
44+
throws SQLException {
45+
String connectionUrl = String.format("jdbc:cloudspanner:/projects/%s/instances/%s/databases/%s",
46+
projectId, instanceId, databaseId);
47+
try (Connection connection = DriverManager.getConnection(
48+
connectionUrl); Statement statement = connection.createStatement()) {
49+
50+
// Partition a query and then execute each partition sequentially.
51+
List<String> partitions = new ArrayList<>();
52+
// This will partition the query and return a result set with the partition IDs encoded as a
53+
// string. Each of these partition IDs can be executed with "RUN PARTITION '<partition-id>'".
54+
System.out.println("Partitioning query 'SELECT SingerId, FirstName, LastName from singers'");
55+
try (ResultSet partitionsResultSet = statement.executeQuery(
56+
"PARTITION SELECT SingerId, FirstName, LastName from Singers")) {
57+
while (partitionsResultSet.next()) {
58+
partitions.add(partitionsResultSet.getString(1));
59+
}
60+
}
61+
System.out.printf("Partition command returned %d partitions\n", partitions.size());
62+
63+
// This executes the partitions serially on the same connection, but each partition can also
64+
// be executed on a different JDBC connection (even on a different host).
65+
for (String partitionId : partitions) {
66+
try (ResultSet resultSet = statement.executeQuery(
67+
String.format("RUN PARTITION '%s'", partitionId))) {
68+
while (resultSet.next()) {
69+
System.out.printf("%s %s %s%n", resultSet.getString(1), resultSet.getString(2),
70+
resultSet.getString(3));
71+
}
72+
}
73+
}
74+
System.out.println("Finished executing all partitions");
75+
}
76+
}
77+
}

‎spanner/jdbc/src/main/java/com/example/spanner/jdbc/PgRunPartitionedQueryExample.java‎

Lines changed: 61 additions & 0 deletions
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.example.spanner.jdbc;
18+
19+
import java.sql.Connection;
20+
import java.sql.DriverManager;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import java.sql.Statement;
24+
25+
public class PgRunPartitionedQueryExample {
26+
27+
public static void main(String[] args) throws SQLException {
28+
runPartitionedQuery();
29+
}
30+
31+
static void runPartitionedQuery() throws SQLException {
32+
// TODO(developer): Replace these variables before running the sample.
33+
String projectId = "my-project";
34+
String instanceId = "my-instance";
35+
String databaseId = "my-database";
36+
runPartitionedQuery(projectId, instanceId, databaseId);
37+
}
38+
39+
// This example shows how to run a query directly as a partitioned query on a JDBC connection.
40+
// The query will be partitioned and each partition will be executed using the same JDBC
41+
// connection. You can set the maximum parallelism that should be used to execute the query with
42+
// the SQL statement 'SET spanner.max_partitioned_parallelism=<parallelism>'.
43+
static void runPartitionedQuery(String projectId, String instanceId, String databaseId)
44+
throws SQLException {
45+
String connectionUrl = String.format("jdbc:cloudspanner:/projects/%s/instances/%s/databases/%s",
46+
projectId, instanceId, databaseId);
47+
try (Connection connection = DriverManager.getConnection(
48+
connectionUrl); Statement statement = connection.createStatement()) {
49+
// Run a query directly as a partitioned query.
50+
// This will execute at most max_partitioned_parallelism partitions in parallel.
51+
statement.execute("set spanner.max_partitioned_parallelism=8");
52+
try (ResultSet resultSet = statement.executeQuery(
53+
"RUN PARTITIONED QUERY SELECT SingerId, FirstName, LastName FROM singers")) {
54+
while (resultSet.next()) {
55+
System.out.printf("%s %s %s%n", resultSet.getString(1), resultSet.getString(2),
56+
resultSet.getString(3));
57+
}
58+
}
59+
}
60+
}
61+
}

‎spanner/jdbc/src/test/java/com/example/spanner/jdbc/BaseJdbcPgExamplesIT.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,11 @@ static class Singer {
112112
this.lastName = lastName;
113113
this.revenues = revenues;
114114
}
115+
116+
@Override
117+
public String toString() {
118+
return String.format("%d %s %s", singerId, firstName, lastName);
119+
}
115120
}
116121

117122
static final List<Singer> TEST_SINGERS =
@@ -140,7 +145,7 @@ public void insertTestData() throws SQLException {
140145
connection
141146
.createStatement()
142147
.execute(
143-
"CREATE TABLE Singers (\n"
148+
"CREATE TABLE IF NOT EXISTS Singers (\n"
144149
+ " SingerId BIGINT NOT NULL PRIMARY KEY,\n"
145150
+ " FirstName VARCHAR(1024),\n"
146151
+ " LastName VARCHAR(1024),\n"

‎spanner/jdbc/src/test/java/com/example/spanner/jdbc/PgPartitionedQueryIT.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
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.example.spanner.jdbc;
18+
19+
import static org.junit.Assert.assertTrue;
20+
21+
import com.google.cloud.ServiceOptions;
22+
import org.junit.Test;
23+
24+
public class PgPartitionedQueryIT extends BaseJdbcPgExamplesIT {
25+
26+
@Override
27+
protected boolean createTestTable() {
28+
return true;
29+
}
30+
31+
@Test
32+
public void testPartitionQuery() {
33+
String out = runExample(() -> PgPartitionQueryExample.partitionQuery(
34+
ServiceOptions.getDefaultProjectId(), instanceId, databaseId));
35+
assertOutputContainsAllSingers(out);
36+
}
37+
38+
@Test
39+
public void testAutoPartitionMode() {
40+
String out = runExample(() -> PgAutoPartitionModeExample.autoPartitionMode(
41+
ServiceOptions.getDefaultProjectId(), instanceId, databaseId));
42+
assertOutputContainsAllSingers(out);
43+
}
44+
45+
@Test
46+
public void testDataBoost() {
47+
String out = runExample(() -> PgDataBoostExample.dataBoost(
48+
ServiceOptions.getDefaultProjectId(), instanceId, databaseId));
49+
assertOutputContainsAllSingers(out);
50+
}
51+
52+
@Test
53+
public void testRunPartitionedQuery() {
54+
String out = runExample(() -> PgRunPartitionedQueryExample.runPartitionedQuery(
55+
ServiceOptions.getDefaultProjectId(), instanceId, databaseId));
56+
assertOutputContainsAllSingers(out);
57+
}
58+
59+
void assertOutputContainsAllSingers(String out) {
60+
for (Singer singer : TEST_SINGERS) {
61+
assertTrue(out + " should contain " + singer.toString(),
62+
out.contains(singer.toString()));
63+
}
64+
}
65+
66+
}

0 commit comments

Comments
 (0)