JobControllerService { Iterable<Pair<Source, Set<Store>>> getSourceToStoreMappings() { Stream<Pair<Source, Store>> distinctPairs = getAllStores().stream() .flatMap( store -> getFeatureSetsForStore(store).stream() .map(f -> f.getSpec().getSource()) .map(source -> Pair.of(source, store))) .distinct(); return groupingStrategy.collectSingleJobInput(distinctPairs); } @Autowired  JobControllerService(
      JobRepository jobRepository,
      CoreServiceGrpc.CoreServiceBlockingStub specService,
      JobManager jobManager,
      FeastProperties feastProperties,
      JobGroupingStrategy groupingStrategy,
      KafkaTemplate<String, FeatureSetSpec> specPublisher); @Scheduled(fixedDelayString = "${feast.jobs.polling_interval_milliseconds}") void Poll(); @Scheduled(fixedDelayString = "${feast.stream.specsOptions.notifyIntervalMilliseconds}") void notifyJobsWhenFeatureSetUpdated(); @KafkaListener( topics = {"${feast.stream.specsOptions.specsAckTopic}"}, containerFactory = "kafkaAckListenerContainerFactory") void listenAckFromJobs(
      ConsumerRecord<String, IngestionJobProto.FeatureSetSpecAck> record); static final String VERSION_LABEL; }
@Test @SneakyThrows public void shouldGroupJobsBySource() { Store store = DataGenerator.createStore( "test", Store.StoreType.REDIS, ImmutableList.of(Triple.of("project1", "*", false))); Source source1 = DataGenerator.createSource("servers:9092", "topic"); Source source2 = DataGenerator.createSource("others.servers:9092", "topic"); FeatureSet featureSet1 = DataGenerator.createFeatureSet(source1, "project1", "fs1"); FeatureSet featureSet2 = DataGenerator.createFeatureSet(source2, "project1", "fs2"); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter( Filter.newBuilder().setFeatureSetName("*").setProject("project1").build()) .build())) .thenReturn( ListFeatureSetsResponse.newBuilder() .addAllFeatureSets(Lists.newArrayList(featureSet1, featureSet2)) .build()); when(specService.listStores(any())) .thenReturn(ListStoresResponse.newBuilder().addStore(store).build()); ArrayList<Pair<Source, Set<Store>>> pairs = Lists.newArrayList(controllerWithConsolidation.getSourceToStoreMappings()); assertThat(pairs, hasSize(2)); assertThat(pairs, hasItem(Pair.of(source1, Sets.newHashSet(store)))); assertThat(pairs, hasItem(Pair.of(source2, Sets.newHashSet(store)))); }
@Test @SneakyThrows public void shouldUseStoreSubscriptionToMapStore() { Store store1 = DataGenerator.createStore( "test", Store.StoreType.REDIS, ImmutableList.of(Triple.of("*", "features1", false))); Store store2 = DataGenerator.createStore( "test", Store.StoreType.REDIS, ImmutableList.of(Triple.of("*", "features2", false))); Source source1 = DataGenerator.createSource("servers:9092", "topic"); Source source2 = DataGenerator.createSource("other.servers:9092", "topic"); FeatureSet featureSet1 = DataGenerator.createFeatureSet(source1, "default", "fs1"); FeatureSet featureSet2 = DataGenerator.createFeatureSet(source2, "default", "fs2"); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter( Filter.newBuilder().setFeatureSetName("features1").setProject("*").build()) .build())) .thenReturn( ListFeatureSetsResponse.newBuilder() .addAllFeatureSets(Lists.newArrayList(featureSet1)) .build()); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter( Filter.newBuilder().setFeatureSetName("features2").setProject("*").build()) .build())) .thenReturn( ListFeatureSetsResponse.newBuilder() .addAllFeatureSets(Lists.newArrayList(featureSet2)) .build()); when(specService.listStores(any())) .thenReturn(ListStoresResponse.newBuilder().addStore(store1).addStore(store2).build()); ArrayList<Pair<Source, Set<Store>>> pairs = Lists.newArrayList(controllerWithConsolidation.getSourceToStoreMappings()); assertThat(pairs, hasSize(2)); assertThat(pairs, hasItem(Pair.of(source1, Sets.newHashSet(store1)))); assertThat(pairs, hasItem(Pair.of(source2, Sets.newHashSet(store2)))); }
@Test @SneakyThrows public void shouldSelectOnlyStoresThatWhitelisted() { Store store1 = DataGenerator.createStore( "normal-store", Store.StoreType.REDIS, ImmutableList.of(Triple.of("project1", "*", false))); Store store2 = DataGenerator.createStore( "blacklisted-store", Store.StoreType.REDIS, ImmutableList.of(Triple.of("project2", "*", false))); Source source1 = DataGenerator.createSource("source-1", "topic"); Source source2 = DataGenerator.createSource("source-2", "topic"); FeatureSet featureSet1 = DataGenerator.createFeatureSet(source1, "default", "fs1"); FeatureSet featureSet2 = DataGenerator.createFeatureSet(source2, "project", "fs3"); when(specService.listStores(any())) .thenReturn(ListStoresResponse.newBuilder().addStore(store1).addStore(store2).build()); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter( Filter.newBuilder().setProject("project1").setFeatureSetName("*").build()) .build())) .thenReturn(ListFeatureSetsResponse.newBuilder().addFeatureSets(featureSet1).build()); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter( Filter.newBuilder().setProject("project2").setFeatureSetName("*").build()) .build())) .thenReturn(ListFeatureSetsResponse.newBuilder().addFeatureSets(featureSet2).build()); ArrayList<Pair<Source, Set<Store>>> pairs = Lists.newArrayList(controllerWithConsolidation.getSourceToStoreMappings()); assertThat(pairs, containsInAnyOrder(Pair.of(source1, ImmutableSet.of(store1)))); }