Java Code Examples for com.google.cloud.bigquery.DatasetId#of()

/** Example of deleting a dataset, even if non-empty. */
// [TARGET delete(DatasetId, DatasetDeleteOption...)]
// [VARIABLE "my_project_id"]
// [VARIABLE "my_dataset_name"]
public boolean deleteDatasetFromId(String projectId, String datasetName) {
  // [START bigquery_delete_dataset]
  DatasetId datasetId = DatasetId.of(projectId, datasetName);
  boolean deleted = bigquery.delete(datasetId, DatasetDeleteOption.deleteContents());
  if (deleted) {
    // the dataset was deleted
  } else {
    // the dataset was not found
  }
  // [END bigquery_delete_dataset]
  return deleted;
}
 

@Override
public StagingTableId provide(EvalContext evalContext) {
  final String location = getDatasetOrThrow().getLocation();

  final TableId stagingTableId = bigQuery().createStagingTableId(tableId, location);
  final DatasetId stagingDatasetId = DatasetId.of(stagingTableId.getProject(), stagingTableId.getDataset());

  if (bigQuery().getDataset(stagingDatasetId) == null) {
    bigQuery().create(DatasetInfo
        .newBuilder(stagingDatasetId)
        .setLocation(location)
        .setDefaultTableLifetime(Duration.ofDays(1).toMillis())
        .build());
    LOG.info("created staging dataset: {}", stagingDatasetId);
  }

  return StagingTableId.of(this, stagingTableId);
}
 

@Override
public Set<String> listTables() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    DatasetId datasetId = DatasetId.of(properties.getDatastoreProperties().projectName.getValue(),
            properties.bqDataset.getValue());
    Page<Table> tables = bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(100));
    Set<String> tablesName = new HashSet<>();
    Iterator<Table> tableIterator = tables.iterateAll().iterator();
    while (tableIterator.hasNext()) {
        tablesName.add(tableIterator.next().getTableId().getTable());
    }
    return tablesName;
}
 

@Test
public void testCreateGetAndDeleteDataset() throws InterruptedException {
  DatasetId datasetId = DatasetId.of(bigquery.getOptions().getProjectId(), OTHER_DATASET);
  Dataset dataset = bigquerySnippets.createDataset(OTHER_DATASET);
  assertNotNull(dataset);
  assertEquals(datasetId, bigquerySnippets.getDataset(OTHER_DATASET).getDatasetId());
  assertNotNull(bigquerySnippets.updateDataset(OTHER_DATASET, "new description"));
  assertEquals(
      "new description",
      bigquerySnippets.getDatasetFromId(datasetId.getProject(), OTHER_DATASET).getDescription());
  Set<DatasetId> datasets =
      Sets.newHashSet(
          Iterators.transform(
              bigquerySnippets.listDatasets().iterateAll().iterator(), TO_DATASET_ID_FUNCTION));
  while (!datasets.contains(datasetId)) {
    Thread.sleep(500);
    datasets =
        Sets.newHashSet(
            Iterators.transform(
                bigquerySnippets.listDatasets().iterateAll().iterator(), TO_DATASET_ID_FUNCTION));
  }
  datasets =
      Sets.newHashSet(
          Iterators.transform(
              bigquerySnippets.listDatasets(datasetId.getProject()).iterateAll().iterator(),
              TO_DATASET_ID_FUNCTION));
  while (!datasets.contains(datasetId)) {
    Thread.sleep(500);
    datasets =
        Sets.newHashSet(
            Iterators.transform(
                bigquerySnippets.listDatasets(datasetId.getProject()).iterateAll().iterator(),
                TO_DATASET_ID_FUNCTION));
  }
  assertTrue(bigquerySnippets.deleteDataset(OTHER_DATASET));
  assertFalse(bigquerySnippets.deleteDatasetFromId(datasetId.getProject(), OTHER_DATASET));
}
 

@Override
DatasetId parse(String... args) throws Exception {
  String message;
  if (args.length == 1) {
    return DatasetId.of(args[0]);
  } else if (args.length > 1) {
    message = "Too many arguments.";
  } else {
    message = "Missing required dataset id.";
  }
  throw new IllegalArgumentException(message);
}
 

/** Example of getting a dataset. */
// [TARGET getDataset(DatasetId, DatasetOption...)]
// [VARIABLE "my_project_id"]
// [VARIABLE "my_dataset_name"]
public Dataset getDatasetFromId(String projectId, String datasetName) {
  // [START bigquery_get_dataset]
  DatasetId datasetId = DatasetId.of(projectId, datasetName);
  Dataset dataset = bigquery.getDataset(datasetId);
  // [END bigquery_get_dataset]
  return dataset;
}
 

