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

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

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

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

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

@BeforeClass
public static void setUpClass() throws ClassNotFoundException, SQLException,
    IOException, URISyntaxException {

  setUpDb(dbSchemaUrl, "sa", "", "tpcds.sql");

  DBI dbi = new DBI(dbSchemaUrl, "sa", "");
  cubeDAO = dbi.onDemand(CubeDAO.class);
}
 

@BeforeClass
public static void setUpClass() throws ClassNotFoundException, SQLException,
    IOException, URISyntaxException {

  setUpDb(dbSchemaUrl, "sa", "", "tpcds.sql");

  DBI dbi = new DBI(dbSchemaUrl, "sa", "");
  dsSetDAO = dbi.onDemand(DSSetDAO.class);
}
 

@BeforeClass
public static void setUpClass() throws ClassNotFoundException, SQLException,
    IOException, URISyntaxException {

  setUpDb(dbSchemaUrl, "sa", "", "tpcds.sql");
  setUpDb(dbSchemaUrl, "sa", "", "tpcds_2.sql");

  DBI dbi = new DBI(dbSchemaUrl, "sa", "");
  viewDAO = dbi.onDemand(ViewDAO.class);
}
 

@BeforeClass
public static void setUpClass() throws ClassNotFoundException, SQLException,
    IOException, URISyntaxException {

  setUpDb(dbSchemaUrl, "sa", "", "tpcds.sql");
  setUpDb(dbSchemaUrl, "sa", "", "tpcds_2.sql");

  DBI dbi = new DBI(dbSchemaUrl, "sa", "");
  jdbcSourceDAO = dbi.onDemand(JdbcSourceDAO.class);
  dbi.define("encryptClass", new AESEncrypt("easy"));
}
 

public int executeCreateView(SqlCreateQuarkView sqlNode) throws SQLException {
  DBI dbi = getDbi();
  List<String> tableNameList = sqlNode.getTableName().names;
  String dataSourceName = tableNameList.get(0);

  ViewDAO viewDAO = dbi.onDemand(ViewDAO.class);

  JdbcSourceDAO jdbcDAO = dbi.onDemand(JdbcSourceDAO.class);
  QuboleDbSourceDAO quboleDAO = dbi.onDemand(QuboleDbSourceDAO.class);
  DataSource dataSource = jdbcDAO.findByName(dataSourceName,
      connection.getDSSet().getId());
  if (dataSource == null) {
    dataSource = quboleDAO.findByName(dataSourceName, connection.getDSSet().getId());
  }

  if (dataSource == null) {
    throw new SQLException("DataSource with name '" + dataSourceName + "' not found");
  }

  SqlPrettyWriter writer = new SqlPrettyWriter(SqlDialect.CALCITE);
  writer.setAlwaysUseParentheses(false);
  writer.setSelectListItemsOnSeparateLines(false);
  writer.setIndentation(0);
  writer.setQuoteAllIdentifiers(true);
  sqlNode.getQuery().unparse(writer, 0, 0);
  final String sql = writer.toString();

  LOG.debug(sql);
  return viewDAO.insert(sqlNode.getName(),
      "No Description", sql,
      0L, dataSource.getId(),
      tableNameList.get(1), tableNameList.get(2),
      connection.getDSSet().getId());
}
 

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

@BeforeMethod
public void setup()
{
    dbi = new DBI("jdbc:h2:mem:test" + System.nanoTime() + ThreadLocalRandom.current().nextLong());
    dummyHandle = dbi.open();
    metadataDao = dbi.onDemand(MetadataDao.class);
    organizerDao = dbi.onDemand(ShardOrganizerDao.class);

    createTablesWithRetry(dbi);
}
 

@Provides
public QueryStore provideQueryStore(DBI dbi)
{
    return dbi.onDemand(QueryStoreDAO.class);
}
 

