Skip to content

Commit

Permalink
feat: Added deadline to gRPC Java client (#4217)
Browse files Browse the repository at this point in the history
Added deadline to gRPC Java Client

Signed-off-by: Jose Acevedo <sharp.acevedo@gmail.com>
  • Loading branch information
acevedosharp committed May 26, 2024
1 parent 55566f0 commit ff429c9
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 3 deletions.
42 changes: 40 additions & 2 deletions java/serving-client/src/main/java/dev/feast/FeastClient.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import feast.proto.serving.ServingServiceGrpc.ServingServiceBlockingStub;
import feast.proto.types.ValueProto;
import io.grpc.CallCredentials;
import io.grpc.Deadline;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
Expand Down Expand Up @@ -63,7 +64,20 @@ public static FeastClient create(String host, int port) {
}

/**
* Create a authenticated client that can access Feast serving with authentication enabled.
* Create a client to access Feast Serving.
*
* @param host hostname or ip address of Feast serving GRPC server
* @param port port number of Feast serving GRPC server
* @param deadline GRPC deadline of Feast serving GRPC server {@link Deadline}
* @return {@link FeastClient}
*/
public static FeastClient create(String host, int port, Deadline deadline) {
// configure client with no security config.
return FeastClient.createSecure(host, port, SecurityConfig.newBuilder().build(), deadline);
}

/**
* Create an authenticated client that can access Feast serving with authentication enabled.
*
* @param host hostname or ip address of Feast serving GRPC server
* @param port port number of Feast serving GRPC server
Expand All @@ -72,6 +86,21 @@ public static FeastClient create(String host, int port) {
* @return {@link FeastClient}
*/
public static FeastClient createSecure(String host, int port, SecurityConfig securityConfig) {
return createSecure(host, port, securityConfig, null);
}

/**
* Create an authenticated client that can access Feast serving with authentication enabled.
*
* @param host hostname or ip address of Feast serving GRPC server
* @param port port number of Feast serving GRPC server
* @param securityConfig security options to configure the Feast client. See {@link
* SecurityConfig} for options.
* @param deadline GRPC deadline of Feast serving GRPC server {@link Deadline}
* @return {@link FeastClient}
*/
public static FeastClient createSecure(
String host, int port, SecurityConfig securityConfig, Deadline deadline) {
// Configure client TLS
ManagedChannel channel = null;
if (securityConfig.isTLSEnabled()) {
Expand All @@ -98,7 +127,7 @@ public static FeastClient createSecure(String host, int port, SecurityConfig sec
channel = ManagedChannelBuilder.forAddress(host, port).usePlaintext().build();
}

return new FeastClient(channel, securityConfig.getCredentials());
return new FeastClient(channel, securityConfig.getCredentials(), Optional.ofNullable(deadline));
}

/**
Expand Down Expand Up @@ -202,6 +231,11 @@ public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> rows, Str
}

protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credentials) {
this(channel, credentials, Optional.empty());
}

protected FeastClient(
ManagedChannel channel, Optional<CallCredentials> credentials, Optional<Deadline> deadline) {
this.channel = channel;
TracingClientInterceptor tracingInterceptor =
TracingClientInterceptor.newBuilder().withTracer(GlobalTracer.get()).build();
Expand All @@ -213,6 +247,10 @@ protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credenti
servingStub = servingStub.withCallCredentials(credentials.get());
}

if (deadline.isPresent()) {
servingStub = servingStub.withDeadline(deadline.get());
}

this.stub = servingStub;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,15 @@
import java.util.HashMap;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

public class FeastClientTest {
private final String AUTH_TOKEN = "test token";
private final Deadline DEADLINE = Deadline.after(2, TimeUnit.SECONDS);

@Rule public GrpcCleanupRule grpcRule;
private AtomicBoolean isAuthenticated;
Expand Down Expand Up @@ -86,7 +88,7 @@ public void setup() throws Exception {
ManagedChannel channel =
this.grpcRule.register(
InProcessChannelBuilder.forName(serverName).directExecutor().build());
this.client = new FeastClient(channel, Optional.empty());
this.client = new FeastClient(channel, Optional.empty(), Optional.of(DEADLINE));
}

@Test
Expand Down

0 comments on commit ff429c9

Please sign in to comment.