Skip to content

Commit e37ebc8

Browse files
author
zhilingc
committed
Add UpdateStores method to core
1 parent 6b5c5ad commit e37ebc8

3 files changed

Lines changed: 96 additions & 23 deletions

File tree

‎core/src/main/java/feast/core/grpc/CoreServiceImpl.java‎

Lines changed: 43 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -29,32 +29,39 @@
2929
import feast.core.CoreServiceProto.GetStoresRequest;
3030
import feast.core.CoreServiceProto.GetStoresRequest.Filter;
3131
import feast.core.CoreServiceProto.GetStoresResponse;
32+
import feast.core.CoreServiceProto.UpdateStoreRequest;
33+
import feast.core.CoreServiceProto.UpdateStoreResponse;
3234
import feast.core.FeatureSetProto.FeatureSetSpec;
3335
import feast.core.StoreProto.Store;
3436
import feast.core.StoreProto.Store.Subscription;
3537
import feast.core.exception.RetrievalException;
3638
import feast.core.service.JobCoordinatorService;
3739
import feast.core.service.SpecService;
38-
import io.grpc.Status;
39-
import io.grpc.StatusRuntimeException;
4040
import io.grpc.stub.StreamObserver;
4141
import java.util.ArrayList;
42+
import java.util.HashSet;
4243
import java.util.List;
44+
import java.util.Set;
4345
import java.util.regex.Pattern;
4446
import java.util.stream.Collectors;
4547
import lombok.extern.slf4j.Slf4j;
4648
import org.lognet.springboot.grpc.GRpcService;
4749
import org.springframework.beans.factory.annotation.Autowired;
4850
import org.springframework.transaction.annotation.Transactional;
4951

50-
/** Implementation of the feast core GRPC service. */
52+
/**
53+
* Implementation of the feast core GRPC service.
54+
*/
5155
@Slf4j
5256
@GRpcService
5357
public class CoreServiceImpl extends CoreServiceImplBase {
5458

55-
@Autowired private SpecService specService;
56-
@Autowired private StatsDClient statsDClient;
57-
@Autowired private JobCoordinatorService jobCoordinatorService;
59+
@Autowired
60+
private SpecService specService;
61+
@Autowired
62+
private StatsDClient statsDClient;
63+
@Autowired
64+
private JobCoordinatorService jobCoordinatorService;
5865

5966
@Override
6067
public void getFeastCoreVersion(
@@ -72,7 +79,7 @@ public void getFeatureSets(
7279
responseObserver.onNext(response);
7380
responseObserver.onCompleted();
7481
} catch (RetrievalException | InvalidProtocolBufferException e) {
75-
responseObserver.onError(getRuntimeException(e));
82+
responseObserver.onError(e);
7683
}
7784
}
7885

@@ -85,7 +92,7 @@ public void getStores(
8592
responseObserver.onNext(response);
8693
responseObserver.onCompleted();
8794
} catch (RetrievalException e) {
88-
responseObserver.onError(getRuntimeException(e));
95+
responseObserver.onError(e);
8996
}
9097
}
9198

@@ -124,17 +131,37 @@ public void applyFeatureSet(
124131
responseObserver.onNext(response);
125132
responseObserver.onCompleted();
126133
} catch (Exception e) {
127-
responseObserver.onError(getRuntimeException(e));
134+
responseObserver.onError(e);
128135
}
129136
}
130137

131-
private StatusRuntimeException getRuntimeException(Exception e) {
132-
return new StatusRuntimeException(
133-
Status.fromCode(Status.Code.INTERNAL).withDescription(e.getMessage()).withCause(e));
134-
}
138+
@Override
139+
public void updateStore(UpdateStoreRequest request,
140+
StreamObserver<UpdateStoreResponse> responseObserver) {
141+
try {
142+
UpdateStoreResponse response = specService.updateStore(request);
143+
responseObserver.onNext(response);
144+
responseObserver.onCompleted();
135145

136-
private StatusRuntimeException getBadRequestException(Exception e) {
137-
return new StatusRuntimeException(
138-
Status.fromCode(Status.Code.OUT_OF_RANGE).withDescription(e.getMessage()).withCause(e));
146+
Set<FeatureSetSpec> featureSetSpecs = new HashSet<>();
147+
Store store = response.getStore();
148+
for (Subscription subscription : store.getSubscriptionsList()) {
149+
featureSetSpecs.addAll(
150+
specService.getFeatureSets(
151+
GetFeatureSetsRequest.Filter.newBuilder()
152+
.setFeatureSetName(subscription.getName())
153+
.setFeatureSetVersion(subscription.getVersion())
154+
.build())
155+
.getFeatureSetsList()
156+
);
157+
}
158+
featureSetSpecs.stream()
159+
.collect(Collectors.groupingBy(FeatureSetSpec::getName))
160+
.entrySet()
161+
.stream()
162+
.forEach(kv -> jobCoordinatorService.startOrUpdateJob(kv.getValue(), store));
163+
} catch (Exception e) {
164+
responseObserver.onError(e);
165+
}
139166
}
140167
}