private int executeCreateDataSource(SqlCreateQuarkDataSource sqlNode) throws SQLException {
  DBI dbi = getDbi();

  Map<String, Object> commonColumns = new HashMap<>();
  Map<String, Object> dbSpecificColumns = new HashMap<>();
  DataSourceDAO dataSourceDAO = dbi.onDemand(DataSourceDAO.class);
  JdbcSourceDAO jdbcSourceDAO = null;
  QuboleDbSourceDAO quboleDbSourceDAO = null;

  int i = 0;
  SqlNodeList rowList = sqlNode.getSourceExpressionList();
  for (SqlNode node : sqlNode.getTargetColumnList()) {
    if (node instanceof SqlIdentifier) {
      switch (((SqlIdentifier) node).getSimple()) {
        case "type":
          commonColumns.put("type", rowList.get(i).toString());
          break;
        case "url":
          commonColumns.put("url", rowList.get(i).toString());
          break;
        case "ds_set_id":
          break;
        case "datasource_type":
          if (rowList.get(i).toString().toUpperCase().equals("JDBC")) {
            jdbcSourceDAO = dbi.onDemand(JdbcSourceDAO.class);
          } else if (rowList.get(i).toString().toUpperCase().equals("QUBOLEDB")) {
            quboleDbSourceDAO = dbi.onDemand(QuboleDbSourceDAO.class);
          } else {
            throw new SQLException("Incorrect argument type <" + rowList.get(i).toString()
                + "> to variable 'datasource_type'");
          }
          commonColumns.put("datasource_type", rowList.get(i).toString());
          break;
        case "username":
          dbSpecificColumns.put("username", rowList.get(i).toString());
          break;
        case "password":
          dbSpecificColumns.put("password", rowList.get(i).toString());
          break;
        case "dbtap_id":
          if (rowList.get(i) instanceof SqlNumericLiteral) {
            dbSpecificColumns.put("dbtap_id",
                ((SqlNumericLiteral) rowList.get(i)).intValue(true));
          } else {
            throw new SQLException("Incorrect argument type to variable"
                + " 'dbtap_id'");
          }
          break;
        case "auth_token":
          dbSpecificColumns.put("auth_token", rowList.get(i).toString());
          break;
        default:
          throw new SQLException("Unknown parameter: " + ((SqlIdentifier) node).getSimple());
      }
      i++;
    }
  }

  Encrypt encrypt;
  if (Boolean.parseBoolean(info.getProperty("encrypt", "false"))) {
    encrypt = new AESEncrypt(info.getProperty("encryptionKey"));
  } else {
    encrypt = new NoopEncrypt();
  }

  if ((jdbcSourceDAO == null && quboleDbSourceDAO == null)
      || (jdbcSourceDAO != null && quboleDbSourceDAO != null)) {
    throw new RuntimeException("Need to pass exact values to create"
        + " data source of type jdbc or quboleDb");
  } else if (jdbcSourceDAO != null) {
    return dataSourceDAO.insertJDBC((String) sqlNode.getIdentifier().getSimple(),
        (String) commonColumns.get("type"),
        (String) commonColumns.get("url"),
        connection.getDSSet().getId(),
        (String) commonColumns.get("datasource_type"),
        jdbcSourceDAO,
        (String) dbSpecificColumns.get("username"),
        (dbSpecificColumns.get("password") == null) ? ""
            : (String) dbSpecificColumns.get("password"),
        encrypt);
  } else {
    return dataSourceDAO.insertQuboleDB((String) sqlNode.getIdentifier().getSimple(),
        (String) commonColumns.get("type"),
        (String) commonColumns.get("url"),
        connection.getDSSet().getId(),
        (String) commonColumns.get("datasource_type"),
        quboleDbSourceDAO,
        (int) dbSpecificColumns.get("dbtap_id"),
        (String) dbSpecificColumns.get("auth_token"),
        encrypt);
  }
}
 

private void executeDeleteOnView(SqlDropQuarkView node) throws SQLException {
  DBI dbi = getDbi();
  ViewDAO viewDAO = dbi.onDemand(ViewDAO.class);
  viewDAO.delete(node.getIdentifier().getSimple(), connection.getDSSet().getId());
}
 

