Examples with TableReference com.google.api.services.bigquery.model.TableReference used on opensource projects

Example 56 with TableReference

use of com.google.api.services.bigquery.model.TableReference in project beam by apache.

the class BigQueryServicesImplTest method testCreateTableRetry.

/**
   * Tests that {@link BigQueryServicesImpl} retries quota rate limited attempts.
   */
@Test
public void testCreateTableRetry() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    Table testTable = new Table().setTableReference(ref);
    // First response is 403 rate limited, second response has valid payload.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(testTable));
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    Table ret = services.tryCreateTable(testTable, new RetryBoundedBackOff(3, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
    assertEquals(testTable, ret);
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
    verifyNotNull(ret.getTableReference());
    expectedLogs.verifyInfo("Quota limit reached when creating table project:dataset.table, " + "retrying up to 5.0 minutes");
}

Example 57 with TableReference

use of com.google.api.services.bigquery.model.TableReference in project beam by apache.

the class BigQueryServicesImplTest method testGetTableSucceeds.

@Test
public void testGetTableSucceeds() throws Exception {
    TableReference tableRef = new TableReference().setProjectId("projectId").setDatasetId("datasetId").setTableId("tableId");
    Table testTable = new Table();
    testTable.setTableReference(tableRef);
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(testTable));
    BigQueryServicesImpl.DatasetServiceImpl datasetService = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    Table table = datasetService.getTable(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
    assertEquals(testTable, table);
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
}

Example 58 with TableReference

use of com.google.api.services.bigquery.model.TableReference in project beam by apache.

the class BigQueryServicesImplTest method testInsertDoesNotRetry.

/**
   * Tests that {@link DatasetServiceImpl#insertAll} does not retry non-rate-limited attempts.
   */
@Test
public void testInsertDoesNotRetry() throws Throwable {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
    rows.add(wrapTableRow(new TableRow()));
    // First response is 403 not-rate-limited, second response has valid payload but should not
    // be invoked.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("actually forbidden", 403))).thenReturn(toStream(new TableDataInsertAllResponse()));
    thrown.expect(GoogleJsonResponseException.class);
    thrown.expectMessage("actually forbidden");
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    try {
        dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null);
        fail();
    } catch (RuntimeException e) {
        verify(response, times(1)).getStatusCode();
        verify(response, times(1)).getContent();
        verify(response, times(1)).getContentType();
        throw e.getCause();
    }
}

Example 59 with TableReference

use of com.google.api.services.bigquery.model.TableReference in project beam by apache.

the class BigQueryServicesImplTest method testGetTableNotFound.

@Test
public void testGetTableNotFound() throws IOException, InterruptedException {
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(404);
    BigQueryServicesImpl.DatasetServiceImpl datasetService = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    TableReference tableRef = new TableReference().setProjectId("projectId").setDatasetId("datasetId").setTableId("tableId");
    Table table = datasetService.getTable(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
    assertNull(table);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

Example 60 with TableReference

use of com.google.api.services.bigquery.model.TableReference in project beam by apache.

the class BigQueryServicesImplTest method testIsTableEmptySucceeds.

@Test
public void testIsTableEmptySucceeds() throws Exception {
    TableReference tableRef = new TableReference().setProjectId("projectId").setDatasetId("datasetId").setTableId("tableId");
    TableDataList testDataList = new TableDataList().setRows(ImmutableList.of(new TableRow()));
    // First response is 403 rate limited, second response has valid payload.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(testDataList));
    BigQueryServicesImpl.DatasetServiceImpl datasetService = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    assertFalse(datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT));
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
}