Java 类io.dropwizard.db.PooledDataSourceFactory 实例源码

@Override
public DBI build(Environment environment, PooledDataSourceFactory configuration, ManagedDataSource dataSource, String name) {
    final DBI dbi = super.build(environment, configuration, dataSource, name);

    dbi.registerArgumentFactory(new ValueArgumentFactory(configuration.getDriverClass()));
    dbi.registerArgumentFactory(new OptionArgumentFactory(configuration.getDriverClass()));
    dbi.registerContainerFactory(new OptionContainerFactory());

    // The order of container factories is important, least specific to most specific
    dbi.registerContainerFactory(new SeqContainerFactory());
    dbi.registerContainerFactory(new IndexedSeqContainerFactory());
    dbi.registerContainerFactory(new SetContainerFactory());
    dbi.registerContainerFactory(new TreeContainerFactory());
    dbi.registerContainerFactory(new ListContainerFactory());
    dbi.registerContainerFactory(new ArrayContainerFactory());
    dbi.registerContainerFactory(new QueueContainerFactory());
    dbi.registerContainerFactory(new StreamContainerFactory());
    dbi.registerContainerFactory(new VectorContainerFactory());

    return dbi;
}

private IDBI connectToDatasource(ShardInfo shardInfo) {
    PooledDataSourceFactory urlReplacedCopy =
            copyAndReplaceURL(dropwizardDSFactory, shardInfo);

    // we cannot re-register using the same name (dropwizard metrics do not merge on name collision)
    String dbShardName = humanFriendlyShardNamer.nameShard(shardInfo);
    DBI jdbi = dropwizardDbiFactory.build(
            dropwizardEnvironment,
            urlReplacedCopy,
            dbShardName);
    if (DBIInitializer != null) {
        DBIInitializer.initialize(jdbi);
    }
    logger.info("New shard connection created: " + dbShardName);

    return jdbi;
}

public Jdbi build(Environment environment,
                  PooledDataSourceFactory configuration,
                  String name) {
    ManagedDataSource dataSource = configuration.build(environment.metrics(), name);
    String validationQuery = configuration.getValidationQuery();
    Jdbi jdbi = Jdbi.create(dataSource);
    jdbi.setTimingCollector(new InstrumentedTimingCollector(environment.metrics(), nameStrategy));
    jdbi.installPlugins();

    environment.lifecycle().manage(dataSource);
    environment.healthChecks().register(name, new JdbiHealthCheck(
            environment.getHealthCheckExecutorService(),
            configuration.getValidationQueryTimeout().orElseGet(() -> Duration.seconds(5)),
            jdbi, validationQuery));

    return jdbi;
}

public SessionFactory build(final RemoteCredentialHibernateBundle<?> bundle,
        final Environment environment,
        final PooledDataSourceFactory dbConfig,
        final ManagedDataSource dataSource,
        final List<Class<?>> entities) {
    final ConnectionProvider provider = this.buildConnectionProvider(dataSource,
            dbConfig.getProperties());
    final SessionFactory factory = this.buildSessionFactory(bundle,
            dbConfig,
            provider,
            dbConfig.getProperties(),
            entities);
    final SessionFactoryManager managedFactory = new SessionFactoryManager(factory, dataSource);
    environment.lifecycle().manage(managedFactory);
    return factory;
}

private void init(final ImmutableList<Class<?>> inEntities) {
    String numShardsEnv = System.getProperty("db.shards", "2");
    int numShards = Integer.parseInt(numShardsEnv);
    shardManager = new ShardManager(numShards);
    for (int i = 0; i < numShards; i++) {
        final int finalI = i;
        shardBundles.add(new HibernateBundle<T>(inEntities, new SessionFactoryFactory()) {
            @Override
            protected String name() {
                return String.format("connectionpool-%d", finalI);
            }

            @Override
            public PooledDataSourceFactory getDataSourceFactory(T t) {
                return getConfig(t).getShards().get(finalI);
            }
        });
    }
}

@Override
public void initialize(final Bootstrap<ProductCatalogConfiguration> bootstrap) {
    bootstrap.addBundle(discoveryBundle);
    bootstrap.addBundle(new MigrationsBundle<ProductCatalogConfiguration>() {

        @Override
        public PooledDataSourceFactory getDataSourceFactory(ProductCatalogConfiguration configuration) {
            return configuration.getDataSourceFactory();
        }
    });
    guiceBundle = GuiceBundle.<ProductCatalogConfiguration> newBuilder().addModule(new ProductCatalogModule())
            .enableAutoConfig(getClass().getPackage().getName()).setConfigClass(ProductCatalogConfiguration.class)
            .build(Stage.PRODUCTION);
    bootstrap.addBundle(guiceBundle);
    // Uncomment below to read the yaml file from Jar
    // bootstrap.setConfigurationSourceProvider(new
    // ResourceConfigurationSourceProvider());
}

