org.skife.jdbi.v2.DBI Java Examples

@Override
protected void configure() {
    bind(jsonMapper).to(ObjectMapper.class);
    //dao
    bind(jdbi.onDemand(AccessCredentialsDao.class)).to(AccessCredentialsDao.class);
    bind(jdbi.onDemand(UserDao.class)).to(UserDao.class);
    bind(jdbi.onDemand(TaskDao.class)).to(TaskDao.class);
    bind(jdbi.onDemand(MonitorDao.class)).to(MonitorDao.class);
    bind(jdbi.onDemand(MonitorTaskDao.class)).to(MonitorTaskDao.class);
    bind(jdbi.onDemand(MonitorNiPingResultDao.class)).to(MonitorNiPingResultDao.class);

    //config
    bind(sapConfiguration).to(SapConfiguration.class);
    bind(jobConfiguration).to(JobConfiguration.class);
    //jdbi
    bind(jdbi).to(DBI.class);
    //service
    bind(AuthService.class).in(Singleton.class).to(AuthService.class);
    bind(MonitorService.class).in(Singleton.class).to(MonitorService.class);
    bind(TaskService.class).in(Singleton.class).to(TaskService.class);
    bind(TaskConverter.class).in(Singleton.class).to(TaskConverter.class);
    bind(MonitorResultService.class).in(Singleton.class).to(MonitorResultService.class);
    //auth
    bind(SapOAuthenticator.class).in(Singleton.class).to(SapOAuthenticator.class);
    bind(SapBasicAuthenticator.class).in(Singleton.class).to(SapBasicAuthenticator.class);
}
 

@Before
public void setupTests() throws IOException {
    final DataSourceFactory dataSourceFactory = new DataSourceFactory();
    dataSourceFactory.setDriverClass("org.h2.Driver");
    dataSourceFactory.setUrl("jdbc:h2:mem:date-time-optional-" + System.currentTimeMillis() + "?user=sa");
    dataSourceFactory.setInitialSize(1);
    final DBI dbi = new DBIFactory().build(env, dataSourceFactory, "test");
    try (Handle h = dbi.open()) {
        h.execute("CREATE TABLE tasks (" +
                "id INT PRIMARY KEY, " +
                "assignee VARCHAR(255) NOT NULL, " +
                "start_date TIMESTAMP, " +
                "end_date TIMESTAMP, " +
                "comments VARCHAR(1024) " +
                ")");
    }
    dao = dbi.onDemand(TaskDao.class);
}
 

@Override
public DynamicAttributes build(Configuration configuration, Environment environment, List<String> scopes)
{
    DBI jdbi = SoaBundle.getFeatures(environment).getNamedRequired(DBI.class, name);

    final JdbiDynamicAttributes dynamicAttributes = new JdbiDynamicAttributes(jdbi, scopes);
    ScheduledExecutorService service = environment.lifecycle().scheduledExecutorService("JdbiDynamicAttributes-%d", true).build();
    Runnable command = new Runnable()
    {
        @Override
        public void run()
        {
            dynamicAttributes.update();
        }
    };
    service.scheduleAtFixedRate(command, refreshPeriodSeconds, refreshPeriodSeconds, TimeUnit.SECONDS);
    return dynamicAttributes;
}
 

private List<DataSource> getDataSourceList(SqlShowDataSources sqlNode) throws SQLException {
  DBI dbi = getDbi();
  JdbcSourceDAO jdbcSourceDAO = dbi.onDemand(JdbcSourceDAO.class);
  QuboleDbSourceDAO quboleDbSourceDAO = dbi.onDemand(QuboleDbSourceDAO.class);

  List<DataSource> dataSources = new ArrayList<>();

  if (sqlNode.getLikePattern() == null) {
    dataSources.addAll(jdbcSourceDAO.findByDSSetId(connection.getDSSet().getId()));
    dataSources.addAll(quboleDbSourceDAO.findByDSSetId(connection.getDSSet().getId()));
  } else {
    dataSources.addAll(jdbcSourceDAO.findLikeName(sqlNode.getLikePattern(),
        connection.getDSSet().getId()));
    dataSources.addAll(quboleDbSourceDAO.findLikeName(sqlNode.getLikePattern(),
        connection.getDSSet().getId()));
  }
  return dataSources;
}
 

