package ai.verta.modeldb.common.connections;

import ai.verta.modeldb.common.authservice.AuthInterceptor;
import ai.verta.modeldb.common.config.Config;
import ai.verta.modeldb.common.config.ServiceUserConfig;
import ai.verta.modeldb.common.futures.InternalFuture;
import ai.verta.modeldb.common.interceptors.MetadataForwarder;
import io.grpc.Context;
import io.grpc.Metadata;
import io.grpc.stub.AbstractStub;
import io.opentracing.contrib.grpc.TracingClientInterceptor;

import java.util.Optional;
import java.util.concurrent.Executor;

public abstract class Connection {
  private final Optional<TracingClientInterceptor> tracingClientInterceptor;

  Connection(Config config) {
    tracingClientInterceptor = config.getTracingClientInterceptor();
  }

  // Place the interceptors in the reverse order of their stacking (the ones attached later will be
  // used first)
  protected <T extends AbstractStub<T>> T attachInterceptors(io.grpc.stub.AbstractStub<T> stub) {
    if (tracingClientInterceptor.isPresent()) {
      stub = stub.withInterceptors(tracingClientInterceptor.get());
    }

    stub = stub.withInterceptors(MetadataForwarder.clientInterceptor());

    return (T) stub;
  }

  public static InternalFuture<Void> setContextCredentials(
      ServiceUserConfig config, Executor executor) {
    return InternalFuture.runAsync(
        () -> {
          inplaceSetContextCredentials(config);
        },
        executor);
  }

  public static void inplaceSetContextCredentials(ServiceUserConfig config) {
    final Metadata authHeaders = new 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> dev_key_hyphen =
        Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
    Metadata.Key<String> source_key = Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

    authHeaders.put(email_key, config.email);
    authHeaders.put(dev_key, config.devKey);
    authHeaders.put(dev_key_hyphen, config.devKey);
    authHeaders.put(source_key, "PythonClient");

    // Force using the ROOT context so that we must lose any context of the current execution
    Context.ROOT.withValue(MetadataForwarder.METADATA_INFO, authHeaders).attach();
  }
}

@@ -529,16 +529,88 @@
 xecutor;
+%0Aimport java.util.function.Consumer;%0Aimport java.util.function.Supplier;
 %0A%0Apublic
@@ -1278,32 +1278,18 @@
 tic 
-InternalFuture%3CVoid%3E set
+%3CT%3E T with
 Cont
@@ -1303,23 +1303,16 @@
 entials(
-%0A      
 ServiceU
@@ -1333,90 +1333,310 @@
 ig, 
-Executor executor)
+Supplier%3CT%3E supplier) %7B%0A    Context previous = inplaceSetContextCredentials(config);%0A%0A    try
  %7B%0A    
+ 
+ 
 return 
-InternalFuture.runAsync(%0A        () -%3E %7B%0A         
+supplier.get();%0A    %7D finally %7B%0A      Context.current().detach(previous);%0A    %7D%0A  %7D%0A%0A  public static void withContextCredentials(ServiceUserConfig config, Runnable run) %7B%0A    Context previous =
  inp
@@ -1674,42 +1674,178 @@
 g);%0A
+%0A
     
-    %7D,%0A        executor);%0A  %7D%0A
+try %7B%0A      run.run();%0A    %7D finally %7B%0A      Context.current().detach(previous);%0A    %7D%0A  %7D%0A%0A  // DO NOT use outside of tests as this causes nasty grpc error messages
 %0A  p
@@ -1857,20 +1857,23 @@
  static 
-void
+Context
  inplace
@@ -2684,24 +2684,31 @@
 ecution%0A    
+return 
 Context.ROOT

Fix grpc context swap error messages (#2227)