feast-dev · pyalex · Apr 9, 2021 · Apr 5, 2021 · Apr 5, 2021 · Apr 5, 2021
diff --git a/common-test/src/main/java/feast/common/it/DataGenerator.java b/common-test/src/main/java/feast/common/it/DataGenerator.java
@@ -249,6 +249,10 @@ public static ValueProto.Value createDoubleValue(double value) {
     return ValueProto.Value.newBuilder().setDoubleVal(value).build();
   }
 
+  public static ValueProto.Value createInt32Value(int value) {
+    return ValueProto.Value.newBuilder().setInt32Val(value).build();
+  }
+
   public static ValueProto.Value createInt64Value(long value) {
     return ValueProto.Value.newBuilder().setInt64Val(value).build();
   }

diff --git a/serving/pom.xml b/serving/pom.xml
@@ -90,6 +90,12 @@
       <version>${project.version}</version>
     </dependency>
 
+    <dependency>
+      <groupId>dev.feast</groupId>
+      <artifactId>feast-storage-connector-cassandra</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+
     <dependency>
       <groupId>dev.feast</groupId>
       <artifactId>feast-common</artifactId>

diff --git a/serving/src/main/java/feast/serving/config/FeastProperties.java b/serving/src/main/java/feast/serving/config/FeastProperties.java
@@ -27,6 +27,7 @@
 import feast.common.auth.credentials.CoreAuthenticationProperties;
 import feast.common.logging.config.LoggingProperties;
 import feast.storage.connectors.bigtable.retriever.BigTableStoreConfig;
+import feast.storage.connectors.cassandra.retriever.CassandraStoreConfig;
 import feast.storage.connectors.redis.retriever.RedisClusterStoreConfig;
 import feast.storage.connectors.redis.retriever.RedisStoreConfig;
 import io.lettuce.core.ReadFrom;
@@ -270,7 +271,7 @@ public void setName(String name) {
     }
 
     /**
-     * Gets the store type. Example are REDIS, REDIS_CLUSTER or BIGTABLE
+     * Gets the store type. Example are REDIS, REDIS_CLUSTER, BIGTABLE or CASSANDRA
      *
      * @return the store type as a String.
      */
@@ -316,6 +317,13 @@ public BigTableStoreConfig getBigtableConfig() {
       return new BigTableStoreConfig(this.config.get("project_id"), this.config.get("instance_id"));
     }
 
+    public CassandraStoreConfig getCassandraConfig() {
+      return new CassandraStoreConfig(
+          this.config.get("connection_string"),
+          this.config.get("data_center"),
+          this.config.get("keyspace"));
+    }
+
     /**
      * Sets the store config. Please protos/feast/core/Store.proto for the specific options for each
      * store.
@@ -329,6 +337,7 @@ public void setConfig(Map<String, String> config) {
 
   public enum StoreType {
     BIGTABLE,
+    CASSANDRA,
     REDIS,
     REDIS_CLUSTER;
   }

diff --git a/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java b/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java
@@ -16,6 +16,8 @@
  */
 package feast.serving.config;
 
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.CqlSessionBuilder;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
 import feast.serving.service.OnlineServingServiceV2;
@@ -24,9 +26,15 @@
 import feast.storage.api.retriever.OnlineRetrieverV2;
 import feast.storage.connectors.bigtable.retriever.BigTableOnlineRetriever;
 import feast.storage.connectors.bigtable.retriever.BigTableStoreConfig;
+import feast.storage.connectors.cassandra.retriever.CassandraOnlineRetriever;
+import feast.storage.connectors.cassandra.retriever.CassandraStoreConfig;
 import feast.storage.connectors.redis.retriever.*;
 import io.opentracing.Tracer;
 import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.ApplicationContext;
