com.google.cloud.bigquery.DatasetId Java Examples

private List<SchemaTableName> listTablesWithTypes(ConnectorSession session, Optional<String> schemaName, TableDefinition.Type... types)
{
    if (schemaName.isPresent() && schemaName.get().equalsIgnoreCase(INFORMATION_SCHEMA)) {
        return ImmutableList.of();
    }
    Set<String> schemaNames = schemaName.map(ImmutableSet::of)
            .orElseGet(() -> ImmutableSet.copyOf(listSchemaNames(session)));

    ImmutableList.Builder<SchemaTableName> tableNames = ImmutableList.builder();
    for (String datasetId : schemaNames) {
        for (Table table : bigQueryClient.listTables(DatasetId.of(projectId, datasetId), types)) {
            tableNames.add(new SchemaTableName(datasetId, table.getTableId().getTable()));
        }
    }
    return tableNames.build();
}
 

/** 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);
}
 

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

@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 shouldReturnTableIdOnJobSuccess() throws InterruptedException {
  when(bigQuery.getDataset(any(DatasetId.class))).thenReturn(dataset);
  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  when(job.waitFor(any(RetryOption.class))).thenReturn(job);
  when(job.getStatus()).thenReturn(mock(JobStatus.class));

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

  final StagingTableId stagingTableId = bigQueryOutput.provide(null);

  final TableId tableId = stagingTableId.publish();

  assertThat(tableId, is(TABLE_ID));
}
 

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

@Override
public void publish(StagingTableId stagingTableId, TableId tableId) {
  stagingTableIds.remove(formatTableIdKey(tableId));

  final DatasetId datasetId = datasetIdOf(tableId);
  publishedTables.computeIfAbsent(datasetId, k -> new ConcurrentSkipListSet<>())
      .add(tableId.getTable());
}
 

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

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

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

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

/**
 * Generating BQ table destination from dataset reference and featuresSet's project and name. If
 * project is undefined "default" would be selected
 *
 * @param dataset {@link DatasetId} reference to bq project and dataset
 * @param featureSetKey Feature Set reference with format <project>/<feature-set-name>
 * @return {@link TableDestination}
 */
public static TableDestination getTableDestination(DatasetId dataset, String featureSetKey) {
  String[] splitName = featureSetKey.split("/");
  String projectName, setName;

  if (splitName.length == 2) {
    projectName = splitName[0];
    setName = splitName[1];
  } else {
    projectName = DEFAULT_PROJECT_NAME;
    setName = splitName[0];
  }

  TimePartitioning timePartitioning =
      new TimePartitioning()
          .setType("DAY")
          .setField(FeatureRowToTableRow.getEventTimestampColumn());

  return new TableDestination(
      String.format(
          "%s:%s.%s_%s",
          dataset.getProject(),
          dataset.getDataset(),
          projectName.replaceAll("-", "_"),
          setName.replaceAll("-", "_")),
      String.format("Feast table for %s", featureSetKey),
      timePartitioning);
}
 

@Override
public void run(BigQuery bigquery, DatasetId datasetId) {
  if (bigquery.delete(datasetId)) {
    System.out.println("Dataset " + datasetId + " was deleted");
  } else {
    System.out.println("Dataset " + datasetId + " not found");
  }
}
 

TableId createDestinationTable(TableId tableId)
{
    String project = viewMaterializationProject.orElse(tableId.getProject());
    String dataset = viewMaterializationDataset.orElse(tableId.getDataset());
    DatasetId datasetId = mapIfNeeded(project, dataset);
    String name = format("_pbc_%s", randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""));
    return TableId.of(datasetId.getProject(), datasetId.getDataset(), name);
}
 

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

Iterable<Table> listTables(DatasetId datasetId, TableDefinition.Type... types)
{
    Set<TableDefinition.Type> allowedTypes = ImmutableSet.copyOf(types);
    DatasetId bigQueryDatasetId = datasetIds.getOrDefault(datasetId, datasetId);
    Iterable<Table> allTables = bigQuery.listTables(bigQueryDatasetId).iterateAll();
    return StreamSupport.stream(allTables.spliterator(), false)
            .filter(table -> allowedTypes.contains(table.getDefinition().getType()))
            .collect(toImmutableList());
}
 

@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
public void run(BigQuery bigquery, DatasetId datasetId) {
  System.out.println("Dataset info: " + bigquery.getDataset(datasetId));
}
 

@Override
public void run(BigQuery bigquery, DatasetId datasetId) {
  for (Table table : bigquery.listTables(datasetId).iterateAll()) {
    System.out.println(table);
  }
}
 

@Override
public void run(BigQuery bigquery, DatasetId datasetId) {
  bigquery.create(DatasetInfo.newBuilder(datasetId).build());
  System.out.println("Created dataset " + datasetId);
}
 

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

@Override
public DatasetId apply(Dataset dataset) {
  return dataset.getDatasetId();
}
 

@Override
public DatasetInfo getDataset(DatasetId datasetId) {
  return Dataset.newBuilder(datasetId)
      .setLocation("test") // TOOD: make mockable?
      .build();
}
 

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

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