Skip to content

Commit fb6d0bf

Browse files
pratickchokhanipratickchokhani
committed
Merge branch 'main' of https://github.com/googleapis/java-spanner into blind-write-release
2 parents ea327fb + 331942f commit fb6d0bf

18 files changed

+3876
-256
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: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
7777
import com.google.cloud.spanner.connection.ConnectionProperty.Context;
7878
import com.google.cloud.spanner.connection.ConnectionState.Type;
79+
import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType;
7980
import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout;
8081
import com.google.cloud.spanner.connection.StatementResult.ResultType;
8182
import com.google.cloud.spanner.connection.UnitOfWork.CallType;
@@ -284,9 +285,17 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
284285
Preconditions.checkNotNull(options);
285286
this.leakedException =
286287
options.isTrackConnectionLeaks() ? new LeakedConnectionException() : null;
288+
StatementExecutorType statementExecutorType;
289+
if (options.getStatementExecutorType() != null) {
290+
statementExecutorType = options.getStatementExecutorType();
291+
} else {
292+
statementExecutorType =
293+
options.isUseVirtualThreads()
294+
? StatementExecutorType.VIRTUAL_THREAD
295+
: StatementExecutorType.DIRECT_EXECUTOR;
296+
}
287297
this.statementExecutor =
288-
new StatementExecutor(
289-
options.isUseVirtualThreads(), options.getStatementExecutionInterceptors());
298+
new StatementExecutor(statementExecutorType, options.getStatementExecutionInterceptors());
290299
this.spannerPool = SpannerPool.INSTANCE;
291300
this.options = options;
292301
this.spanner = spannerPool.getSpanner(options, this);
@@ -330,7 +339,11 @@ && getDialect() == Dialect.POSTGRESQL
330339
this.leakedException =
331340
options.isTrackConnectionLeaks() ? new LeakedConnectionException() : null;
332341
this.statementExecutor =
333-
new StatementExecutor(options.isUseVirtualThreads(), Collections.emptyList());
342+
new StatementExecutor(
343+
options.isUseVirtualThreads()
344+
? StatementExecutorType.VIRTUAL_THREAD
345+
: StatementExecutorType.DIRECT_EXECUTOR,
346+
Collections.emptyList());
334347
this.spannerPool = Preconditions.checkNotNull(spannerPool);
335348
this.options = Preconditions.checkNotNull(options);
336349
this.spanner = spannerPool.getSpanner(options, this);
@@ -929,8 +942,12 @@ public boolean isRetryAbortsInternally() {
929942

930943
@Override
931944
public void setRetryAbortsInternally(boolean retryAbortsInternally) {
945+
setRetryAbortsInternally(retryAbortsInternally, /* local = */ false);
946+
}
947+
948+
void setRetryAbortsInternally(boolean retryAbortsInternally, boolean local) {
932949
checkSetRetryAbortsInternallyAvailable();
933-
setConnectionPropertyValue(RETRY_ABORTS_INTERNALLY, retryAbortsInternally);
950+
setConnectionPropertyValue(RETRY_ABORTS_INTERNALLY, retryAbortsInternally, local);
934951
}
935952

936953
@Override

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
import com.google.cloud.spanner.SpannerException;
7171
import com.google.cloud.spanner.SpannerExceptionFactory;
7272
import com.google.cloud.spanner.SpannerOptions;
73+
import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType;
7374
import com.google.common.annotations.VisibleForTesting;
7475
import com.google.common.base.Preconditions;
7576
import com.google.common.base.Strings;
@@ -614,6 +615,7 @@ public static class Builder {
614615
new HashMap<>();
615616
private String uri;
616617
private Credentials credentials;
618+
private StatementExecutorType statementExecutorType;
617619
private SessionPoolOptions sessionPoolOptions;
618620
private List<StatementExecutionInterceptor> statementExecutionInterceptors =
619621
Collections.emptyList();
@@ -777,6 +779,11 @@ Builder setCredentials(Credentials credentials) {
777779
return this;
778780
}
779781

782+
Builder setStatementExecutorType(StatementExecutorType statementExecutorType) {
783+
this.statementExecutorType = statementExecutorType;
784+
return this;
785+
}
786+
780787
public Builder setOpenTelemetry(OpenTelemetry openTelemetry) {
781788
this.openTelemetry = openTelemetry;
782789
return this;
@@ -814,6 +821,7 @@ public static Builder newBuilder() {
814821
private final String instanceId;
815822
private final String databaseName;
816823
private final Credentials credentials;
824+
private final StatementExecutorType statementExecutorType;
817825
private final SessionPoolOptions sessionPoolOptions;
818826

819827
private final OpenTelemetry openTelemetry;
@@ -834,6 +842,7 @@ private ConnectionOptions(Builder builder) {
834842
ConnectionPropertyValue<Boolean> value = cast(connectionPropertyValues.get(LENIENT.getKey()));
835843
this.warnings = checkValidProperties(value != null && value.getValue(), uri);
836844
this.fixedCredentials = builder.credentials;
845+
this.statementExecutorType = builder.statementExecutorType;
837846

838847
this.openTelemetry = builder.openTelemetry;
839848
this.statementExecutionInterceptors =
@@ -1105,6 +1114,10 @@ CredentialsProvider getCredentialsProvider() {
11051114
return getInitialConnectionPropertyValue(CREDENTIALS_PROVIDER);
11061115
}
11071116

1117+
StatementExecutorType getStatementExecutorType() {
1118+
return this.statementExecutorType;
1119+
}
1120+
11081121
/** The {@link SessionPoolOptions} of this {@link ConnectionOptions}. */
11091122
public SessionPoolOptions getSessionPoolOptions() {
11101123
return sessionPoolOptions;

‎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/java/com/google/cloud/spanner/connection/StatementExecutor.java‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,10 @@ java.time.Duration asDuration() {
146146
ThreadFactoryUtil.createVirtualOrPlatformDaemonThreadFactory("connection-executor", false);
147147

148148
/** Creates an {@link ExecutorService} for a {@link StatementExecutor}. */
149-
private static ListeningExecutorService createExecutorService(boolean useVirtualThreads) {
149+
private static ListeningExecutorService createExecutorService(StatementExecutorType type) {
150+
if (type == StatementExecutorType.DIRECT_EXECUTOR) {
151+
return MoreExecutors.newDirectExecutorService();
152+
}
150153
return MoreExecutors.listeningDecorator(
151154
Context.taskWrapping(
152155
new ThreadPoolExecutor(
@@ -155,7 +158,7 @@ private static ListeningExecutorService createExecutorService(boolean useVirtual
155158
0L,
156159
TimeUnit.MILLISECONDS,
157160
new LinkedBlockingQueue<>(),
158-
useVirtualThreads
161+
type == StatementExecutorType.VIRTUAL_THREAD
159162
? DEFAULT_VIRTUAL_THREAD_FACTORY
160163
: DEFAULT_DAEMON_THREAD_FACTORY)));
161164
}
@@ -168,13 +171,23 @@ private static ListeningExecutorService createExecutorService(boolean useVirtual
168171
*/
169172
private final List<StatementExecutionInterceptor> interceptors;
170173

174+
enum StatementExecutorType {
175+
PLATFORM_THREAD,
176+
VIRTUAL_THREAD,
177+
DIRECT_EXECUTOR,
178+
}
179+
171180
@VisibleForTesting
172181
StatementExecutor() {
173-
this(DEFAULT_USE_VIRTUAL_THREADS, Collections.emptyList());
182+
this(
183+
DEFAULT_USE_VIRTUAL_THREADS
184+
? StatementExecutorType.VIRTUAL_THREAD
185+
: StatementExecutorType.PLATFORM_THREAD,
186+
Collections.emptyList());
174187
}
175188

176-
StatementExecutor(boolean useVirtualThreads, List<StatementExecutionInterceptor> interceptors) {
177-
this.executor = createExecutorService(useVirtualThreads);
189+
StatementExecutor(StatementExecutorType type, List<StatementExecutionInterceptor> interceptors) {
190+
this.executor = createExecutorService(type);
178191
this.interceptors = Collections.unmodifiableList(interceptors);
179192
}
180193

‎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/AbstractMockServerTest.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,7 @@ ITConnection createConnection(
283283
ConnectionOptions.newBuilder()
284284
.setUri(getBaseUrl() + additionalUrlOptions)
285285
.setStatementExecutionInterceptors(interceptors);
286+
configureConnectionOptions(builder);
286287
ConnectionOptions options = builder.build();
287288
ITConnection connection = createITConnection(options);
288289
for (TransactionRetryListener listener : transactionRetryListeners) {
@@ -291,6 +292,11 @@ ITConnection createConnection(
291292
return connection;
292293
}
293294

295+
protected ConnectionOptions.Builder configureConnectionOptions(
296+
ConnectionOptions.Builder builder) {
297+
return builder;
298+
}
299+
294300
protected String getBaseUrl() {
295301
return String.format(
296302
"cloudspanner://localhost:%d/projects/proj/instances/inst/databases/db?usePlainText=true;autocommit=false;retryAbortsInternally=true",

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@
3232
import com.google.cloud.spanner.Options;
3333
import com.google.cloud.spanner.SpannerExceptionFactory;
3434
import com.google.cloud.spanner.Statement;
35+
import com.google.cloud.spanner.connection.ConnectionOptions.Builder;
3536
import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection;
37+
import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType;
3638
import com.google.common.collect.Collections2;
3739
import com.google.common.collect.ImmutableList;
3840
import com.google.common.collect.Lists;
@@ -127,6 +129,11 @@ ITConnection createConnection(TransactionRetryListener listener) {
127129
return connection;
128130
}
129131

132+
@Override
133+
protected Builder configureConnectionOptions(Builder builder) {
134+
return builder.setStatementExecutorType(StatementExecutorType.PLATFORM_THREAD);
135+
}
136+
130137
@Test
131138
public void testSingleQueryAborted() {
132139
RetryCounter counter = new RetryCounter();

0 commit comments

Comments
 (0)