Skip to content

Commit f2ad888

Browse files
chore: add addArtifactId() to ConnectorRegistry (#508)
1 parent f74bab6 commit f2ad888

11 files changed

+113
-30
lines changed

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java

+3-4
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ class AlloyDBAdminClientFactory {
3232
private static final String DEFAULT_ENDPOINT = "alloydb.googleapis.com:443";
3333

3434
static AlloyDBAdminClient create(
35-
FixedCredentialsProvider credentialsProvider, ConnectorConfig config) throws IOException {
35+
FixedCredentialsProvider credentialsProvider, ConnectorConfig config, String userAgents)
36+
throws IOException {
3637
AlloyDBAdminSettings.Builder settingsBuilder = AlloyDBAdminSettings.newBuilder();
3738

3839
String endpoint = config.getAdminServiceEndpoint();
@@ -41,9 +42,7 @@ static AlloyDBAdminClient create(
4142
}
4243

4344
Map<String, String> headers =
44-
ImmutableMap.<String, String>builder()
45-
.put("user-agent", "alloydb-java-connector/" + Version.VERSION)
46-
.build();
45+
ImmutableMap.<String, String>builder().put("user-agent", userAgents).build();
4746

4847
settingsBuilder
4948
.setEndpoint(endpoint)

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionInfoRepositoryFactory.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
package com.google.cloud.alloydb;
1818

19-
/** Factory interface for creating SQLAdmin clients to interact with AlloyDB Admin API. */
19+
/** Factory interface for creating AlloyDBAdminClient to interact with AlloyDB Admin API. */
2020
public interface ConnectionInfoRepositoryFactory {
2121

2222
ConnectionInfoRepository create(CredentialFactory credentialFactory, ConnectorConfig config);

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionSocket.java

+5-3
Original file line numberDiff line numberDiff line change
@@ -60,23 +60,25 @@ class ConnectionSocket {
6060
private static final String X_509 = "X.509";
6161
private static final String ROOT_CA_CERT = "rootCaCert";
6262
private static final String CLIENT_CERT = "clientCert";
63-
private static final String USER_AGENT = "alloydb-java-connector/" + Version.VERSION;
6463
private static final int IO_TIMEOUT_MS = 30000;
6564
private static final int SERVER_SIDE_PROXY_PORT = 5433;
6665
private final ConnectionInfo connectionInfo;
6766
private final ConnectionConfig connectionConfig;
6867
private final KeyPair clientConnectorKeyPair;
6968
private final AccessTokenSupplier accessTokenSupplier;
69+
private final String userAgents;
7070

7171
ConnectionSocket(
7272
ConnectionInfo connectionInfo,
7373
ConnectionConfig connectionConfig,
7474
KeyPair clientConnectorKeyPair,
75-
AccessTokenSupplier accessTokenSupplier) {
75+
AccessTokenSupplier accessTokenSupplier,
76+
String userAgents) {
7677
this.connectionInfo = connectionInfo;
7778
this.connectionConfig = connectionConfig;
7879
this.clientConnectorKeyPair = clientConnectorKeyPair;
7980
this.accessTokenSupplier = accessTokenSupplier;
81+
this.userAgents = userAgents;
8082
}
8183

8284
Socket connect() throws IOException {
@@ -234,7 +236,7 @@ private void metadataExchange(SSLSocket socket) throws IOException {
234236
MetadataExchangeRequest.newBuilder()
235237
.setAuthType(authType)
236238
.setOauth2Token(tokenValue)
237-
.setUserAgent(USER_AGENT)
239+
.setUserAgent(userAgents)
238240
.build();
239241

240242
// Write data to the server.

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java

+10-4
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ class Connector {
3636
private final ConcurrentHashMap<ConnectionConfig, ConnectionInfoCache> instances;
3737
private final ConnectorConfig config;
3838
private final AccessTokenSupplier accessTokenSupplier;
39+
private final String userAgents;
3940

4041
Connector(
4142
ConnectorConfig config,
@@ -44,14 +45,16 @@ class Connector {
4445
KeyPair clientConnectorKeyPair,
4546
ConnectionInfoCacheFactory connectionInfoCacheFactory,
4647
ConcurrentHashMap<ConnectionConfig, ConnectionInfoCache> instances,
47-
AccessTokenSupplier accessTokenSupplier) {
48+
AccessTokenSupplier accessTokenSupplier,
49+
String userAgents) {
4850
this.config = config;
4951
this.executor = executor;
5052
this.connectionInfoRepo = connectionInfoRepo;
5153
this.clientConnectorKeyPair = clientConnectorKeyPair;
5254
this.connectionInfoCacheFactory = connectionInfoCacheFactory;
5355
this.instances = instances;
5456
this.accessTokenSupplier = accessTokenSupplier;
57+
this.userAgents = userAgents;
5558
}
5659

5760
public ConnectorConfig getConfig() {
@@ -70,7 +73,8 @@ Socket connect(ConnectionConfig config) throws IOException {
7073

7174
try {
7275
ConnectionSocket socket =
73-
new ConnectionSocket(connectionInfo, config, clientConnectorKeyPair, accessTokenSupplier);
76+
new ConnectionSocket(
77+
connectionInfo, config, clientConnectorKeyPair, accessTokenSupplier, userAgents);
7478
return socket.connect();
7579
} catch (IOException e) {
7680
logger.debug(
@@ -124,7 +128,8 @@ public boolean equals(Object o) {
124128
&& Objects.equal(clientConnectorKeyPair, that.clientConnectorKeyPair)
125129
&& Objects.equal(connectionInfoCacheFactory, that.connectionInfoCacheFactory)
126130
&& Objects.equal(instances, that.instances)
127-
&& Objects.equal(accessTokenSupplier, that.accessTokenSupplier);
131+
&& Objects.equal(accessTokenSupplier, that.accessTokenSupplier)
132+
&& Objects.equal(userAgents, that.userAgents);
128133
}
129134

130135
@Override
@@ -136,6 +141,7 @@ public int hashCode() {
136141
clientConnectorKeyPair,
137142
connectionInfoCacheFactory,
138143
instances,
139-
accessTokenSupplier);
144+
accessTokenSupplier,
145+
userAgents);
140146
}
141147
}

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java

+10
Original file line numberDiff line numberDiff line change
@@ -53,4 +53,14 @@ public static void reset() {
5353
public static void shutdown() {
5454
InternalConnectorRegistry.INSTANCE.shutdownInstance();
5555
}
56+
57+
/**
58+
* Adds an external application name to the user agent string for tracking. This is known to be
59+
* used by the spring-cloud-gcp project.
60+
*
61+
* @throws IllegalStateException if the AlloyDB Admin client has already been initialized
62+
*/
63+
public static void addArtifactId(String artifactId) {
64+
InternalConnectorRegistry.INSTANCE.addArtifactId(artifactId);
65+
}
5666
}

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoRepositoryFactory.java

+7-3
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,15 @@
2020
import com.google.common.util.concurrent.ListeningScheduledExecutorService;
2121
import java.io.IOException;
2222

23-
/** Factory for creating a SQLAdmin client that interacts with the real AlloyDB Admin API. */
23+
/** Factory for creating a AlloyDBAdminClient that interacts with the real AlloyDB Admin API. */
2424
class DefaultConnectionInfoRepositoryFactory implements ConnectionInfoRepositoryFactory {
2525
private final ListeningScheduledExecutorService executor;
26+
private final String userAgents;
2627

27-
DefaultConnectionInfoRepositoryFactory(ListeningScheduledExecutorService executor) {
28+
DefaultConnectionInfoRepositoryFactory(
29+
ListeningScheduledExecutorService executor, String userAgents) {
2830
this.executor = executor;
31+
this.userAgents = userAgents;
2932
}
3033

3134
@Override
@@ -34,7 +37,8 @@ public DefaultConnectionInfoRepository create(
3437

3538
AlloyDBAdminClient alloyDBAdminClient;
3639
try {
37-
alloyDBAdminClient = AlloyDBAdminClientFactory.create(credentialFactory.create(), config);
40+
alloyDBAdminClient =
41+
AlloyDBAdminClientFactory.create(credentialFactory.create(), config, userAgents);
3842
return new DefaultConnectionInfoRepository(executor, alloyDBAdminClient);
3943
} catch (IOException e) {
4044
throw new RuntimeException(e);

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/InternalConnectorRegistry.java

+28-2
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import java.io.Closeable;
2424
import java.io.IOException;
2525
import java.net.Socket;
26+
import java.util.ArrayList;
27+
import java.util.List;
2628
import java.util.concurrent.ConcurrentHashMap;
2729
import java.util.concurrent.Executors;
2830

@@ -53,6 +55,11 @@ enum InternalConnectorRegistry implements Closeable {
5355
@GuardedBy("shutdownGuard")
5456
private boolean shutdown = false;
5557

58+
@SuppressWarnings("ImmutableEnumChecker")
59+
private List<String> userAgents = new ArrayList<>();
60+
61+
private static final String USER_AGENT = "alloydb-java-connector/" + Version.VERSION;
62+
5663
InternalConnectorRegistry() {
5764
// During refresh, each instance consumes 2 threads from the thread pool. By using 8 threads,
5865
// there should be enough free threads so that there will not be a deadlock. Most users
@@ -62,6 +69,7 @@ enum InternalConnectorRegistry implements Closeable {
6269
this.unnamedConnectors = new ConcurrentHashMap<>();
6370
this.namedConnectors = new ConcurrentHashMap<>();
6471
this.credentialFactoryProvider = new CredentialFactoryProvider();
72+
this.addArtifactId(USER_AGENT);
6573
}
6674

6775
/** Test use only: Set a new CredentialFactoryProvider */
@@ -154,6 +162,23 @@ public void shutdownInstance() {
154162
}
155163
}
156164

165+
/**
166+
* Sets the default string which is appended to the AlloyDB Admin API client User-Agent header.
167+
*/
168+
public void addArtifactId(String artifactId) {
169+
if (!userAgents.contains(artifactId)) {
170+
userAgents.add(artifactId);
171+
}
172+
}
173+
174+
/**
175+
* Returns the default string which is appended to the AlloyDB Admin API client User-Agent header.
176+
*/
177+
@VisibleForTesting
178+
String getUserAgents() {
179+
return String.join(" ", userAgents);
180+
}
181+
157182
private Connector getConnector(ConnectionConfig config) {
158183
return unnamedConnectors.computeIfAbsent(
159184
config.getConnectorConfig(), k -> createConnector(config.getConnectorConfig()));
@@ -163,7 +188,7 @@ private Connector createConnector(ConnectorConfig config) {
163188
CredentialFactory instanceCredentialFactory =
164189
credentialFactoryProvider.getInstanceCredentialFactory(config);
165190
DefaultConnectionInfoRepositoryFactory connectionInfoRepositoryFactory =
166-
new DefaultConnectionInfoRepositoryFactory(executor);
191+
new DefaultConnectionInfoRepositoryFactory(executor, getUserAgents());
167192
DefaultConnectionInfoRepository connectionInfoRepository =
168193
connectionInfoRepositoryFactory.create(instanceCredentialFactory, config);
169194
AccessTokenSupplier accessTokenSupplier =
@@ -176,7 +201,8 @@ private Connector createConnector(ConnectorConfig config) {
176201
RsaKeyPairGenerator.generateKeyPair(),
177202
new DefaultConnectionInfoCacheFactory(),
178203
new ConcurrentHashMap<>(),
179-
accessTokenSupplier);
204+
accessTokenSupplier,
205+
getUserAgents());
180206
}
181207

182208
private Connector getNamedConnector(String name) {

alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectorTest.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ public class ConnectorTest {
4747
private static final String ERROR_MESSAGE_PERMISSION_DENIED =
4848
"Location not found or access is unauthorized.";
4949
private static final String ERROR_MESSAGE_INTERNAL = "Internal Error";
50+
private static final String USER_AGENT = "unit tests";
5051

5152
ListeningScheduledExecutorService defaultExecutor;
5253

@@ -160,7 +161,8 @@ private Connector newConnector(ConnectorConfig config, MockAlloyDBAdminGrpc mock
160161
TestCertificates.INSTANCE.getClientKey(),
161162
new DefaultConnectionInfoCacheFactory(),
162163
new ConcurrentHashMap<>(),
163-
accessTokenSupplier);
164+
accessTokenSupplier,
165+
USER_AGENT);
164166
}
165167

166168
private String readLine(Socket socket) throws IOException {

alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITConnectorTest.java

+36-11
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public class ITConnectorTest {
4444
private ConnectionInfoRepository connectionInfoRepo;
4545
private CredentialFactoryProvider credentialFactoryProvider;
4646
private AccessTokenSupplier accessTokenSupplier;
47+
private static final String USER_AGENT = "integration tests";
4748

4849
@Before
4950
public void setUp() throws IOException {
@@ -54,7 +55,8 @@ public void setUp() throws IOException {
5455
credentialFactoryProvider = new CredentialFactoryProvider();
5556
CredentialFactory instanceCredentialFactory =
5657
credentialFactoryProvider.getInstanceCredentialFactory(config);
57-
connectionInfoRepositoryFactory = new DefaultConnectionInfoRepositoryFactory(executor);
58+
connectionInfoRepositoryFactory =
59+
new DefaultConnectionInfoRepositoryFactory(executor, USER_AGENT);
5860
connectionInfoRepo = connectionInfoRepositoryFactory.create(instanceCredentialFactory, config);
5961
accessTokenSupplier = new DefaultAccessTokenSupplier(instanceCredentialFactory);
6062
}
@@ -87,7 +89,8 @@ public void testConnect_createsSocketConnection() throws IOException {
8789
RsaKeyPairGenerator.generateKeyPair(),
8890
new DefaultConnectionInfoCacheFactory(),
8991
new ConcurrentHashMap<>(),
90-
accessTokenSupplier);
92+
accessTokenSupplier,
93+
USER_AGENT);
9194

9295
socket = (SSLSocket) connector.connect(config);
9396

@@ -132,7 +135,8 @@ public void testConnect_whenTlsHandshakeFails()
132135
clientConnectorKeyPair,
133136
connectionInfoCacheFactory,
134137
new ConcurrentHashMap<>(),
135-
accessTokenSupplier);
138+
accessTokenSupplier,
139+
USER_AGENT);
136140
socket = (SSLSocket) connector.connect(config);
137141
} catch (ConnectException ignore) {
138142
// The socket connect will fail because it's trying to connect to localhost with TLS certs.
@@ -174,7 +178,8 @@ public void testEquals() {
174178
clientConnectorKeyPair,
175179
connectionInfoCacheFactory,
176180
new ConcurrentHashMap<>(),
177-
accessTokenSupplier);
181+
accessTokenSupplier,
182+
USER_AGENT);
178183

179184
assertThat(a)
180185
.isNotEqualTo(
@@ -185,7 +190,8 @@ public void testEquals() {
185190
clientConnectorKeyPair,
186191
connectionInfoCacheFactory,
187192
new ConcurrentHashMap<>(),
188-
accessTokenSupplier));
193+
accessTokenSupplier,
194+
USER_AGENT));
189195

190196
assertThat(a)
191197
.isNotEqualTo(
@@ -196,7 +202,8 @@ public void testEquals() {
196202
clientConnectorKeyPair,
197203
connectionInfoCacheFactory,
198204
new ConcurrentHashMap<>(),
199-
accessTokenSupplier));
205+
accessTokenSupplier,
206+
USER_AGENT));
200207

201208
assertThat(a)
202209
.isNotEqualTo(
@@ -207,7 +214,8 @@ public void testEquals() {
207214
clientConnectorKeyPair,
208215
connectionInfoCacheFactory,
209216
new ConcurrentHashMap<>(),
210-
accessTokenSupplier));
217+
accessTokenSupplier,
218+
USER_AGENT));
211219

212220
assertThat(a)
213221
.isNotEqualTo(
@@ -218,7 +226,8 @@ public void testEquals() {
218226
RsaKeyPairGenerator.generateKeyPair(), // Different
219227
connectionInfoCacheFactory,
220228
new ConcurrentHashMap<>(),
221-
accessTokenSupplier));
229+
accessTokenSupplier,
230+
USER_AGENT));
222231

223232
assertThat(a)
224233
.isNotEqualTo(
@@ -229,7 +238,8 @@ public void testEquals() {
229238
clientConnectorKeyPair,
230239
new DefaultConnectionInfoCacheFactory(), // Different
231240
new ConcurrentHashMap<>(),
232-
accessTokenSupplier));
241+
accessTokenSupplier,
242+
USER_AGENT));
233243

234244
assertThat(a)
235245
.isNotEqualTo(
@@ -240,6 +250,19 @@ public void testEquals() {
240250
clientConnectorKeyPair,
241251
connectionInfoCacheFactory,
242252
new ConcurrentHashMap<>(),
253+
null, // Different
254+
USER_AGENT));
255+
256+
assertThat(a)
257+
.isNotEqualTo(
258+
new Connector(
259+
config,
260+
executor,
261+
connectionInfoRepo,
262+
clientConnectorKeyPair,
263+
connectionInfoCacheFactory,
264+
new ConcurrentHashMap<>(),
265+
accessTokenSupplier,
243266
null)); // Different
244267
}
245268

@@ -259,7 +282,8 @@ public void testHashCode() {
259282
clientConnectorKeyPair,
260283
connectionInfoCacheFactory,
261284
instances,
262-
accessTokenSupplier);
285+
accessTokenSupplier,
286+
USER_AGENT);
263287

264288
assertThat(a.hashCode())
265289
.isEqualTo(
@@ -270,6 +294,7 @@ public void testHashCode() {
270294
clientConnectorKeyPair,
271295
connectionInfoCacheFactory,
272296
instances,
273-
accessTokenSupplier));
297+
accessTokenSupplier,
298+
USER_AGENT));
274299
}
275300
}

0 commit comments

Comments
 (0)