public int executeAlterDataSource(SqlAlterQuarkDataSource sqlNode) throws SQLException {
  DBI dbi = getDbi();
  DataSourceDAO dataSourceDAO = dbi.onDemand(DataSourceDAO.class);
  JdbcSourceDAO jdbcDAO = dbi.onDemand(JdbcSourceDAO.class);
  QuboleDbSourceDAO quboleDAO = dbi.onDemand(QuboleDbSourceDAO.class);
  DataSource dataSource = jdbcDAO.findByName(sqlNode.getIdentifier().getSimple(),
      connection.getDSSet().getId());
  if (dataSource == null) {
    dataSource = quboleDAO.findByName(sqlNode.getIdentifier().getSimple(),
        connection.getDSSet().getId());
  }
  if (dataSource == null) {
    return 0;
  }
  SqlNodeList rowList = sqlNode.getSourceExpressionList();
  int i = 0;
  for (SqlNode node : sqlNode.getTargetColumnList()) {
    if (node instanceof SqlIdentifier) {
      switch (((SqlIdentifier) node).getSimple()) {
        case "name":
          dataSource.setName(rowList.get(i).toString());
          break;
        case "type":
          dataSource.setType(rowList.get(i).toString());
          break;
        case "url":
          dataSource.setUrl(rowList.get(i).toString());
          break;
        case "ds_set_id":
          break;
        case "datasource_type":
          dataSource.setDatasourceType(rowList.get(i).toString());
          break;
        case "username":
          if (dataSource instanceof JdbcSource) {
            ((JdbcSource) dataSource)
                .setUsername(rowList.get(i).toString());
          }
          break;
        case "password":
          if (dataSource instanceof JdbcSource) {
            ((JdbcSource) dataSource)
                .setPassword(rowList.get(i).toString());
          }
          break;
        case "dbtap_id":
          if (dataSource instanceof QuboleDbSource) {
            if (rowList.get(i) instanceof SqlNumericLiteral) {
              ((QuboleDbSource) dataSource).setDbTapId(
                  ((SqlNumericLiteral) rowList.get(i)).intValue(true));
            } else {
              throw new SQLException("Incorrect argument type to variable"
                  + " 'dbtap_id'");
            }
          }
          break;
        case "auth_token":
          if (dataSource instanceof QuboleDbSource) {
            ((QuboleDbSource) dataSource)
                .setAuthToken(rowList.get(i).toString());
          }
          break;
        default:
          throw new SQLException("Unknown parameter: " + ((SqlIdentifier) node).getSimple());
      }
      i++;
    }
  }

  Encrypt encrypt;
  if (Boolean.parseBoolean(info.getProperty("encrypt", "false"))) {
    encrypt = new AESEncrypt(info.getProperty("encryptionKey"));
  } else {
    encrypt = new NoopEncrypt();
  }
  if (dataSource instanceof JdbcSource) {
    return jdbcDAO.update((JdbcSource) dataSource, dataSourceDAO, encrypt);
  } else {
    return quboleDAO.update((QuboleDbSource) dataSource, dataSourceDAO, encrypt);
  }
}
 

public JdbiDynamicAttributes(DBI jdbi, List<String> scopes)
{
    container = new StandardAttributesContainer(scopes);
    dao = jdbi.onDemand(AttributeEntityDao.class);
}
 

/**
 * Creates list of QuarkSchema
 *
 * @param info A JSON string which contains database credentials
 * @return
 * @throws QuarkException
 */
