fix: Azure blob storage support in Java feature server (feast-dev#2319)

- Add azure blob storage support in java feature server - Fix S3 integration test to work without a real AWS account - Add GCS mock to integration tests to be able to run them without a real google cloud account - Adding dependency management in maven for libraries with older incompatible versions as transitive dependencies
szaboferee · Mar 15, 2024 · 7aab7d7 · 7aab7d7
1 parent 34cabfb
commit 7aab7d7
Show file tree

Hide file tree

Showing 10 changed files with 335 additions and 52 deletions.
diff --git a/java/CONTRIBUTING.md b/java/CONTRIBUTING.md
@@ -50,7 +50,7 @@ Automatically format the code to conform the style guide by:
 
 ```sh
 # formats all code in the feast-java repository
-mvn spotless:apply
+make format-java
 ```
 
 > If you're using IntelliJ, you can import these [code style settings](https://github.com/google/styleguide/blob/gh-pages/intellij-java-google-style.xml)
@@ -66,7 +66,7 @@ Run all Unit tests:
 make test-java
 ```
 
-Run all Integration tests (note: this also runs GCS + S3 based tests which should fail):
+Run all Integration tests:
 ```
 make test-java-integration
 ```

diff --git a/java/pom.xml b/java/pom.xml
@@ -68,6 +68,8 @@
         <google.auth.library.oauth2.http.version>0.21.0</google.auth.library.oauth2.http.version>
         <auto.value.version>1.6.6</auto.value.version>
         <guava.version>30.1-jre</guava.version>
+        <reactor.version>3.4.34</reactor.version>
+        <netty.version>4.1.101.Final</netty.version>
 
         <license.content><![CDATA[
 /*
@@ -186,6 +188,49 @@
                 <version>${javax.validation.version}</version>
             </dependency>
 
+            <dependency>
+                <groupId>com.fasterxml.jackson.core</groupId>
+                <artifactId>jackson-core</artifactId>
+                <version>${jackson.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>com.fasterxml.jackson.core</groupId>
+                <artifactId>jackson-databind</artifactId>
+                <version>${jackson.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>com.fasterxml.jackson.core</groupId>
+                <artifactId>jackson-annotations</artifactId>
+                <version>${jackson.version}</version>
+            </dependency>
+
+            <dependency>
+                <groupId>io.netty</groupId>
+                <artifactId>netty-common</artifactId>
+                <version>${netty.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>io.netty</groupId>
+                <artifactId>netty-buffer</artifactId>
+                <version>${netty.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>io.netty</groupId>
+                <artifactId>netty-handler</artifactId>
+                <version>${netty.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>io.netty</groupId>
+                <artifactId>netty-transport</artifactId>
+                <version>${netty.version}</version>
+            </dependency>
+
+            <dependency>
+                <groupId>io.projectreactor</groupId>
+                <artifactId>reactor-core</artifactId>
+                <version>${reactor.version}</version>
+            </dependency>
+
             <dependency>
                 <groupId>org.junit.platform</groupId>
                 <artifactId>junit-platform-engine</artifactId>
@@ -246,7 +291,7 @@
                 <configuration>
                     <java>
                         <licenseHeader>
-                          <content>${license.content}</content>
+                            <content>${license.content}</content>
                         </licenseHeader>
                         <googleJavaFormat>
                             <version>1.7</version>
@@ -264,15 +309,15 @@
                     </scala>
                 </configuration>
                 <executions>
-                  <!-- Move check to fail faster, but after compilation. Default is verify phase -->
-                  <execution>
-                      <id>spotless-check</id>
-                      <phase>process-test-classes</phase>
-                      <goals>
-                          <goal>check</goal>
-                      </goals>
-                  </execution>
-              </executions>
+                    <!-- Move check to fail faster, but after compilation. Default is verify phase -->
+                    <execution>
+                        <id>spotless-check</id>
+                        <phase>process-test-classes</phase>
+                        <goals>
+                            <goal>check</goal>
+                        </goals>
+                    </execution>
+                </executions>
             </plugin>
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>

diff --git a/java/serving/.gitignore b/java/serving/.gitignore
@@ -34,4 +34,7 @@ feast-serving.jar
 /.nb-gradle/
 
 ## Feast Temporary Files ##
-/temp/
+/temp/
+
+## Generated test data ##
+**/*.parquet
diff --git a/java/serving/pom.xml b/java/serving/pom.xml
@@ -16,8 +16,8 @@
   ~
   -->
 <project xmlns="http://maven.apache.org/POM/4.0.0"
-    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
   <modelVersion>4.0.0</modelVersion>
 
   <parent>
@@ -121,6 +121,19 @@
       <version>5.0.1</version>
     </dependency>
 
+    <!-- azure blob storage -->
+    <!-- https://mvnrepository.com/artifact/com.azure/azure-storage-blob -->
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-storage-blob</artifactId>
+      <version>12.25.2</version>
+    </dependency>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-identity</artifactId>
+      <version>1.11.3</version>
+    </dependency>
+
     <!-- TODO: SLF4J is being used via Lombok, but also jog4j - pick one -->
     <dependency>
       <groupId>org.slf4j</groupId>
@@ -356,11 +369,11 @@
       <version>2.7.4</version>
       <scope>test</scope>
     </dependency>
-      <dependency>
-          <groupId>io.lettuce</groupId>
-          <artifactId>lettuce-core</artifactId>
-          <version>6.0.2.RELEASE</version>
-      </dependency>
+    <dependency>
+      <groupId>io.lettuce</groupId>
+      <artifactId>lettuce-core</artifactId>
+      <version>6.0.2.RELEASE</version>
+    </dependency>
     <dependency>
       <groupId>org.apache.commons</groupId>
       <artifactId>commons-lang3</artifactId>

diff --git a/java/serving/src/main/java/feast/serving/registry/AzureRegistryFile.java b/java/serving/src/main/java/feast/serving/registry/AzureRegistryFile.java
@@ -0,0 +1,57 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright 2018-2021 The Feast Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package feast.serving.registry;
+
+import com.azure.storage.blob.BlobClient;
+import com.azure.storage.blob.BlobServiceClient;
+import com.google.protobuf.InvalidProtocolBufferException;
+import feast.proto.core.RegistryProto;
+import java.util.Objects;
+import java.util.Optional;
+
+public class AzureRegistryFile implements RegistryFile {
+  private final BlobClient blobClient;
+  private String lastKnownETag;
+
+  public AzureRegistryFile(BlobServiceClient blobServiceClient, String url) {
+    String[] split = url.replace("az://", "").split("/");
+    String objectPath = String.join("/", java.util.Arrays.copyOfRange(split, 1, split.length));
+    this.blobClient = blobServiceClient.getBlobContainerClient(split[0]).getBlobClient(objectPath);
+  }
+
+  @Override
+  public RegistryProto.Registry getContent() {
+    try {
+      return RegistryProto.Registry.parseFrom(blobClient.downloadContent().toBytes());
+    } catch (InvalidProtocolBufferException e) {
+      throw new RuntimeException(
+          String.format(
+              "Couldn't read remote registry: %s. Error: %s",
+              blobClient.getBlobUrl(), e.getMessage()));
+    }
+  }
+
+  @Override
+  public Optional<RegistryProto.Registry> getContentIfModified() {
+    String eTag = blobClient.getProperties().getETag();
+    if (Objects.equals(eTag, this.lastKnownETag)) {
+      return Optional.empty();
+    } else this.lastKnownETag = eTag;
+
+    return Optional.of(getContent());
+  }
+}
diff --git a/java/serving/src/main/java/feast/serving/service/config/ApplicationProperties.java b/java/serving/src/main/java/feast/serving/service/config/ApplicationProperties.java
@@ -95,6 +95,7 @@ public static class FeastProperties {
     private String gcpProject;
     private String awsRegion;
     private String transformationServiceEndpoint;
+    private String azureStorageAccount;
 
     public String getRegistry() {
       return registry;
@@ -205,6 +206,14 @@ public String getTransformationServiceEndpoint() {
     public void setTransformationServiceEndpoint(String transformationServiceEndpoint) {
       this.transformationServiceEndpoint = transformationServiceEndpoint;
     }
+
+    public String getAzureStorageAccount() {
+      return azureStorageAccount;
+    }
+
+    public void setAzureStorageAccount(String azureStorageAccount) {
+      this.azureStorageAccount = azureStorageAccount;
+    }
   }
 
   /** Store configuration class for database that this Feast Serving uses. */

diff --git a/java/serving/src/main/java/feast/serving/service/config/RegistryConfigModule.java b/java/serving/src/main/java/feast/serving/service/config/RegistryConfigModule.java
@@ -18,6 +18,9 @@
 
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3ClientBuilder;
+import com.azure.identity.DefaultAzureCredentialBuilder;
+import com.azure.storage.blob.BlobServiceClient;
+import com.azure.storage.blob.BlobServiceClientBuilder;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.StorageOptions;
 import com.google.inject.AbstractModule;
@@ -43,11 +46,27 @@ public AmazonS3 awsStorage(ApplicationProperties applicationProperties) {
         .build();
   }
 
+  @Provides
+  public BlobServiceClient azureStorage(ApplicationProperties applicationProperties) {
+
+    BlobServiceClient blobServiceClient =
+        new BlobServiceClientBuilder()
+            .endpoint(
+                String.format(
+                    "https://%s.blob.core.windows.net",
+                    applicationProperties.getFeast().getAzureStorageAccount()))
+            .credential(new DefaultAzureCredentialBuilder().build())
+            .buildClient();
+
+    return blobServiceClient;
+  }
+
   @Provides
   RegistryFile registryFile(
       ApplicationProperties applicationProperties,
       Provider<Storage> storageProvider,
-      Provider<AmazonS3> amazonS3Provider) {
+      Provider<AmazonS3> amazonS3Provider,
+      Provider<BlobServiceClient> azureProvider) {
 
     String registryPath = applicationProperties.getFeast().getRegistry();
     Optional<String> scheme = Optional.ofNullable(URI.create(registryPath).getScheme());
@@ -57,6 +76,8 @@ RegistryFile registryFile(
         return new GSRegistryFile(storageProvider.get(), registryPath);
       case "s3":
         return new S3RegistryFile(amazonS3Provider.get(), registryPath);
+      case "az":
+        return new AzureRegistryFile(azureProvider.get(), registryPath);
       case "":
       case "file":
         return new LocalRegistryFile(registryPath);