Java Code Examples for org.skife.jdbi.v2.DBI#registerMapper()

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

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

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

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

    NodeManager nodeManager = new TestingNodeManager();
    NodeSupplier nodeSupplier = nodeManager::getWorkerNodes;
    shardManager = createShardManager(dbi, nodeSupplier, systemTicker());
    metadata = new RaptorMetadata(dbi, shardManager);
}
 

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

    metadata = new RaptorMetadata(dbi, createShardManager(dbi));

    metadataDao = dbi.onDemand(MetadataDao.class);
    shardManager = createShardManager(dbi);
}
 

@BeforeMethod
public void setup()
{
    TypeManager typeManager = new InternalTypeManager(createTestMetadataManager());
    DBI dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    dbi.registerMapper(new TableColumn.Mapper(typeManager));
    dummyHandle = dbi.open();
    metadataDao = dbi.onDemand(MetadataDao.class);
    createTablesWithRetry(dbi);
    dataDir = Files.createTempDir();

    RaptorConnectorId connectorId = new RaptorConnectorId("test");
    NodeManager nodeManager = new TestingNodeManager();
    NodeSupplier nodeSupplier = nodeManager::getWorkerNodes;
    ShardManager shardManager = createShardManager(dbi);
    StorageManager storageManager = createOrcStorageManager(dbi, dataDir);
    StorageManagerConfig config = new StorageManagerConfig();
    connector = new RaptorConnector(
            new LifeCycleManager(ImmutableList.of(), null),
            new TestingNodeManager(),
            new RaptorMetadataFactory(dbi, shardManager),
            new RaptorSplitManager(connectorId, nodeSupplier, shardManager, false),
            new RaptorPageSourceProvider(storageManager),
            new RaptorPageSinkProvider(storageManager,
                    new PagesIndexPageSorter(new PagesIndex.TestingFactory(false)),
                    new TemporalFunction(DateTimeZone.forID("America/Los_Angeles")),
                    config),
            new RaptorNodePartitioningProvider(nodeSupplier),
            new RaptorSessionProperties(config),
            new RaptorTableProperties(typeManager),
            ImmutableSet.of(),
            new AllowAllAccessControl(),
            dbi);
}
 

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

@BeforeMethod
public void setup()
        throws Exception
{
    DBI dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    dbi.registerMapper(new TableColumn.Mapper(new InternalTypeManager(createTestMetadataManager())));
    dummyHandle = dbi.open();
    createTablesWithRetry(dbi);
    temporary = createTempDir();
    AssignmentLimiter assignmentLimiter = new AssignmentLimiter(ImmutableSet::of, systemTicker(), new MetadataConfig());
    shardManager = new DatabaseShardManager(dbi, new DaoSupplier<>(dbi, ShardDao.class), ImmutableSet::of, assignmentLimiter, systemTicker(), new Duration(0, MINUTES));
    TestingNodeManager nodeManager = new TestingNodeManager();
    NodeSupplier nodeSupplier = nodeManager::getWorkerNodes;

    String nodeName = UUID.randomUUID().toString();
    nodeManager.addNode(new InternalNode(nodeName, new URI("http://127.0.0.1/"), NodeVersion.UNKNOWN, false));

    RaptorConnectorId connectorId = new RaptorConnectorId("raptor");
    metadata = new RaptorMetadata(dbi, shardManager);

    metadata.createTable(SESSION, TEST_TABLE, false);
    tableHandle = metadata.getTableHandle(SESSION, TEST_TABLE.getTable());

    List<ShardInfo> shards = ImmutableList.<ShardInfo>builder()
            .add(shardInfo(UUID.randomUUID(), nodeName))
            .add(shardInfo(UUID.randomUUID(), nodeName))
            .add(shardInfo(UUID.randomUUID(), nodeName))
            .add(shardInfo(UUID.randomUUID(), nodeName))
            .build();

    tableId = ((RaptorTableHandle) tableHandle).getTableId();

    List<ColumnInfo> columns = metadata.getColumnHandles(SESSION, tableHandle).values().stream()
            .map(RaptorColumnHandle.class::cast)
            .map(ColumnInfo::fromHandle)
            .collect(toList());

    long transactionId = shardManager.beginTransaction();
    shardManager.commitShards(transactionId, tableId, columns, shards, Optional.empty(), 0);

    raptorSplitManager = new RaptorSplitManager(connectorId, nodeSupplier, shardManager, false);
}
 

