org.finos.legend.connection.ConnectionFactory Maven / Gradle / Ivy
// Copyright 2023 Goldman Sachs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package org.finos.legend.connection;
import org.eclipse.collections.api.factory.Lists;
import org.finos.legend.engine.protocol.pure.v1.packageableElement.connection.AuthenticationConfiguration;
import org.finos.legend.engine.protocol.pure.v1.packageableElement.connection.ConnectionSpecification;
import org.finos.legend.engine.shared.core.identity.Credential;
import org.finos.legend.engine.shared.core.identity.Identity;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
public class ConnectionFactory
{
private final LegendEnvironment environment;
private final Map credentialBuildersIndex = new LinkedHashMap<>();
private final Map connectionBuildersIndex = new LinkedHashMap<>();
private ConnectionFactory(LegendEnvironment environment, List credentialBuilders, List connectionBuilders)
{
this.environment = Objects.requireNonNull(environment, "environment is missing");
for (ConnectionBuilder builder : connectionBuilders)
{
this.connectionBuildersIndex.put(new ConnectionBuilder.Key(builder.getConnectionSpecificationType(), builder.getCredentialType()), builder);
}
for (CredentialBuilder builder : credentialBuilders)
{
this.credentialBuildersIndex.put(new CredentialBuilder.Key(builder.getAuthenticationConfigurationType(), builder.getInputCredentialType(), builder.getOutputCredentialType()), builder);
}
}
public LegendEnvironment getEnvironment()
{
return environment;
}
public Authenticator getAuthenticator(Identity identity, Connection connection, AuthenticationConfiguration authenticationConfiguration)
{
AuthenticationMechanismType authenticationMechanismType = Objects.requireNonNull(connection.getAuthenticationMechanism(authenticationConfiguration.getClass()), String.format("Connection '%s' is not compatible with authentication configuration type '%s'. Supported configuration type(s):\n%s",
connection.getIdentifier(),
authenticationConfiguration.getClass().getSimpleName(),
connection.getAuthenticationConfigurationTypes().collect(configType -> "- " + configType.getSimpleName()).makeString("\n")
));
return this.getAuthenticator(identity, connection, authenticationMechanismType, authenticationConfiguration);
}
private Authenticator getAuthenticator(Identity identity, Connection connection, AuthenticationMechanismType authenticationMechanismType, AuthenticationConfiguration authenticationConfiguration)
{
AuthenticationFlowResolver.ResolutionResult result = AuthenticationFlowResolver.run(this.credentialBuildersIndex, this.connectionBuildersIndex, identity, authenticationMechanismType, authenticationConfiguration, connection.getConnectionSpecification());
if (result == null)
{
throw new RuntimeException(String.format("No authentication flow for connection '%s' can be resolved for the specified identity (authentication configuration: %s, connection specification: %s)",
connection.getIdentifier(),
authenticationConfiguration.getClass().getSimpleName(),
connection.getConnectionSpecification().getClass().getSimpleName()
));
}
return new Authenticator(connection, authenticationMechanismType, authenticationConfiguration, result.sourceCredentialType, result.targetCredentialType, result.flow, connectionBuildersIndex.get(new ConnectionBuilder.Key(connection.getConnectionSpecification().getClass(), result.targetCredentialType)), this.environment);
}
public Authenticator getAuthenticator(Identity identity, Connection connection)
{
Authenticator authenticator = null;
for (AuthenticationMechanismType authenticationMechanismType : connection.getAuthenticationMechanisms())
{
AuthenticationMechanism authenticationMechanism = connection.getAuthenticationMechanism(authenticationMechanismType);
AuthenticationConfiguration authenticationConfiguration = connection.getAuthenticationConfiguration();
if (authenticationConfiguration != null)
{
AuthenticationFlowResolver.ResolutionResult result = AuthenticationFlowResolver.run(this.credentialBuildersIndex, this.connectionBuildersIndex, identity, authenticationMechanismType, authenticationConfiguration, connection.getConnectionSpecification());
if (result != null)
{
authenticator = new Authenticator(connection, authenticationMechanismType, authenticationConfiguration, result.sourceCredentialType, result.targetCredentialType, result.flow, connectionBuildersIndex.get(new ConnectionBuilder.Key(connection.getConnectionSpecification().getClass(), result.targetCredentialType)), this.environment);
break;
}
}
}
if (authenticator == null)
{
throw new RuntimeException(String.format("No authentication flow for connection '%s' can be resolved for the specified identity. Try specifying another authentication configuration. Supported configuration type(s):\n%s",
connection.getIdentifier(),
connection.getAuthenticationConfigurationTypes().collect(configType -> "- " + configType.getSimpleName() + " (" + connection.getAuthenticationMechanism(configType).getIdentifier() + ")").makeString("\n")
));
}
return authenticator;
}
private static class AuthenticationFlowResolver
{
private final Map credentialBuildersIndex = new HashMap<>();
private final Set nodes = new HashSet<>();
private final Map> edges = new HashMap<>();
private final FlowNode startNode;
private final FlowNode endNode;
/**
* This constructor sets up the authentication flow (directed non-cyclic) graph to help with flow resolution
*
* The start node is the identity and the end node is the connection
* The immediately adjacent nodes to the start node are credential nodes
* The remaining nodes are credential-type nodes
*
* The edges coming out from the start node correspond to credentials that the identity comes with
* The edges going to end node correspond to available connection builders
* The remaining edges correspond to available credential builders
*
* NOTE:
* - Since some credential builders do not require a specific input credential type, we added a generic `Credential` node
* to Identity (start node)
* - We want to differentiate credential and credential-type nodes because we want to account for (short-circuit) cases where
* no resolution is needed: some credentials that belong to the identity is enough to build the connection (e.g. Kerberos).
* We want to be very explicit about this case, we don't want this behavior to be generic for all types of credentials; for example,
* just because an identity comes with a username-password credential, does not mean this credential is appropriate to be used to
* connect to a database which supports username-password authentication mechanism, unless this intention is explicitly stated.
*
* With this setup, we can use a basic graph search algorithm (e.g. BFS) to resolve the shortest path to build a connection
*/
private AuthenticationFlowResolver(Map credentialBuildersIndex, Map connectionBuildersIndex, Identity identity, AuthenticationConfiguration authenticationConfiguration, AuthenticationMechanismType authenticationMechanismType, ConnectionSpecification connectionSpecification)
{
// add start node (i.e. identity node)
this.startNode = new FlowNode(identity);
// add identity's credential nodes
identity.getCredentials().forEach(cred -> this.processEdge(this.startNode, new FlowNode(identity, cred.getClass())));
// add special `Credential` node for catch-all credential builders
this.processEdge(this.startNode, new FlowNode(identity, Credential.class));
// process credential builders
credentialBuildersIndex.values().stream()
.filter(builder -> builder.getAuthenticationConfigurationType().equals(authenticationConfiguration.getClass()))
.forEach(builder ->
{
if (!(builder.getInputCredentialType().equals(builder.getOutputCredentialType())))
{
this.processEdge(new FlowNode(builder.getInputCredentialType()), new FlowNode(builder.getOutputCredentialType()));
}
this.processEdge(new FlowNode(identity, builder.getInputCredentialType()), new FlowNode(builder.getOutputCredentialType()));
this.credentialBuildersIndex.put(createCredentialBuilderKey(builder.getInputCredentialType().getSimpleName(), builder.getOutputCredentialType().getSimpleName()), builder);
});
// add end node (i.e. connection node)
this.endNode = new FlowNode(connectionSpecification);
// process connection builders
connectionBuildersIndex.values().stream()
.filter(builder -> builder.getConnectionSpecificationType().equals(connectionSpecification.getClass()))
.forEach(builder -> this.processEdge(new FlowNode(builder.getCredentialType()), this.endNode));
}
static String createCredentialBuilderKey(String inputCredentialType, String outputCredentialType)
{
return inputCredentialType + "__" + outputCredentialType;
}
private void processEdge(FlowNode node, FlowNode adjacentNode)
{
this.nodes.add(node);
this.nodes.add(adjacentNode);
if (this.edges.get(node.id) != null)
{
Set adjacentNodes = this.edges.get(node.id);
adjacentNodes.add(adjacentNode);
}
else
{
LinkedHashSet adjacentNodes = new LinkedHashSet<>();
adjacentNodes.add(adjacentNode);
this.edges.put(node.id, adjacentNodes);
}
}
/**
* Resolves the authentication flow in order to build a connection for a specified identity
*/
public static ResolutionResult run(Map credentialBuildersIndex, Map connectionBuildersIndex, Identity identity, AuthenticationMechanismType authenticationMechanismType, AuthenticationConfiguration authenticationConfiguration, ConnectionSpecification connectionSpecification)
{
// using BFS algo to search for the shortest (non-cyclic) path
AuthenticationFlowResolver state = new AuthenticationFlowResolver(credentialBuildersIndex, connectionBuildersIndex, identity, authenticationConfiguration, authenticationMechanismType, connectionSpecification);
boolean found = false;
Set visitedNodes = new HashSet<>(); // Create a set to keep track of visited vertices
Deque queue = new ArrayDeque<>();
queue.add(state.startNode);
Map distances = new HashMap<>();
Map previousNodes = new HashMap<>();
state.nodes.forEach(node -> distances.put(node.id, Integer.MAX_VALUE));
distances.put(state.startNode.id, 0);
while (!queue.isEmpty())
{
if (found)
{
break;
}
FlowNode node = queue.removeFirst();
visitedNodes.add(node);
if (state.edges.get(node.id) != null)
{
for (FlowNode adjNode : state.edges.get(node.id))
{
if (!visitedNodes.contains(adjNode))
{
distances.put(adjNode.id, distances.get(node.id) + 1);
previousNodes.put(adjNode.id, node);
queue.addLast(adjNode);
if (adjNode.equals(state.endNode))
{
found = true;
break;
}
}
}
}
}
if (!found)
{
return null;
}
// resolve the path
LinkedList nodes = new LinkedList<>();
FlowNode currentNode = previousNodes.get(connectionSpecification.getClass().getSimpleName());
while (!state.startNode.equals(currentNode))
{
nodes.addFirst(currentNode);
currentNode = previousNodes.get(currentNode.id);
}
if (nodes.size() < 2)
{
throw new IllegalStateException("Can't resolve connection authentication flow for specified identity: invalid flow state found!");
}
List flow = new ArrayList<>();
for (int i = 0; i < nodes.size() - 1; i++)
{
flow.add(Objects.requireNonNull(
state.credentialBuildersIndex.get(createCredentialBuilderKey(nodes.get(i).credentialType.getSimpleName(), nodes.get(i + 1).credentialType.getSimpleName())),
String.format("Can't find a matching credential builder (input: %s, output: %s)", nodes.get(i).credentialType.getSimpleName(), nodes.get(i + 1).credentialType.getSimpleName()
)));
}
return new ResolutionResult(flow, nodes.get(0).credentialType, nodes.get(nodes.size() - 1).credentialType);
}
private static class FlowNode
{
private static final String IDENTITY_NODE_ID = "__identity__";
private static final String IDENTITY_CREDENTIAL_NODE_ID_PREFIX = "identity__";
public final String id;
public final Class credentialType;
public FlowNode(Identity identity)
{
this.id = IDENTITY_NODE_ID;
this.credentialType = null;
}
public FlowNode(Identity identity, Class credentialType)
{
this.id = IDENTITY_CREDENTIAL_NODE_ID_PREFIX + credentialType.getSimpleName();
this.credentialType = credentialType;
}
public FlowNode(Class credentialType)
{
this.id = credentialType.getSimpleName();
this.credentialType = credentialType;
}
public FlowNode(ConnectionSpecification connectionSpecification)
{
this.id = connectionSpecification.getClass().getSimpleName();
this.credentialType = null;
}
@Override
public boolean equals(Object o)
{
if (this == o)
{
return true;
}
if (o == null || getClass() != o.getClass())
{
return false;
}
FlowNode that = (FlowNode) o;
return this.id.equals(that.id);
}
@Override
public int hashCode()
{
return Objects.hash(this.id);
}
}
private static class ResolutionResult
{
public final List flow;
public final Class sourceCredentialType;
public final Class targetCredentialType;
public ResolutionResult(List flow, Class sourceCredentialType, Class targetCredentialType)
{
this.flow = flow;
this.sourceCredentialType = sourceCredentialType;
this.targetCredentialType = targetCredentialType;
}
}
}
public T getConnection(Identity identity, Connection connection, AuthenticationConfiguration authenticationConfiguration) throws Exception
{
return this.getConnection(identity, this.getAuthenticator(identity, connection, authenticationConfiguration));
}
public T getConnection(Identity identity, Connection connection) throws Exception
{
return this.getConnection(identity, this.getAuthenticator(identity, connection));
}
public T getConnection(Identity identity, Authenticator authenticator) throws Exception
{
ConnectionBuilder flow = (ConnectionBuilder) authenticator.getConnectionBuilder();
return flow.getConnection(authenticator.getConnection().getConnectionSpecification(), flow.getAuthenticatorCompatible(authenticator), identity);
}
public static Builder builder()
{
return new Builder();
}
public static class Builder
{
private LegendEnvironment environment;
private final List credentialBuilders = Lists.mutable.empty();
private final List connectionBuilders = Lists.mutable.empty();
private Builder()
{
}
public Builder environment(LegendEnvironment environment)
{
this.environment = environment;
return this;
}
public Builder credentialBuilders(List credentialBuilders)
{
this.credentialBuilders.addAll(credentialBuilders);
return this;
}
public Builder credentialBuilders(CredentialBuilder... credentialBuilders)
{
this.credentialBuilders.addAll(Lists.mutable.with(credentialBuilders));
return this;
}
public Builder credentialBuilder(CredentialBuilder credentialBuilder)
{
this.credentialBuilders.add(credentialBuilder);
return this;
}
public Builder connectionBuilders(List connectionBuilders)
{
this.connectionBuilders.addAll(connectionBuilders);
return this;
}
public Builder connectionBuilders(ConnectionBuilder... connectionBuilders)
{
this.connectionBuilders.addAll(Lists.mutable.with(connectionBuilders));
return this;
}
public Builder connectionBuilder(ConnectionBuilder connectionBuilder)
{
this.connectionBuilders.add(connectionBuilder);
return this;
}
public ConnectionFactory build()
{
for (ConnectionBuilder connectionBuilder : connectionBuilders)
{
ConnectionManager connectionManager = connectionBuilder.getConnectionManager();
if (connectionManager != null)
{
connectionManager.initialize(environment);
}
}
return new ConnectionFactory(
this.environment,
this.credentialBuilders,
this.connectionBuilders
);
}
}
}