public QuarkFactoryResult create(Properties info) throws QuarkException {
  try {
    Connection connection = new Connection(info);
    connection.runFlyWay();

    DSSet dsSet = connection.getDSSet();
    DBI dbi = connection.getDbi();

    JdbcSourceDAO jdbcSourceDAO = dbi.onDemand(JdbcSourceDAO.class);
    QuboleDbSourceDAO quboleDbSourceDAO = dbi.onDemand(QuboleDbSourceDAO.class);
    CubeDAO cubeDAO = dbi.onDemand(CubeDAO.class);
    ViewDAO viewDAO = dbi.onDemand(ViewDAO.class);
    MeasureDAO measureDAO = dbi.onDemand(MeasureDAO.class);
    DimensionDAO dimensionDAO = dbi.onDemand(DimensionDAO.class);

    long dsSetId = dsSet.getId();
    long defaultDataSourceId = dsSet.getDefaultDatasourceId();

    List<JdbcSource> jdbcSources = jdbcSourceDAO.findByDSSetId(dsSetId);
    List<QuboleDbSource> quboleDbSources = quboleDbSourceDAO.findByDSSetId(dsSetId);

    List<DataSource> dataSources = new ArrayList<>();
    dataSources.addAll(jdbcSources);
    dataSources.addAll(quboleDbSources);

    ImmutableList.Builder<com.qubole.quark.planner.DataSourceSchema> schemaList =
        new ImmutableList.Builder<>();

    com.qubole.quark.planner.DataSourceSchema defaultSchema = null;

    for (DataSource dataSource : dataSources) {
      com.qubole.quark.planner.DataSourceSchema dataSourceSchema = new DataSourceSchema(
          dataSource.getProperties(defaultDataSourceId));
      if (dataSource.getId() == defaultDataSourceId) {
        defaultSchema = dataSourceSchema;
      }
      schemaList.add(dataSourceSchema);
    }

    RelSchema relSchema = getRelSchema(viewDAO, cubeDAO, measureDAO, dimensionDAO, dsSetId);

    return new QuarkFactoryResult(schemaList.build(), relSchema, defaultSchema);
  } catch (Exception se) {
    LOG.error(se.getMessage());
    throw new QuarkException(se);
  }
}
 

@Override
public void run(RufusConfiguration conf, Environment env) throws Exception {
    final DBIFactory factory = new DBIFactory();
    final DBI jdbi = factory.build(env, conf.getDataSourceFactory(), DB_SOURCE);

    final UserDao userDao = jdbi.onDemand(UserDao.class);
    final ArticleDao articleDao = jdbi.onDemand(ArticleDao.class);

    final FeedProcessorImpl processor = FeedProcessorImpl.newInstance(articleDao);
    final FeedParser parser = new FeedParser(articleDao, processor);

    final JwtConsumer jwtConsumer = new JwtConsumerBuilder()
        .setAllowedClockSkewInSeconds(30)
        .setRequireExpirationTime()
        .setRequireSubject()
        .setVerificationKey(new HmacKey(VERIFICATION_KEY))
        .setRelaxVerificationKeyValidation()
        .build();
    final CachingJwtAuthenticator<User> cachingJwtAuthenticator = new CachingJwtAuthenticator<>(
        env.metrics(),
        new JwtAuthenticator(userDao),
        conf.getAuthenticationCachePolicy()
    );

    env.jersey().register(new ArticleResource(userDao, articleDao, processor, parser));
    env.jersey().register(
        new UserResource(
            new BasicAuthenticator(userDao),
            new TokenGenerator(VERIFICATION_KEY),
            userDao,
            articleDao
        )
    );

    //route source
    env.jersey().setUrlPattern(ROOT_PATH);

    env.jersey().register(new AuthValueFactoryProvider.Binder<>(User.class));
    env.jersey().register(new AuthDynamicFeature(
        new JwtAuthFilter.Builder<User>()
            .setJwtConsumer(jwtConsumer)
            .setRealm(REALM)
            .setPrefix(BEARER)
            .setAuthenticator(cachingJwtAuthenticator)
            .buildAuthFilter()
    ));
}
 

@BeforeClass
public static void setUpClass() throws ClassNotFoundException, SQLException,
    IOException, URISyntaxException {

  setUpDb(dbSchemaUrl, "sa", "", "tpcds.sql");


  DBI dbi = new DBI(dbSchemaUrl, "sa", "");
  measureDAO = dbi.onDemand(MeasureDAO.class);
}