package ai.verta.modeldb.common.connections;

import ai.verta.modeldb.common.CommonConstants;
import ai.verta.modeldb.common.CommonMessages;
import ai.verta.modeldb.common.authservice.AuthInterceptor;
import ai.verta.modeldb.common.config.Config;
import ai.verta.modeldb.common.exceptions.UnavailableException;
import ai.verta.modeldb.common.interceptors.MetadataForwarder;
import ai.verta.uac.AuthzServiceGrpc;
import ai.verta.uac.CollaboratorServiceGrpc;
import ai.verta.uac.RoleServiceGrpc;
import ai.verta.uac.UACServiceGrpc;
import ai.verta.uac.WorkspaceServiceGrpc;
import io.grpc.*;
import io.grpc.stub.AbstractStub;
import io.grpc.stub.MetadataUtils;
import io.opentracing.contrib.grpc.TracingClientInterceptor;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.Optional;

public class UAC extends Connection {
  private static final Logger LOGGER = LogManager.getLogger(UAC.class);

  private final ManagedChannel authServiceChannel;

  private final CollaboratorServiceGrpc.CollaboratorServiceFutureStub collaboratorServiceFutureStub;
  private final UACServiceGrpc.UACServiceFutureStub uacServiceFutureStub;
  private final WorkspaceServiceGrpc.WorkspaceServiceFutureStub workspaceServiceFutureStub;
  private final AuthzServiceGrpc.AuthzServiceFutureStub authzServiceFutureStub;
  private final RoleServiceGrpc.RoleServiceFutureStub roleServiceFutureStub;

  public static UAC FromConfig(Config config) {
    if (!config.hasAuth()) return null;
    else return new UAC(config);
  }

  private UAC(Config config) {
    this(
        config.authService.host,
        config.authService.port,
        config);
  }

  public UAC(
      String host, Integer port, Config config) {
    super(config);
    LOGGER.trace(CommonMessages.HOST_PORT_INFO_STR, host, port);
    if (host != null && port != null) { // AuthService not available.
      authServiceChannel =
          ManagedChannelBuilder.forTarget(host + CommonConstants.STRING_COLON + port)
              .usePlaintext()
              .build();
    } else {
      throw new UnavailableException(
          "Host OR Port not found for contacting authentication service");
    }

    collaboratorServiceFutureStub = CollaboratorServiceGrpc.newFutureStub(authServiceChannel);
    uacServiceFutureStub = UACServiceGrpc.newFutureStub(authServiceChannel);
    workspaceServiceFutureStub = WorkspaceServiceGrpc.newFutureStub(authServiceChannel);
    authzServiceFutureStub = AuthzServiceGrpc.newFutureStub(authServiceChannel);
    roleServiceFutureStub = RoleServiceGrpc.newFutureStub(authServiceChannel);
  }

  public CollaboratorServiceGrpc.CollaboratorServiceFutureStub getCollaboratorService() {
    return attachInterceptors(collaboratorServiceFutureStub);
  }

  public UACServiceGrpc.UACServiceFutureStub getUACService() {
    return attachInterceptors(uacServiceFutureStub);
  }

  public WorkspaceServiceGrpc.WorkspaceServiceFutureStub getWorkspaceService() {
    return attachInterceptors(workspaceServiceFutureStub);
  }

  public AuthzServiceGrpc.AuthzServiceFutureStub getAuthzService() {
    return attachInterceptors(authzServiceFutureStub);
  }

  public RoleServiceGrpc.RoleServiceFutureStub getRoleService() {
    return attachInterceptors(roleServiceFutureStub);
  }
}

@@ -138,68 +138,8 @@
 es;%0A
-import ai.verta.modeldb.common.authservice.AuthInterceptor;%0A
 impo
@@ -264,399 +264,31 @@
 rta.
-modeldb.common.interceptors.MetadataForwarder;%0Aimport ai.verta.uac.AuthzServiceGrpc;%0Aimport ai.verta.uac.CollaboratorServiceGrpc;%0Aimport ai.verta.uac.RoleServiceGrpc;%0Aimport ai.verta.uac.UACServiceGrpc;%0Aimport ai.verta.uac.WorkspaceServiceGrpc;%0Aimport io.grpc.*;%0Aimport io.grpc.stub.AbstractStub;%0Aimport io.grpc.stub.MetadataUtils;%0Aimport io.opentracing.contrib.grpc.TracingClientInterceptor
+uac.*;%0Aimport io.grpc.*
 ;%0Aim
@@ -374,36 +374,8 @@
 r;%0A%0A
-import java.util.Optional;%0A%0A
 publ
@@ -956,16 +956,117 @@
 ureStub;
+%0A  private final OrganizationServiceGrpc.OrganizationServiceFutureStub organizationServiceFutureStub;
 %0A%0A  publ
@@ -2253,16 +2253,111 @@
 annel);%0A
+    organizationServiceFutureStub = OrganizationServiceGrpc.newFutureStub(authServiceChannel);%0A
   %7D%0A%0A  p
@@ -2897,32 +2897,190 @@
 tureStub);%0A  %7D%0A%0A
+  public OrganizationServiceGrpc.OrganizationServiceFutureStub getOrganizationService( ) %7B%0A    return attachInterceptors(organizationServiceFutureStub);%0A  %7D%0A%0A
   public RoleSer

Allow registry to access organization service so we can create and delete orgs for tests (#2347)