connman: cleanup bootstraper for new connection framework

finos · Oct 11, 2023 · 1e2d9d1 · 1e2d9d1
1 parent 8f0f31f
commit 1e2d9d1
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 19 deletions.
diff --git a/...gine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java b/...gine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java
@@ -45,17 +45,18 @@
 import org.finos.legend.connection.AuthenticationMechanismConfiguration;
 import org.finos.legend.connection.ConnectionFactory;
 import org.finos.legend.connection.DatabaseType;
-import org.finos.legend.connection.impl.DefaultStoreInstanceProvider;
 import org.finos.legend.connection.LegendEnvironment;
 import org.finos.legend.connection.RelationalDatabaseStoreSupport;
 import org.finos.legend.connection.StoreInstanceProvider;
+import org.finos.legend.connection.impl.DefaultStoreInstanceProvider;
 import org.finos.legend.connection.impl.EncryptedPrivateKeyPairAuthenticationConfiguration;
+import org.finos.legend.connection.impl.HACKY__SnowflakeConnectionAdapter;
 import org.finos.legend.connection.impl.KerberosCredentialExtractor;
 import org.finos.legend.connection.impl.KeyPairCredentialBuilder;
 import org.finos.legend.connection.impl.SnowflakeConnectionBuilder;
+import org.finos.legend.connection.impl.StaticJDBCConnectionBuilder;
 import org.finos.legend.connection.impl.UserPasswordAuthenticationConfiguration;
 import org.finos.legend.connection.impl.UserPasswordCredentialBuilder;
-import org.finos.legend.connection.impl.StaticJDBCConnectionBuilder;
 import org.finos.legend.connection.protocol.AuthenticationMechanismType;
 import org.finos.legend.engine.api.analytics.DataSpaceAnalytics;
 import org.finos.legend.engine.api.analytics.DiagramAnalytics;
@@ -279,6 +280,9 @@ public void run(T serverConfiguration, Environment environment)
             relationalExecution.setFlowProviderConfiguration(new LegendDefaultDatabaseAuthenticationFlowProviderConfiguration());
         }
         relationalExecution.setConnectionFactory(this.setupConnectionFactory(serverConfiguration.vaults));
+        relationalExecution.setRelationalDatabaseConnectionAdapters(Lists.mutable.of(
+                new HACKY__SnowflakeConnectionAdapter.WithKeyPair()
+        ));
 
         relationalStoreExecutor = (RelationalStoreExecutor) Relational.build(serverConfiguration.relationalexecution);
 