@Test
public void testMapToDbObject() throws Exception {
    DBI dbi = new DBI(DbHelper.getHsqlDataSource());
    dbi.registerMapper(new SfmResultSetMapperFactory());
    Handle handle = dbi.open();
    try {
        DbObject dbObject = handle.createQuery(DbHelper.TEST_DB_OBJECT_QUERY).mapTo(DbObject.class).first();
        DbHelper.assertDbObjectMapping(dbObject);

        SfmBindTest.SfmBindExample attach = handle.attach(SfmBindTest.SfmBindExample.class);
        attach.insert(DbObject.newInstance());
        assertTrue(handle.createQuery("select * from TEST_DB_OBJECT").mapTo(DbObject.class).list().size() > 1);
    } finally {
        handle.close();
    }
}
 

@Override
public void configure(Binder binder)
{
    binder.bind(DatabaseConfig.class).toProvider(DatabaseConfigProvider.class).in(Scopes.SINGLETON);
    binder.bind(DataSource.class).toProvider(DataSourceProvider.class).in(Scopes.SINGLETON);
    binder.bind(AutoMigrator.class);
    binder.bind(DBI.class).toProvider(DbiProvider.class);  // don't make this singleton because DBI.registerMapper is called for each StoreManager
    binder.bind(TransactionManager.class).to(ThreadLocalTransactionManager.class).in(Scopes.SINGLETON);
    binder.bind(ConfigMapper.class).in(Scopes.SINGLETON);
    binder.bind(DatabaseMigrator.class).in(Scopes.SINGLETON);
    binder.bind(ProjectStoreManager.class).to(DatabaseProjectStoreManager.class).in(Scopes.SINGLETON);
    binder.bind(QueueSettingStoreManager.class).to(DatabaseQueueSettingStoreManager.class).in(Scopes.SINGLETON);
    binder.bind(SessionStoreManager.class).to(DatabaseSessionStoreManager.class).in(Scopes.SINGLETON);
    binder.bind(ScheduleStoreManager.class).to(DatabaseScheduleStoreManager.class).in(Scopes.SINGLETON);
    if (withTaskQueueServer) {
        binder.bind(DatabaseTaskQueueConfig.class).in(Scopes.SINGLETON);
        binder.bind(DatabaseTaskQueueServer.class).in(Scopes.SINGLETON);
    }
}
 

@Before
public void setUp() throws SQLException, IOException {
  Server.createTcpServer().start();

  DBI dbi = new DBI(DB_URL);
  Handle handle = dbi.open();
  Connection conn = handle.getConnection();

  // to suppress output of ScriptRunner
  PrintStream tmp = new PrintStream(new OutputStream() {
    @Override
    public void write(int buff) throws IOException {
      // do nothing
    }
  });
  PrintStream console = System.out;
  System.setOut(tmp);

  String cwd = Paths.get("").toAbsolutePath().toString();
  String path = cwd + "/../src/test/resources/db/postgres/V17_0_0__multi_instance.sql";
  ScriptRunner scriptExecutor = new ScriptRunner(conn, false, true);
  Reader reader = new BufferedReader(new FileReader(path));
  scriptExecutor.runScript(reader);

  System.setOut(console);
}
 

@BeforeMethod
public void setup()
{
    dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    dbi.registerMapper(new Distribution.Mapper(new InternalTypeManager(createTestMetadataManager())));
    dummyHandle = dbi.open();
    createTablesWithRetry(dbi);

    metadataDao = dbi.onDemand(MetadataDao.class);
    nodeManager = new TestingNodeManager(AVAILABLE_WORKERS.stream()
            .map(TestBucketBalancer::createTestingNode)
            .collect(Collectors.toList()));

    NodeSupplier nodeSupplier = nodeManager::getWorkerNodes;
    shardManager = createShardManager(dbi, nodeSupplier);
    balancer = new BucketBalancer(nodeSupplier, shardManager, true, new Duration(1, DAYS), true, true, "test");
}
 

@BeforeMethod
public void setup()
{
    temporary = createTempDir();
    File directory = new File(temporary, "data");
    File backupDirectory = new File(temporary, "backup");
    backupStore = new FileBackupStore(backupDirectory);
    backupStore.start();
    storageService = new FileStorageService(directory);
    storageService.start();

    IDBI dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    dummyHandle = dbi.open();
    createTablesWithRetry(dbi);
    ShardManager shardManager = createShardManager(dbi);
    recoveryManager = createShardRecoveryManager(storageService, Optional.of(backupStore), shardManager);
}
 