‎core/src/main/java/feast/core/service/SpecService.java‎

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,10 @@
2626
import feast.core.CoreServiceProto.GetStoresRequest;
2727
import feast.core.CoreServiceProto.GetStoresResponse;
2828
import feast.core.CoreServiceProto.GetStoresResponse.Builder;
29+
import feast.core.CoreServiceProto.UpdateStoreRequest;
30+
import feast.core.CoreServiceProto.UpdateStoreResponse;
2931
import feast.core.FeatureSetProto.FeatureSetSpec;
32+
import feast.core.StoreProto;
3033
import feast.core.dao.FeatureSetRepository;
3134
import feast.core.dao.StoreRepository;
3235
import feast.core.exception.RetrievalException;
@@ -134,7 +137,10 @@ public GetStoresResponse getStores(GetStoresRequest.Filter filter) {
134137
.addStore(store.toProto())
135138
.build();
136139
} catch (InvalidProtocolBufferException e) {
137-
throw new RetrievalException("Unable to retrieve stores", e);
140+
throw io.grpc.Status.NOT_FOUND
141+
.withDescription("Unable to retrieve stores")
142+
.withCause(e)
143+
.asRuntimeException();
138144
}
139145
}
140146

@@ -188,6 +194,32 @@ public ApplyFeatureSetResponse applyFeatureSet(FeatureSetSpec newFeatureSetSpec)
188194
.build();
189195
}
190196

197+
/**
198+
* UpdateStore updates the repository with the new given store.
199+
*
200+
* @param updateStoreRequest containing the new store definition
201+
* @return UpdateStoreResponse containing the new store definition
202+
* @throws InvalidProtocolBufferException
203+
*/
204+
public UpdateStoreResponse updateStore(UpdateStoreRequest updateStoreRequest)
205+
throws InvalidProtocolBufferException {
206+
StoreProto.Store newStoreProto = updateStoreRequest.getStore();
207+
Store existingStore = storeRepository.findById(newStoreProto.getName()).orElse(null);
208+
209+
// Do nothing if no change
210+
if (existingStore != null && existingStore.toProto().equals(newStoreProto)) {
211+
return UpdateStoreResponse.newBuilder()
212+
.setStore(updateStoreRequest.getStore())
213+
.build();
214+
}
215+
216+
Store newStore = Store.fromProto(newStoreProto);
217+
storeRepository.save(newStore);
218+
return UpdateStoreResponse.newBuilder()
219+
.setStore(updateStoreRequest.getStore())
220+
.build();
221+
}
222+
191223
private Predicate<? super FeatureSet> getVersionFilter(String versionFilter) {
192224
if (versionFilter.equals("")) {
193225
return v -> true;
@@ -196,11 +228,12 @@ private Predicate<? super FeatureSet> getVersionFilter(String versionFilter) {
196228
match.find();
197229

198230
if (!match.matches()) {
199-
throw new RetrievalException(
200-
String.format(
231+
throw io.grpc.Status.INVALID_ARGUMENT
232+
.withDescription(String.format(
201233
"Invalid version string '%s' provided. Version string may either "
202234
+ "be a fixed version, e.g. 10, or contain a comparator, e.g. >10.",
203-
versionFilter));
235+
versionFilter))
236+
.asRuntimeException();
204237
}
205238

206239
int versionNumber = Integer.valueOf(match.group("version"));
@@ -217,11 +250,12 @@ private Predicate<? super FeatureSet> getVersionFilter(String versionFilter) {
217250
case "":
218251
return v -> v.getVersion() == versionNumber;
219252
default:
220-
throw new RetrievalException(
221-
String.format(
253+
throw io.grpc.Status.INVALID_ARGUMENT
254+
.withDescription(String.format(
222255
"Invalid comparator '%s' provided. Version string may either "
223256
+ "be a fixed version, e.g. 10, or contain a comparator, e.g. >10.",
224-
comparator));
257+
comparator))
258+
.asRuntimeException();
225259
}
226260
}
227261

‎protos/feast/core/CoreService.proto‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ service CoreService {
3939

4040
// Idempotent creation of feature set. Will not create a new feature set if schema does not change
4141
rpc ApplyFeatureSet (ApplyFeatureSetRequest) returns (ApplyFeatureSetResponse);
42+
43+
// UpdateStore updates core with the configuration of the store. If the changes are valid,
44+
// core will return the given store configuration in response.
45+
rpc UpdateStore(UpdateStoreRequest) returns (UpdateStoreResponse);
4246
}
4347

4448
// Retrieves details for all versions of a specific feature set
@@ -94,3 +98,11 @@ message GetFeastCoreVersionRequest {}
9498
message GetFeastCoreVersionResponse {
9599
string version = 1;
96100
}
101+
102+
message UpdateStoreRequest {
103+
feast.core.Store store = 1;
104+
}
105+
106+
message UpdateStoreResponse {
107+
feast.core.Store store = 1;
108+
}

0 commit comments

Comments
 (0)