@Override
public void initialize(final Bootstrap<ProductReviewConfiguration> bootstrap) {
    bootstrap.addBundle(discoveryBundle);
    bootstrap.addBundle(new MigrationsBundle<ProductReviewConfiguration>() {

        @Override
        public PooledDataSourceFactory getDataSourceFactory(ProductReviewConfiguration configuration) {
            return configuration.getDataSourceFactory();
        }
    });
    guiceBundle = GuiceBundle.<ProductReviewConfiguration> newBuilder().addModule(new ProductReviewModule())
            .enableAutoConfig(getClass().getPackage().getName()).setConfigClass(ProductReviewConfiguration.class)
            .build(Stage.PRODUCTION);
    bootstrap.addBundle(guiceBundle);
    // Uncomment below to read the yaml file from Jar
    // bootstrap.setConfigurationSourceProvider(new
    // ResourceConfigurationSourceProvider());
}

@Override
public final void run(T configuration, Environment environment) throws Exception {
    final PooledDataSourceFactory dbConfig = getDataSourceFactory(configuration);

    this.entityManagerFactory = entityManagerFactoryFactory.build(this, environment, dbConfig, entities, name());
    this.entityManagerContext = new EntityManagerContext(entityManagerFactory);
    this.sharedEntityManager = sharedEntityManagerFactory.build(entityManagerContext);

    registerUnitOfWorkListerIfAbsent(environment).registerEntityManagerFactory(name(), entityManagerFactory);
    environment.healthChecks().register(name(),
            new EntityManagerFactoryHealthCheck(
                    environment.getHealthCheckExecutorService(),
                    dbConfig.getValidationQueryTimeout().orElse(Duration.seconds(5)),
                    entityManagerFactory,
                    dbConfig.getValidationQuery()));
}

@Test
@SuppressWarnings("unchecked")
public void testFindEntityClassesFromDirectory() {
    String packageWithEntities = "com.scottescue.dropwizard.entitymanager.entity.fake.entities.pckg";
    ScanningEntityManagerBundle bundle = new ScanningEntityManagerBundle(packageWithEntities) {
        @Override
        public void run(Object o, Environment environment) throws Exception {
        }

        @Override
        public PooledDataSourceFactory getDataSourceFactory(Configuration configuration) {
            return null;
        }
    };

    assertThat(bundle.getEntities()).containsOnly(
            FakeEntity1.class,
            DeepFakeEntity.class,
            DeeperFakeEntity.class);
}

@Test
@SuppressWarnings("unchecked")
public void testFindEntityClassesFromMultipleDirectories() {
    String root = "com.scottescue.dropwizard.entitymanager.entity.fake.entities.";

    ScanningEntityManagerBundle bundle = new ScanningEntityManagerBundle(
            root.concat("pckg"), root.concat("pckg2"), root.concat("pckg3")) {
        @Override
        public void run(Object o, Environment environment) throws Exception {
        }

        @Override
        public PooledDataSourceFactory getDataSourceFactory(Configuration configuration) {
            return null;
        }
    };

    assertThat(bundle.getEntities()).containsOnly(
            FakeEntity1.class,
            FakeEntity2.class,
            FakeEntity3.class,
            DeepFakeEntity.class,
            DeeperFakeEntity.class);
}

@Override
public DBI build(Environment environment,
                 PooledDataSourceFactory configuration,
                 ManagedDataSource dataSource,
                 String name) {
    final DBI dbi = super.build(environment, configuration, dataSource, name);

    dbi.registerArgumentFactory(new OptionalArgumentFactory(configuration.getDriverClass()));
    dbi.registerContainerFactory(new OptionalContainerFactory());
    dbi.registerArgumentFactory(new LocalDateArgumentFactory());
    dbi.registerArgumentFactory(new OptionalLocalDateArgumentFactory());
    dbi.registerArgumentFactory(new LocalDateTimeArgumentFactory());
    dbi.registerArgumentFactory(new OptionalLocalDateTimeArgumentFactory());
    dbi.registerMapper(new LocalDateMapper());
    dbi.registerMapper(new LocalDateTimeMapper());

    final Optional<TimeZone> tz = Optional.ofNullable(databaseTimeZone().orNull());
    dbi.registerArgumentFactory(new InstantArgumentFactory(tz));
    dbi.registerArgumentFactory(new OptionalInstantArgumentFactory(tz));
    dbi.registerMapper(new InstantMapper(tz));


    return dbi;
}

