Examples with ConnectionKey org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey used on opensource projects

Example 1 with ConnectionKey

use of org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey in project legend-engine by finos.

the class TestDatasourceCreation method userAcquiresConcurrentConnectionsToDifferentDbs.

@Test
public void userAcquiresConcurrentConnectionsToDifferentDbs() throws Exception {
    ConcurrentMutableMap datasourceSpecifications = null;
    Identity identity = IdentityFactoryProvider.getInstance().makeIdentityForTesting("testuser1");
    // User gets connection to db1
    RelationalDatabaseConnection database1 = buildStaticDatabaseSpec("127.0.0.1", server.getPort(), "db2");
    ConnectionKey connectionKey = this.connectionManagerSelector.generateKeyFromDatabaseConnection(database1);
    String pool1 = this.connectionStateManager.poolNameFor(identity, connectionKey);
    this.connectionManagerSelector.getDatabaseConnection(identity, database1);
    // We have a single data source
    datasourceSpecifications = ConnectionPoolTestUtils.getDataSourceSpecifications();
    assertEquals(1, datasourceSpecifications.size());
    // We have a single data source for user1
    DataSourceWithStatistics datasource1 = this.getDatasourceByPool(pool1);
    // User gets another connection to db2
    RelationalDatabaseConnection database2 = buildStaticDatabaseSpec("127.0.0.1", server.getPort(), "db3");
    ConnectionKey key2 = this.connectionManagerSelector.generateKeyFromDatabaseConnection(database2);
    String pool2 = this.connectionStateManager.poolNameFor(identity, key2);
    this.connectionManagerSelector.getDatabaseConnection(identity, database2);
    // We now have 2 data sources one per database
    datasourceSpecifications = ConnectionPoolTestUtils.getDataSourceSpecifications();
    assertEquals(2, datasourceSpecifications.size());
    // We have a single data source for user2
    DataSourceWithStatistics datasource2 = this.getDatasourceByPool(pool2);
    assertNotSame("found same datasource when distinct datasources was expected", datasource1, datasource2);
}

Example 2 with ConnectionKey

use of org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey in project legend-engine by finos.

the class TestDatasourceCreation method multipleUsersAcquireConnectionsToDifferentDatabases.

@Test
public void multipleUsersAcquireConnectionsToDifferentDatabases() throws Exception {
    ConcurrentMutableMap datasourceSpecifications = null;
    Identity identity1 = IdentityFactoryProvider.getInstance().makeIdentityForTesting("testuser1");
    // User gets connection to db1
    RelationalDatabaseConnection database1 = buildStaticDatabaseSpec("127.0.0.1", server.getPort(), "db4");
    ConnectionKey connectionKey = this.connectionManagerSelector.generateKeyFromDatabaseConnection(database1);
    String key1 = this.connectionStateManager.poolNameFor(identity1, connectionKey);
    this.connectionManagerSelector.getDatabaseConnection(identity1, database1);
    DataSourceSpecification ds = builStaticDataSourceSpecification("127.0.0.1", server.getPort(), "db4");
    assertEquals(ds.getConnectionKey(), connectionKey);
    // We have a single data source
    datasourceSpecifications = ConnectionPoolTestUtils.getDataSourceSpecifications();
    assertEquals(1, datasourceSpecifications.size());
    // We have a single data source for user1
    DataSourceWithStatistics datasource1 = this.getDatasourceByPool(key1);
    Identity identity2 = IdentityFactoryProvider.getInstance().makeIdentityForTesting("testuser2");
    // User gets another connection to db2
    RelationalDatabaseConnection database2 = buildStaticDatabaseSpec("127.0.0.1", server.getPort(), "db5");
    ConnectionKey key2 = this.connectionManagerSelector.generateKeyFromDatabaseConnection(database2);
    String pool2 = this.connectionStateManager.poolNameFor(identity2, key2);
    this.connectionManagerSelector.getDatabaseConnection(identity2, database2);
    // We now have 2 data sources one per database + user
    datasourceSpecifications = ConnectionPoolTestUtils.getDataSourceSpecifications();
    assertEquals(2, datasourceSpecifications.size());
    // We have a single data source for user2
    DataSourceWithStatistics datasource2 = this.getDatasourceByPool(pool2);
    assertNotSame("found same datasource when distinct datasources was expected", datasource1, datasource2);
}

Example 3 with ConnectionKey

use of org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey in project legend-engine by finos.

the class TestLocalH2ConnectionCreation method multipleUsersAcquireConnectionsToSameDatabase.

@Test
public void multipleUsersAcquireConnectionsToSameDatabase() throws Exception {
    Identity identity1 = IdentityFactoryProvider.getInstance().makeIdentityForTesting("identity1");
    RelationalDatabaseConnection db1 = this.buildLocalH2DatasourceSpec();
    Connection db1Conn1 = this.connectionManagerSelector.getDatabaseConnection(identity1, db1);
    String db1Conn1Name = h2Name(db1Conn1);
    String poolName1 = this.connectionStateManager.poolNameFor(identity1, this.connectionManagerSelector.generateKeyFromDatabaseConnection(db1));
    Identity identity2 = IdentityFactoryProvider.getInstance().makeIdentityForTesting("identity2");
    RelationalDatabaseConnection db2 = this.buildLocalH2DatasourceSpec();
    Connection db2Conn1 = this.connectionManagerSelector.getDatabaseConnection(identity2, db2);
    String db2Conn1Name = h2Name(db2Conn1);
    String poolName2 = this.connectionStateManager.poolNameFor(identity2, this.connectionManagerSelector.generateKeyFromDatabaseConnection(db2));
    // Connections are distinct
    assertNotSame(db1Conn1Name, db2Conn1Name);
    assertNotSame(poolName1, poolName2);
    // connections keys are same
    ConnectionKey key1 = this.connectionManagerSelector.generateKeyFromDatabaseConnection(db1);
    Assert.assertNotNull(key1);
    ConnectionKey key2 = this.connectionManagerSelector.generateKeyFromDatabaseConnection(db2);
    Assert.assertNotNull(key2);
    Assert.assertEquals(key1, key2);
    Assert.assertNotNull(this.connectionStateManager.get(poolName1));
    Assert.assertNotNull(this.connectionStateManager.get(poolName2));
    Assert.assertNotSame(this.connectionStateManager.get(poolName1), this.connectionStateManager.get(poolName2));
    H2TestUtils.closeProperly(db1Conn1, db2Conn1);
}

