2626import feast .proto .serving .ServingServiceGrpc .ServingServiceBlockingStub ;
2727import feast .proto .types .ValueProto ;
2828import io .grpc .CallCredentials ;
29+ import io .grpc .Deadline ;
2930import io .grpc .ManagedChannel ;
3031import io .grpc .ManagedChannelBuilder ;
3132import io .grpc .netty .shaded .io .grpc .netty .GrpcSslContexts ;
@@ -57,9 +58,9 @@ public class FeastClient implements AutoCloseable {
5758 * @param port port number of Feast serving GRPC server
5859 * @return {@link FeastClient}
5960 */
60- public static FeastClient create (String host , int port ) {
61+ public static FeastClient create (String host , int port , Deadline deadline ) {
6162 // configure client with no security config.
62- return FeastClient .createSecure (host , port , SecurityConfig .newBuilder ().build ());
63+ return FeastClient .createSecure (host , port , SecurityConfig .newBuilder ().build (), deadline );
6364 }
6465
6566 /**
@@ -71,7 +72,7 @@ public static FeastClient create(String host, int port) {
7172 * SecurityConfig} for options.
7273 * @return {@link FeastClient}
7374 */
74- public static FeastClient createSecure (String host , int port , SecurityConfig securityConfig ) {
75+ public static FeastClient createSecure (String host , int port , SecurityConfig securityConfig , Deadline deadline ) {
7576 // Configure client TLS
7677 ManagedChannel channel = null ;
7778 if (securityConfig .isTLSEnabled ()) {
@@ -98,7 +99,7 @@ public static FeastClient createSecure(String host, int port, SecurityConfig sec
9899 channel = ManagedChannelBuilder .forAddress (host , port ).usePlaintext ().build ();
99100 }
100101
101- return new FeastClient (channel , securityConfig .getCredentials ());
102+ return new FeastClient (channel , securityConfig .getCredentials (), deadline );
102103 }
103104
104105 /**
@@ -201,7 +202,7 @@ public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> rows, Str
201202 return getOnlineFeatures (featureRefs , rows );
202203 }
203204
204- protected FeastClient (ManagedChannel channel , Optional <CallCredentials > credentials ) {
205+ protected FeastClient (ManagedChannel channel , Optional <CallCredentials > credentials , Deadline deadline ) {
205206 this .channel = channel ;
206207 TracingClientInterceptor tracingInterceptor =
207208 TracingClientInterceptor .newBuilder ().withTracer (GlobalTracer .get ()).build ();
@@ -213,6 +214,8 @@ protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credenti
213214 servingStub = servingStub .withCallCredentials (credentials .get ());
214215 }
215216
217+ servingStub = servingStub .withDeadline (deadline );
218+
216219 this .stub = servingStub ;
217220 }
218221
0 commit comments