How to use

Dataset

in

com.google.cloud.bigquery

@Test
public void testGetDataset() {
 EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS))
   .andReturn(DATASET_INFO_WITH_PROJECT.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery = options.getService();
 Dataset dataset = bigquery.getDataset(DATASET);
 assertEquals(
   new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset);
}

@Test
public void testDeleteTrue() {
 initializeExpectedDataset(1);
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.delete(DATASET_INFO.getDatasetId())).andReturn(true);
 replay(bigquery);
 initializeDataset();
 assertTrue(dataset.delete());
}

@Test
public void testGetDataset() {
 Dataset dataset = bigquery.getDataset(DATASET);
 assertEquals(bigquery.getOptions().getProjectId(), dataset.getDatasetId().getProject());
 assertEquals(DATASET, dataset.getDatasetId().getDataset());
 assertEquals(DESCRIPTION, dataset.getDescription());
 assertEquals(LABELS, dataset.getLabels());
 assertNotNull(dataset.getAcl());
 assertNotNull(dataset.getEtag());
 assertNotNull(dataset.getGeneratedId());
 assertNotNull(dataset.getLastModified());
 assertNotNull(dataset.getSelfLink());
}

@Test
public void testToAndFromPb() {
 initializeExpectedDataset(4);
 replay(bigquery);
 compareDataset(
   expectedDataset, Dataset.fromPb(serviceMockReturnsOptions, expectedDataset.toPb()));
}

@Test
public void testListDatasetsWithProjects() {
 bigquery = options.getService();
 ImmutableList<Dataset> datasetList =
   ImmutableList.of(
     new Dataset(
       bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO.setProjectId(OTHER_PROJECT))));
 Tuple<String, Iterable<com.google.api.services.bigquery.model.Dataset>> result =
   Tuple.of(CURSOR, Iterables.transform(datasetList, DatasetInfo.TO_PB_FUNCTION));
 EasyMock.expect(bigqueryRpcMock.listDatasets(OTHER_PROJECT, EMPTY_RPC_OPTIONS))
   .andReturn(result);
 EasyMock.replay(bigqueryRpcMock);
 Page<Dataset> page = bigquery.listDatasets(OTHER_PROJECT);
 assertEquals(CURSOR, page.getNextPageToken());
 assertArrayEquals(
   datasetList.toArray(), Iterables.toArray(page.getValues(), DatasetInfo.class));
}

@Test
public void testGetNull() throws Exception {
 initializeExpectedDataset(1);
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.getTable(TABLE_INFO1.getTableId())).andReturn(null);
 replay(bigquery);
 initializeDataset();
 assertNull(dataset.get(TABLE_INFO1.getTableId().getTable()));
}

@Test
public void testGetDatasetWithSelectedFields() {
 Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance();
 EasyMock.expect(bigqueryRpcMock.getDataset(eq(PROJECT), eq(DATASET), capture(capturedOptions)))
   .andReturn(DATASET_INFO_WITH_PROJECT.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery = options.getService();
 Dataset dataset = bigquery.getDataset(DATASET, DATASET_OPTION_FIELDS);
 String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.getRpcOption());
 assertTrue(selector.contains("datasetReference"));
 assertTrue(selector.contains("access"));
 assertTrue(selector.contains("etag"));
 assertEquals(28, selector.length());
 assertEquals(
   new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset);
}

@Test
public void testReload() throws Exception {
 initializeExpectedDataset(4);
 DatasetInfo updatedInfo = DATASET_INFO.toBuilder().setDescription("Description").build();
 Dataset expectedDataset =
   new Dataset(serviceMockReturnsOptions, new DatasetInfo.BuilderImpl(updatedInfo));
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.getDataset(DATASET_INFO.getDatasetId().getDataset()))
   .andReturn(expectedDataset);
 replay(bigquery);
 initializeDataset();
 Dataset updatedDataset = dataset.reload();
 compareDataset(expectedDataset, updatedDataset);
}

