org.skife.jdbi.v2.Handle Java Examples

@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);
}
 

@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);
}
 

/**
 * Postgres does not allow to read from the large object after the connection has been closed.
 */
private InputStream doReadPostgres(String path) {
    try (Handle h = dbi.open()) {
        h.getConnection().setAutoCommit(false);

        List<Map<String, Object>> res = h.select("SELECT data FROM filestore WHERE path=?", path);

        Optional<Long> oid = res.stream()
            .map(row -> row.get("data"))
            .map(Long.class::cast)
            .findFirst();

        if (oid.isPresent()) {
            LargeObjectManager lobj = getPostgresConnection(h.getConnection()).getLargeObjectAPI();
            LargeObject obj = lobj.open(oid.get(), LargeObjectManager.READ);
            return new HandleCloserInputStream(h, obj.getInputStream());
        }

        return null;
    } catch (SQLException e) {
        throw DaoException.launderThrowable(e);
    }
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    if (context.isPostgres()) {
        handle.update("alter table tasks" +
                " add column started_at timestamp with time zone");
    }
    else {
        handle.update("alter table tasks" +
                " add column started_at timestamp");
    }

    if (context.isPostgres()) {
        handle.update("create index tasks_on_state_and_started_at on tasks (state, started_at, id asc) where started_at is not null");
    } else {
        handle.update("create index tasks_on_state_and_started_at on tasks (state, started_at, id asc)");
    }
}
 

private boolean tableExists(Handle h, String tableName) {
    try {
        String tableToCheck = tableName;
        boolean caseSensitive = this.databaseKind == DatabaseKind.PostgreSQL;
        if (!caseSensitive) {
            tableToCheck = tableName.toUpperCase(Locale.ROOT);
        }
        DatabaseMetaData metaData = h.getConnection().getMetaData();

        try (ResultSet rs = metaData.getTables(null, null, tableToCheck, null)) {
            while (rs.next()) {
                String foundTable = rs.getString("TABLE_NAME");
                if (tableToCheck.equalsIgnoreCase(foundTable)) {
                    return true;
                }
            }
        }
        return false;
    } catch (SQLException ex) {
        throw DaoException.launderThrowable("Cannot check if the table " + tableName + " already exists", ex);
    }
}
 

@Override
public Optional<RepairSegment> getNextFreeSegmentInRange(UUID runId, Optional<RingRange> range) {
  if (range.isPresent()) {
    RepairSegment result;
    try (Handle h = jdbi.open()) {
      IStoragePostgreSql storage = getPostgresStorage(h);
      if (!range.get().isWrapping()) {
        result = storage.getNextFreeRepairSegmentInNonWrappingRange(
            UuidUtil.toSequenceId(runId), range.get().getStart(), range.get().getEnd());
      } else {
        result = storage.getNextFreeRepairSegmentInWrappingRange(
            UuidUtil.toSequenceId(runId), range.get().getStart(), range.get().getEnd());
      }
    }
    return Optional.ofNullable(result);
  } else {
    return getNextFreeSegment(runId);
  }
}
 

@Override
public RegisterEmployeeOutput callRegisterEmployee(String name, String surname, String email, BigDecimal salary, String departmentName, String companyName) {
    try (Handle h = dbi.open()) {
        return h.createQuery("SELECT employee_id employeePid, department_id departmentPid, company_id companyPid FROM register_employee(" +
                "  :name, \n" +
                "  :surname, \n" +
                "  :email, \n" +
                "  :salary, \n" +
                "  :departmentName, \n" +
                "  :companyName\n" +
                ")")
                .bind("name", name)
                .bind("surname", surname)
                .bind("email", email)
                .bind("salary", salary)
                .bind("departmentName", departmentName)
                .bind("companyName", companyName)
                .map(RegisterEmployeeOutput.class)
                .first();
    }
}
 

@Override
public void storeNodeMetrics(UUID runId, NodeMetrics nodeMetrics) {
  if (null != jdbi) {
    try (Handle h = jdbi.open()) {
      getPostgresStorage(h).storeNodeMetrics(
          UuidUtil.toSequenceId(runId),
          nodeMetrics.getNode(),
          nodeMetrics.getCluster(),
          nodeMetrics.getDatacenter(),
          nodeMetrics.isRequested(),
          nodeMetrics.getPendingCompactions(),
          nodeMetrics.hasRepairRunning(),
          nodeMetrics.getActiveAnticompactions()
      );
    }
  }
}
 

