Skip to content

Commit 6c8c37c

Browse files
vkediavkedia
authored
Adds support for attaching labels to cloud spanner sessions (#2606)
* Adds support for attaching labels to cloud spanner sessions * Fixed style issues * Fixed style issues Author: Vikas Kedia <[email protected]> * Few more style fixes and null check * Adds tests for NPE * Doc fix
1 parent 09aeb93 commit 6c8c37c

7 files changed

Lines changed: 76 additions & 11 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ Session createSession(final DatabaseId db) throws SpannerException {
221221
new Callable<com.google.spanner.v1.Session>() {
222222
@Override
223223
public com.google.spanner.v1.Session call() throws Exception {
224-
return rpc.createSession(db.getName(), options);
224+
return rpc.createSession(db.getName(), getOptions().getSessionLabels(), options);
225225
}
226226
});
227227
return new SessionImpl(session.getName(), options);

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.google.common.base.MoreObjects;
2828
import com.google.common.base.Preconditions;
2929
import com.google.common.collect.ImmutableList;
30+
import com.google.common.collect.ImmutableMap;
3031
import com.google.common.collect.ImmutableSet;
3132
import io.grpc.ClientInterceptor;
3233
import io.grpc.ManagedChannel;
@@ -36,6 +37,7 @@
3637
import java.net.MalformedURLException;
3738
import java.net.URL;
3839
import java.util.List;
40+
import java.util.Map;
3941
import java.util.Set;
4042
import javax.net.ssl.SSLException;
4143

@@ -75,7 +77,8 @@ public ServiceRpc create(SpannerOptions options) {
7577
private final int prefetchChunks;
7678
private final int numChannels;
7779
private final String userAgent;
78-
80+
private final ImmutableMap<String, String> sessionLabels;
81+
7982
private SpannerOptions(Builder builder) {
8083
super(SpannerFactory.class, SpannerRpcFactory.class, builder, new SpannerDefaults());
8184
numChannels = builder.numChannels;
@@ -94,6 +97,7 @@ private SpannerOptions(Builder builder) {
9497
? builder.sessionPoolOptions
9598
: SessionPoolOptions.newBuilder().build();
9699
prefetchChunks = builder.prefetchChunks;
100+
sessionLabels = builder.sessionLabels;
97101
}
98102

99103
/** Builder for {@link SpannerOptions} instances. */
@@ -108,6 +112,7 @@ public static class Builder
108112
private int prefetchChunks = DEFAULT_PREFETCH_CHUNKS;
109113
private SessionPoolOptions sessionPoolOptions;
110114
private String userAgentPrefix;
115+
private ImmutableMap<String, String> sessionLabels;
111116

112117
private Builder() {}
113118

@@ -117,6 +122,7 @@ private Builder() {}
117122
this.sessionPoolOptions = options.sessionPoolOptions;
118123
this.prefetchChunks = options.prefetchChunks;
119124
this.userAgentPrefix = options.userAgent;
125+
this.sessionLabels = options.sessionLabels;
120126
}
121127

122128
@Override
@@ -152,6 +158,23 @@ public Builder setSessionPoolOption(SessionPoolOptions sessionPoolOptions) {
152158
return this;
153159
}
154160

161+
/**
162+
* Sets the labels to add to all Sessions created in this client.
163+
*
164+
* @param sessionLabels Map from label key to label value. Label key and value cannot be null.
165+
* For more information on valid syntax see
166+
* <a href="https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1#google.spanner.v1.Session">
167+
* api docs </a>.
168+
*/
169+
public Builder setSessionLabels(Map<String, String> sessionLabels) {
170+
Preconditions.checkNotNull(sessionLabels, "Session labels map cannot be null");
171+
for (String value : sessionLabels.values()) {
172+
Preconditions.checkNotNull(value, "Null values are not allowed in the labels map.");
173+
}
174+
this.sessionLabels = ImmutableMap.copyOf(sessionLabels);
175+
return this;
176+
}
177+
155178
/**
156179
* Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
157180
* PartialResultSet} chunks for each read and query. The data size of each chunk depends on the
@@ -205,6 +228,10 @@ public SessionPoolOptions getSessionPoolOptions() {
205228
return sessionPoolOptions;
206229
}
207230

231+
public Map<String, String> getSessionLabels() {
232+
return sessionLabels;
233+
}
234+
208235
public int getPrefetchChunks() {
209236
return prefetchChunks;
210237
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GrpcSpannerRpc.java‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import com.google.spanner.v1.PartialResultSet;
6262
import com.google.spanner.v1.ReadRequest;
6363
import com.google.spanner.v1.RollbackRequest;
64+
import com.google.spanner.v1.Session;
6465
import com.google.spanner.v1.SpannerGrpc;
6566
import com.google.spanner.v1.Transaction;
6667
import io.grpc.CallCredentials;
@@ -325,14 +326,18 @@ public Operation getOperation(String name) throws SpannerException {
325326
}
326327

327328
@Override
328-
public com.google.spanner.v1.Session createSession(
329-
String databaseName, @Nullable Map<Option, ?> options) {
330-
CreateSessionRequest request =
331-
CreateSessionRequest.newBuilder().setDatabase(databaseName).build();
329+
public Session createSession(
330+
String databaseName, @Nullable Map<String, String> labels, @Nullable Map<Option, ?> options) {
331+
CreateSessionRequest.Builder request =
332+
CreateSessionRequest.newBuilder().setDatabase(databaseName);
333+
if (labels != null && !labels.isEmpty()) {
334+
Session.Builder session = Session.newBuilder().putAllLabels(labels);
335+
request.setSession(session);
336+
}
332337
return get(
333338
doUnaryCall(
334339
SpannerGrpc.METHOD_CREATE_SESSION,
335-
request,
340+
request.build(),
336341
databaseName,
337342
Option.CHANNEL_HINT.getLong(options)));
338343
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ Operation updateDatabaseDdl(
188188
/** Retrieves a long running operation. */
189189
Operation getOperation(String name) throws SpannerException;
190190

191-
Session createSession(String databaseName, @Nullable Map<Option, ?> options)
191+
Session createSession(String databaseName, @Nullable Map<String, String> labels, @Nullable Map<Option, ?> options)
192192
throws SpannerException;
193193

194194
void deleteSession(String sessionName, @Nullable Map<Option, ?> options) throws SpannerException;

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,8 @@ public void setUp() {
7575
DatabaseId db = DatabaseId.of(dbName);
7676

7777
Session sessionProto = Session.newBuilder().setName(sessionName).build();
78-
Mockito.when(rpc.createSession(Mockito.eq(dbName), optionsCaptor.capture()))
78+
Mockito.when(rpc.createSession(Mockito.eq(dbName),
79+
Mockito.anyMapOf(String.class, String.class), optionsCaptor.capture()))
7980
.thenReturn(sessionProto);
8081
session = spanner.createSession(db);
8182
// We expect the same options, "options", on all calls on "session".

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import static com.google.common.truth.Truth.assertThat;
2020

2121
import com.google.cloud.spanner.spi.v1.SpannerRpc;
22+
23+
import java.util.HashMap;
2224
import java.util.Map;
2325
import org.junit.Before;
2426
import org.junit.Test;
@@ -47,13 +49,20 @@ public void setUp() {
4749

4850
@Test
4951
public void createAndCloseSession() {
52+
Map<String, String> labels = new HashMap<>();
53+
labels.put("env", "dev");
54+
Mockito.when(spannerOptions.getSessionLabels()).thenReturn(labels);
5055
String dbName = "projects/p1/instances/i1/databases/d1";
5156
String sessionName = dbName + "/sessions/s1";
5257
DatabaseId db = DatabaseId.of(dbName);
5358

5459
com.google.spanner.v1.Session sessionProto =
55-
com.google.spanner.v1.Session.newBuilder().setName(sessionName).build();
56-
Mockito.when(rpc.createSession(Mockito.eq(dbName), options.capture())).thenReturn(sessionProto);
60+
com.google.spanner.v1.Session.newBuilder()
61+
.setName(sessionName)
62+
.putAllLabels(labels)
63+
.build();
64+
Mockito.when(rpc.createSession(Mockito.eq(dbName), Mockito.eq(labels), options.capture()))
65+
.thenReturn(sessionProto);
5766
Session session = impl.createSession(db);
5867
assertThat(session.getName()).isEqualTo(sessionName);
5968

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@
1818

1919
import static com.google.common.truth.Truth.assertThat;
2020

21+
import java.util.HashMap;
22+
import java.util.Map;
23+
2124
import com.google.cloud.TransportOptions;
25+
2226
import io.grpc.ManagedChannel;
2327
import io.grpc.netty.NettyChannelBuilder;
2428
import org.junit.Rule;
@@ -54,27 +58,46 @@ public void defaultBuilder() {
5458
.build();
5559
assertThat(options.getHost()).isEqualTo("https://spanner.googleapis.com");
5660
assertThat(options.getPrefetchChunks()).isEqualTo(4);
61+
assertThat(options.getSessionLabels()).isNull();
5762
}
5863

5964
@Test
6065
public void builder() {
6166
String host = "http://localhost:8000/";
6267
String projectId = "test-project";
68+
Map<String, String> labels = new HashMap<>();
69+
labels.put("env", "dev");
6370
SpannerOptions options =
6471
SpannerOptions.newBuilder()
6572
.setRpcChannelFactory(new TestChannelFactory())
6673
.setHost(host)
6774
.setProjectId(projectId)
6875
.setPrefetchChunks(2)
76+
.setSessionLabels(labels)
6977
.build();
7078
assertThat(options.getHost()).isEqualTo(host);
7179
assertThat(options.getProjectId()).isEqualTo(projectId);
7280
assertThat(options.getPrefetchChunks()).isEqualTo(2);
81+
assertThat(options.getSessionLabels()).containsExactlyEntriesIn(labels);
7382
}
7483

7584
@Test
7685
public void testInvalidTransport() {
7786
thrown.expect(IllegalArgumentException.class);
7887
SpannerOptions.newBuilder().setTransportOptions(Mockito.mock(TransportOptions.class));
7988
}
89+
90+
@Test
91+
public void testInvalidSessionLabels() {
92+
thrown.expect(NullPointerException.class);
93+
Map<String, String> labels = new HashMap<>();
94+
labels.put("env", null);
95+
SpannerOptions.newBuilder().setSessionLabels(labels);
96+
}
97+
98+
@Test
99+
public void testNullSessionLabels() {
100+
thrown.expect(NullPointerException.class);
101+
SpannerOptions.newBuilder().setSessionLabels(null);
102+
}
80103
}

0 commit comments

Comments
 (0)