Example 4 with ConnectionKey

use of org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey in project legend-engine by finos.

the class DataSourceSpecification method getConnection.

protected Connection getConnection(IdentityState identityState, Supplier<DataSource> dataSourcePoolBuilder) {
    try (Scope scope = GlobalTracer.get().buildSpan("Get Connection").startActive(true)) {
        ConnectionKey connectionKey = this.getConnectionKey();
        // Logs and traces -----
        String principal = identityState.getIdentity().getName();
        scope.span().setTag("Principal", principal);
        scope.span().setTag("DataSourceSpecification", this.toString());
        LOGGER.info("Get Connection as [{}] for datasource [{}]", principal, connectionKey.shortId());
        // ---------------------
        try {
            DataSourceWithStatistics dataSourceWithStatistics = this.connectionStateManager.getDataSourceForIdentityIfAbsentBuild(identityState, this, dataSourcePoolBuilder);
            // Logs and traces and stats -----
            String poolName = dataSourceWithStatistics.getPoolName();
            scope.span().setTag("Pool", poolName);
            int requests = dataSourceWithStatistics.requestConnection();
            LOGGER.info("Principal [{}] has requested [{}] connections for pool [{}]", principal, requests, poolName);
            return authenticationStrategy.getConnection(dataSourceWithStatistics, identityState.getIdentity());
        } catch (ConnectionException ce) {
            LOGGER.error("ConnectionException  {{}} : pool stats [{}] ", principal, connectionStateManager.getPoolStatisticsAsJSON(poolNameFor(identityState.getIdentity())));
            LOGGER.error("ConnectionException ", ce);
            throw ce;
        }
    }
}

Example 5 with ConnectionKey

use of org.finos.legend.engine.plan.execution.stores.relational.connection.ConnectionKey in project legend-engine by finos.

the class ConnectionStateManager method getDataSourceForIdentityIfAbsentBuild.

public DataSourceWithStatistics getDataSourceForIdentityIfAbsentBuild(IdentityState identityState, DataSourceSpecification dataSourceSpecification, Supplier<DataSource> dataSourceBuilder) {
    String principal = identityState.getIdentity().getName();
    String poolName = poolNameFor(identityState.getIdentity(), dataSourceSpecification.getConnectionKey());
    ConnectionKey connectionKey = dataSourceSpecification.getConnectionKey();
    // why do we need getIfAbsentPut?  the first ever pool creation request will create a new Hikari Data Source
    // because we have configured hikari to fail fast a new connection will be created.
    // This will invoke the DriverWrapper connect method, for this method to create that test connection we need to pass minimal state
    Function0<DataSourceWithStatistics> dsSupplier = () -> new DataSourceWithStatistics(poolName, identityState, dataSourceSpecification);
    DataSource dataSource = this.connectionPools.getIfAbsentPut(poolName, dsSupplier).getDataSource();
    // not for every pool request
    if (dataSource == null) {
        synchronized (poolLockManager.getLock(poolName)) {
            dataSource = this.connectionPools.getIfAbsentPut(poolName, dsSupplier).getDataSource();
            if (dataSource == null) {
                LOGGER.info("Pool not found for [{}] for datasource [{}], creating one", principal, connectionKey.shortId());
                try {
                    DataSourceWithStatistics dataSourceWithStatistics = new DataSourceWithStatistics(poolName, dataSourceBuilder.get(), identityState, dataSourceSpecification);
                    this.connectionPools.put(poolName, dataSourceWithStatistics);
                    LOGGER.info("Pool created for [{}] for datasource [{}], name {}", principal, connectionKey.shortId(), poolName);
                } catch (Exception e) {
                    LOGGER.error("Error creating pool {} {}", poolName, e);
                    throw new RuntimeException(e);
                }
            }
        }
    }
    // eventually kerberos credentials will expire(ie invalid) so we need to recreate the pool with the latest subject from the incoming request
    if (!this.connectionPools.get(poolName).getIdentityState().isValid()) {
        LOGGER.info("Pool [{}] for datasource [{}] does not have a valid identity state", principal, connectionKey.shortId());
        synchronized (poolLockManager.getLock(poolName)) {
            DataSourceWithStatistics dataSourceWithStatistics = this.connectionPools.get(poolName);
            if (!dataSourceWithStatistics.getIdentityState().isValid()) {
                DataSourceWithStatistics newDataSourceWithStatistics = new DataSourceWithStatistics(poolName, dataSourceBuilder.get(), identityState, dataSourceSpecification);
                this.connectionPools.put(poolName, newDataSourceWithStatistics);
                LOGGER.info("DataSource re-created for [{}] for datasource [{}], name {}", principal, connectionKey.shortId(), poolName);
                dataSourceWithStatistics.close();
            }
        }
    }
    LOGGER.info("Pool found for [{}] in datasource [{}] : pool Name [{}]", principal, connectionKey.shortId(), poolName);
    return this.connectionPools.get(poolName);
}