@Test
public void testCreateTable() throws Exception {
 initializeExpectedDataset(2);
 Table expectedTable =
   new Table(serviceMockReturnsOptions, new TableInfo.BuilderImpl(TABLE_INFO1));
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.create(TABLE_INFO1)).andReturn(expectedTable);
 replay(bigquery);
 initializeDataset();
 Table table = dataset.create(TABLE_INFO1.getTableId().getTable(), TABLE_DEFINITION);
 assertEquals(expectedTable, table);
}

@Test
public void testExists_True() throws Exception {
 initializeExpectedDataset(1);
 BigQuery.DatasetOption[] expectedOptions = {BigQuery.DatasetOption.fields()};
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.getDataset(DATASET_INFO.getDatasetId(), expectedOptions))
   .andReturn(expectedDataset);
 replay(bigquery);
 initializeDataset();
 assertTrue(dataset.exists());
}

@Test
public void testList() throws Exception {
 initializeExpectedDataset(4);
 List<Table> tableResults =
   ImmutableList.of(
     new Table(serviceMockReturnsOptions, new Table.BuilderImpl(TABLE_INFO1)),
     new Table(serviceMockReturnsOptions, new Table.BuilderImpl(TABLE_INFO2)),
     new Table(serviceMockReturnsOptions, new Table.BuilderImpl(TABLE_INFO3)));
 PageImpl<Table> expectedPage = new PageImpl<>(null, "c", tableResults);
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.listTables(DATASET_INFO.getDatasetId())).andReturn(expectedPage);
 replay(bigquery);
 initializeDataset();
 Page<Table> tablePage = dataset.list();
 assertArrayEquals(
   tableResults.toArray(), Iterables.toArray(tablePage.getValues(), Table.class));
 assertEquals(expectedPage.getNextPageToken(), tablePage.getNextPageToken());
}

private void initializeExpectedDataset(int optionsCalls) {
 expect(serviceMockReturnsOptions.getOptions()).andReturn(mockOptions).times(optionsCalls);
 replay(serviceMockReturnsOptions);
 bigquery = createStrictMock(BigQuery.class);
 expectedDataset = new Dataset(serviceMockReturnsOptions, new Dataset.BuilderImpl(DATASET_INFO));
}

@Test
public void testUpdateDatasetWithSelectedFields() {
 Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance();
 DatasetInfo updatedDatasetInfo =
   DATASET_INFO.toBuilder().setDescription("newDescription").build();
 DatasetInfo updatedDatasetInfoWithProject =
   DATASET_INFO_WITH_PROJECT.toBuilder().setDescription("newDescription").build();
 EasyMock.expect(
     bigqueryRpcMock.patch(
       eq(updatedDatasetInfoWithProject.toPb()), capture(capturedOptions)))
   .andReturn(updatedDatasetInfoWithProject.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery = options.getService();
 Dataset dataset = bigquery.update(updatedDatasetInfo, DATASET_OPTION_FIELDS);
 String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.getRpcOption());
 assertTrue(selector.contains("datasetReference"));
 assertTrue(selector.contains("access"));
 assertTrue(selector.contains("etag"));
 assertEquals(28, selector.length());
 assertEquals(
   new Dataset(bigquery, new DatasetInfo.BuilderImpl(updatedDatasetInfoWithProject)), dataset);
}

@Test
public void testRetryableException() {
 EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS))
   .andThrow(new BigQueryException(500, "InternalError"))
   .andReturn(DATASET_INFO_WITH_PROJECT.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery =
   options
     .toBuilder()
     .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
     .build()
     .getService();
 Dataset dataset = bigquery.getDataset(DATASET);
 assertEquals(
   new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset);
}

