Skip to content

Commit fe905c1

Browse files
igorbernstein2igorbernstein2
committed
03. Surface: SampleRowKeys
Adds the surface to wrap SampleRowKeys RPC. The surface: - hides the fact that the RPC is streaming - wraps the request as a simple string table id - wraps the response in KeyOffset
1 parent dfa1972 commit fe905c1

7 files changed

Lines changed: 259 additions & 0 deletions

File tree

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,14 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2;
1717

18+
import com.google.api.core.ApiFuture;
1819
import com.google.api.core.InternalApi;
20+
import com.google.api.gax.rpc.UnaryCallable;
1921
import com.google.bigtable.admin.v2.InstanceName;
2022
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
23+
import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
2124
import java.io.IOException;
25+
import java.util.List;
2226

2327
/**
2428
* Client for reading from and writing to existing Bigtable tables.
@@ -112,4 +116,46 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
112116
public void close() throws Exception {
113117
stub.close();
114118
}
119+
120+
/**
121+
* Convenience method to asynchronously return a sample of row keys in the table. The returned row
122+
* keys will delimit contiguous sections of the table of approximately equal size, which can be
123+
* used to break up the data for distributed tasks like mapreduces. The returned callable object
124+
* allows for customization of api invocation.
125+
*
126+
* <p>Sample code:
127+
*
128+
* <pre>{@code
129+
* InstanceName instanceName = InstanceName.of("[PROJECT]", "[INSTANCE]");
130+
* try (BigtableClient bigtableClient = BigtableClient.create(instanceName)) {
131+
* ApiFuture<List<KeyOffset>> keyOffsets = bigtableClient.sampleRowKeys("[TABLE]");
132+
* }
133+
* }</pre>
134+
*/
135+
public ApiFuture<List<KeyOffset>> sampleRowKeys(String tableId) {
136+
return sampleRowKeysCallable().futureCall(tableId);
137+
}
138+
139+
/**
140+
* Returns a sample of row keys in the table. The returned row keys will delimit contiguous
141+
* sections of the table of approximately equal size, which can be used to break up the data for
142+
* distributed tasks like mapreduces. The returned callable object allows for customization of api
143+
* invocation.
144+
*
145+
* <p>Sample code:
146+
*
147+
* <pre>{@code
148+
* InstanceName instanceName = InstanceName.of("[PROJECT]", "[INSTANCE]");
149+
* try (BigtableClient bigtableClient = BigtableClient.create(instanceName)) {
150+
* // Synchronous invocation
151+
* List<KeyOffset> keyOffsets = bigtableClient.sampleRowKeysCallable().call("[TABLE]");
152+
*
153+
* // Asynchronous invocation
154+
* ApiFuture<List<KeyOffset>> keyOffsets = bigtableClient.sampleRowKeysCallable().futureCall("[TABLE]");
155+
* }
156+
* }</pre>
157+
*/
158+
public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
159+
return stub.sampleRowKeysCallable();
160+
}
115161
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataSettings.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,12 @@
1616
package com.google.cloud.bigtable.data.v2;
1717

1818
import com.google.api.gax.rpc.ClientSettings;
19+
import com.google.api.gax.rpc.UnaryCallSettings;
1920
import com.google.bigtable.admin.v2.InstanceName;
2021
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
22+
import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
2123
import java.io.IOException;
24+
import java.util.List;
2225
import javax.annotation.Nonnull;
2326