private void configureDataSourceFactory(
        final C configuration,
        final Environment environment,
        final String name,
        final PooledDataSourceFactory dataSourceFactory
) throws RuntimeException {

    try {
        final JooqFactory jooqFactory = getJooqFactory(configuration);
        final Configuration cfg = jooqFactory.build(environment, dataSourceFactory, name);
        final JooqHealthCheck healthCheck = new JooqHealthCheck(cfg, dataSourceFactory.getValidationQuery());

        environment.healthChecks().register(name, healthCheck);

        jooqFactoryConfigurationMap.put(name, cfg);
    } catch (ClassNotFoundException e) {
        throw new RuntimeException(e);
    }
}

public void initialize(Environment environment, PooledDataSourceFactory dataSourceFactory, String name) {
   dataSource = dataSourceFactory.build(environment.metrics(), name);
   bCryptPasswordEncryptor = BCryptPasswordEncryptor.newInstance(12);
   environment.lifecycle().manage(dataSource);
   environment.healthChecks().register(name,
                                       new DataSourceHealthCheck(environment.getHealthCheckExecutorService(),
                                                                 dataSourceFactory.getValidationQueryTimeout()
                                                                       .orElse(Duration.seconds(5)),
                                                                 dataSource,
                                                                 dataSourceFactory.getValidationQuery()));
}

@Override
public void run(final BenchConfiguration config,
                final Environment environment) {
    final DBIFactory factory = new DBIFactory();
    final Optional<PooledDataSourceFactory> tomcatFactory = config.getTomcatFactory().map(x -> x);

    config.getHikariFactory().ifPresent(x -> x.setHealthCheckRegistry(environment.healthChecks()));
    final Optional<PooledDataSourceFactory> hikariFactory = config.getHikariFactory().map(x -> x);
    final PooledDataSourceFactory datasource = tomcatFactory.orElse(hikariFactory.orElse(null));
    final DBI jdbi = factory.build(environment, datasource, "postgresql");
    jdbi.registerMapper(new QuestionMapper());
    final QuestionQuery dao = jdbi.onDemand(QuestionQuery.class);
    environment.jersey().register(new QuestionResource(dao));
}

private PooledDataSourceFactory copyAndReplaceURL(
        DataSourceFactory dsFactory,
        ShardInfo shardInfo) {

    DataSourceFactory copy = new CopiedDataSourceFactory(dsFactory);
    copy.setUrl(shardInfo.url());
    return copy;
}

public SessionFactory build(final RemoteCredentialHibernateBundle<?> bundle,
        final Environment environment,
        final PooledDataSourceFactory dbConfig,
        final List<Class<?>> entities,
        final String name) {
    final ManagedDataSource dataSource = dbConfig.build(environment.metrics(), name);
    return this.build(bundle, environment, dbConfig, dataSource, entities);
}

private SessionFactory buildSessionFactory(final RemoteCredentialHibernateBundle<?> bundle,
        final PooledDataSourceFactory dbConfig,
        final ConnectionProvider connectionProvider,
        final Map<String, String> properties,
        final List<Class<?>> entities) {
    final Configuration configuration = new Configuration();
    configuration.setProperty(AvailableSettings.CURRENT_SESSION_CONTEXT_CLASS, "managed");
    configuration.setProperty(AvailableSettings.USE_SQL_COMMENTS,
            Boolean.toString(dbConfig.isAutoCommentsEnabled()));
    configuration.setProperty(AvailableSettings.USE_GET_GENERATED_KEYS, "true");
    configuration.setProperty(AvailableSettings.GENERATE_STATISTICS, "true");
    configuration.setProperty(AvailableSettings.USE_REFLECTION_OPTIMIZER, "true");
    configuration.setProperty(AvailableSettings.ORDER_UPDATES, "true");
    configuration.setProperty(AvailableSettings.ORDER_INSERTS, "true");
    configuration.setProperty(AvailableSettings.USE_NEW_ID_GENERATOR_MAPPINGS, "true");
    configuration.setProperty("jadira.usertype.autoRegisterUserTypes", "true");
    for (final Map.Entry<String, String> property : properties.entrySet()) {
        configuration.setProperty(property.getKey(), property.getValue());
    }

    this.addAnnotatedClasses(configuration, entities);

    final ServiceRegistry registry = new StandardServiceRegistryBuilder()
            .addService(ConnectionProvider.class, connectionProvider)
            .applySettings(properties)
            .build();

    this.configure(configuration, registry);

    return configuration.buildSessionFactory(registry);
}