@@ -77,6 +85,32 @@ public ServingServiceV2 servingServiceV2(
         OnlineRetrieverV2 bigtableRetriever = new BigTableOnlineRetriever(bigtableClient);
         servingService = new OnlineServingServiceV2(bigtableRetriever, specService, tracer);
         break;
+      case CASSANDRA:
+        CassandraStoreConfig config = feastProperties.getActiveStore().getCassandraConfig();
+        String connectionString = config.getConnectionString();
+        String dataCenter = config.getDataCenter();
+        String keySpace = config.getKeySpace();
+
+        List<InetSocketAddress> contactPoints =
+            Arrays.stream(connectionString.split(","))
+                .map(String::trim)
+                .map(cs -> cs.split(":"))
+                .map(
+                    hostPort -> {
+                      int port = hostPort.length > 1 ? Integer.parseInt(hostPort[1]) : 9042;
+                      return new InetSocketAddress(hostPort[0], port);
+                    })
+                .collect(Collectors.toList());
+
+        CqlSession session =
+            new CqlSessionBuilder()
+                .addContactPoints(contactPoints)
+                .withLocalDatacenter(dataCenter)
+                .withKeyspace(keySpace)
+                .build();
+        OnlineRetrieverV2 cassandraRetriever = new CassandraOnlineRetriever(session);
+        servingService = new OnlineServingServiceV2(cassandraRetriever, specService, tracer);
+        break;
     }
 
     return servingService;

diff --git a/serving/src/main/resources/application.yml b/serving/src/main/resources/application.yml
@@ -59,6 +59,12 @@ feast:
       config:
         project_id: <gcp_project>
         instance_id: <gcp_bigtable_instance>
+    - name: cassandra
+      type: CASSANDRA
+      config:
+        connection_string: localhost:9042
+        data_center: datacenter1
+        keyspace: feast
   tracing:
     # If true, Feast will provide tracing data (using OpenTracing API) for various RPC method calls
     # which can be useful to debug performance issues and perform benchmarking

diff --git a/serving/src/test/java/feast/serving/it/BaseAuthIT.java b/serving/src/test/java/feast/serving/it/BaseAuthIT.java
@@ -54,6 +54,13 @@ public class BaseAuthIT {
   static final String BIGTABLE = "bigtable_1";
   static final int BIGTABLE_PORT = 8086;
 
+  static final String CASSANDRA = "cassandra_1";
+  static final int CASSANDRA_PORT = 9042;
+  static final String CASSANDRA_DATACENTER = "datacenter1";
+  static final String CASSANDRA_KEYSPACE = "feast";
+  static final String CASSANDRA_SCHEMA_TABLE = "feast_schema_reference";
+  static final String CASSANDRA_ENTITY_KEY = "key";
+
   static final int FEAST_CORE_PORT = 6565;
 
   @DynamicPropertySource
@@ -72,14 +79,40 @@ static void properties(DynamicPropertyRegistry registry) {
           }
         });
     registry.add("feast.stores[0].config.port", () -> REDIS_PORT);
-    registry.add("feast.stores[0].subscriptions[0].name", () -> "*");
-    registry.add("feast.stores[0].subscriptions[0].project", () -> "*");
 
     registry.add("feast.stores[1].name", () -> "bigtable");
     registry.add("feast.stores[1].type", () -> "BIGTABLE");
     registry.add("feast.stores[1].config.project_id", () -> "test-project");
     registry.add("feast.stores[1].config.instance_id", () -> "test-instance");
 
+    registry.add("feast.stores[2].name", () -> "cassandra");
+    registry.add("feast.stores[2].type", () -> "CASSANDRA");
+    registry.add(
+        "feast.stores[2].config.host",
+        () -> {
+          try {
+            return InetAddress.getLocalHost().getHostAddress();
+          } catch (UnknownHostException e) {
+            e.printStackTrace();
+            return "";
+          }
+        });
+
+    registry.add(
+        "feast.stores[2].config.connection_string",
+        () -> {
+          String hostAddress = "";
+          try {
+            hostAddress = InetAddress.getLocalHost().getHostAddress();
+          } catch (UnknownHostException e) {
+            e.printStackTrace();
+          }
+
+          return String.format("%s:%s", hostAddress, CASSANDRA_PORT);
+        });
+    registry.add("feast.stores[2].config.data_center", () -> CASSANDRA_DATACENTER);
+    registry.add("feast.stores[2].config.keyspace", () -> CASSANDRA_KEYSPACE);
+
     registry.add("feast.core-authentication.options.oauth_url", () -> TOKEN_URL);
     registry.add("feast.core-authentication.options.grant_type", () -> GRANT_TYPE);
     registry.add("feast.core-authentication.options.client_id", () -> CLIENT_ID);