Address review comments

Author: piyush-zlai

.github/workflows/test_scala_no_spark.yaml

@@ -66,3 +66,8 @@ jobs:
         run: |
           export SBT_OPTS="-Xmx8G -Xms2G"
           sbt "++ 2.12.18 hub/test"
+
+      - name: Run service tests
+        run: |
+          export SBT_OPTS="-Xmx8G -Xms2G"
+          sbt "++ 2.12.18 service/test"

service/README.md

@@ -28,7 +28,7 @@ This command will write out a file in the target/scala-2.12 sub-directory.
 
 We can now use this to start up the feature service:
 ```bash
-~/workspace/chronon $ java -jar service/target/scala-2.12/service-*.jar run ai.chronon.service.WebServiceVerticle \
+~/workspace/chronon $ java -jar service/target/scala-2.12/service-*.jar run ai.chronon.service.FetcherVerticle \
 -Dserver.port=9000 -conf service/src/main/resources/example_config.json
 ...
 14:39:26.626 [vert.x-eventloop-thread-1] INFO  a.chronon.service.WebServiceVerticle - HTTP server started on port 9000
@@ -46,7 +46,7 @@ Some examples to curl the webservice:
 ```bash
 $ curl 'http://localhost:9000/ping'
 $ curl 'http://localhost:9000/config'
-$ curl -X POST   'http://localhost:9000/v1/features/join/quickstart%2Ftraining_set.v2'   -H 'Content-Type: application/json'   -d '[{"user_id": "5"}]'
+$ curl -X POST   'http://localhost:9000/v1/fetch/join/quickstart%2Ftraining_set.v2'   -H 'Content-Type: application/json'   -d '[{"user_id": "5"}]'
 ```
 
 ## Metrics
@@ -77,9 +77,9 @@ StatsD Metric: ai.zipline.join.fetch.group_by_request.count 1|c|#null,null,accur
 
 ## Features Lookup Response Structure
 
-The /v1/features/join and /v1/features/groupby endpoints are bulkGet endpoints (against a single GroupBy or Join). Users can request multiple lookups, for example:
+The /v1/fetch/join and /v1/fetch/groupby endpoints are bulkGet endpoints (against a single GroupBy or Join). Users can request multiple lookups, for example:
 ```bash