@Override
public final void run(final T configuration, final Environment environment) throws Exception {
    this.environment = environment;

    final PooledDataSourceFactory dbConfig = this.getDataSourceFactory(configuration);
    this.registerUnitOfWorkListerIfAbsent(environment).registerBundle(this.name(), this);
    environment.healthChecks().register(this.name(),
            new SessionFactoryHealthCheck(
                    environment.getHealthCheckExecutorService(),
                    dbConfig.getValidationQueryTimeout().or(Duration.seconds(5)),
                    this.sessionFactory.get(),
                    dbConfig.getValidationQuery()));

    this.dataSourceFactory = (RemoteCredentialDataSourceFactory) this
            .getDataSourceFactory(configuration);
    // If the feature is disabled we don't need to create the client and
    // retrieve the credentials.
    if (this.dataSourceFactory.isRetrieveCredentials()) {
        try {
            this.client = this.createCredentialClient();
            this.credential = this.client.getCredential();
        } catch (NoSuchAlgorithmException | InvalidKeySpecException | IOException
                | InvalidKeyException | NoSuchPaddingException | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new RuntimeException("Failed to initialize credential storage client.", e);
        }
    }
    this.createDataSourceAndSessionFactory(this.environment.metrics());
    // The scheduled credential retrieval is useless if the feature is
    // disabled.
    if (this.dataSourceFactory.isRetrieveCredentials()) {
        this.scheduleCredentialRetrieval(this.dataSourceFactory);
    }
}

@Override
public void initialize(Bootstrap<JdbiAppConfiguration> bootstrap) {

    bootstrap.addBundle(GuiceBundle.builder()
            .enableAutoConfig(JdbiApplication.class.getPackage().getName())
            .bundles(JdbiBundle.<JdbiAppConfiguration>forDatabase((conf, env) -> conf.getDatabase()))
            .build());
    // used for manual run to init db
    bootstrap.addBundle(new FlywayBundle<JdbiAppConfiguration>() {
        @Override
        public PooledDataSourceFactory getDataSourceFactory(JdbiAppConfiguration configuration) {
            return configuration.getDatabase();
        }
    });
}

@Override
public void initialize(Bootstrap<AppConfiguration> bootstrap) {
    bootstrap.addBundle(new AssetsBundle("/ui", "/", "index.html"));
    bootstrap.addBundle(hibernate);
    bootstrap.addBundle(new MigrationsBundle<AppConfiguration>() {

        @Override
        public PooledDataSourceFactory getDataSourceFactory(AppConfiguration configuration) {
            return configuration.getDataSourceFactory();
        }
    });
}

EntityManagerFactory build(EntityManagerBundle<?> bundle,
                            Environment environment,
                            PooledDataSourceFactory dbConfig,
                            List<Class<?>> entities,
                            String name) {
    final ManagedDataSource dataSource = dbConfig.build(environment.metrics(), name);
    return build(bundle, environment, dbConfig, dataSource, entities);
}

EntityManagerFactory build(EntityManagerBundle<?> bundle,
                            Environment environment,
                            PooledDataSourceFactory dbConfig,
                            ManagedDataSource dataSource,
                            List<Class<?>> entities) {
    final EntityManagerFactory factory = buildSessionFactory(bundle,
            dbConfig,
            dataSource,
            dbConfig.getProperties(),
            entities);
    final EntityManagerFactoryManager managedFactory = new EntityManagerFactoryManager(factory, dataSource);
    environment.lifecycle().manage(managedFactory);
    return factory;
}