/** Example of listing the tables in a dataset. */
// [TARGET listTables(DatasetId, TableListOption...)]
// [VARIABLE "my_project_id"]
// [VARIABLE "my_dataset_name"]
public Page<Table> listTablesFromId(String projectId, String datasetName) {
  // [START bigquery_list_tables]
  DatasetId datasetId = DatasetId.of(projectId, datasetName);
  Page<Table> tables = bigquery.listTables(datasetId, TableListOption.pageSize(100));
  for (Table table : tables.iterateAll()) {
    // do something with the table
  }
  // [END bigquery_list_tables]
  return tables;
}
 

@AfterClass
public static void cleanDatasetAndTable() {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    for (String dataset : datasets) {
        DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, dataset);
        bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
    }
}
 

@BeforeClass
public static void initDatasetAndTable() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    for (String dataset : datasets) {
        DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, dataset);
        bigquery.create(DatasetInfo.of(datasetId));
    }

    for (String table : tables) {
        TableDefinition tableDefinition =
                StandardTableDefinition.of(Schema.of(Field.of("test", LegacySQLTypeName.STRING)));
        TableId tableId = TableId.of(BigQueryTestConstants.PROJECT, datasets.get(0), table);
        bigquery.create(TableInfo.of(tableId, tableDefinition));
    }
}
 

@Test
public void shouldCreateStagingDatasetIfDoesNotExist() {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(dataset);

  final BigQueryOutput bigQueryOutput = BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID);

  bigQueryOutput.provide(null);

  verify(bigQuery).create(any(DatasetInfo.class));
  final DatasetInfo createdStagingDataset = datasetInfoCaptor.getValue();
  final DatasetId expectedDatasetId = DatasetId.of(DATASET_ID.getProject(), "_incoming_" + LOCATION);
  assertThat(createdStagingDataset.getDatasetId(), is(expectedDatasetId));
  assertThat(createdStagingDataset.getLocation(), is(LOCATION));
  assertThat(createdStagingDataset.getDefaultTableLifetime(), is(Duration.ofDays(1).toMillis()));
}
 

private DatasetInfo getDatasetOrThrow() {
  final DatasetId datasetId = DatasetId.of(tableId.getProject(), tableId.getDataset());

  final DatasetInfo dataset = bigQuery().getDataset(datasetId);

  if (dataset == null) {
    LOG.error("Could not find dataset {}", datasetId);
    throw new IllegalArgumentException(
        "Dataset does not exist. Please create it before attempting to write to it.");
  }

  return dataset;
}
 

@After
public void tearDown() {
  System.setOut(null);
  bout.reset();
  DatasetId datasetId = DatasetId.of(bigquery.getOptions().getProjectId(), datasetName);
  bigquery.delete(datasetId, DatasetDeleteOption.deleteContents());
}
 

@After
public void tearDown() {
  String consoleOutput = bout.toString();
  System.setOut(null);
  deleteObjects();
  DatasetId datasetId = DatasetId.of(bigquery.getOptions().getProjectId(), datasetName);
  bigquery.delete(datasetId, DatasetDeleteOption.deleteContents());
}
 

private Dataset addDataSetMappingIfNeeded(Dataset dataset)
{
    DatasetId bigQueryDatasetId = dataset.getDatasetId();
    DatasetId prestoDatasetId = DatasetId.of(bigQueryDatasetId.getProject(), bigQueryDatasetId.getDataset().toLowerCase(ENGLISH));
    datasetIds.putIfAbsent(prestoDatasetId, bigQueryDatasetId);
    return dataset;
}
 

public static void createDataset(String dataset) {
    BigQuery bq = getBigquery();
    DatasetId datasetId = DatasetId.of(dataset);
    logger.warn("Creating test dataset: {}", datasetId);
    bq.create(DatasetInfo.of(datasetId));
}
 

private static DatasetId datasetIdOf(TableId tableId) {
  return DatasetId.of(tableId.getProject(), tableId.getDataset());
}
 

@BeforeClass
public static void initDataset() {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, datasetName);
    bigquery.create(DatasetInfo.of(datasetId));
}
 

@AfterClass
public static void cleanDataset() {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, datasetName);
    bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
}
 

private DatasetId mapIfNeeded(String project, String dataset)
{
    DatasetId datasetId = DatasetId.of(project, dataset);
    return datasetIds.getOrDefault(datasetId, datasetId);
}
 

DatasetId toDatasetId(TableId tableId)
{
    return DatasetId.of(tableId.getProject(), tableId.getDataset());
}