-$ curl -X POST   'http://localhost:9000/v1/features/join/quickstart%2Ftraining_set.v2'   -H 'Content-Type: application/json'   -d '[{"user_id": "5"}, {"user_id": "7"}]'
+$ curl -X POST   'http://localhost:9000/v1/fetch/join/quickstart%2Ftraining_set.v2'   -H 'Content-Type: application/json'   -d '[{"user_id": "5"}, {"user_id": "7"}]'
 ```
 
 The response status is 4xx (in case of errors parsing the incoming json request payload), 5xx (internal error like the KV store being unreachable) or 200 (some / all successful lookups).

service/src/main/java/ai/chronon/service/ApiProvider.java

@@ -23,14 +23,12 @@ public class ApiProvider {
     private static final Logger logger = LoggerFactory.getLogger(ApiProvider.class);
 
     public static Api buildApi(ConfigStore configStore) throws Exception {
-        Optional<String> maybeJarPath = configStore.getOnlineJar();
-        Optional<String> maybeClass = configStore.getOnlineClass();
-        if (!(maybeJarPath.isPresent() && maybeClass.isPresent())) {
-            throw new IllegalArgumentException("Both 'online.jar' and 'online.class' configs must be set.");
-        }
+        configStore.validateOnlineApiConfig();
+
+        // we've already validated and confirmed these are present
+        String jarPath = configStore.getOnlineJar().get();
+        String className = configStore.getOnlineClass().get();
 
-        String jarPath = maybeJarPath.get();
-        String className = maybeClass.get();
         File jarFile = new File(jarPath);
         if (!jarFile.exists()) {
             throw new IllegalArgumentException("JAR file does not exist: " + jarPath);

service/src/main/java/ai/chronon/service/ConfigStore.java

@@ -64,6 +64,12 @@ public Optional<String> getOnlineClass() {
         return Optional.ofNullable(jsonConfig.getString(ONLINE_CLASS));
     }
 
+    public void validateOnlineApiConfig() {
+        if (!(getOnlineJar().isPresent() && getOnlineClass().isPresent())) {
+            throw new IllegalArgumentException("Both 'online.jar' and 'online.class' configs must be set.");
+        }
+    }
+
     public Map<String, String> getOnlineApiProps() {
         JsonObject apiProps = jsonConfig.getJsonObject(ONLINE_API_PROPS);
         if (apiProps == null) {

service/src/main/java/ai/chronon/service/WebServiceVerticle.java renamed to service/src/main/java/ai/chronon/service/FetcherVerticle.java

@@ -1,7 +1,7 @@
 package ai.chronon.service;
 
 import ai.chronon.online.Api;
-import ai.chronon.service.handlers.FeaturesRouter;
+import ai.chronon.service.handlers.FetcherRouter;
 import io.vertx.core.AbstractVerticle;
 import io.vertx.core.Promise;
 import io.vertx.core.http.HttpServer;
@@ -11,12 +11,12 @@
 import org.slf4j.LoggerFactory;
 
 /**
- * Entry point for the Chronon webservice. We wire up our API routes and configure and launch our HTTP service here.
+ * Entry point for the Chronon fetcher endpoints. We wire up our API routes and configure and launch our HTTP service here.
  * We choose to use just 1 verticle for now as it allows us to keep things simple and we don't need to scale /
  * independently deploy different endpoint routes.
  */
-public class WebServiceVerticle extends AbstractVerticle {
-    private static final Logger logger = LoggerFactory.getLogger(WebServiceVerticle.class);
+public class FetcherVerticle extends AbstractVerticle {
+    private static final Logger logger = LoggerFactory.getLogger(FetcherVerticle.class);
 
     private HttpServer server;
 
@@ -32,7 +32,7 @@ protected void startHttpServer(int port, String configJsonString, Api api, Promi
         // Define routes
 
         // Set up sub-routes for the various feature retrieval apis
-        router.route("/v1/features/*").subRouter(FeaturesRouter.createFeaturesRoutes(vertx, api));
+        router.route("/v1/fetch/*").subRouter(FetcherRouter.createFetchRoutes(vertx, api));
 
         // Health check route
         router.get("/ping").handler(ctx -> {

service/src/main/java/ai/chronon/service/handlers/FeaturesRouter.java

@@ -19,13 +19,15 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static ai.chronon.service.model.GetFeaturesResponse.Result.Status.Failure;
 import static ai.chronon.service.model.GetFeaturesResponse.Result.Status.Success;
 
 /**
- * Concrete implementation of the GetFeatures endpoints. Supports loading groupBys and joins.
+ * Concrete implementation of the Chronon fetcher endpoints. Supports loading groupBys and joins.
  * Some notes on this:
  * We currently support bulkGet lookups against a single groupBy / join. Attempts to lookup n different GroupBys / Joins
  * need to be split up into n different requests.
@@ -36,39 +38,23 @@
  * As an example:
  * { results: [ {"status": "Success", "features": ...}, {"status": "Failure", "error": ...} ] }
  */
-public class FeaturesHandler implements Handler<RoutingContext> {
+public class FetchHandler implements Handler<RoutingContext> {
 
-    public enum EntityType {
-        GroupBy,
-        Join
-    }
-
-    // PoJo to simplify transforming responses from the Fetcher to the final results form we return
-    private static class EntityKeyToValues {
-        public EntityKeyToValues(Map<String, Object> keys, JTry<Map<String, Object>> values) {
-            this.entityKeys = keys;
-            this.features = values;
-        }
-
-        public Map<String, Object> entityKeys;
-        public JTry<Map<String, Object>> features;
-    }
-
-    private static final Logger logger = LoggerFactory.getLogger(FeaturesHandler.class);
+    private static final Logger logger = LoggerFactory.getLogger(FetchHandler.class);
     private static final ObjectMapper objectMapper = new ObjectMapper();
 
-    private final EntityType entityType;
     private final JavaFetcher fetcher;
+    private final BiFunction<JavaFetcher, List<JavaRequest>, CompletableFuture<List<JavaResponse>>> fetchFunction;
 
-    public FeaturesHandler(EntityType entityType, JavaFetcher fetcher) {
-        this.entityType = entityType;
+    public FetchHandler(JavaFetcher fetcher, BiFunction<JavaFetcher, List<JavaRequest>, CompletableFuture<List<JavaResponse>>> fetchFunction) {
         this.fetcher = fetcher;
+        this.fetchFunction = fetchFunction;
     }
 
     @Override
     public void handle(RoutingContext ctx) {
         String entityName = ctx.pathParam("name");
-        logger.debug("Retrieving {} - {}", entityType.name(), entityName);
+        logger.debug("Retrieving {}", entityName);
         JTry<List<JavaRequest>> maybeRequest = parseJavaRequest(entityName, ctx.body());
         if (! maybeRequest.isSuccess()) {
             logger.error("Unable to parse request body", maybeRequest.getException());
@@ -81,27 +67,19 @@ public void handle(RoutingContext ctx) {
         }
 
         List<JavaRequest> requests = maybeRequest.getValue();
-        CompletableFuture<List<JavaResponse>> resultsJavaFuture =
-                entityType.equals(EntityType.GroupBy) ? fetcher.fetchGroupBys(requests) : fetcher.fetchJoin(requests);
+        CompletableFuture<List<JavaResponse>> resultsJavaFuture = fetchFunction.apply(fetcher, requests);
         // wrap the Java future we get in a Vert.x Future to not block the worker thread
-        Future<List<EntityKeyToValues>> maybeFeatureResponses =
+        Future<List<GetFeaturesResponse.Result>> maybeFeatureResponses =
                 Future.fromCompletionStage(resultsJavaFuture)
-                      .map(result -> result.stream().map(FeaturesHandler::responseToPoJo)
+                      .map(result -> result.stream().map(FetchHandler::responseToPoJo)
                       .collect(Collectors.toList()));
 
         maybeFeatureResponses.onSuccess(resultList -> {
             // as this is a bulkGet request, we might have some successful and some failed responses
             // we return the responses in the same order as they come in and mark them as successful / failed based
             // on the lookups
             GetFeaturesResponse.Builder responseBuilder = GetFeaturesResponse.builder();
-            List<GetFeaturesResponse.Result> results = resultList.stream().map(resultsPojo -> {
-                if (resultsPojo.features.isSuccess()) {
-                    return GetFeaturesResponse.Result.builder().status(Success).entityKeys(resultsPojo.entityKeys).features(resultsPojo.features.getValue()).build();
-                } else {
-                    return GetFeaturesResponse.Result.builder().status(Failure).entityKeys(resultsPojo.entityKeys).error(resultsPojo.features.getException().getMessage()).build();
-                }
-            }).collect(Collectors.toList());
-            responseBuilder.results(results);
+            responseBuilder.results(resultList);
 
             ctx.response()
                         .setStatusCode(200)
@@ -118,8 +96,22 @@ public void handle(RoutingContext ctx) {
         });
     }
 
-    public static EntityKeyToValues responseToPoJo(JavaResponse response) {
-        return new EntityKeyToValues(response.request.keys, response.values);
+    public static GetFeaturesResponse.Result responseToPoJo(JavaResponse response) {
+        if (response.values.isSuccess()) {
+            return GetFeaturesResponse.Result
+                    .builder()
+                    .status(Success)
+                    .entityKeys(response.request.keys)
+                    .features(response.values.getValue())
+                    .build();
+        } else {
+            return GetFeaturesResponse.Result
+                    .builder()
+                    .status(Failure)
+                    .entityKeys(response.request.keys)
+                    .error(response.values.getException().getMessage())
+                    .build();
+        }
     }
 
     public static JTry<List<JavaRequest>> parseJavaRequest(String name, RequestBody body) {

service/src/main/java/ai/chronon/service/handlers/FeaturesHandler.java renamed to service/src/main/java/ai/chronon/service/handlers/FetchHandler.java

@@ -0,0 +1,40 @@
+package ai.chronon.service.handlers;
+
+import ai.chronon.online.*;
+import io.vertx.core.Vertx;
+import io.vertx.ext.web.Router;
+import io.vertx.ext.web.handler.BodyHandler;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
+
+// Configures the routes for our get features endpoints
+// We support bulkGets of groupBys and bulkGets of joins
+public class FetcherRouter {
+
+    public static class GroupByFetcherFunction implements BiFunction<JavaFetcher, List<JavaRequest>, CompletableFuture<List<JavaResponse>>> {
+        @Override
+        public CompletableFuture<List<JavaResponse>> apply(JavaFetcher fetcher, List<JavaRequest> requests) {
+            return fetcher.fetchGroupBys(requests);
+        }
+    }
+
+    public static class JoinFetcherFunction implements BiFunction<JavaFetcher, List<JavaRequest>, CompletableFuture<List<JavaResponse>>> {
+        @Override
+        public CompletableFuture<List<JavaResponse>> apply(JavaFetcher fetcher, List<JavaRequest> requests) {
+            return fetcher.fetchJoin(requests);
+        }
+    }
+
+    public static Router createFetchRoutes(Vertx vertx, Api api) {
+        Router router = Router.router(vertx);
+        router.route().handler(BodyHandler.create());
+        JavaFetcher fetcher = api.buildJavaFetcher("feature-service", false);
+
+        router.post("/groupby/:name").handler(new FetchHandler(fetcher, new GroupByFetcherFunction()));
+        router.post("/join/:name").handler(new FetchHandler(fetcher, new JoinFetcherFunction()));
+
+        return router;
+    }
+}

service/src/main/java/ai/chronon/service/handlers/FetcherRouter.java

@@ -7,7 +7,7 @@
 import java.util.Map;
 
 /**
- * PoJo capturing the response we return back as part of /v1/features/groupby and /v1/features/join endpoints
+ * PoJo capturing the response we return back as part of /v1/fetch/groupby and /v1/fetch/join endpoints
  * when the individual bulkGet lookups were either all successful or partially successful.
  */
 @JsonInclude(JsonInclude.Include.NON_NULL)

service/src/main/java/ai/chronon/service/model/GetFeaturesResponse.java

@@ -11,7 +11,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-public class FeaturesHandlerJsonSerDeTest {
+public class FetchHandlerJsonSerDeTest {
 
     @Test
     public void testParsingOfSimpleJavaRequests() {
@@ -20,7 +20,7 @@ public void testParsingOfSimpleJavaRequests() {
         when(mockRequestBody.asString()).thenReturn(mockRequest);
 
         String groupByName = "my_groupby.1";
-        JTry<List<JavaRequest>> maybeRequest = FeaturesHandler.parseJavaRequest(groupByName, mockRequestBody);
+        JTry<List<JavaRequest>> maybeRequest = FetchHandler.parseJavaRequest(groupByName, mockRequestBody);
         assertTrue(maybeRequest.isSuccess());
         List<JavaRequest> reqs = maybeRequest.getValue();
         assertEquals(1, reqs.size());
@@ -38,7 +38,7 @@ public void testParsingInvalidRequest() {
         when(mockRequestBody.asString()).thenReturn(mockRequest);
 
         String groupByName = "my_groupby.1";
-        JTry<List<JavaRequest>> maybeRequest = FeaturesHandler.parseJavaRequest(groupByName, mockRequestBody);
+        JTry<List<JavaRequest>> maybeRequest = FetchHandler.parseJavaRequest(groupByName, mockRequestBody);
         assertFalse(maybeRequest.isSuccess());
         assertNotNull(maybeRequest.getException());
     }
@@ -50,7 +50,7 @@ public void testParsingOneValidAndInvalidRequest() {
         when(mockRequestBody.asString()).thenReturn(mockRequest);
 
         String groupByName = "my_groupby.1";
-        JTry<List<JavaRequest>> maybeRequest = FeaturesHandler.parseJavaRequest(groupByName, mockRequestBody);
+        JTry<List<JavaRequest>> maybeRequest = FetchHandler.parseJavaRequest(groupByName, mockRequestBody);
         assertFalse(maybeRequest.isSuccess());
         assertNotNull(maybeRequest.getException());
     }

service/src/test/java/ai/chronon/service/handlers/FeaturesHandlerJsonSerDeTest.java renamed to service/src/test/java/ai/chronon/service/handlers/FetchHandlerJsonSerDeTest.java

@@ -31,7 +31,7 @@
 import static org.mockito.Mockito.when;
 
 @RunWith(VertxUnitRunner.class)
-public class FeaturesHandlerTest {
+public class FetchHandlerTest {
 
     @Mock
     private JavaFetcher mockFetcher;
@@ -45,7 +45,7 @@ public class FeaturesHandlerTest {
     @Mock
     RequestBody requestBody;
 
-    private FeaturesHandler handler;
+    private FetchHandler handler;
     private Vertx vertx;
 
     private static final String TEST_GROUP_BY = "test_groupby.v1";
@@ -54,7 +54,8 @@ public class FeaturesHandlerTest {
     public void setUp(TestContext context) {
         MockitoAnnotations.openMocks(this);
         vertx = Vertx.vertx();
-        handler = new FeaturesHandler(FeaturesHandler.EntityType.Join, mockFetcher);
+
+        handler = new FetchHandler(mockFetcher, new FetcherRouter.JoinFetcherFunction());
 
         // Set up common routing context behavior
         when(routingContext.response()).thenReturn(response);