Skip to content

Commit 1654c2a

Browse files
authored
test: add tag tests (#1445)
1 parent 8be2f4c commit 1654c2a

File tree

2 files changed

+210
-2
lines changed

2 files changed

+210
-2
lines changed

src/test/java/com/google/cloud/spanner/jdbc/SavepointMockServerTest.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import com.google.cloud.spanner.connection.AbstractMockServerTest;
2929
import com.google.cloud.spanner.connection.RandomResultSetGenerator;
3030
import com.google.cloud.spanner.connection.SavepointSupport;
31-
import com.google.cloud.spanner.connection.SpannerPool;
3231
import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcAbortedDueToConcurrentModificationException;
3332
import com.google.common.base.Strings;
3433
import com.google.protobuf.AbstractMessage;
@@ -83,7 +82,7 @@ public void clearRequests() {
8382
private String createUrl() {
8483
return String.format(
8584
"jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true;autoCommit=false",
86-
getPort(), "proj", "inst", "db");
85+
getPort(), "proj", "inst", "db" + (dialect == Dialect.POSTGRESQL ? "pg" : ""));
8786
}
8887

8988
private Connection createConnection() throws SQLException {
Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
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 static org.junit.Assert.assertArrayEquals;
20+
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertTrue;
22+
23+
import com.google.cloud.spanner.Dialect;
24+
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
25+
import com.google.cloud.spanner.connection.AbstractMockServerTest;
26+
import com.google.spanner.v1.ExecuteBatchDmlRequest;
27+
import com.google.spanner.v1.ExecuteSqlRequest;
28+
import java.sql.Connection;
29+
import java.sql.DriverManager;
30+
import java.sql.ResultSet;
31+
import java.sql.SQLException;
32+
import java.sql.Statement;
33+
import org.junit.After;
34+
import org.junit.Before;
35+
import org.junit.Test;
36+
import org.junit.runner.RunWith;
37+
import org.junit.runners.Parameterized;
38+
import org.junit.runners.Parameterized.Parameter;
39+
import org.junit.runners.Parameterized.Parameters;
40+
41+
@RunWith(Parameterized.class)
42+
public class TagMockServerTest extends AbstractMockServerTest {
43+
@Parameter public Dialect dialect;
44+
45+
@Parameters(name = "dialect = {0}")
46+
public static Object[] data() {
47+
return Dialect.values();
48+
}
49+
50+
@Before
51+
public void setupDialect() {
52+
mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
53+
}
54+
55+
@After
56+
public void clearRequests() {
57+
mockSpanner.clearRequests();
58+
}
59+
60+
private String getVariablePrefix() {
61+
return dialect == Dialect.POSTGRESQL ? "spanner." : "";
62+
}
63+
64+
private String createUrl() {
65+
return String.format(
66+
"jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true;autoCommit=false",
67+
getPort(), "proj", "inst", "db" + (dialect == Dialect.POSTGRESQL ? "pg" : ""));
68+
}
69+
70+
private Connection createConnection() throws SQLException {
71+
return DriverManager.getConnection(createUrl());
72+
}
73+
74+
@Test
75+
public void testStatementTagForQuery() throws SQLException {
76+
try (Connection connection = createConnection()) {
77+
connection
78+
.createStatement()
79+
.execute(String.format("set %sstatement_tag='my-tag'", getVariablePrefix()));
80+
try (ResultSet resultSet =
81+
connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
82+
assertTrue(resultSet.next());
83+
}
84+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
85+
ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
86+
assertEquals("my-tag", request.getRequestOptions().getRequestTag());
87+
88+
// Verify that the tag is cleared after having been used.
89+
mockSpanner.clearRequests();
90+
try (ResultSet resultSet =
91+
connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
92+
assertTrue(resultSet.next());
93+
}
94+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
95+
request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
96+
assertEquals("", request.getRequestOptions().getRequestTag());
97+
}
98+
}
99+
100+
@Test
101+
public void testTransactionTagForQuery() throws SQLException {
102+
try (Connection connection = createConnection()) {
103+
connection
104+
.createStatement()
105+
.execute(String.format("set %stransaction_tag='my-tag'", getVariablePrefix()));
106+
try (ResultSet resultSet =
107+
connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
108+
assertTrue(resultSet.next());
109+
}
110+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
111+
ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
112+
assertEquals("my-tag", request.getRequestOptions().getTransactionTag());
113+
114+
// Verify that the tag is used for the entire transaction.
115+
mockSpanner.clearRequests();
116+
try (ResultSet resultSet =
117+
connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
118+
assertTrue(resultSet.next());
119+
}
120+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
121+
request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
122+
assertEquals("my-tag", request.getRequestOptions().getTransactionTag());
123+
124+
// Verify that committing the transaction will remove the transaction tag.
125+
connection.commit();
126+
127+
mockSpanner.clearRequests();
128+
try (ResultSet resultSet =
129+
connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
130+
assertTrue(resultSet.next());
131+
}
132+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
133+
request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
134+
assertEquals("", request.getRequestOptions().getTransactionTag());
135+
}
136+
}
137+
138+
@Test
139+
public void testStatementTagForBatchDml() throws SQLException {
140+
try (Connection connection = createConnection()) {
141+
connection
142+
.createStatement()
143+
.execute(String.format("set %sstatement_tag='my-tag'", getVariablePrefix()));
144+
145+
try (Statement statement = connection.createStatement()) {
146+
statement.addBatch(INSERT_STATEMENT.getSql());
147+
statement.addBatch(INSERT_STATEMENT.getSql());
148+
assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
149+
}
150+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
151+
ExecuteBatchDmlRequest request =
152+
mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
153+
assertEquals("my-tag", request.getRequestOptions().getRequestTag());
154+
155+
// Verify that the tag is cleared after having been used.
156+
mockSpanner.clearRequests();
157+
try (Statement statement = connection.createStatement()) {
158+
statement.addBatch(INSERT_STATEMENT.getSql());
159+
statement.addBatch(INSERT_STATEMENT.getSql());
160+
assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
161+
}
162+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
163+
request = mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
164+
assertEquals("", request.getRequestOptions().getRequestTag());
165+
}
166+
}
167+
168+
@Test
169+
public void testTransactionTagForBatchDml() throws SQLException {
170+
try (Connection connection = createConnection()) {
171+
connection
172+
.createStatement()
173+
.execute(String.format("set %stransaction_tag='my-tag'", getVariablePrefix()));
174+
175+
try (Statement statement = connection.createStatement()) {
176+
statement.addBatch(INSERT_STATEMENT.getSql());
177+
statement.addBatch(INSERT_STATEMENT.getSql());
178+
assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
179+
}
180+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
181+
ExecuteBatchDmlRequest request =
182+
mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
183+
assertEquals("my-tag", request.getRequestOptions().getTransactionTag());
184+
185+
// Verify that the tag is used for the entire transaction.
186+
mockSpanner.clearRequests();
187+
try (Statement statement = connection.createStatement()) {
188+
statement.addBatch(INSERT_STATEMENT.getSql());
189+
statement.addBatch(INSERT_STATEMENT.getSql());
190+
assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
191+
}
192+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
193+
request = mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
194+
assertEquals("my-tag", request.getRequestOptions().getTransactionTag());
195+
196+
// Verify that committing the transaction will remove the transaction tag.
197+
connection.commit();
198+
mockSpanner.clearRequests();
199+
try (Statement statement = connection.createStatement()) {
200+
statement.addBatch(INSERT_STATEMENT.getSql());
201+
statement.addBatch(INSERT_STATEMENT.getSql());
202+
assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
203+
}
204+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
205+
request = mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
206+
assertEquals("", request.getRequestOptions().getTransactionTag());
207+
}
208+
}
209+
}

0 commit comments

Comments
 (0)