private EntityManagerFactory buildSessionFactory(EntityManagerBundle<?> bundle,
                                                 PooledDataSourceFactory dbConfig,
                                                 ManagedDataSource dataSource,
                                                 Map<String, String> properties,
                                                 List<Class<?>> entities) {

    PersistenceUnitInfoImpl persistenceUnitInfo = new PersistenceUnitInfoImpl(bundle.name(), dataSource);

    persistenceUnitInfo.setProperty(AvailableSettings.CURRENT_SESSION_CONTEXT_CLASS, "managed");
    persistenceUnitInfo.setProperty(AvailableSettings.USE_SQL_COMMENTS,
            Boolean.toString(dbConfig.isAutoCommentsEnabled()));
    persistenceUnitInfo.setProperty(AvailableSettings.USE_GET_GENERATED_KEYS, "true");
    persistenceUnitInfo.setProperty(AvailableSettings.GENERATE_STATISTICS, "true");
    persistenceUnitInfo.setProperty(AvailableSettings.USE_REFLECTION_OPTIMIZER, "true");
    persistenceUnitInfo.setProperty(AvailableSettings.ORDER_UPDATES, "true");
    persistenceUnitInfo.setProperty(AvailableSettings.ORDER_INSERTS, "true");
    persistenceUnitInfo.setProperty(AvailableSettings.USE_NEW_ID_GENERATOR_MAPPINGS, "true");
    persistenceUnitInfo.setProperty("jadira.usertype.autoRegisterUserTypes", "true");
    for (Map.Entry<String, String> property : properties.entrySet()) {
        persistenceUnitInfo.setProperty(property.getKey(), property.getValue());
    }

    addAnnotatedClasses(persistenceUnitInfo, entities);
    bundle.configure(persistenceUnitInfo);

    configure(persistenceUnitInfo);

    return new HibernatePersistenceProvider().createContainerEntityManagerFactory(persistenceUnitInfo, null);
}

@Test
public void addsAllEntities() {
    EntityManagerBundle<Configuration> customBundle = new EntityManagerBundle<Configuration>(Person.class, FakeEntity1.class) {
        @Override
        public PooledDataSourceFactory getDataSourceFactory(Configuration configuration) {
            return dbConfig;
        }
    };
    assertThat(customBundle.getEntities()).containsExactly(Person.class, FakeEntity1.class);
}

@Override
protected void run(final Bootstrap<T> bootstrap, final Namespace namespace, final T configuration) throws Exception {
    final PooledDataSourceFactory datasourceFactory = databaseConfiguration.getDataSourceFactory(configuration);
    final FlywayFactory flywayFactory = flywayConfiguration.getFlywayFactory(configuration);
    final Flyway flyway = flywayFactory.build(datasourceFactory.build(bootstrap.getMetricRegistry(), "Flyway"));

    try {
        run(namespace, flyway);
    } catch (FlywayException e) {
        LOG.error("Error while running database command", e);
        throw e;
    }
}

public Configuration build(Environment environment, PooledDataSourceFactory factory, String name) throws ClassNotFoundException {
    final Settings settings = buildSettings();
    final ManagedDataSource dataSource = factory.build(environment.metrics(), name);
    final SQLDialect dialect = determineDialect(factory, dataSource);
    final ConnectionProvider connectionProvider = new DataSourceConnectionProvider(dataSource);
    final Configuration config = new DefaultConfiguration()
            .set(settings)
            .set(dialect)
            .set(connectionProvider);

    environment.lifecycle().manage(dataSource);

    return config;
}

private SQLDialect determineDialect(PooledDataSourceFactory dataSourceFactory, ManagedDataSource dataSource) {
    // If a dialect was specified, great!
    if (getDialect().isPresent()) {
        return dialect.get();
    }

    return JDBCUtils.dialect(dataSourceFactory.getUrl());
}

@Override
public PooledDataSourceFactory getDataSourceFactory(HbnAppConfiguration configuration) {
    return configuration.getDataSourceFactory();
}

@Override
public PooledDataSourceFactory getDataSourceFactory(MappedJsonConfiguration configuration) {
    DataSourceFactory dataSourceFactory = configureDataSourceFactory(configuration);
    dataSourceFactory.setCheckConnectionWhileIdle(false);
    return dataSourceFactory;
}

EntityManagerFactory build(EntityManagerBundle<?> bundle,
                            Environment environment,
                            PooledDataSourceFactory dbConfig,
                            List<Class<?>> entities) {
    return build(bundle, environment, dbConfig, entities, DEFAULT_NAME);
}

@Override
public PooledDataSourceFactory getDataSourceFactory(AbstractIntegrationTest.TestConfiguration configuration) {
    return configuration.getDataSource();
}

@Override
public PooledDataSourceFactory getDataSourceFactory(T configuration) {
    return configuration.getHibernate().getDataSourceFactory(configuration);
}

@Test
@Order(order = 8)
public void getDatasourceFactory() {
    PooledDataSourceFactory dataSourceFactory = RobeHibernateBundle.getInstance().getDataSourceFactory(new TestConfig());
    Assert.assertNotNull(dataSourceFactory);
}