connection manager: cleanups (#2332)

* connman: address review in #2323 * connman: tidy up connection builder interface * connman: improve immutability * connman: cleanup ConnectionFactory * connman: cleanup
finos · Oct 2, 2023 · 2eb7c83 · 2eb7c83
1 parent d948bd9
commit 2eb7c83
Show file tree

Hide file tree

Showing 38 changed files with 473 additions and 538 deletions.
diff --git a/legend-engine-config/legend-engine-connection-integration-tests/pom.xml b/legend-engine-config/legend-engine-connection-integration-tests/pom.xml
@@ -74,26 +74,7 @@
         </dependency>
         <!-- RUNTIME -->
 
-        <!-- ECLIPSE COLLECTIONS -->
-<!--        <dependency>-->
-<!--            <groupId>org.eclipse.collections</groupId>-->
-<!--            <artifactId>eclipse-collections-api</artifactId>-->
-<!--        </dependency>-->
-        <!-- ECLIPSE COLLECTIONS -->
-
-        <!-- Authentication -->
-<!--        <dependency>-->
-<!--            <groupId>org.pac4j</groupId>-->
-<!--            <artifactId>pac4j-core</artifactId>-->
-<!--        </dependency>-->
-        <!-- Authentication -->
-
         <!-- TEST -->
-        <dependency>
-            <groupId>junit</groupId>
-            <artifactId>junit</artifactId>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-api</artifactId>

diff --git a/.../src/test/java/org/finos/legend/engine/connection/test/AbstractConnectionFactoryTest.java b/.../src/test/java/org/finos/legend/engine/connection/test/AbstractConnectionFactoryTest.java
@@ -20,9 +20,10 @@
 import org.finos.legend.connection.Authenticator;
 import org.finos.legend.connection.ConnectionFactory;
 import org.finos.legend.connection.DatabaseType;
-import org.finos.legend.connection.EnvironmentConfiguration;
 import org.finos.legend.connection.IdentityFactory;
 import org.finos.legend.connection.IdentitySpecification;
+import org.finos.legend.connection.InstrumentedStoreInstanceProvider;
+import org.finos.legend.connection.LegendEnvironment;
 import org.finos.legend.connection.RelationalDatabaseStoreSupport;
 import org.finos.legend.connection.StoreInstance;
 import org.finos.legend.connection.impl.KerberosCredentialExtractor;
@@ -41,16 +42,17 @@ public abstract class AbstractConnectionFactoryTest<T>
 {
     protected static final String TEST_STORE_INSTANCE_NAME = "test-store";
 
-    protected EnvironmentConfiguration environmentConfiguration;
+    protected LegendEnvironment environment;
     protected IdentityFactory identityFactory;
+    protected InstrumentedStoreInstanceProvider storeInstanceProvider;
     protected ConnectionFactory connectionFactory;
 
     @BeforeEach
     public void initialize()
     {
         this.setup();
 
-        EnvironmentConfiguration.Builder environmentConfigurationBuilder = new EnvironmentConfiguration.Builder()
+        LegendEnvironment.Builder environmentBuilder = new LegendEnvironment.Builder()
                 .withVaults(
                         new SystemPropertiesCredentialVault(),
                         new EnvironmentCredentialVault()
@@ -82,15 +84,16 @@ public void initialize()
         CredentialVault credentialVault = this.getCredentialVault();
         if (credentialVault != null)
         {
-            environmentConfigurationBuilder.withVault(credentialVault);
+            environmentBuilder.withVault(credentialVault);
         }
 
-        this.environmentConfiguration = environmentConfigurationBuilder.build();
+        this.environment = environmentBuilder.build();
 
-        this.identityFactory = new IdentityFactory.Builder(environmentConfiguration)
+        this.identityFactory = new IdentityFactory.Builder(this.environment)
                 .build();
 
-        this.connectionFactory = new ConnectionFactory.Builder(environmentConfiguration)
+        this.storeInstanceProvider = new InstrumentedStoreInstanceProvider();
+        this.connectionFactory = new ConnectionFactory.Builder(this.environment, this.storeInstanceProvider)
                 .withCredentialBuilders(
                         new KerberosCredentialExtractor(),
                         new UserPasswordCredentialBuilder(),
@@ -129,12 +132,12 @@ public CredentialVault getCredentialVault()
     @Test
     public void runTest() throws Exception
     {
-        this.connectionFactory.injectStoreInstance(this.getStoreInstance());
+        this.storeInstanceProvider.injectStoreInstance(this.getStoreInstance());
         Identity identity = this.getIdentity();
         AuthenticationConfiguration authenticationConfiguration = this.getAuthenticationConfiguration();
 
         Authenticator authenticator = this.connectionFactory.getAuthenticator(identity, TEST_STORE_INSTANCE_NAME, authenticationConfiguration);
-        T connection = this.connectionFactory.getConnection(authenticator);
+        T connection = this.connectionFactory.getConnection(identity, authenticator);
 
         this.runTestWithConnection(connection);
         System.out.println("Successfully established and checked connection!");

diff --git a/...n-tests/src/test/java/org/finos/legend/engine/connection/test/TestPostgresConnection.java b/...n-tests/src/test/java/org/finos/legend/engine/connection/test/TestPostgresConnection.java
@@ -30,7 +30,7 @@
 import java.sql.Statement;
 import java.util.Properties;
 
-import static org.junit.Assume.assumeTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 public class TestPostgresConnection
 {
@@ -48,7 +48,7 @@ public void setup()
             }
             catch (Exception e)
             {
-                assumeTrue("Can't start PostgreSQLContainer", false);
+                assumeTrue(false, "Can't start PostgreSQLContainer");
             }
         }
 
@@ -77,7 +77,7 @@ public StoreInstance getStoreInstance()
                     this.postgresContainer.getPort(),
                     this.postgresContainer.getDatabaseName()
             );
-            return new StoreInstance.Builder(this.environmentConfiguration)
+            return new StoreInstance.Builder(this.environment)
                     .withIdentifier(TEST_STORE_INSTANCE_NAME)
                     .withStoreSupportIdentifier("Postgres")
                     .withAuthenticationMechanisms(

diff --git a/...-tests/src/test/java/org/finos/legend/engine/connection/test/TestSnowflakeConnection.java b/...-tests/src/test/java/org/finos/legend/engine/connection/test/TestSnowflakeConnection.java
@@ -29,14 +29,14 @@
 import java.sql.Statement;
 import java.util.Properties;
 
-import static org.junit.Assume.assumeTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 public class TestSnowflakeConnection
 {
     public static class WithKeyPair extends AbstractConnectionFactoryTest<Connection>
     {
-        private static final String TEST_SNOWFLAKE_PK = "TEST_SNOWFLAKE_PK";
-        private static final String TEST_SNOWFLAKE_PK_PASSPHRASE = "TEST_SNOWFLAKE_PK_PASSPHRASE";
+        private static final String CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK = "CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK";
+        private static final String CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK_PASSPHRASE = "CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK_PASSPHRASE";
         private String snowflakePrivateKey;
         private String snowflakePassPhrase;
 
@@ -45,12 +45,12 @@ public void setup()
         {
             try
             {
-                this.snowflakePrivateKey = this.environmentConfiguration.lookupVaultSecret(new EnvironmentCredentialVaultSecret(TEST_SNOWFLAKE_PK), null);
-                this.snowflakePassPhrase = this.environmentConfiguration.lookupVaultSecret(new EnvironmentCredentialVaultSecret(TEST_SNOWFLAKE_PK_PASSPHRASE), null);
+                this.snowflakePrivateKey = this.environment.lookupVaultSecret(new EnvironmentCredentialVaultSecret(CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK), null);
+                this.snowflakePassPhrase = this.environment.lookupVaultSecret(new EnvironmentCredentialVaultSecret(CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK_PASSPHRASE), null);
             }
             catch (Exception e)
             {
-                assumeTrue("Can't retrieve Snowflake test instance key-pair info (TEST_SNOWFLAKE_PK, TEST_SNOWFLAKE_PK_PASSPHRASE)", false);
+                assumeTrue(false, String.format("Can't retrieve Snowflake connection key-pair info (%s, %s environment variables are expected)", CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK, CONNECTION_INTEGRATION_TEST__SNOWFLAKE_PK_PASSPHRASE));
             }
         }
 
@@ -79,7 +79,7 @@ public StoreInstance getStoreInstance()
             connectionSpecification.region = "us-east-2";
             connectionSpecification.cloudType = "aws";
             connectionSpecification.role = "SUMMIT_DEV";
-            return new StoreInstance.Builder(this.environmentConfiguration)
+            return new StoreInstance.Builder(this.environment)
                     .withIdentifier(TEST_STORE_INSTANCE_NAME)
                     .withStoreSupportIdentifier("Snowflake")
                     .withAuthenticationMechanisms(

diff --git a/...tion-factory/src/main/java/org/finos/legend/connection/AuthenticationMechanismLoader.java b/...tion-factory/src/main/java/org/finos/legend/connection/AuthenticationMechanismLoader.java
diff --git a/...on-factory/src/main/java/org/finos/legend/connection/AuthenticationMechanismProvider.java b/...on-factory/src/main/java/org/finos/legend/connection/AuthenticationMechanismProvider.java
diff --git a/...ntication-connection-factory/src/main/java/org/finos/legend/connection/Authenticator.java b/...ntication-connection-factory/src/main/java/org/finos/legend/connection/Authenticator.java
@@ -14,6 +14,8 @@
 
 package org.finos.legend.connection;
 
+import org.eclipse.collections.api.factory.Lists;
+import org.eclipse.collections.api.list.ImmutableList;
 import org.finos.legend.connection.protocol.AuthenticationConfiguration;
 import org.finos.legend.engine.shared.core.identity.Credential;
 import org.finos.legend.engine.shared.core.identity.Identity;
@@ -23,31 +25,28 @@
 
 public class Authenticator
 {
-    private final Identity identity;
     private final StoreInstance storeInstance;
     private final AuthenticationConfiguration authenticationConfiguration;
     private final Class<? extends Credential> sourceCredentialType;
-
-    private final List<CredentialBuilder> credentialBuilders;
+    private final ImmutableList<CredentialBuilder> credentialBuilders;
     private final ConnectionBuilder connectionBuilder;
 
-    public Authenticator(Identity identity, StoreInstance storeInstance, AuthenticationConfiguration authenticationConfiguration, Class<? extends Credential> sourceCredentialType, List<CredentialBuilder> credentialBuilders, ConnectionBuilder connectionBuilder)
+    public Authenticator(StoreInstance storeInstance, AuthenticationConfiguration authenticationConfiguration, Class<? extends Credential> sourceCredentialType, List<CredentialBuilder> credentialBuilders, ConnectionBuilder connectionBuilder)
     {
-        this.identity = identity;
         this.storeInstance = storeInstance;
         this.authenticationConfiguration = authenticationConfiguration;
         this.sourceCredentialType = sourceCredentialType;
-        this.credentialBuilders = credentialBuilders;
+        this.credentialBuilders = Lists.immutable.withAll(credentialBuilders);
         this.connectionBuilder = connectionBuilder;
     }
 
-    public Credential makeCredential(EnvironmentConfiguration configuration) throws Exception
+    public Credential makeCredential(Identity identity, LegendEnvironment environment) throws Exception
     {
         Credential credential = null;
         // no need to resolve the source credential if the flow starts with generic `Credential` node
         if (!this.sourceCredentialType.equals(Credential.class))
         {
-            Optional<Credential> credentialOptional = this.identity.getCredential((Class<Credential>) this.sourceCredentialType);
+            Optional<Credential> credentialOptional = identity.getCredential((Class<Credential>) this.sourceCredentialType);
             if (!credentialOptional.isPresent())
             {
                 throw new RuntimeException(String.format("Can't resolve source credential of type '%s' from the specified identity", this.sourceCredentialType.getSimpleName()));
@@ -59,7 +58,7 @@ public Credential makeCredential(EnvironmentConfiguration configuration) throws
         }
         for (CredentialBuilder credentialBuilder : this.credentialBuilders)
         {
-            credential = credentialBuilder.makeCredential(this.identity, this.authenticationConfiguration, credential, configuration);
+            credential = credentialBuilder.makeCredential(identity, this.authenticationConfiguration, credential, environment);
         }
         return credential;
     }
@@ -79,7 +78,7 @@ public Class<? extends Credential> getSourceCredentialType()
         return sourceCredentialType;
     }
 
-    public List<CredentialBuilder> getCredentialBuilders()
+    public ImmutableList<CredentialBuilder> getCredentialBuilders()
     {
         return credentialBuilders;
     }

diff --git a/...ation-connection-factory/src/main/java/org/finos/legend/connection/ConnectionBuilder.java b/...ation-connection-factory/src/main/java/org/finos/legend/connection/ConnectionBuilder.java
@@ -21,9 +21,9 @@
 import java.lang.reflect.Type;
 import java.util.Objects;
 
-public abstract class ConnectionBuilder<T, CRED extends Credential, SPEC extends ConnectionSpecification>
+public abstract class ConnectionBuilder<CONNECTION, CRED extends Credential, SPEC extends ConnectionSpecification>
 {
-    public abstract T getConnection(CRED credential, SPEC connectionSpecification, StoreInstance storeInstance) throws Exception;
+    public abstract CONNECTION getConnection(StoreInstance storeInstance, CRED credential) throws Exception;
 
     public Class<? extends Credential> getCredentialType()
     {
@@ -42,6 +42,11 @@ private Type[] actualTypeArguments()
         return parameterizedType.getActualTypeArguments();
     }
 
+    protected SPEC getCompatibleConnectionSpecification(StoreInstance storeInstance)
+    {
+        return (SPEC) storeInstance.getConnectionSpecification(this.getConnectionSpecificationType());
+    }
+
     public static class Key
     {
         private final Class<? extends ConnectionSpecification> connectionSpecificationType;