2427
/**
@@ -64,6 +67,11 @@ public String getAppProfileId() {
6467
return getTypedStubSettings().getAppProfileId();
6568
}
6669

70+
/** Returns the object with the settings used for calls to sampleRowKeys. */
71+
public UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings() {
72+
return getTypedStubSettings().sampleRowKeysSettings();
73+
}
74+
6775
@SuppressWarnings("unchecked")
6876
EnhancedBigtableStubSettings getTypedStubSettings() {
6977
return (EnhancedBigtableStubSettings) getStubSettings();
@@ -124,6 +132,11 @@ public String getAppProfileId() {
124132
return getTypedStubSettings().getAppProfileId();
125133
}
126134

135+
/** Returns the builder for the settings used for calls to SampleRowKeysSettings. */
136+
public UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings() {
137+
return getTypedStubSettings().sampleRowKeysSettings();
138+
}
139+
127140
@SuppressWarnings("unchecked")
128141
private EnhancedBigtableStubSettings.Builder getTypedStubSettings() {
129142
return (EnhancedBigtableStubSettings.Builder) getStubSettings();

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,15 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.stub;
1717

18+
import com.google.api.core.ApiFuture;
1819
import com.google.api.core.InternalApi;
20+
import com.google.api.gax.rpc.ApiCallContext;
1921
import com.google.api.gax.rpc.ClientContext;
22+
import com.google.api.gax.rpc.UnaryCallable;
2023
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
24+
import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
2125
import java.io.IOException;
26+
import java.util.List;
2227

2328
/**
2429
* The core client that converts method calls to RPCs.
@@ -39,6 +44,8 @@ public class EnhancedBigtableStub implements AutoCloseable {
3944
private final ClientContext clientContext;
4045
private final RequestContext requestContext;
4146

47+
private final UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable;
48+
4249
public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
4350
throws IOException {
4451
// Configure the base settings
@@ -63,6 +70,8 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
6370
this.stub = stub;
6471
this.requestContext =
6572
RequestContext.create(settings.getInstanceName(), settings.getAppProfileId());
73+
74+
sampleRowKeysCallable = createSampleRowKeysCallable();
6675
}
6776

6877
@Override
@@ -71,8 +80,31 @@ public void close() throws Exception {
7180
}
7281

7382
// <editor-fold desc="Callable creators">
83+
84+
/**
85+
* Creates a callable chain to handle SampleRowKeys RPcs. The chain will:
86+
*
87+
* <ul>
88+
* <li>Convert a table id to a {@link com.google.bigtable.v2.SampleRowKeysRequest}.
89+
* <li>Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}.
90+
* <li>Spool responses into a list.
91+
* <li>Retry on failure.
92+
* <li>Convert the responses into {@link KeyOffset}s.
93+
* </ul>
94+
*/
95+
private UnaryCallable<String, List<KeyOffset>> createSampleRowKeysCallable() {
96+
return new UnaryCallable<String, List<KeyOffset>>() {
97+
@Override
98+
public ApiFuture<List<KeyOffset>> futureCall(String request, ApiCallContext context) {
99+
throw new UnsupportedOperationException("todo");
100+
}
101+
};
102+
}
74103
// </editor-fold>
75104

76105
// <editor-fold desc="Callable accessors">
106+
public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
107+
return sampleRowKeysCallable;
108+
}
77109
// </editor-fold>
78110
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,16 @@
1616
package com.google.cloud.bigtable.data.v2.stub;
1717

1818
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
19+
import com.google.api.gax.retrying.RetrySettings;
20+
import com.google.api.gax.rpc.StatusCode.Code;
1921
import com.google.api.gax.rpc.StubSettings;
2022
import com.google.api.gax.rpc.UnaryCallSettings;
2123
import com.google.bigtable.admin.v2.InstanceName;
24+
import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
2225
import com.google.common.base.Preconditions;
26+
import java.util.List;
2327
import javax.annotation.Nonnull;
28+
import org.threeten.bp.Duration;
2429

2530
/**
2631
* Settings class to configure an instance of {@link EnhancedBigtableStub}.
@@ -53,12 +58,15 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
5358
private final InstanceName instanceName;
5459
private final String appProfileId;
5560

61+
private final UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings;
62+
5663
private EnhancedBigtableStubSettings(Builder builder) {
5764
super(builder);
5865
instanceName = builder.instanceName;
5966
appProfileId = builder.appProfileId;
6067

6168
// Per method settings.
69+
sampleRowKeysSettings = builder.sampleRowKeysSettings.build();
6270
}
6371

6472
public static Builder newBuilder() {
@@ -75,6 +83,11 @@ public String getAppProfileId() {
7583
return appProfileId;
7684
}
7785

86+
/** Returns the object with the settings used for calls to SampleRowKeys. */
87+
public UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings() {
88+
return sampleRowKeysSettings;
89+
}
90+
7891
/** Returns a builder containing all the values of this settings class. */
7992
public Builder toBuilder() {
8093
return new Builder(this);
@@ -85,6 +98,8 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
8598
private InstanceName instanceName;
8699
private String appProfileId;
87100

101+
private final UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings;
102+
88103
/**
89104
* Initializes a new Builder with sane defaults for all settings.
90105
*
@@ -109,6 +124,22 @@ private Builder() {
109124
.build());
110125

111126
// Per-method settings using baseSettings for defaults.
127+
128+
sampleRowKeysSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
129+
/* TODO: copy retryCodes & retrySettings from baseSettings.sampleRowKeysSettings once it exists in GAPIC */
130+
sampleRowKeysSettings
131+
.setRetryableCodes(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED)
132+
.setRetrySettings(
133+
RetrySettings.newBuilder()
134+
.setMaxAttempts(10)
135+
.setTotalTimeout(Duration.ofMinutes(1))
136+
.setInitialRetryDelay(Duration.ofMillis(100))
137+
.setRetryDelayMultiplier(1.3)
138+
.setMaxRetryDelay(Duration.ofMinutes(1))
139+
.setInitialRpcTimeout(Duration.ofSeconds(20))
140+
.setRpcTimeoutMultiplier(1)
141+
.setMaxRpcTimeout(Duration.ofSeconds(20))
142+
.build());
112143
}
113144

114145
private Builder(EnhancedBigtableStubSettings settings) {
@@ -117,6 +148,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
117148
appProfileId = settings.appProfileId;
118149

119150
// Per method settings.
151+
sampleRowKeysSettings = settings.sampleRowKeysSettings.toBuilder();
120152
}
121153

122154
// <editor-fold desc="Private Helpers">
@@ -164,6 +196,11 @@ public String getAppProfileId() {
164196
return appProfileId;
165197
}
166198

199+
/** Returns the builder for the settings used for calls to SampleRowKeysSettings. */
200+
public UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings() {
201+
return sampleRowKeysSettings;
202+
}
203+
167204
@SuppressWarnings("unchecked")
168205
public EnhancedBigtableStubSettings build() {
169206
Preconditions.checkState(instanceName != null, "InstanceName must be set");

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/wrappers/KeyOffset.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright 2018 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.wrappers;
17+
18+
import com.google.api.core.InternalApi;
19+
import com.google.auto.value.AutoValue;
20+
import com.google.protobuf.ByteString;
21+
22+
/** Represents the offset of a row key in a table. */
23+
@AutoValue
24+
public abstract class KeyOffset {
25+
@InternalApi
26+
public static KeyOffset create(ByteString key, long offsetBytes) {
27+
return new AutoValue_KeyOffset(key, offsetBytes);
28+
}
29+
30+
/**
31+
* Sorted streamed sequence of sample row keys in the table. The table might have contents before
32+
* the first row key in the list and after the last one, but a key containing the empty string
33+
* indicates "end of table" and will be the last response given, if present.
34+
*
35+
* <p>Note that row keys in this list may not have ever been written to or read from, and users
36+
* should therefore not make any assumptions about the row key structure that are specific to
37+
* their use case.
38+
*/
39+
public abstract ByteString key();
40+
41+
/**
42+
* Approximate total storage space used by all rows in the table which precede {@link #key()}.
43+
* Buffering the contents of all rows between two subsequent samples would require space roughly
44+
* equal to the difference in their {@link #offsetBytes()} fields.
45+
*/
46+
public abstract long offsetBytes();
47+
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTest.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,12 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2;
1717

18+
import static com.google.common.truth.Truth.assertThat;
19+
20+
import com.google.api.gax.rpc.UnaryCallable;
1821
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
22+
import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
23+
import java.util.List;
1924
import org.junit.Before;
2025
import org.junit.Test;
2126
import org.junit.runner.RunWith;
@@ -26,11 +31,13 @@
2631
@RunWith(MockitoJUnitRunner.class)
2732
public class BigtableDataClientTest {
2833
@Mock private EnhancedBigtableStub mockStub;
34+
@Mock private UnaryCallable<String, List<KeyOffset>> mockSampleRowKeysCallable;
2935

3036
private BigtableDataClient bigtableDataClient;
3137

3238
@Before
3339
public void setUp() {
40+
Mockito.when(mockStub.sampleRowKeysCallable()).thenReturn(mockSampleRowKeysCallable);
3441
bigtableDataClient = new BigtableDataClient(mockStub);
3542
}
3643

@@ -39,4 +46,15 @@ public void proxyCloseTest() throws Exception {
3946
bigtableDataClient.close();
4047
Mockito.verify(mockStub).close();
4148
}
49+
50+
@Test
51+
public void proxySampleRowKeysCallableTest() {
52+
assertThat(bigtableDataClient.sampleRowKeysCallable()).isSameAs(mockSampleRowKeysCallable);
53+
}
54+
55+
@Test
56+
public void proxySampleRowKeysTest() {
57+
bigtableDataClient.sampleRowKeys("fake-table");
58+
Mockito.verify(mockSampleRowKeysCallable).futureCall("fake-table");
59+
}
4260
}

0 commit comments

Comments
 (0)