Skip to content

Commit

Permalink
refactor: consolidate the declaration of the common grpc context auth…
Browse files Browse the repository at this point in the history 
… keys (#3876)
  • Loading branch information
@jkwatson-verta
jkwatson-verta committed Jun 12, 2023
1 parent ec3be71 commit 3a80bba
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 27 deletions.
9 changes: 5 additions & 4 deletions backend/common/src/main/java/ai/verta/modeldb/common/authservice/AuthInterceptor.java
View file
@@ -1,5 +1,6 @@
package ai.verta.modeldb.common.authservice;

import ai.verta.modeldb.common.connections.Connection;
import io.grpc.Context;
import io.grpc.Contexts;
import io.grpc.ForwardingServerCallListener;
Expand Down Expand Up @@ -34,14 +35,14 @@ public <R, S> Listener<R> interceptCall(
|| methodName.equals("grpc.health.v1.Health/Check")
|| methodName.equals("grpc.health.v1.Health/Watch"))) {
// validate empty headers from user request
var emailKey = Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
var devKeyUnderscore = Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
var devKey = Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
var emailKey = Connection.EMAIL_GRPC_METADATA_KEY;
var devKeyUnderscore = Connection.DEV_KEY_GRPC_METADATA_KEY;
var devKey = Connection.DEV_KEY_WITH_HYPHEN_GRPC_METADATA_KEY;
var bearerAccessToken =
Metadata.Key.of("bearer_access_token", Metadata.ASCII_STRING_MARSHALLER);
var sessionId = Metadata.Key.of("sessionId", Metadata.ASCII_STRING_MARSHALLER);
var sessionIdSig = Metadata.Key.of("sessionIdSig", Metadata.ASCII_STRING_MARSHALLER);
var sourceKey = Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);
var sourceKey = Connection.SOURCE_GRPC_METADATA_KEY;

var parameterMissing = false;
if (!requestHeaders.containsKey(sourceKey)) {
Expand Down
13 changes: 4 additions & 9 deletions backend/common/src/main/java/ai/verta/modeldb/common/authservice/AuthServiceChannel.java
View file
Expand Up @@ -52,15 +52,10 @@ public AuthServiceChannel(Config config, Optional<ClientInterceptor> tracingClie

private Metadata getServiceUserMetadataHeaders() {
var requestHeaders = new Metadata();
var emailKey = Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
var devKey = Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
var devKeyHyphen = Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
var sourceKey = Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

requestHeaders.put(emailKey, this.serviceUserEmail);
requestHeaders.put(devKey, this.serviceUserDevKey);
requestHeaders.put(devKeyHyphen, this.serviceUserDevKey);
requestHeaders.put(sourceKey, "PythonClient");
requestHeaders.put(Connection.EMAIL_GRPC_METADATA_KEY, this.serviceUserEmail);
requestHeaders.put(Connection.DEV_KEY_GRPC_METADATA_KEY, this.serviceUserDevKey);
requestHeaders.put(Connection.DEV_KEY_WITH_HYPHEN_GRPC_METADATA_KEY, this.serviceUserDevKey);
requestHeaders.put(Connection.SOURCE_GRPC_METADATA_KEY, "PythonClient");
return requestHeaders;
}

Expand Down
21 changes: 13 additions & 8 deletions backend/common/src/main/java/ai/verta/modeldb/common/connections/Connection.java
View file
Expand Up @@ -10,6 +10,15 @@
import java.util.function.Supplier;

public abstract class Connection {
public static final Metadata.Key<String> EMAIL_GRPC_METADATA_KEY =
Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
public static final Metadata.Key<String> DEV_KEY_GRPC_METADATA_KEY =
Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
public static final Metadata.Key<String> DEV_KEY_WITH_HYPHEN_GRPC_METADATA_KEY =
Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
public static final Metadata.Key<String> SOURCE_GRPC_METADATA_KEY =
Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

private final Optional<ClientInterceptor> tracingClientInterceptor;

protected Connection(Optional<ClientInterceptor> tracingClientInterceptor) {
Expand Down Expand Up @@ -64,15 +73,11 @@ protected Optional<ClientInterceptor> getTracingClientInterceptor() {

protected static Metadata getServiceUserMetadata(ServiceUserConfig config) {
var requestHeaders = new Metadata();
var emailKey = Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
var devKey = Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
var devKeyHyphen = Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
var sourceKey = Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

requestHeaders.put(emailKey, config.getEmail());
requestHeaders.put(devKey, config.getDevKey());
requestHeaders.put(devKeyHyphen, config.getDevKey());
requestHeaders.put(sourceKey, "PythonClient");
requestHeaders.put(EMAIL_GRPC_METADATA_KEY, config.getEmail());
requestHeaders.put(DEV_KEY_GRPC_METADATA_KEY, config.getDevKey());
requestHeaders.put(DEV_KEY_WITH_HYPHEN_GRPC_METADATA_KEY, config.getDevKey());
requestHeaders.put(SOURCE_GRPC_METADATA_KEY, "PythonClient");
return requestHeaders;
}
}
10 changes: 4 additions & 6 deletions backend/server/src/test/java/ai/verta/modeldb/AuthClientInterceptor.java
View file
@@ -1,6 +1,7 @@
package ai.verta.modeldb;

import ai.verta.modeldb.common.config.ServiceUserConfig;
import ai.verta.modeldb.common.connections.Connection;
import ai.verta.modeldb.config.TestConfig;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.grpc.*;
Expand Down Expand Up @@ -83,12 +84,9 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
@Override
public void start(Listener<RespT> responseListener, Metadata headers) {
// TODO: Here set request metadata
Metadata.Key<String> email_key =
Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
Metadata.Key<String> dev_key =
Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
Metadata.Key<String> source_key =
Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);
Metadata.Key<String> email_key = Connection.EMAIL_GRPC_METADATA_KEY;
Metadata.Key<String> dev_key = Connection.DEV_KEY_GRPC_METADATA_KEY;
Metadata.Key<String> source_key = Connection.SOURCE_GRPC_METADATA_KEY;

headers.put(email_key, clientEmail);
headers.put(dev_key, clientDevKey);
Expand Down

2 comments on commit 3a80bba

@github-actions
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docker Tag: main-2023-06-12T17-49-50--3a80bba

@github-actions
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Total coverage (common): 11.24
Total coverage (server): 60.49

Changed Files coverage (common): coverage 0
Changed Files coverage (server): 100

Please sign in to comment.