Skip to content

Commit 734f521

Browse files
olavloiteolavloite
authored
fix: show statements failed in pgx (#629)
* fix: show statements failed in pgx SHOW variable_name statements failed in pgx, because pgx tries to describe all statements (including SHOW). This failed as these statements do not have any metadata, as they are not returned by the backend. This fix also adds automatic detection of the pgx driver based on the first Parse message that we receive. pgx uses a very specific name for the prepared statements that it creates that can be used for (relatively) safe detection. * fix: add safeguard against unknown parameter
1 parent 715f97f commit 734f521

21 files changed

+246
-39
lines changed

‎src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java‎

Lines changed: 57 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,11 @@
5252
import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
5353
import com.google.cloud.spanner.pgadapter.wireoutput.TerminateResponse;
5454
import com.google.cloud.spanner.pgadapter.wireprotocol.BootstrapMessage;
55+
import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
5556
import com.google.cloud.spanner.pgadapter.wireprotocol.SSLMessage;
5657
import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
5758
import com.google.common.annotations.VisibleForTesting;
59+
import com.google.common.base.Strings;
5860
import com.google.common.cache.Cache;
5961
import com.google.common.cache.CacheBuilder;
6062
import com.google.common.collect.ImmutableList;
@@ -69,6 +71,7 @@
6971
import java.time.Duration;
7072
import java.util.HashMap;
7173
import java.util.HashSet;
74+
import java.util.Locale;
7275
import java.util.Map;
7376
import java.util.Map.Entry;
7477
import java.util.Properties;
@@ -704,6 +707,14 @@ public WellKnownClient getWellKnownClient() {
704707

705708
public void setWellKnownClient(WellKnownClient wellKnownClient) {
706709
this.wellKnownClient = wellKnownClient;
710+
if (this.wellKnownClient != WellKnownClient.UNSPECIFIED) {
711+
logger.log(
712+
Level.INFO,
713+
() ->
714+
String.format(
715+
"Well-known client %s detected for connection %d.",
716+
this.wellKnownClient, getConnectionId()));
717+
}
707718
}
708719

709720
/**
@@ -713,23 +724,58 @@ public void setWellKnownClient(WellKnownClient wellKnownClient) {
713724
* executed.
714725
*/
715726
public void maybeDetermineWellKnownClient(Statement statement) {
716-
if (!this.hasDeterminedClientUsingQuery
717-
&& this.wellKnownClient == WellKnownClient.UNSPECIFIED
718-
&& getServer().getOptions().shouldAutoDetectClient()) {
719-
this.wellKnownClient = ClientAutoDetector.detectClient(ImmutableList.of(statement));
720-
if (this.wellKnownClient != WellKnownClient.UNSPECIFIED) {
721-
logger.log(
722-
Level.INFO,
723-
() ->
724-
String.format(
725-
"Well-known client %s detected for connection %d.",
726-
this.wellKnownClient, getConnectionId()));
727+
if (!this.hasDeterminedClientUsingQuery) {
728+
if (this.wellKnownClient == WellKnownClient.UNSPECIFIED
729+
&& getServer().getOptions().shouldAutoDetectClient()) {
730+
setWellKnownClient(ClientAutoDetector.detectClient(ImmutableList.of(statement)));
727731
}
732+
maybeSetApplicationName();
728733
}
729734
// Make sure that we only try to detect the client once.
730735
this.hasDeterminedClientUsingQuery = true;
731736
}
732737

738+
/**
739+
* This is called by the extended query protocol {@link
740+
* com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage} to give the connection the
741+
* opportunity to determine the client that is connected based on the data in the (first) parse
742+
* messages.
743+
*/
744+
public void maybeDetermineWellKnownClient(ParseMessage parseMessage) {
745+
if (!this.hasDeterminedClientUsingQuery) {
746+
if (this.wellKnownClient == WellKnownClient.UNSPECIFIED
747+
&& getServer().getOptions().shouldAutoDetectClient()) {
748+
setWellKnownClient(ClientAutoDetector.detectClient(parseMessage));
749+
}
750+
maybeSetApplicationName();
751+
}
752+
// Make sure that we only try to detect the client once.
753+
this.hasDeterminedClientUsingQuery = true;
754+
}
755+
756+
private void maybeSetApplicationName() {
757+
try {
758+
if (this.wellKnownClient != WellKnownClient.UNSPECIFIED
759+
&& getExtendedQueryProtocolHandler() != null
760+
&& Strings.isNullOrEmpty(
761+
getExtendedQueryProtocolHandler()
762+
.getBackendConnection()
763+
.getSessionState()
764+
.get(null, "application_name")
765+
.getSetting())) {
766+
getExtendedQueryProtocolHandler()
767+
.getBackendConnection()
768+
.getSessionState()
769+
.set(null, "application_name", wellKnownClient.name().toLowerCase(Locale.ENGLISH));
770+
getExtendedQueryProtocolHandler().getBackendConnection().getSessionState().commit();
771+
}
772+
} catch (Throwable ignore) {
773+
// Safeguard against a theoretical situation that 'application_name' has been removed from
774+
// the list of settings. Just ignore this situation, as the only consequence is that the
775+
// 'application_name' setting has not been set.
776+
}
777+
}
778+
733779
/** Status of a {@link ConnectionHandler} */
734780
public enum ConnectionStatus {
735781
UNAUTHENTICATED,

‎src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,7 @@ public String toString() {
358358
return String.format("ProxyServer[port: %d]", getLocalPort());
359359
}
360360

361-
ConcurrentLinkedQueue<WireMessage> getDebugMessages() {
361+
public ConcurrentLinkedQueue<WireMessage> getDebugMessages() {
362362
return debugMessages;
363363
}
364364

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
import com.google.cloud.spanner.PartitionOptions;
3131
import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
3232
import com.google.cloud.spanner.ResultSet;
33-
import com.google.cloud.spanner.ResultSets;
3433
import com.google.cloud.spanner.Spanner;
3534
import com.google.cloud.spanner.SpannerBatchUpdateException;
3635
import com.google.cloud.spanner.SpannerException;
@@ -1313,7 +1312,7 @@ public ClientSideStatementType getClientSideStatementType() {
13131312

13141313
@Override
13151314
public ResultSet getResultSet() {
1316-
return ResultSets.forRows(
1315+
return ClientSideResultSet.forRows(
13171316
Type.struct(StructField.of("partition", Type.bytes())),
13181317
partitions.stream()
13191318
.map(

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/ClientSideResultSet.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Copyright 2023 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package com.google.cloud.spanner.pgadapter.statements;
16+
17+
import com.google.api.core.InternalApi;
18+
import com.google.cloud.spanner.ForwardingResultSet;
19+
import com.google.cloud.spanner.ResultSet;
20+
import com.google.cloud.spanner.ResultSets;
21+
import com.google.cloud.spanner.Struct;
22+
import com.google.cloud.spanner.Type;
23+
import com.google.spanner.v1.ResultSetMetadata;
24+
import com.google.spanner.v1.ResultSetStats;
25+
26+
/** Wrapper class for query results that are handled directly in PGAdapter. */
27+
@InternalApi
28+
public class ClientSideResultSet extends ForwardingResultSet {
29+
public static ResultSet forRows(Type type, Iterable<Struct> rows) {
30+
return new ClientSideResultSet(ResultSets.forRows(type, rows), type);
31+
}
32+
33+
private final Type type;
34+
35+
private ClientSideResultSet(ResultSet delegate, Type type) {
36+
super(delegate);
37+
this.type = type;
38+
}
39+
40+
@Override
41+
public Type getType() {
42+
return type;
43+
}
44+
45+
@Override
46+
public ResultSetStats getStats() {
47+
return ResultSetStats.getDefaultInstance();
48+
}
49+
50+
@Override
51+
public ResultSetMetadata getMetadata() {
52+
return ResultSetMetadata.getDefaultInstance();
53+
}
54+
}

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/SessionStatementParser.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import com.google.api.client.util.Strings;
1818
import com.google.api.core.InternalApi;
1919
import com.google.cloud.spanner.ErrorCode;
20-
import com.google.cloud.spanner.ResultSets;
2120
import com.google.cloud.spanner.SpannerExceptionFactory;
2221
import com.google.cloud.spanner.Struct;
2322
import com.google.cloud.spanner.Type;
@@ -186,7 +185,7 @@ static ShowStatement createShowAll() {
186185
public StatementResult execute(SessionState sessionState) {
187186
if (name != null) {
188187
return new QueryResult(
189-
ResultSets.forRows(
188+
ClientSideResultSet.forRows(
190189
Type.struct(StructField.of(getKey(), Type.string())),
191190
ImmutableList.of(
192191
Struct.newBuilder()
@@ -195,7 +194,7 @@ public StatementResult execute(SessionState sessionState) {
195194
.build())));
196195
}
197196
return new QueryResult(
198-
ResultSets.forRows(
197+
ClientSideResultSet.forRows(
199198
Type.struct(
200199
StructField.of("name", Type.string()),
201200
StructField.of("setting", Type.string()),

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/local/DjangoGetTableNamesStatement.java‎

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

1717
import com.google.api.core.InternalApi;
1818
import com.google.cloud.spanner.ResultSet;
19-
import com.google.cloud.spanner.ResultSets;
2019
import com.google.cloud.spanner.Type;
2120
import com.google.cloud.spanner.Type.StructField;
2221
import com.google.cloud.spanner.connection.StatementResult;
2322
import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
2423
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
24+
import com.google.cloud.spanner.pgadapter.statements.ClientSideResultSet;
2525
import com.google.common.collect.ImmutableList;
2626

2727
/*
@@ -60,7 +60,7 @@ public String[] getSql() {
6060
@Override
6161
public StatementResult execute(BackendConnection backendConnection) {
6262
ResultSet resultSet =
63-
ResultSets.forRows(
63+
ClientSideResultSet.forRows(
6464
Type.struct(
6565
StructField.of("relname", Type.string()), StructField.of("case", Type.string())),
6666
ImmutableList.of());

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/local/ListDatabasesStatement.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import com.google.cloud.spanner.Dialect;
2121
import com.google.cloud.spanner.InstanceId;
2222
import com.google.cloud.spanner.ResultSet;
23-
import com.google.cloud.spanner.ResultSets;
2423
import com.google.cloud.spanner.Spanner;
2524
import com.google.cloud.spanner.Struct;
2625
import com.google.cloud.spanner.Type;
@@ -31,6 +30,7 @@
3130
import com.google.cloud.spanner.pgadapter.ConnectionHandler;
3231
import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
3332
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
33+
import com.google.cloud.spanner.pgadapter.statements.ClientSideResultSet;
3434
import com.google.common.base.Preconditions;
3535
import com.google.common.collect.ImmutableList;
3636
import java.util.Comparator;
@@ -74,7 +74,7 @@ public StatementResult execute(BackendConnection backendConnection) {
7474
InstanceId defaultInstanceId =
7575
connectionHandler.getServer().getOptions().getDefaultInstanceId();
7676
ResultSet resultSet =
77-
ResultSets.forRows(
77+
ClientSideResultSet.forRows(
7878
Type.struct(
7979
StructField.of("Name", Type.string()),
8080
StructField.of("Owner", Type.string()),

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/local/SelectCurrentCatalogStatement.java‎

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

1717
import com.google.api.core.InternalApi;
1818
import com.google.cloud.spanner.ResultSet;
19-
import com.google.cloud.spanner.ResultSets;
2019
import com.google.cloud.spanner.Struct;
2120
import com.google.cloud.spanner.Type;
2221
import com.google.cloud.spanner.Type.StructField;
2322
import com.google.cloud.spanner.connection.StatementResult;
2423
import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
2524
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
25+
import com.google.cloud.spanner.pgadapter.statements.ClientSideResultSet;
2626
import com.google.common.collect.ImmutableList;
2727

2828
@InternalApi
@@ -52,7 +52,7 @@ public String[] getSql() {
5252
@Override
5353
public StatementResult execute(BackendConnection backendConnection) {
5454
ResultSet resultSet =
55-
ResultSets.forRows(
55+
ClientSideResultSet.forRows(
5656
Type.struct(StructField.of("current_catalog", Type.string())),
5757
ImmutableList.of(
5858
Struct.newBuilder()

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/local/SelectCurrentDatabaseStatement.java‎

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

1717
import com.google.api.core.InternalApi;
1818
import com.google.cloud.spanner.ResultSet;
19-
import com.google.cloud.spanner.ResultSets;
2019
import com.google.cloud.spanner.Struct;
2120
import com.google.cloud.spanner.Type;
2221
import com.google.cloud.spanner.Type.StructField;
2322
import com.google.cloud.spanner.connection.StatementResult;
2423
import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
2524
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
25+
import com.google.cloud.spanner.pgadapter.statements.ClientSideResultSet;
2626
import com.google.common.collect.ImmutableList;
2727

2828
@InternalApi
@@ -53,7 +53,7 @@ public String[] getSql() {
5353
@Override
5454
public StatementResult execute(BackendConnection backendConnection) {
5555
ResultSet resultSet =
56-
ResultSets.forRows(
56+
ClientSideResultSet.forRows(
5757
Type.struct(StructField.of("current_database", Type.string())),
5858
ImmutableList.of(
5959
Struct.newBuilder()

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/local/SelectCurrentSchemaStatement.java‎

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

1717
import com.google.api.core.InternalApi;
1818
import com.google.cloud.spanner.ResultSet;
19-
import com.google.cloud.spanner.ResultSets;
2019
import com.google.cloud.spanner.Struct;
2120
import com.google.cloud.spanner.Type;
2221
import com.google.cloud.spanner.Type.StructField;
2322
import com.google.cloud.spanner.connection.StatementResult;
2423
import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
2524
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
25+
import com.google.cloud.spanner.pgadapter.statements.ClientSideResultSet;
2626
import com.google.common.collect.ImmutableList;
2727

2828
@InternalApi
@@ -50,7 +50,7 @@ public String[] getSql() {
5050
@Override
5151
public StatementResult execute(BackendConnection backendConnection) {
5252
ResultSet resultSet =
53-
ResultSets.forRows(
53+
ClientSideResultSet.forRows(
5454
Type.struct(StructField.of("current_schema", Type.string())),
5555
ImmutableList.of(
5656
Struct.newBuilder()

0 commit comments

Comments
 (0)