diff --git a/...src/main/java/org/finos/legend/connection/HACKY__RelationalDatabaseConnectionAdapter.java b/...src/main/java/org/finos/legend/connection/HACKY__RelationalDatabaseConnectionAdapter.java
@@ -16,6 +16,7 @@
 
 import org.finos.legend.connection.protocol.AuthenticationConfiguration;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.shared.core.identity.Identity;
 
 /**
  * NOTE: this is hacky way of us to realize the relational database connection that we can support
@@ -24,7 +25,7 @@
  */
 public interface HACKY__RelationalDatabaseConnectionAdapter
 {
-    ConnectionFactoryMaterial adapt(RelationalDatabaseConnection relationalDatabaseConnection, LegendEnvironment environment);
+    ConnectionFactoryMaterial adapt(RelationalDatabaseConnection relationalDatabaseConnection, Identity identity, LegendEnvironment environment);
 
     class ConnectionFactoryMaterial
     {

diff --git a/...ion/src/main/java/org/finos/legend/connection/impl/HACKY__SnowflakeConnectionAdapter.java b/...ion/src/main/java/org/finos/legend/connection/impl/HACKY__SnowflakeConnectionAdapter.java
@@ -23,13 +23,14 @@
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.authentication.SnowflakePublicAuthenticationStrategy;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.specification.SnowflakeDatasourceSpecification;
+import org.finos.legend.engine.shared.core.identity.Identity;
 
 public class HACKY__SnowflakeConnectionAdapter
 {
     public static class WithKeyPair implements HACKY__RelationalDatabaseConnectionAdapter
     {
         @Override
-        public ConnectionFactoryMaterial adapt(RelationalDatabaseConnection relationalDatabaseConnection, LegendEnvironment environment)
+        public ConnectionFactoryMaterial adapt(RelationalDatabaseConnection relationalDatabaseConnection, Identity identity, LegendEnvironment environment)
         {
             if (
                     DatabaseType.Snowflake.equals(relationalDatabaseConnection.databaseType) &&

diff --git a/...ain/java/org/finos/legend/engine/plan/execution/stores/relational/RelationalExecutor.java b/...ain/java/org/finos/legend/engine/plan/execution/stores/relational/RelationalExecutor.java
@@ -107,7 +107,7 @@ public RelationalExecutor(TemporaryTestDbConfiguration temporarytestdb, Relation
     public RelationalExecutor(TemporaryTestDbConfiguration temporarytestdb, RelationalExecutionConfiguration relationalExecutionConfiguration, Optional<DatabaseAuthenticationFlowProvider> flowProviderHolder)
     {
         this.flowProviderHolder = flowProviderHolder;
-        this.connectionManager = new ConnectionManagerSelector(temporarytestdb, relationalExecutionConfiguration.oauthProfiles, flowProviderHolder, relationalExecutionConfiguration.getConnectionFactory());
+        this.connectionManager = new ConnectionManagerSelector(temporarytestdb, relationalExecutionConfiguration.oauthProfiles, flowProviderHolder, relationalExecutionConfiguration.getConnectionFactory(), relationalExecutionConfiguration.getRelationalDatabaseConnectionAdapters(), false);
         this.relationalExecutionConfiguration = relationalExecutionConfiguration;
         this.resultInterpreterExtensions = Iterate.addAllTo(ResultInterpreterExtensionLoader.extensions(), Lists.mutable.empty()).collect(ResultInterpreterExtension::additionalResultBuilder);
     }

diff --git a/...gend/engine/plan/execution/stores/relational/config/RelationalExecutionConfiguration.java b/...gend/engine/plan/execution/stores/relational/config/RelationalExecutionConfiguration.java
@@ -18,6 +18,7 @@
 import org.eclipse.collections.api.factory.Lists;
 import org.finos.legend.authentication.credentialprovider.CredentialProviderProvider;
 import org.finos.legend.connection.ConnectionFactory;
+import org.finos.legend.connection.HACKY__RelationalDatabaseConnectionAdapter;
 import org.finos.legend.engine.authentication.provider.DatabaseAuthenticationFlowProvider;
 import org.finos.legend.engine.authentication.provider.DatabaseAuthenticationFlowProviderConfiguration;
 import org.finos.legend.engine.plan.execution.stores.StoreExecutorConfiguration;
@@ -37,6 +38,7 @@ public class RelationalExecutionConfiguration implements StoreExecutorConfigurat
     @JsonProperty
     private RelationalGraphFetchExecutionConfig relationalGraphFetchExecutionConfig;
     private ConnectionFactory connectionFactory;
+    private List<HACKY__RelationalDatabaseConnectionAdapter> relationalDatabaseConnectionAdapters = Lists.mutable.empty();
 
     @Override
     public StoreType getStoreType()
@@ -78,6 +80,11 @@ public ConnectionFactory getConnectionFactory()
         return connectionFactory;
     }
 
+    public List<HACKY__RelationalDatabaseConnectionAdapter> getRelationalDatabaseConnectionAdapters()
+    {
+        return relationalDatabaseConnectionAdapters;
+    }
+
     public void setCredentialProviderProvider(CredentialProviderProvider credentialProviderProvider)
     {
         this.credentialProviderProvider = credentialProviderProvider;
@@ -98,6 +105,11 @@ public void setConnectionFactory(ConnectionFactory connectionFactory)
         this.connectionFactory = connectionFactory;
     }
 
+    public void setRelationalDatabaseConnectionAdapters(List<HACKY__RelationalDatabaseConnectionAdapter> adapters)
+    {
+        this.relationalDatabaseConnectionAdapters = adapters;
+    }
+
     public static Builder newInstance()
     {
         return new Builder();
@@ -118,6 +130,7 @@ public static class Builder
         private CredentialProviderProvider credentialProviderProvider;
         private RelationalGraphFetchExecutionConfig relationalGraphFetchExecutionConfig;
         private ConnectionFactory connectionFactory;
+        private final List<HACKY__RelationalDatabaseConnectionAdapter> relationalDatabaseConnectionAdapters = Lists.mutable.empty();
 
         public Builder withTempPath(String tempPath)
         {
@@ -170,6 +183,12 @@ public Builder withConnectionFactory(ConnectionFactory connectionFactory)
             return this;
         }
 
+        public Builder withRelationalDatabaseConnectionAdapters(List<HACKY__RelationalDatabaseConnectionAdapter> adapters)
+        {
+            this.relationalDatabaseConnectionAdapters.addAll(adapters);
+            return this;
+        }
+
         public RelationalExecutionConfiguration build()
         {
             RelationalExecutionConfiguration relationalExecutionConfiguration = new RelationalExecutionConfiguration();
@@ -181,6 +200,7 @@ public RelationalExecutionConfiguration build()
             relationalExecutionConfiguration.credentialProviderProvider = credentialProviderProvider;
             relationalExecutionConfiguration.relationalGraphFetchExecutionConfig = relationalGraphFetchExecutionConfig;
             relationalExecutionConfiguration.connectionFactory = connectionFactory;
+            relationalExecutionConfiguration.relationalDatabaseConnectionAdapters = relationalDatabaseConnectionAdapters;
             return relationalExecutionConfiguration;
         }
     }

diff --git a/...engine/plan/execution/stores/relational/connection/manager/ConnectionManagerSelector.java b/...engine/plan/execution/stores/relational/connection/manager/ConnectionManagerSelector.java
@@ -68,18 +68,7 @@ public ConnectionManagerSelector(TemporaryTestDbConfiguration temporaryTestDb, L
         this.connectionFactory = null;
     }
 
-    public ConnectionManagerSelector(TemporaryTestDbConfiguration temporaryTestDb, List<OAuthProfile> oauthProfiles, Optional<DatabaseAuthenticationFlowProvider> flowProviderHolder, ConnectionFactory connectionFactory)
-    {
-        MutableList<ConnectionManagerExtension> extensions = Iterate.addAllTo(ServiceLoader.load(ConnectionManagerExtension.class), Lists.mutable.empty());
-        this.connectionManagers = Lists.mutable.<ConnectionManager>with(
-                new RelationalConnectionManager(temporaryTestDb.port, oauthProfiles, flowProviderHolder)
-        ).withAll(extensions.collect(e -> e.getExtensionManager(temporaryTestDb.port, oauthProfiles)));
-        this.flowProviderHolder = flowProviderHolder;
-        this.connectionFactory = connectionFactory;
-        this.relationalDatabaseConnectionAdapters.addAll(Lists.mutable.withAll(ServiceLoader.load(HACKY__RelationalDatabaseConnectionAdapter.class)));
-    }
-
-    public ConnectionManagerSelector(TemporaryTestDbConfiguration temporaryTestDb, List<OAuthProfile> oauthProfiles, Optional<DatabaseAuthenticationFlowProvider> flowProviderHolder, ConnectionFactory connectionFactory, List<HACKY__RelationalDatabaseConnectionAdapter> relationalDatabaseConnectionAdapters)
+    public ConnectionManagerSelector(TemporaryTestDbConfiguration temporaryTestDb, List<OAuthProfile> oauthProfiles, Optional<DatabaseAuthenticationFlowProvider> flowProviderHolder, ConnectionFactory connectionFactory, List<HACKY__RelationalDatabaseConnectionAdapter> relationalDatabaseConnectionAdapters, boolean enableNewConnectionFrameworkByDefault)
     {
         MutableList<ConnectionManagerExtension> extensions = Iterate.addAllTo(ServiceLoader.load(ConnectionManagerExtension.class), Lists.mutable.empty());
         this.connectionManagers = Lists.mutable.<ConnectionManager>with(
@@ -89,7 +78,7 @@ public ConnectionManagerSelector(TemporaryTestDbConfiguration temporaryTestDb, L
 
         this.connectionFactory = connectionFactory;
         this.relationalDatabaseConnectionAdapters.addAll(relationalDatabaseConnectionAdapters);
-        this.enableNewConnectionFramework = true;
+        this.enableNewConnectionFramework = enableNewConnectionFrameworkByDefault;
     }
 
     public Optional<DatabaseAuthenticationFlowProvider> getFlowProviderHolder()
@@ -155,7 +144,7 @@ public Connection getDatabaseConnectionImpl(Identity identity, DatabaseConnectio
                     HACKY__RelationalDatabaseConnectionAdapter.ConnectionFactoryMaterial connectionFactoryMaterial = null;
                     for (HACKY__RelationalDatabaseConnectionAdapter adapter : this.relationalDatabaseConnectionAdapters)
                     {
-                        connectionFactoryMaterial = adapter.adapt(relationalDatabaseConnection, this.connectionFactory.getEnvironment());
+                        connectionFactoryMaterial = adapter.adapt(relationalDatabaseConnection, identity, this.connectionFactory.getEnvironment());
                         if (connectionFactoryMaterial != null)
                         {
                             break;