@Override
public void releaseLead(UUID leaderId) {
  if (null != jdbi) {
    try (Handle h = jdbi.open()) {
      int rowsDeleted = getPostgresStorage(h).releaseLead(
          leaderId,
          reaperInstanceId
      );
      if (rowsDeleted == 1) {
        LOG.debug("Released lead on segment {}", leaderId);
      } else {
        LOG.error("Could not release lead on segment {}", leaderId);
      }
    }
  }
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    if (context.isPostgres()) {
        handle.update("alter table projects" +
                " add column deleted_at timestamp with time zone");
        handle.update("alter table projects" +
                " add column deleted_name text");
        handle.update("alter table projects" +
                " alter column name drop not null");
    }
    else {
        handle.update("alter table projects" +
                " add column deleted_at timestamp");
        handle.update("alter table projects" +
                " add column deleted_name varchar(255)");
        handle.update("alter table projects" +
                " alter column name drop not null");
    }
}
 

@Override
public Set<UUID> getExistingShardUuids(long tableId, Set<UUID> shardUuids)
{
    try (Handle handle = dbi.open()) {
        String args = Joiner.on(",").join(nCopies(shardUuids.size(), "?"));
        String selectShards = format(
                "SELECT shard_uuid FROM %s WHERE shard_uuid IN (%s)",
                shardIndexTable(tableId), args);

        ImmutableSet.Builder<UUID> existingShards = ImmutableSet.builder();
        try (PreparedStatement statement = handle.getConnection().prepareStatement(selectShards)) {
            bindUuids(statement, shardUuids);
            try (ResultSet rs = statement.executeQuery()) {
                while (rs.next()) {
                    existingShards.add(uuidFromBytes(rs.getBytes("shard_uuid")));
                }
            }
        }
        return existingShards.build();
    }
    catch (SQLException e) {
        throw new RuntimeException(e);
    }
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    if (context.isPostgres()) {
        handle.update("alter table sessions" +
                " add column last_attempt_created_at timestamp with time zone");
        handle.update(
                "update sessions" +
                " set last_attempt_created_at = session_attempts.created_at" +
                " from session_attempts" +
                " where session_attempts.id = sessions.last_attempt_id");
    }
    else {
        handle.update("alter table sessions" +
                " add column last_attempt_created_at timestamp");
        handle.update("update sessions" +
                " set last_attempt_created_at = (" +
                    " select created_at from session_attempts" +
                    " where session_attempts.id = sessions.last_attempt_id" +
                ")");
    }

    handle.update("create index sessions_on_project_id_and_workflow_name_and_last_attempt_created_at on sessions (project_id, workflow_name, last_attempt_created_at desc)");
}
 

@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);
}
 

@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 Collection<RepairSchedule> getRepairSchedulesForCluster(String clusterName) {
  Collection<RepairSchedule> result;
  try (Handle h = jdbi.open()) {
    result = getPostgresStorage(h).getRepairSchedulesForCluster(clusterName);
  }
  return result;
}
 

@Override
public SortedSet<UUID> getRepairRunIdsForCluster(String clusterName) {
  SortedSet<UUID> result = Sets.newTreeSet(Collections.reverseOrder());
  try (Handle h = jdbi.open()) {
    for (Long l : getPostgresStorage(h).getRepairRunIdsForCluster(clusterName)) {
      result.add(UuidUtil.fromSequenceId(l));
    }
  }
  return result;
}
 

@Override
public Collection<RepairSegment> getSegmentsWithState(UUID runId, RepairSegment.State segmentState) {
  Collection<RepairSegment> result;
  try (Handle h = jdbi.open()) {
    result = getPostgresStorage(h).getRepairSegmentsForRunWithState(UuidUtil.toSequenceId(runId), segmentState);
  }
  return result;
}
 

private Optional<RepairSegment> getNextFreeSegment(UUID runId) {
  RepairSegment result;
  try (Handle h = jdbi.open()) {
    result = getPostgresStorage(h).getNextFreeRepairSegment(UuidUtil.toSequenceId(runId));
  }
  return Optional.ofNullable(result);
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    if (context.isPostgres()) {
        handle.update("create index concurrently sessions_on_project_id_and_workflow_name_desc on sessions (project_id, workflow_name, id DESC)");
    } else {
        handle.update("create index sessions_on_project_id_and_workflow_name_desc on sessions (project_id, workflow_name, id DESC)");
    }
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    handle.update("alter table task_state_details" +
            " add column reset_store_params text");
    handle.update("alter table resuming_tasks" +
            " add column reset_store_params text");
}
 

@Override
public Project findProject(Integer pid) {
    try (Handle h = dbi.open()) {
        return h.createQuery("SELECT pid, name, datestarted FROM project WHERE pid = :pid")
                .bind("pid", pid)
                .map((index, r, ctx) -> {
                    Project project = new Project();
                    project.setPid(r.getInt("pid"));
                    project.setName(r.getString("name"));
                    project.setDate(r.getDate("datestarted").toLocalDate());
                    return project;
                })
                .first();
    }
}
 

