Search in sources :

Example 1 with RetryBoundedBackOff

use of com.google.cloud.hadoop.util.RetryBoundedBackOff in project beam by apache.

the class BigQueryServicesImplTest method testCreateTableSucceedsAlreadyExists.

/**
   * Tests that table creation succeeds when the table already exists.
   */
@Test
public void testCreateTableSucceedsAlreadyExists() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    TableSchema schema = new TableSchema().setFields(ImmutableList.of(new TableFieldSchema().setName("column1").setType("String"), new TableFieldSchema().setName("column2").setType("Integer")));
    Table testTable = new Table().setTableReference(ref).setSchema(schema);
    // 409 means already exists
    when(response.getStatusCode()).thenReturn(409);
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    Table ret = services.tryCreateTable(testTable, new RetryBoundedBackOff(0, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
    assertNull(ret);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}
Also used : RetryBoundedBackOff(com.google.cloud.hadoop.util.RetryBoundedBackOff) TableReference(com.google.api.services.bigquery.model.TableReference) Table(com.google.api.services.bigquery.model.Table) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) TableSchema(com.google.api.services.bigquery.model.TableSchema) TableFieldSchema(com.google.api.services.bigquery.model.TableFieldSchema) Test(org.junit.Test)

Example 2 with RetryBoundedBackOff

use of com.google.cloud.hadoop.util.RetryBoundedBackOff in project beam by apache.

the class BigQueryServicesImplTest method testCreateTableDoesNotRetry.

/**
   * Tests that {@link BigQueryServicesImpl} does not retry non-rate-limited attempts.
   */
@Test
public void testCreateTableDoesNotRetry() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    Table testTable = new Table().setTableReference(ref);
    // 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(testTable));
    thrown.expect(GoogleJsonResponseException.class);
    thrown.expectMessage("actually forbidden");
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    try {
        services.tryCreateTable(testTable, new RetryBoundedBackOff(3, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
        fail();
    } catch (IOException e) {
        verify(response, times(1)).getStatusCode();
        verify(response, times(1)).getContent();
        verify(response, times(1)).getContentType();
        throw e;
    }
}
Also used : RetryBoundedBackOff(com.google.cloud.hadoop.util.RetryBoundedBackOff) TableReference(com.google.api.services.bigquery.model.TableReference) Table(com.google.api.services.bigquery.model.Table) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) IOException(java.io.IOException) Test(org.junit.Test)

Example 3 with RetryBoundedBackOff

use of com.google.cloud.hadoop.util.RetryBoundedBackOff in project beam by apache.

the class BigQueryServicesImplTest method testCreateTableSucceeds.

@Test
public void testCreateTableSucceeds() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    Table testTable = new Table().setTableReference(ref);
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(testTable));
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    Table ret = services.tryCreateTable(testTable, new RetryBoundedBackOff(0, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
    assertEquals(testTable, ret);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}
Also used : RetryBoundedBackOff(com.google.cloud.hadoop.util.RetryBoundedBackOff) TableReference(com.google.api.services.bigquery.model.TableReference) Table(com.google.api.services.bigquery.model.Table) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) Test(org.junit.Test)

Example 4 with RetryBoundedBackOff

use of com.google.cloud.hadoop.util.RetryBoundedBackOff 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");
}
Also used : RetryBoundedBackOff(com.google.cloud.hadoop.util.RetryBoundedBackOff) TableReference(com.google.api.services.bigquery.model.TableReference) Table(com.google.api.services.bigquery.model.Table) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) Test(org.junit.Test)

Aggregations

Table (com.google.api.services.bigquery.model.Table)4 TableReference (com.google.api.services.bigquery.model.TableReference)4 RetryBoundedBackOff (com.google.cloud.hadoop.util.RetryBoundedBackOff)4 DatasetServiceImpl (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl)4 Test (org.junit.Test)4 TableFieldSchema (com.google.api.services.bigquery.model.TableFieldSchema)1 TableSchema (com.google.api.services.bigquery.model.TableSchema)1 IOException (java.io.IOException)1