@Test
public void testCreateDataset() {
 DatasetInfo datasetInfo = DATASET_INFO.setProjectId(OTHER_PROJECT);
 EasyMock.expect(bigqueryRpcMock.create(datasetInfo.toPb(), EMPTY_RPC_OPTIONS))
   .andReturn(datasetInfo.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery = options.getService();
 Dataset dataset = bigquery.create(datasetInfo);
 assertEquals(new Dataset(bigquery, new DatasetInfo.BuilderImpl(datasetInfo)), dataset);
}

@Test
public void testUpdateDataset() {
 DatasetInfo updatedDatasetInfo =
   DATASET_INFO
     .setProjectId(OTHER_PROJECT)
     .toBuilder()
     .setDescription("newDescription")
     .build();
 EasyMock.expect(bigqueryRpcMock.patch(updatedDatasetInfo.toPb(), EMPTY_RPC_OPTIONS))
   .andReturn(updatedDatasetInfo.toPb());
 EasyMock.replay(bigqueryRpcMock);
 bigquery = options.getService();
 Dataset dataset = bigquery.update(updatedDatasetInfo);
 assertEquals(new Dataset(bigquery, new DatasetInfo.BuilderImpl(updatedDatasetInfo)), dataset);
}

@Test
public void testReloadNull() throws Exception {
 initializeExpectedDataset(1);
 expect(bigquery.getOptions()).andReturn(mockOptions);
 expect(bigquery.getDataset(DATASET_INFO.getDatasetId().getDataset())).andReturn(null);
 replay(bigquery);
 initializeDataset();
 assertNull(dataset.reload());
}

@Test
public void testBuilder() {
 initializeExpectedDataset(2);
 replay(bigquery);
 Dataset builtDataset =
   new Dataset.Builder(serviceMockReturnsOptions, DATASET_ID)
     .setLabels(LABELS)
     .build();
 assertEquals(DATASET_ID, builtDataset.getDatasetId());
 assertEquals(ACCESS_RULES, builtDataset.getAcl());
 assertEquals(CREATION_TIME, builtDataset.getCreationTime());
 assertEquals(DEFAULT_TABLE_EXPIRATION, builtDataset.getDefaultTableLifetime());
 assertEquals(DESCRIPTION, builtDataset.getDescription());
 assertEquals(ETAG, builtDataset.getEtag());
 assertEquals(FRIENDLY_NAME, builtDataset.getFriendlyName());
 assertEquals(GENERATED_ID, builtDataset.getGeneratedId());
 assertEquals(LAST_MODIFIED, builtDataset.getLastModified());
 assertEquals(LOCATION, builtDataset.getLocation());
 assertEquals(SELF_LINK, builtDataset.getSelfLink());
 assertEquals(LABELS, builtDataset.getLabels());

@Test
public void testUpdateDatasetWithSelectedFields() {
 Dataset dataset =
   bigquery.create(
     DatasetInfo.newBuilder(OTHER_DATASET).setDescription("Some Description").build());
 assertNotNull(dataset);
 assertEquals(bigquery.getOptions().getProjectId(), dataset.getDatasetId().getProject());
 assertEquals(OTHER_DATASET, dataset.getDatasetId().getDataset());
 assertEquals("Some Description", dataset.getDescription());
 Dataset updatedDataset =
   bigquery.update(
     dataset.toBuilder().setDescription("Updated Description").build(),
     DatasetOption.fields(DatasetField.DESCRIPTION));
 assertEquals("Updated Description", updatedDataset.getDescription());
 assertNull(updatedDataset.getCreationTime());
 assertNull(updatedDataset.getDefaultTableLifetime());
 assertNull(updatedDataset.getAcl());
 assertNull(updatedDataset.getEtag());
 assertNull(updatedDataset.getFriendlyName());
 assertNull(updatedDataset.getGeneratedId());
 assertNull(updatedDataset.getLastModified());
 assertNull(updatedDataset.getLocation());
 assertNull(updatedDataset.getSelfLink());
 assertTrue(dataset.delete());
}

  bigquery.create(
    DatasetInfo.newBuilder("locationset_" + UUID.randomUUID().toString().replace("-", "_"))
      .setLocation(location)
      .build());
try {
 TableId tableId = TableId.of(dataset.getDatasetId().getDataset(), "sometable");
 Schema schema = Schema.of(Field.of("name", LegacySQLTypeName.STRING));
 TableDefinition tableDef = StandardTableDefinition.of(schema);
 Table table = bigquery.create(TableInfo.newBuilder(tableId, tableDef).build());
    bigquery.create(
      JobInfo.of(
        JobId.newBuilder().setLocation(location).build(),
 bigquery.delete(dataset.getDatasetId(), DatasetDeleteOption.deleteContents());