private Sha1HashId insertMetricDefinition(Handle h, MetricDefinitionAndTenantId mdtid) {
  final String region = ""; // TODO We currently don't have region
  final String definitionIdStringToHash =
      trunc(mdtid.metricDefinition.name, MAX_COLUMN_LENGTH)
          + trunc(mdtid.tenantId, MAX_COLUMN_LENGTH) + trunc(region, MAX_COLUMN_LENGTH);
  final byte[] id = DigestUtils.sha(definitionIdStringToHash);
  h.insert("insert into metric_definition(id, name, tenant_id) values (?, ?, ?) " +
           "on duplicate key update id=id", id, mdtid.metricDefinition.name, mdtid.tenantId);
  return new Sha1HashId(id);
}
 

@Override
public boolean isStorageConnected() {
  String currentDate = null;
  if (null != jdbi) {
    try (Handle h = jdbi.open()) {
      currentDate = getPostgresStorage(h).getCurrentDate();
    }
  }
  return null != currentDate && !currentDate.trim().isEmpty();
}
 

@Override
public void updateState(String id, AlarmState state, long msTimestamp) {

  try (final Handle h = db.open()) {
    String timestamp  = formatDateFromMillis(msTimestamp);
    h.createStatement("update alarm set state = :state, state_updated_at = :timestamp, updated_at = :timestamp where id = :id")
        .bind("id", id).bind("timestamp", timestamp).bind("state", state.toString()).execute();
  }
}
 

@Override
public Cluster getCluster(String clusterName) {
  try (Handle h = jdbi.open()) {
    Cluster result = getPostgresStorage(h).getCluster(clusterName);
    if (null != result) {
      return result;
    }
  }
  throw new IllegalArgumentException("no such cluster: " + clusterName);
}
 

@Override
public void migrate(Handle handle, MigrationContext context)
{
    // DatabaseProjectStoreManager.PgDao.getLatestActiveWorkflowDefinitions uses these indexes.
    handle.update("create index workflow_definitions_on_revision_id_and_id on workflow_definitions (revision_id, id)");
    if (context.isPostgres()) {
        handle.update("create index projects_on_site_id_and_id on projects (site_id, id) where deleted_at is null");
    }
    else {
        handle.update("create index projects_on_site_id_and_id on projects (site_id, id)");
    }
}
 

@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);
}
 

@Override
public Employee findEmployee(Integer pid) {
    try (Handle h = dbi.open()) {
        return h.createQuery("SELECT pid, name, surname, email, department_pid departmentPid, salary FROM employee WHERE pid = :pid")
                .bind("pid", pid)
                .map(Employee.class)
                .first();
    }
}
 

@Override
public List<ProjectsWithCostsGreaterThanOutput> getProjectsWithCostsGreaterThan(int totalCostBoundary) {
    try (Handle h = dbi.open()) {
        String query;
        query = "WITH project_info AS (\n" +
                "    SELECT project.pid project_pid, project.name project_name, salary monthly_cost, company.name company_name\n" +
                "    FROM project\n" +
                "      JOIN projectemployee ON project.pid = projectemployee.project_pid\n" +
                "      JOIN employee ON projectemployee.employee_pid = employee.pid\n" +
                "      LEFT JOIN department ON employee.department_pid = department.pid\n" +
                "      LEFT JOIN company ON department.company_pid = company.pid\n" +
                "),\n" +
                "project_cost AS (\n" +
                "    SELECT project_pid, sum(monthly_cost) total_cost\n" +
                "    FROM project_info GROUP BY project_pid\n" +
                ")\n" +
                "SELECT project_name projectName, total_cost totalCost, company_name companyName, sum(monthly_cost) companyCost FROM project_info\n" +
                "  JOIN project_cost USING (project_pid)\n" +
                "WHERE total_cost > :totalCostBoundary\n" +
                "GROUP BY project_name, total_cost, company_name\n" +
                "ORDER BY company_name";

        return h.createQuery(query)
                .bind("totalCostBoundary", totalCostBoundary)
                .map(ProjectsWithCostsGreaterThanOutput.class)
                .list();
    }
}
 

@Test
public void testReleaseLead() {
  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 leaderIdForSelf = UUID.randomUUID();
  UUID leaderIdForOther = UUID.randomUUID();

  storage.takeLead(leaderIdForSelf);
  storage.takeLead(leaderIdForOther);

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

  handle.createStatement("UPDATE leader SET reaper_instance_id = 0 WHERE leader_id = :id")
      .bind("id", UuidUtil.toSequenceId(leaderIdForOther))
      .execute();

  // test that releaseLead succeeds for entry where instance_id = self
  storage.releaseLead(leaderIdForSelf);
  fetchedLeaderIds = storage.getLeaders();
  Assertions.assertThat(fetchedLeaderIds.size()).isEqualTo(1);

  // test that releaseLead fails for entry where instance_id != self
  storage.releaseLead(leaderIdForOther);
  fetchedLeaderIds = storage.getLeaders();
  Assertions.assertThat(fetchedLeaderIds.size()).isEqualTo(1);
}