@Before
public void setupTests() throws IOException {
    final DataSourceFactory dataSourceFactory = new DataSourceFactory();
    dataSourceFactory.setDriverClass("org.h2.Driver");
    dataSourceFactory.setUrl("jdbc:h2:mem:date-time-optional-" + System.currentTimeMillis() + "?user=sa");
    dataSourceFactory.setInitialSize(1);
    final DBI dbi = new DBIFactory().build(env, dataSourceFactory, "test");
    try (Handle h = dbi.open()) {
        h.execute("CREATE TABLE tasks (" +
                "id INT PRIMARY KEY, " +
                "assignee VARCHAR(255) NOT NULL, " +
                "start_date TIMESTAMP, " +
                "end_date TIMESTAMP, " +
                "comments VARCHAR(1024) " +
                ")");
    }
    dao = dbi.onDemand(TaskDao.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;
}
 

@Override
public void main()
        throws Exception
{
    checkArgs();
    DatabaseConfig dbConfig = DatabaseConfig.convertFrom(buildConfig());
    try (DataSourceProvider dsp = new DataSourceProvider(dbConfig)) {
        DBI dbi = new DBI(dsp.get());
        DatabaseMigrator migrator = new DatabaseMigrator(dbi, dbConfig);
        switch (subCommand) {
            case RUN:
                runMigrate(migrator);
                break;
            case CHECK:
                checkMigrate(migrator);
                break;
            default:
                throw new RuntimeException("No command");
        }
    }
}
 

public SqlJsonDB(DBI dbi, EventBus bus, Collection<Index> indexes) {
    this.dbi = dbi;
    this.bus = bus;
    this.indexes = indexes;

    for (Index index : indexes) {
        this.indexPaths.add(index.getPath()+"/#"+index.getField());
    }

    // Lets find out the type of DB we are working with.
    try (Handle handle = dbi.open()) {
        String dbName = handle.getConnection().getMetaData().getDatabaseProductName();
        databaseKind = DatabaseKind.valueOf(dbName);
    } catch (SQLException e) {
        throw new IllegalStateException(e);
    }
}
 

private void executeDeleteOnDataSource(SqlDropQuarkDataSource node) throws SQLException {
  DBI dbi = getDbi();
  DataSourceDAO dataSourceDAO = dbi.onDemand(DataSourceDAO.class);
  JdbcSourceDAO jdbcDao = dbi.onDemand(JdbcSourceDAO.class);
  QuboleDbSourceDAO quboleDao = dbi.onDemand(QuboleDbSourceDAO.class);
  DataSource jdbcSource = jdbcDao.findByName(node.getIdentifier().getSimple(),
      connection.getDSSet().getId());
  if (jdbcSource != null) {
    jdbcDao.delete(jdbcSource.getId());
    dataSourceDAO.delete(jdbcSource.getId());
  } else {
    DataSource quboleSource = quboleDao.findByName(node.getIdentifier().getSimple(),
        connection.getDSSet().getId());
    if (quboleSource != null) {
      jdbcDao.delete(quboleSource.getId());
      dataSourceDAO.delete(quboleSource.getId());
    }
  }
}
 

@BeforeMethod
public void setup()
{
    this.dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    this.dbi.registerMapper(new TableColumn.Mapper(new InternalTypeManager(createTestMetadataManager())));
    this.dummyHandle = dbi.open();
    createTablesWithRetry(dbi);
    this.metadata = new RaptorMetadata(dbi, createShardManager(dbi));

    // Create table
    ConnectorTableMetadata table = tableMetadataBuilder(DEFAULT_TEST_ORDERS)
            .column("orderkey", BIGINT)
            .column("orderdate", DATE)
            .property("temporal_column", "orderdate")
            .build();
    createTable(table);
}
 

@Singleton
@Provides
public DBI provideDBI(ObjectMapper objectMapper)
        throws ClassNotFoundException
{
    final DBIFactory factory = new DBIFactory();
    final DBI dbi = factory.build(environment, config.getDataSourceFactory(), provideDbType().name());
    dbi.registerMapper(new TableRow.TableRowMapper(objectMapper));
    dbi.registerMapper(new QueryStoreMapper(objectMapper));
    dbi.registerArgumentFactory(new UUIDArgumentFactory());
    dbi.registerArgumentFactory(new URIArgumentFactory());

    return dbi;
}
 

@Inject
public AutoMigrator(DataSource ds, DatabaseConfig config)
{
    if (config.getAutoMigrate()) {
        this.migrator = new DatabaseMigrator(new DBI(ds), config);
    }
}
 

public static DatabaseFactory setupDatabase(boolean autoAutoCommit)
{
    DatabaseConfig config = getEnvironmentDatabaseConfig();
    DataSourceProvider dsp = new DataSourceProvider(config);

    DBI dbi = new DBI(dsp.get());
    TransactionManager tm = new ThreadLocalTransactionManager(dsp.get(), autoAutoCommit);
    // FIXME
    new DatabaseMigrator(dbi, config).migrate();

    cleanDatabase(config.getType(), dbi);

    return new DatabaseFactory(tm, dsp, config);
}
 

@VisibleForTesting
public PostgresStorage(UUID reaperInstanceId,
                       DBI jdbi,
                       int leaderTimeoutInMinutes,
                       int reaperTimeoutInMinutes,
                       int metricsTimeoutInMinutes,
                       int nodeOperationsTimeoutInMinutes) {
  this.reaperInstanceId = reaperInstanceId;
  this.jdbi = jdbi;
  leaderTimeout = Duration.ofMinutes(leaderTimeoutInMinutes);
  reaperTimeout = Duration.ofMinutes(reaperTimeoutInMinutes);
  metricsTimeout = Duration.ofMinutes(metricsTimeoutInMinutes);
  nodeOperationsTimeout = Duration.ofMinutes(nodeOperationsTimeoutInMinutes);
}
 

protected <T> void bindSqlType(final Class<T> sqlType) {
  final Provider<DBI> dbProvider = getProvider(DBI.class);
  bind(sqlType).toProvider(new Provider<T>() {
    @Override
    public T get() {
      return dbProvider.get().onDemand(sqlType);
    }
  });
}
 

/**
 * Check the behavior for upgrading existing database
 * @throws Exception
 */
@Test
public void checkDatabaseUpgrade()
{
    assumeTrue(server.isRemoteDatabase());

    try {
        server.setupDatabase();
        DataSource ds = server.getTestDBDataSource();
        DBI dbi = new DBI(ds);
        DatabaseMigrator migrator = new DatabaseMigrator(dbi, server.getRemoteTestDatabaseConfig());
        MigrationContext context = new MigrationContext(migrator.getDatabaseType());

        //Prepare for tables
        try (Handle handle = dbi.open()) {
            migrator.createSchemaMigrationsTable(handle, context);
            Migration m = new Migration_20151204221156_CreateTables();
            migrator.applyMigration(m, handle, context);
        }

        //Apply rest migrations when digdag server start
        server.start(true);
    }
    catch (Exception e) {
        e.printStackTrace();
        fail(e.toString());
    }
}
 

private DBI getDbi() throws SQLException {
  try {
    return this.connection.getDbi();
  } catch (QuarkException e) {
    throw new SQLException(e);
  }
}
 

@Test
public void testRenewLead() throws InterruptedException {
  DBI dbi = new DBI(DB_URL);
  UUID reaperInstanceId = UUID.randomUUID();
  PostgresStorage storage = new PostgresStorage(reaperInstanceId, dbi);
  Assertions.assertThat(storage.isStorageConnected()).isTrue();

  Handle handle = dbi.open();
  handle.execute("DELETE from leader");

  UUID leaderId = UUID.randomUUID();
  int sleepTime = 3;

  final Instant initialTime = Instant.now();
  storage.takeLead(leaderId);

  // sleep 3 seconds, then renew lead
  TimeUnit.SECONDS.sleep(sleepTime);
  Assertions.assertThat(storage.renewLead(leaderId)).isTrue();

  Instant hbTime = handle.createQuery("SELECT last_heartbeat FROM leader")
      .mapTo(Timestamp.class)
      .first()
      .toInstant();

  Duration between = Duration.between(initialTime, hbTime);
  Assertions.assertThat(between.getSeconds()).isGreaterThanOrEqualTo(sleepTime);
}
 

@Test
public void testNoLeaders() {
  DBI dbi = new DBI(DB_URL);
  UUID reaperInstanceId = UUID.randomUUID();
  PostgresStorage storage = new PostgresStorage(reaperInstanceId, dbi);
  Assertions.assertThat(storage.isStorageConnected()).isTrue();

  Handle handle = dbi.open();
  handle.execute("DELETE from leader");

  List<UUID> fetchedLeaderIds = storage.getLeaders();
  Assertions.assertThat(fetchedLeaderIds.size()).isEqualTo(0);
}
 

@Test
public void testTakeLead() {
  DBI dbi = new DBI(DB_URL);
  UUID reaperInstanceId = UUID.randomUUID();
  PostgresStorage storage = new PostgresStorage(reaperInstanceId, dbi);
  Assertions.assertThat(storage.isStorageConnected()).isTrue();

  Handle handle = dbi.open();
  handle.execute("DELETE from leader");

  int numEntries = 5;
  Set<UUID> leaderIds = new HashSet<>();
  for (int i = 0; i < numEntries; i++) {
    UUID msbLeaderId = UuidUtil.fromSequenceId(UuidUtil.toSequenceId(UUID.randomUUID()));
    leaderIds.add(msbLeaderId);
  }

  // insert all five leader entries
  for (UUID leaderId : leaderIds) {
    boolean result = storage.takeLead(leaderId);
    Assertions.assertThat(result).isEqualTo(true);
  }

  // make sure fetched leaders has all the inserted leaders
  List<UUID> fetchedLeaderIds = storage.getLeaders();
  for (UUID fetchedLeaderId : fetchedLeaderIds) {
    Assertions.assertThat(leaderIds.contains(fetchedLeaderId)).isTrue();
  }
}
 

@Test
public void testDeleteOldNodeMetrics() throws InterruptedException {
  System.out.println("Testing metrics timeout (this will take a minute)...");
  DBI dbi = new DBI(DB_URL);
  UUID reaperInstanceId = UUID.randomUUID();
  PostgresStorage storage = new PostgresStorage(reaperInstanceId, dbi, 1, 1, 1, 1);
  Assertions.assertThat(storage.isStorageConnected()).isTrue();

  Handle handle = dbi.open();
  handle.execute("DELETE from node_metrics_v1");

  UUID runId = UUID.randomUUID();
  NodeMetrics originalNm = NodeMetrics.builder()
      .withNode("fake_node")
      .withCluster("fake_cluster")
      .withDatacenter("NYDC")
      .withHasRepairRunning(true)
      .withPendingCompactions(4)
      .withActiveAnticompactions(1)
      .build();
  storage.storeNodeMetrics(runId, originalNm);

  // first delete attempt shouldn't do anything because the entry hasn't passed its expiration time
  storage.purgeNodeMetrics();
  int numMetrics = handle.createQuery("SELECT COUNT(*) FROM node_metrics_v1")
      .mapTo(Integer.class)
      .first();
  Assertions.assertThat(numMetrics).isEqualTo(1);

  TimeUnit.SECONDS.sleep(61);

  // second delete attempt should work because entry has passed its expiration time
  storage.purgeNodeMetrics();
  numMetrics = handle.createQuery("SELECT COUNT(*) FROM node_metrics_v1")
      .mapTo(Integer.class)
      .first();
  Assertions.assertThat(numMetrics).isEqualTo(0);
}
 

@SuppressWarnings("SpringJavaAutowiringInspection")
@Bean
public DBI jdbiFactory(DataSource dataSource) {
    // note that for JDBI we have to wrap datasource with TransactionAwareDataSourceProxy otherwise JDBI won't respect
    // transactions created by spring
    TransactionAwareDataSourceProxy transactionAwareDataSourceProxy = new TransactionAwareDataSourceProxy(dataSource);
    DBI dbi = new DBI(transactionAwareDataSourceProxy);
    dbi.setSQLLog(new SLF4JLog());  // to enable SLF4J logging
    return dbi;
}
 

@BeforeClass
protected void setupClass() throws Exception {
  // See class comment
  db = new DBI("jdbc:mysql://192.168.10.6/mon", "monapi", "password");
  handle = db.open();
  dao = new AlarmDAOImpl(db);
}
 

@ForMetadata
@Singleton
@Provides
public IDBI createDBI(@ForMetadata ConnectionFactory connectionFactory, TypeManager typeManager)
{
    DBI dbi = new DBI(connectionFactory);
    dbi.registerMapper(new TableColumn.Mapper(typeManager));
    dbi.registerMapper(new Distribution.Mapper(typeManager));
    createTablesWithRetry(dbi);
    return dbi;
}
 

@Override
@SuppressWarnings("unchecked")
protected void configure() {
  DBI db = mock(DBI.class);
  when(db.onDemand(any(Class.class))).thenReturn(dao);
  bind(DBI.class).toInstance(db);
}