Skip to content

Commit 331942f

Browse files
olavloiteolavloite
authored
feat: support 'set local' for retry_aborts_internally (#3532)
Adds support for `set local retry_aborts_internally=true|false` in the Connection API. This change also adds the parsing infrastructure that is needed to support `set local` for all connection variables. Support for this will be added to other connection variables in follow-up pull requests.
1 parent 46464fc commit 331942f

File tree

11 files changed

+3740
-246
lines changed

11 files changed

+3740
-246
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementSetExecutor.java

Lines changed: 50 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,18 @@
1616

1717
package com.google.cloud.spanner.connection;
1818

19+
import com.google.cloud.Tuple;
1920
import com.google.cloud.spanner.ErrorCode;
2021
import com.google.cloud.spanner.SpannerExceptionFactory;
2122
import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
2223
import com.google.cloud.spanner.connection.ClientSideStatementImpl.CompileException;
2324
import com.google.common.base.Preconditions;
25+
import com.google.common.cache.Cache;
26+
import com.google.common.cache.CacheBuilder;
27+
import com.google.common.util.concurrent.UncheckedExecutionException;
2428
import java.lang.reflect.Constructor;
2529
import java.lang.reflect.Method;
30+
import java.util.concurrent.ExecutionException;
2631
import java.util.regex.Matcher;
2732
import java.util.regex.Pattern;
2833

@@ -31,8 +36,10 @@
3136
* AUTOCOMMIT=TRUE.
3237
*/
3338
class ClientSideStatementSetExecutor<T> implements ClientSideStatementExecutor {
39+
private final Cache<String, Tuple<T, Boolean>> cache;
3440
private final ClientSideStatementImpl statement;
3541
private final Method method;
42+
private final boolean supportsLocal;
3643
private final ClientSideStatementValueConverter<T> converter;
3744
private final Pattern allowedValuesPattern;
3845

@@ -46,12 +53,18 @@ class ClientSideStatementSetExecutor<T> implements ClientSideStatementExecutor {
4653
@SuppressWarnings("unchecked")
4754
ClientSideStatementSetExecutor(ClientSideStatementImpl statement) throws CompileException {
4855
Preconditions.checkNotNull(statement.getSetStatement());
56+
this.cache =
57+
CacheBuilder.newBuilder()
58+
.maximumSize(25)
59+
// Set the concurrency level to 1, as we don't expect many concurrent updates.
60+
.concurrencyLevel(1)
61+
.build();
4962
try {
5063
this.statement = statement;
5164
this.allowedValuesPattern =
5265
Pattern.compile(
5366
String.format(
54-
"(?is)\\A\\s*set\\s+%s\\s*%s\\s*%s\\s*\\z",
67+
"(?is)\\A\\s*set\\s+((?:local|session)\\s+)?%s\\s*%s\\s*%s\\s*\\z",
5568
statement.getSetStatement().getPropertyName(),
5669
statement.getSetStatement().getSeparator(),
5770
statement.getSetStatement().getAllowedValues()));
@@ -64,9 +77,21 @@ class ClientSideStatementSetExecutor<T> implements ClientSideStatementExecutor {
6477
Constructor<ClientSideStatementValueConverter<T>> constructor =
6578
converterClass.getConstructor(String.class);
6679
this.converter = constructor.newInstance(statement.getSetStatement().getAllowedValues());
67-
this.method =
68-
ConnectionStatementExecutor.class.getDeclaredMethod(
69-
statement.getMethodName(), converter.getParameterClass());
80+
Method method;
81+
boolean supportsLocal;
82+
try {
83+
method =
84+
ConnectionStatementExecutor.class.getDeclaredMethod(
85+
statement.getMethodName(), converter.getParameterClass());
86+
supportsLocal = false;
87+
} catch (NoSuchMethodException ignore) {
88+
method =
89+
ConnectionStatementExecutor.class.getDeclaredMethod(
90+
statement.getMethodName(), converter.getParameterClass(), Boolean.class);
91+
supportsLocal = true;
92+
}
93+
this.method = method;
94+
this.supportsLocal = supportsLocal;
7095
} catch (Exception e) {
7196
throw new CompileException(e, statement);
7297
}
@@ -75,17 +100,29 @@ class ClientSideStatementSetExecutor<T> implements ClientSideStatementExecutor {
75100
@Override
76101
public StatementResult execute(ConnectionStatementExecutor connection, ParsedStatement statement)
77102
throws Exception {
78-
return (StatementResult)
79-
method.invoke(connection, getParameterValue(statement.getSqlWithoutComments()));
103+
Tuple<T, Boolean> value;
104+
try {
105+
value =
106+
this.cache.get(
107+
statement.getSqlWithoutComments(),
108+
() -> getParameterValue(statement.getSqlWithoutComments()));
109+
} catch (ExecutionException | UncheckedExecutionException executionException) {
110+
throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
111+
}
112+
if (this.supportsLocal) {
113+
return (StatementResult) method.invoke(connection, value.x(), value.y());
114+
}
115+
return (StatementResult) method.invoke(connection, value.x());
80116
}
81117

82-
T getParameterValue(String sql) {
118+
Tuple<T, Boolean> getParameterValue(String sql) {
83119
Matcher matcher = allowedValuesPattern.matcher(sql);
84-
if (matcher.find() && matcher.groupCount() >= 1) {
85-
String value = matcher.group(1);
120+
if (matcher.find() && matcher.groupCount() >= 2) {
121+
boolean local = matcher.group(1) != null && "local".equalsIgnoreCase(matcher.group(1).trim());
122+
String value = matcher.group(2);
86123
T res = converter.convert(value);
87124
if (res != null) {
88-
return res;
125+
return Tuple.of(res, local);
89126
}
90127
throw SpannerExceptionFactory.newSpannerException(
91128
ErrorCode.INVALID_ARGUMENT,
@@ -94,8 +131,9 @@ T getParameterValue(String sql) {
94131
this.statement.getSetStatement().getPropertyName(), value));
95132
} else {
96133
Matcher invalidMatcher = this.statement.getPattern().matcher(sql);
97-
if (invalidMatcher.find() && invalidMatcher.groupCount() == 1) {
98-
String invalidValue = invalidMatcher.group(1);
134+
int valueGroup = this.supportsLocal ? 2 : 1;
135+
if (invalidMatcher.find() && invalidMatcher.groupCount() == valueGroup) {
136+
String invalidValue = invalidMatcher.group(valueGroup);
99137
throw SpannerExceptionFactory.newSpannerException(
100138
ErrorCode.INVALID_ARGUMENT,
101139
String.format(

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -942,8 +942,12 @@ public boolean isRetryAbortsInternally() {
942942

943943
@Override
944944
public void setRetryAbortsInternally(boolean retryAbortsInternally) {
945+
setRetryAbortsInternally(retryAbortsInternally, /* local = */ false);
946+
}
947+
948+
void setRetryAbortsInternally(boolean retryAbortsInternally, boolean local) {
945949
checkSetRetryAbortsInternallyAvailable();
946-
setConnectionPropertyValue(RETRY_ABORTS_INTERNALLY, retryAbortsInternally);
950+
setConnectionPropertyValue(RETRY_ABORTS_INTERNALLY, retryAbortsInternally, local);
947951
}
948952

949953
@Override

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ interface ConnectionStatementExecutor {
4444

4545
StatementResult statementShowReadOnly();
4646

47-
StatementResult statementSetRetryAbortsInternally(Boolean retryAbortsInternally);
47+
StatementResult statementSetRetryAbortsInternally(Boolean retryAbortsInternally, Boolean local);
4848

4949
StatementResult statementShowRetryAbortsInternally();
5050

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -181,9 +181,10 @@ public StatementResult statementShowReadOnly() {
181181
}
182182

183183
@Override
184-
public StatementResult statementSetRetryAbortsInternally(Boolean retryAbortsInternally) {
184+
public StatementResult statementSetRetryAbortsInternally(
185+
Boolean retryAbortsInternally, Boolean local) {
185186
Preconditions.checkNotNull(retryAbortsInternally);
186-
getConnection().setRetryAbortsInternally(retryAbortsInternally);
187+
getConnection().setRetryAbortsInternally(retryAbortsInternally, local);
187188
return noResult(SET_RETRY_ABORTS_INTERNALLY);
188189
}
189190

‎google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -355,13 +355,18 @@
355355
}
356356
},
357357
{
358-
"name": "SET RETRY_ABORTS_INTERNALLY = TRUE|FALSE",
358+
"name": "SET [LOCAL] RETRY_ABORTS_INTERNALLY = TRUE|FALSE",
359359
"executorName": "ClientSideStatementSetExecutor",
360360
"resultType": "NO_RESULT",
361361
"statementType": "SET_RETRY_ABORTS_INTERNALLY",
362-
"regex": "(?is)\\A\\s*set\\s+retry_aborts_internally\\s*(?:=)\\s*(.*)\\z",
362+
"regex": "(?is)\\A\\s*set\\s+(local\\s+)?retry_aborts_internally\\s*(?:=)\\s*(.*)\\z",
363363
"method": "statementSetRetryAbortsInternally",
364-
"exampleStatements": ["set retry_aborts_internally = true", "set retry_aborts_internally = false"],
364+
"exampleStatements": [
365+
"set retry_aborts_internally = true",
366+
"set retry_aborts_internally = false",
367+
"set local retry_aborts_internally = true",
368+
"set local retry_aborts_internally = false"
369+
],
365370
"examplePrerequisiteStatements": ["set readonly = false", "set autocommit = false"],
366371
"setStatement": {
367372
"propertyName": "RETRY_ABORTS_INTERNALLY",

‎google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -404,13 +404,26 @@
404404
}
405405
},
406406
{
407-
"name": "SET SPANNER.RETRY_ABORTS_INTERNALLY =|TO TRUE|FALSE",
407+
"name": "SET [SESSION|LOCAL] SPANNER.RETRY_ABORTS_INTERNALLY =|TO TRUE|FALSE",
408408
"executorName": "ClientSideStatementSetExecutor",
409409
"resultType": "NO_RESULT",
410410
"statementType": "SET_RETRY_ABORTS_INTERNALLY",
411-
"regex": "(?is)\\A\\s*set\\s+spanner\\.retry_aborts_internally(?:\\s*=\\s*|\\s+to\\s+)(.*)\\z",
411+
"regex": "(?is)\\A\\s*set\\s+((?:session|local)\\s+)?spanner\\.retry_aborts_internally(?:\\s*=\\s*|\\s+to\\s+)(.*)\\z",
412412
"method": "statementSetRetryAbortsInternally",
413-
"exampleStatements": ["set spanner.retry_aborts_internally = true", "set spanner.retry_aborts_internally = false", "set spanner.retry_aborts_internally to true", "set spanner.retry_aborts_internally to false"],
413+
"exampleStatements": [
414+
"set spanner.retry_aborts_internally = true",
415+
"set spanner.retry_aborts_internally = false",
416+
"set spanner.retry_aborts_internally to true",
417+
"set spanner.retry_aborts_internally to false",
418+
"set local spanner.retry_aborts_internally = true",
419+
"set local spanner.retry_aborts_internally = false",
420+
"set local spanner.retry_aborts_internally to true",
421+
"set local spanner.retry_aborts_internally to false",
422+
"set session spanner.retry_aborts_internally = true",
423+
"set session spanner.retry_aborts_internally = false",
424+
"set session spanner.retry_aborts_internally to true",
425+
"set session spanner.retry_aborts_internally to false"
426+
],
414427
"examplePrerequisiteStatements": ["set spanner.readonly = false", "set autocommit = false"],
415428
"setStatement": {
416429
"propertyName": "SPANNER.RETRY_ABORTS_INTERNALLY",

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionStateMockServerTest.java

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,15 @@
1919
import static com.google.cloud.spanner.connection.ConnectionProperties.CONNECTION_STATE_TYPE;
2020
import static org.junit.Assert.assertEquals;
2121
import static org.junit.Assert.assertFalse;
22+
import static org.junit.Assert.assertThrows;
2223
import static org.junit.Assert.assertTrue;
24+
import static org.junit.Assume.assumeTrue;
2325

2426
import com.google.cloud.spanner.Dialect;
27+
import com.google.cloud.spanner.ErrorCode;
2528
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
29+
import com.google.cloud.spanner.SpannerException;
30+
import com.google.cloud.spanner.Statement;
2631
import com.google.cloud.spanner.connection.ConnectionState.Type;
2732
import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection;
2833
import org.junit.After;
@@ -84,6 +89,10 @@ ITConnection createConnection(ConnectionState.Type type) {
8489
return createConnection(";" + CONNECTION_STATE_TYPE.getKey() + "=" + type.name());
8590
}
8691

92+
String getPrefix() {
93+
return dialect == Dialect.POSTGRESQL ? "SPANNER." : "";
94+
}
95+
8796
@Test
8897
public void testConnectionStateType() {
8998
try (Connection connection = createConnection()) {
@@ -228,4 +237,58 @@ public void testLocalChangeIsLostAfterTransaction() {
228237
}
229238
}
230239
}
240+
241+
@Test
242+
public void testSetLocalWithSqlStatement() {
243+
try (Connection connection = createConnection()) {
244+
connection.setAutocommit(false);
245+
246+
assertTrue(connection.isRetryAbortsInternally());
247+
connection.execute(
248+
Statement.of(String.format("set local %sretry_aborts_internally=false", getPrefix())));
249+
assertFalse(connection.isRetryAbortsInternally());
250+
connection.commit();
251+
assertTrue(connection.isRetryAbortsInternally());
252+
}
253+
}
254+
255+
@Test
256+
public void testSetSessionWithSqlStatement() {
257+
assumeTrue("Only PostgreSQL supports the 'session' keyword", dialect == Dialect.POSTGRESQL);
258+
259+
try (Connection connection = createConnection()) {
260+
connection.setAutocommit(false);
261+
262+
assertTrue(connection.isRetryAbortsInternally());
263+
connection.execute(
264+
Statement.of(String.format("set session %sretry_aborts_internally=false", getPrefix())));
265+
assertFalse(connection.isRetryAbortsInternally());
266+
connection.commit();
267+
assertFalse(connection.isRetryAbortsInternally());
268+
}
269+
}
270+
271+
@Test
272+
public void testSetLocalInvalidValue() {
273+
try (Connection connection = createConnection()) {
274+
connection.setAutocommit(false);
275+
276+
assertTrue(connection.isRetryAbortsInternally());
277+
SpannerException exception =
278+
assertThrows(
279+
SpannerException.class,
280+
() ->
281+
connection.execute(
282+
Statement.of(
283+
String.format("set local %sretry_aborts_internally=foo", getPrefix()))));
284+
assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
285+
assertTrue(
286+
exception.getMessage(),
287+
exception
288+
.getMessage()
289+
.endsWith(
290+
String.format("Unknown value for %sRETRY_ABORTS_INTERNALLY: foo", getPrefix())));
291+
assertTrue(connection.isRetryAbortsInternally());
292+
}
293+
}
231294
}

0 commit comments

Comments
 (0)