Skip to content

Commit ffdd1f1

Browse files
mrzzyZhu Zhanyan
mrzzy
and
Zhu Zhanyan
authored
Fix Feast Serving not registering its store in Feast Core (#641)
* Fixed missing call in serving's CachedSpecService to register store in core. * Add unit test to check that cachedSpecService will register store with Core. Co-authored-by: Zhu Zhanyan <zhu.zhanyan@gojek.com>
1 parent 0816cd4 commit ffdd1f1

3 files changed

Lines changed: 32 additions & 7 deletions

File tree

‎serving/src/main/java/feast/serving/specs/CachedSpecService.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@
4747
import org.apache.commons.lang3.tuple.Pair;
4848
import org.slf4j.Logger;
4949

50-
/** In-memory cache of specs. */
50+
/** In-memory cache of specs hosted in Feast Core. */
5151
public class CachedSpecService {
5252

5353
private static final int MAX_SPEC_COUNT = 1000;
@@ -76,7 +76,7 @@ public class CachedSpecService {
7676

7777
public CachedSpecService(CoreSpecService coreService, StoreProto.Store store) {
7878
this.coreService = coreService;
79-
this.store = store;
79+
this.store = coreService.registerStore(store);
8080

8181
Map<String, FeatureSetSpec> featureSets = getFeatureSetMap();
8282
featureToFeatureSetMapping =

‎serving/src/main/java/feast/serving/specs/CoreSpecService.java‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,12 @@
2323
import feast.core.CoreServiceProto.ListFeatureSetsResponse;
2424
import feast.core.CoreServiceProto.UpdateStoreRequest;
2525
import feast.core.CoreServiceProto.UpdateStoreResponse;
26+
import feast.core.StoreProto.Store;
2627
import io.grpc.ManagedChannel;
2728
import io.grpc.ManagedChannelBuilder;
2829
import org.slf4j.Logger;
2930

30-
/** Client for spec retrieval from core. */
31+
/** Client for interfacing with specs in Feast Core. */
3132
public class CoreSpecService {
3233

3334
private static final Logger log = org.slf4j.LoggerFactory.getLogger(CoreSpecService.class);
@@ -50,4 +51,24 @@ public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest ListFeatur
5051
public UpdateStoreResponse updateStore(UpdateStoreRequest updateStoreRequest) {
5152
return blockingStub.updateStore(updateStoreRequest);
5253
}
54+
55+
/**
56+
* Register the given store entry in Feast Core. If store already exists in Feast Core, updates
57+
* the store entry in feast core.
58+
*
59+
* @param store entry to register/update in Feast Core.
60+
* @return The register/updated store entry
61+
*/
62+
public Store registerStore(Store store) {
63+
UpdateStoreRequest request = UpdateStoreRequest.newBuilder().setStore(store).build();
64+
try {
65+
UpdateStoreResponse updateStoreResponse = this.updateStore(request);
66+
if (!updateStoreResponse.getStore().equals(store)) {
67+
throw new RuntimeException("Core store config not matching current store config");
68+
}
69+
return updateStoreResponse.getStore();
70+
} catch (Exception e) {
71+
throw new RuntimeException("Unable to update store configuration", e);
72+
}
73+
}
5374
}

‎serving/src/test/java/feast/serving/service/CachedSpecServiceTest.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,14 @@
1919
import static org.hamcrest.CoreMatchers.equalTo;
2020
import static org.hamcrest.Matchers.containsInAnyOrder;
2121
import static org.junit.Assert.assertThat;
22+
import static org.mockito.Mockito.times;
23+
import static org.mockito.Mockito.verify;
2224
import static org.mockito.Mockito.when;
2325
import static org.mockito.MockitoAnnotations.initMocks;
2426

2527
import com.google.common.collect.Lists;
2628
import feast.core.CoreServiceProto.ListFeatureSetsRequest;
2729
import feast.core.CoreServiceProto.ListFeatureSetsResponse;
28-
import feast.core.CoreServiceProto.UpdateStoreRequest;
29-
import feast.core.CoreServiceProto.UpdateStoreResponse;
3030
import feast.core.FeatureSetProto;
3131
import feast.core.FeatureSetProto.FeatureSetSpec;
3232
import feast.core.FeatureSetProto.FeatureSpec;
@@ -82,8 +82,7 @@ public void setUp() {
8282
.build())
8383
.build();
8484

85-
when(coreService.updateStore(UpdateStoreRequest.newBuilder().setStore(store).build()))
86-
.thenReturn(UpdateStoreResponse.newBuilder().setStore(store).build());
85+
when(coreService.registerStore(store)).thenReturn(store);
8786

8887
featureSetSpecs = new LinkedHashMap<>();
8988
featureSetSpecs.put(
@@ -143,6 +142,11 @@ public void setUp() {
143142
cachedSpecService = new CachedSpecService(coreService, store);
144143
}
145144

145+
@Test
146+
public void shouldRegisterStoreWithCore() {
147+
verify(coreService, times(1)).registerStore(cachedSpecService.getStore());
148+
}
149+
146150
@Test
147151
public void shouldPopulateAndReturnStore() {
148152
cachedSpecService.populateCache();

0 commit comments

Comments
 (0)