@Override
public Handle getHandle(ConfigMapper configMapper)
{
    if (state != State.ACTIVE) {
        throw new IllegalStateException("Transaction is already " + state.name().toLowerCase(ENGLISH));
    }

    if (handle == null) {
        DBI dbi = new DBI(ds);
        ConfigKeyListMapper cklm = new ConfigKeyListMapper();
        dbi.registerMapper(new DatabaseProjectStoreManager.StoredProjectMapper(configMapper));
        dbi.registerMapper(new DatabaseProjectStoreManager.StoredProjectWithRevisionMapper(configMapper));
        dbi.registerMapper(new DatabaseProjectStoreManager.StoredRevisionMapper(configMapper));
        dbi.registerMapper(new DatabaseProjectStoreManager.StoredWorkflowDefinitionMapper(configMapper));
        dbi.registerMapper(new DatabaseProjectStoreManager.StoredWorkflowDefinitionWithProjectMapper(configMapper));
        dbi.registerMapper(new DatabaseProjectStoreManager.WorkflowConfigMapper());
        dbi.registerMapper(new DatabaseProjectStoreManager.IdNameMapper());
        dbi.registerMapper(new DatabaseProjectStoreManager.ScheduleStatusMapper());
        dbi.registerMapper(new DatabaseQueueSettingStoreManager.StoredQueueSettingMapper(configMapper));
        dbi.registerMapper(new DatabaseScheduleStoreManager.StoredScheduleMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredTaskMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.ArchivedTaskMapper(cklm, configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.ResumingTaskMapper(cklm, configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredSessionMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredSessionWithLastAttemptMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredSessionAttemptMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredSessionAttemptWithSessionMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.TaskStateSummaryMapper());
        dbi.registerMapper(new DatabaseSessionStoreManager.TaskAttemptSummaryMapper());
        dbi.registerMapper(new DatabaseSessionStoreManager.SessionAttemptSummaryMapper());
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredSessionMonitorMapper(configMapper));
        dbi.registerMapper(new DatabaseSessionStoreManager.StoredDelayedSessionAttemptMapper());
        dbi.registerMapper(new DatabaseSessionStoreManager.TaskRelationMapper());
        dbi.registerMapper(new DatabaseSessionStoreManager.InstantMapper());
        dbi.registerMapper(new DatabaseSecretStore.ScopedSecretMapper());
        dbi.registerMapper(new DatabaseTaskQueueServer.ImmutableTaskQueueLockMapper());

        dbi.registerArgumentFactory(configMapper.getArgumentFactory());
        handle = dbi.open();

        try {
            handle.getConnection().setAutoCommit(autoAutoCommit);
        }
        catch (SQLException ex) {
            throw new TransactionFailedException("Failed to set auto commit: " + autoAutoCommit, ex);
        }
        if (!autoAutoCommit) {
            handle.begin();
        }
    }
    return handle;
}
 

/**
 * Construct a new DBI wrapper for a JDBC datasource using an optional secondary mapper. The secondary mapper comes
 * into play when the default mapper is unable to directly map a value that the SQL query returned to a property in
 * the bean class you have specified. In this case you give it a lambda to call that will insert the value correctly
 * into the bean object.
 * @param dataSource JDBC datasource to be wrapped.
 * @param secondaryMapper A lambda that is invoked when a data element can't be directly mapped to the bean is
 *                        extracted. The lambda takes as its arguments the bean that is being extracted to and
 *                        a Map.Entry that represents the (key, value) pair of the data element that could not be
 *                        directly mapped to the bean. If set to null, unrecognized properties will be ignored.
 * @return A fully constructed DBI object.
 */
public static <T> DBI newDBI(DataSource dataSource, @Nullable BiConsumer<T, Map.Entry<String, String>> secondaryMapper) {
    DBI dbi = new DBI(dataSource);
    dbi.setStatementRewriter(new HashPrefixStatementRewriter());
    dbi.registerMapper(new EtlBeanMapperFactory(secondaryMapper));
    dbi.registerArgumentFactory(new PostgresStringArrayArgumentFactory());

    return dbi;
}