org.apache.beam.sdk.extensions.gcp.util.Transport Java Examples

@Test
public void testMatchingAgainstRequestType() throws Exception {
  GcsOptions options = PipelineOptionsFactory.as(GcsOptions.class);
  options.setGcpCredential(new TestCredential());
  options.setGoogleApiTrace(
      new GoogleApiTracer()
          .addTraceFor(
              Transport.newStorageClient(options)
                  .build()
                  .objects()
                  .get("aProjectId", "aObjectId"),
              "TraceDestination"));

  Storage.Objects.Get getRequest =
      Transport.newStorageClient(options).build().objects().get("testBucketId", "testObjectId");
  assertEquals("TraceDestination", getRequest.get("$trace"));

  Storage.Objects.List listRequest =
      Transport.newStorageClient(options).build().objects().list("testProjectId");
  assertNull(listRequest.get("$trace"));
}
 

@Test
public void testCloudServiceCall() throws Exception {
  WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);

  when(request.execute()).thenReturn(generateMockResponse(workItem));

  WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);

  assertEquals(Optional.of(workItem), client.getWorkItem());

  LeaseWorkItemRequest actualRequest =
      Transport.getJsonFactory()
          .fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
  assertEquals(WORKER_ID, actualRequest.getWorkerId());
  assertEquals(
      ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"),
      actualRequest.getWorkerCapabilities());
  assertEquals(
      ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"),
      actualRequest.getWorkItemTypes());
  assertEquals("1234", DataflowWorkerLoggingMDC.getWorkId());
}
 

@Test
public void testCloudServiceCallNoWorkPresent() throws Exception {
  // If there's no work the service should return an empty work item.
  WorkItem workItem = new WorkItem();

  when(request.execute()).thenReturn(generateMockResponse(workItem));

  WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);

  assertEquals(Optional.absent(), client.getWorkItem());

  LeaseWorkItemRequest actualRequest =
      Transport.getJsonFactory()
          .fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
  assertEquals(WORKER_ID, actualRequest.getWorkerId());
  assertEquals(
      ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"),
      actualRequest.getWorkerCapabilities());
  assertEquals(
      ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"),
      actualRequest.getWorkItemTypes());
}
 

@Test
public void testCloudServiceCallNoWorkId() throws Exception {
  // If there's no work the service should return an empty work item.
  WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
  workItem.setId(null);

  when(request.execute()).thenReturn(generateMockResponse(workItem));

  WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);

  assertEquals(Optional.absent(), client.getWorkItem());

  LeaseWorkItemRequest actualRequest =
      Transport.getJsonFactory()
          .fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
  assertEquals(WORKER_ID, actualRequest.getWorkerId());
  assertEquals(
      ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"),
      actualRequest.getWorkerCapabilities());
  assertEquals(
      ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"),
      actualRequest.getWorkItemTypes());
}
 

/**
 * Determines the serialized size (in bytes) of the {@link GenericJson} object that will be
 * serialized and sent to the Google Cloud Dataflow service API.
 *
 * <p>Uses only constant memory.
 */
public static long computeSerializedSizeBytes(GenericJson object) throws IOException {
  JsonFactory factory = object.getFactory();
  if (factory == null) {
    factory = Transport.getJsonFactory();
  }

  CountingOutputStream stream = new CountingOutputStream(ByteStreams.nullOutputStream());
  JsonGenerator generator = null;
  try {
    generator = factory.createJsonGenerator(stream, StandardCharsets.UTF_8);
    generator.serialize(object);
    generator.close(); // also closes the stream.
  } finally {
    if (generator != null) {
      generator.close();
    }
  }
  return stream.getCount();
}
 

@Test
public void testCloudServiceCallNoWorkItem() throws Exception {
  when(request.execute()).thenReturn(generateMockResponse());

  WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);

  assertEquals(Optional.absent(), client.getWorkItem());

  LeaseWorkItemRequest actualRequest =
      Transport.getJsonFactory()
          .fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
  assertEquals(WORKER_ID, actualRequest.getWorkerId());
  assertEquals(
      ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"),
      actualRequest.getWorkerCapabilities());
  assertEquals(
      ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"),
      actualRequest.getWorkItemTypes());
}
 

@Before
public void setUp() {
  MockitoAnnotations.initMocks(this);

  // Set up the MockHttpRequest for future inspection
  request =
      new MockLowLevelHttpRequest() {
        @Override
        public LowLevelHttpResponse execute() throws IOException {
          return response;
        }
      };

  // A mock transport that lets us mock the API responses.
  MockHttpTransport transport =
      new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build();

  // A sample BigQuery API client that uses default JsonFactory and RetryHttpInitializer.
  bigquery =
      new Bigquery.Builder(
              transport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())
          .build();
}
 

@Override
public PubsubClient newClient(
    @Nullable String timestampAttribute, @Nullable String idAttribute, PubsubOptions options)
    throws IOException {
  Pubsub pubsub =
      new Builder(
              Transport.getTransport(),
              Transport.getJsonFactory(),
              chainHttpRequestInitializer(
                  options.getGcpCredential(),
                  // Do not log 404. It clutters the output and is possibly even required by the
                  // caller.
                  new RetryHttpRequestInitializer(ImmutableList.of(404))))
          .setRootUrl(options.getPubsubRootUrl())
          .setApplicationName(options.getAppName())
          .setGoogleClientRequestInitializer(options.getGoogleApiTrace())
          .build();
  return new PubsubJsonClient(timestampAttribute, idAttribute, pubsub);
}
 

/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
  RetryHttpRequestInitializer httpRequestInitializer =
      new RetryHttpRequestInitializer(ImmutableList.of(404));
  httpRequestInitializer.setCustomErrors(createBigQueryClientCustomErrors());
  httpRequestInitializer.setWriteTimeout(options.getHTTPWriteTimeout());
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              httpRequestInitializer))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

private WorkItemServiceState generateServiceState(
    @Nullable Position suggestedStopPosition, long millisToNextUpdate) {
  WorkItemServiceState responseState = new WorkItemServiceState();
  responseState.setFactory(Transport.getJsonFactory());
  responseState.setLeaseExpireTime(
      toCloudTime(new Instant(clock.currentTimeMillis() + LEASE_MS)));
  responseState.setReportStatusInterval(toCloudDuration(Duration.millis(millisToNextUpdate)));

  if (suggestedStopPosition != null) {
    responseState.setSplitRequest(
        ReaderTestUtils.approximateSplitRequestAtPosition(suggestedStopPosition));
  }

  HotKeyDetection hotKeyDetection = new HotKeyDetection();
  hotKeyDetection.setUserStepName(STEP_ID);
  hotKeyDetection.setHotKeyAge(toCloudDuration(HOT_KEY_AGE));
  responseState.setHotKeyDetection(hotKeyDetection);

  return responseState;
}
 

@Test
public void testMatchingAgainstClient() throws Exception {
  GcsOptions options = PipelineOptionsFactory.as(GcsOptions.class);
  options.setGcpCredential(new TestCredential());
  options.setGoogleApiTrace(
      new GoogleApiTracer()
          .addTraceFor(Transport.newStorageClient(options).build(), "TraceDestination"));

  Storage.Objects.Get getRequest =
      Transport.newStorageClient(options).build().objects().get("testBucketId", "testObjectId");
  assertEquals("TraceDestination", getRequest.get("$trace"));

  Delete deleteRequest =
      GcpOptions.GcpTempLocationFactory.newCloudResourceManagerClient(
              options.as(CloudResourceManagerOptions.class))
          .build()
          .projects()
          .delete("testProjectId");
  assertNull(deleteRequest.get("$trace"));
}
 

@Test
public void testMatchingAllCalls() throws Exception {
  String[] args = new String[] {STORAGE_TRACE};
  GcsOptions options = PipelineOptionsFactory.fromArgs(args).as(GcsOptions.class);
  options.setGcpCredential(new TestCredential());

  assertNotNull(options.getGoogleApiTrace());

  Storage.Objects.Get getRequest =
      Transport.newStorageClient(options).build().objects().get("testBucketId", "testObjectId");
  assertEquals("TraceDestination", getRequest.get("$trace"));

  Storage.Objects.List listRequest =
      Transport.newStorageClient(options).build().objects().list("testProjectId");
  assertEquals("TraceDestination", listRequest.get("$trace"));
}
 

@Test
public void testWithMultipleTraces() throws Exception {
  String[] args = new String[] {STORAGE_GET_AND_LIST_TRACE};
  GcsOptions options = PipelineOptionsFactory.fromArgs(args).as(GcsOptions.class);
  options.setGcpCredential(new TestCredential());

  assertNotNull(options.getGoogleApiTrace());

  Storage.Objects.Get getRequest =
      Transport.newStorageClient(options).build().objects().get("testBucketId", "testObjectId");
  assertEquals("GetTraceDestination", getRequest.get("$trace"));

  Storage.Objects.List listRequest =
      Transport.newStorageClient(options).build().objects().list("testProjectId");
  assertEquals("ListTraceDestination", listRequest.get("$trace"));
}
 

/**
 * Returns a CloudResourceManager client builder using the specified {@link
 * CloudResourceManagerOptions}.
 */
@VisibleForTesting
static CloudResourceManager.Builder newCloudResourceManagerClient(
    CloudResourceManagerOptions options) {
  Credentials credentials = options.getGcpCredential();
  if (credentials == null) {
    NullCredentialInitializer.throwNullCredentialException();
  }
  return new CloudResourceManager.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              credentials,
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

@Test
public void testRead() {
  InstructionOutput readOutput = createInstructionOutput("Read.out");
  ParallelInstruction read = createParallelInstruction("Read", readOutput);
  read.setRead(new ReadInstruction());

  MapTask mapTask = new MapTask();
  mapTask.setInstructions(ImmutableList.of(read));
  mapTask.setFactory(Transport.getJsonFactory());

  Network<Node, Edge> network =
      new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
  assertNetworkProperties(network);
  assertEquals(2, network.nodes().size());
  assertEquals(1, network.edges().size());

  ParallelInstructionNode readNode = get(network, read);
  InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
  assertEquals(readOutput, readOutputNode.getInstructionOutput());
}
 

/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

/**
 * Returns a {@link MapTask} with the provided {@code instructions} and default values everywhere
 * else.
 */
private MapTask defaultMapTask(List<ParallelInstruction> instructions) {
  MapTask mapTask =
      new MapTask()
          .setStageName(DEFAULT_MAP_STAGE_NAME)
          .setSystemName(DEFAULT_MAP_SYSTEM_NAME)
          .setInstructions(instructions);
  mapTask.setFactory(Transport.getJsonFactory());
  return mapTask;
}
 

private WorkItem createWorkItem(String projectId, String jobId) {
  WorkItem workItem = new WorkItem();
  workItem.setFactory(Transport.getJsonFactory());
  workItem.setProjectId(projectId);
  workItem.setJobId(jobId);

  // We need to set a work id because otherwise the client will treat the response as
  // indicating no work is available.
  workItem.setId(1234L);
  return workItem;
}
 

@Test
public void testWrite() {
  InstructionOutput readOutput = createInstructionOutput("Read.out");
  ParallelInstruction read = createParallelInstruction("Read", readOutput);
  read.setRead(new ReadInstruction());

  WriteInstruction writeInstruction = new WriteInstruction();
  writeInstruction.setInput(createInstructionInput(0, 0)); // Read.out
  ParallelInstruction write = createParallelInstruction("Write");
  write.setWrite(writeInstruction);

  MapTask mapTask = new MapTask();
  mapTask.setInstructions(ImmutableList.of(read, write));
  mapTask.setFactory(Transport.getJsonFactory());

  Network<Node, Edge> network =
      new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
  assertNetworkProperties(network);
  assertEquals(3, network.nodes().size());
  assertEquals(2, network.edges().size());

  ParallelInstructionNode readNode = get(network, read);
  InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
  assertEquals(readOutput, readOutputNode.getInstructionOutput());

  ParallelInstructionNode writeNode = getOnlySuccessor(network, readOutputNode);
  assertNotNull(writeNode);
}
 

private LowLevelHttpResponse generateMockResponse(WorkItem... workItems) throws Exception {
  MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
  response.setContentType(Json.MEDIA_TYPE);
  LeaseWorkItemResponse lease = new LeaseWorkItemResponse();
  lease.setWorkItems(Lists.newArrayList(workItems));
  // N.B. Setting the factory is necessary in order to get valid JSON.
  lease.setFactory(Transport.getJsonFactory());
  response.setContent(lease.toPrettyString());
  return response;
}
 

@Test
public void testExecutionContextPlumbing() throws Exception {
  List<ParallelInstruction> instructions =
      Arrays.asList(
          createReadInstruction("Read", ReaderFactoryTest.SingletonTestReaderFactory.class),
          createParDoInstruction(0, 0, "DoFn1", "DoFnUserName"),
          createParDoInstruction(1, 0, "DoFnWithContext", "DoFnWithContextUserName"));

  MapTask mapTask = new MapTask();
  mapTask.setStageName(STAGE);
  mapTask.setInstructions(instructions);
  mapTask.setFactory(Transport.getJsonFactory());

  BatchModeExecutionContext context =
      BatchModeExecutionContext.forTesting(options, counterSet, "testStage");

  try (DataflowMapTaskExecutor executor =
      mapTaskExecutorFactory.create(
          null /* beamFnControlClientHandler */,
          null /* beamFnDataService */,
          null /* beamFnStateService */,
          null,
          mapTaskToNetwork.apply(mapTask),
          options,
          STAGE,
          readerRegistry,
          sinkRegistry,
          context,
          counterSet,
          idGenerator)) {
    executor.execute();
  }

  List<String> stepNames = new ArrayList<>();
  for (BatchModeExecutionContext.StepContext stepContext : context.getAllStepContexts()) {
    stepNames.add(stepContext.getNameContext().systemName());
  }
  assertThat(stepNames, hasItems("DoFn1", "DoFnWithContext"));
}
 

@Test
public void testParDo() {
  InstructionOutput readOutput = createInstructionOutput("Read.out");
  ParallelInstruction read = createParallelInstruction("Read", readOutput);
  read.setRead(new ReadInstruction());

  MultiOutputInfo parDoMultiOutput = createMultiOutputInfo("output");
  ParDoInstruction parDoInstruction = new ParDoInstruction();
  parDoInstruction.setInput(createInstructionInput(0, 0)); // Read.out
  parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput));
  InstructionOutput parDoOutput = createInstructionOutput("ParDo.out");
  ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput);
  parDo.setParDo(parDoInstruction);

  MapTask mapTask = new MapTask();
  mapTask.setInstructions(ImmutableList.of(read, parDo));
  mapTask.setFactory(Transport.getJsonFactory());

  Network<Node, Edge> network =
      new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
  assertNetworkProperties(network);
  assertEquals(4, network.nodes().size());
  assertEquals(3, network.edges().size());

  ParallelInstructionNode readNode = get(network, read);
  InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
  assertEquals(readOutput, readOutputNode.getInstructionOutput());

  ParallelInstructionNode parDoNode = getOnlySuccessor(network, readOutputNode);
  InstructionOutputNode parDoOutputNode = getOnlySuccessor(network, parDoNode);
  assertEquals(parDoOutput, parDoOutputNode.getInstructionOutput());

  assertEquals(
      parDoMultiOutput,
      ((MultiOutputInfoEdge)
              Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutputNode)))
          .getMultiOutputInfo());
}
 

@Test
public void testParallelEdgeFlatten() {
  //                  /---\
  // Read --> Read.out --> Flatten
  //                  \---/
  InstructionOutput readOutput = createInstructionOutput("Read.out");
  ParallelInstruction read = createParallelInstruction("Read", readOutput);
  read.setRead(new ReadInstruction());

  FlattenInstruction flattenInstruction = new FlattenInstruction();
  flattenInstruction.setInputs(
      ImmutableList.of(
          createInstructionInput(0, 0), // Read.out
          createInstructionInput(0, 0), // Read.out
          createInstructionInput(0, 0))); // Read.out
  InstructionOutput flattenOutput = createInstructionOutput("Flatten.out");
  ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput);
  flatten.setFlatten(flattenInstruction);

  MapTask mapTask = new MapTask();
  mapTask.setInstructions(ImmutableList.of(read, flatten));
  mapTask.setFactory(Transport.getJsonFactory());

  Network<Node, Edge> network =
      new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
  assertNetworkProperties(network);
  assertEquals(4, network.nodes().size());
  assertEquals(5, network.edges().size());

  ParallelInstructionNode readNode = get(network, read);
  InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
  assertEquals(readOutput, readOutputNode.getInstructionOutput());

  ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputNode);
  // Assert that the three parallel edges are maintained
  assertEquals(3, network.edgesConnecting(readOutputNode, flattenNode).size());

  InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
  assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
}
 

@Before
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  when(transport.buildRequest(anyString(), anyString())).thenReturn(request);
  doCallRealMethod().when(request).getContentAsString();

  Dataflow service = new Dataflow(transport, Transport.getJsonFactory(), null);
  pipelineOptions = PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class);
  pipelineOptions.setProject(PROJECT_ID);
  pipelineOptions.setJobId(JOB_ID);
  pipelineOptions.setWorkerId(WORKER_ID);
  pipelineOptions.setGcpCredential(new TestCredential());
  pipelineOptions.setDataflowClient(service);
  pipelineOptions.setRegion("us-central1");
}
 

@Before
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  when(transport.buildRequest(anyString(), anyString())).thenReturn(request);
  doCallRealMethod().when(request).getContentAsString();

  Dataflow service = new Dataflow(transport, Transport.getJsonFactory(), null);
  pipelineOptions = PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class);
  pipelineOptions.setProject(PROJECT_ID);
  pipelineOptions.setJobId(JOB_ID);
  pipelineOptions.setWorkerId(WORKER_ID);
  pipelineOptions.setGcpCredential(new TestCredential());
  pipelineOptions.setDataflowClient(service);
}
 

@Test
public void testPartialGroupByKey() {
  // Read --> PGBK --> Write
  InstructionOutput readOutput = createInstructionOutput("Read.out");
  ParallelInstruction read = createParallelInstruction("Read", readOutput);
  read.setRead(new ReadInstruction());

  PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction();
  pgbkInstruction.setInput(createInstructionInput(0, 0)); // Read.out
  InstructionOutput pgbkOutput = createInstructionOutput("PGBK.out");
  ParallelInstruction pgbk = createParallelInstruction("PGBK", pgbkOutput);
  pgbk.setPartialGroupByKey(pgbkInstruction);

  WriteInstruction writeInstruction = new WriteInstruction();
  writeInstruction.setInput(createInstructionInput(1, 0)); // PGBK.out
  ParallelInstruction write = createParallelInstruction("Write");
  write.setWrite(writeInstruction);

  MapTask mapTask = new MapTask();
  mapTask.setInstructions(ImmutableList.of(read, pgbk, write));
  mapTask.setFactory(Transport.getJsonFactory());

  Network<Node, Edge> network =
      new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
  assertNetworkProperties(network);
  assertEquals(5, network.nodes().size());
  assertEquals(4, network.edges().size());

  ParallelInstructionNode readNode = get(network, read);
  InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
  assertEquals(readOutput, readOutputNode.getInstructionOutput());

  ParallelInstructionNode pgbkNode = getOnlySuccessor(network, readOutputNode);
  InstructionOutputNode pgbkOutputNode = getOnlySuccessor(network, pgbkNode);
  assertEquals(pgbkOutput, pgbkOutputNode.getInstructionOutput());

  ParallelInstructionNode writeNode = getOnlySuccessor(network, pgbkOutputNode);
  assertNotNull(write);
}
 

private JobMetrics buildJobMetrics(List<MetricUpdate> metricList) {
  JobMetrics jobMetrics = new JobMetrics();
  jobMetrics.setMetrics(metricList);
  // N.B. Setting the factory is necessary in order to get valid JSON.
  jobMetrics.setFactory(Transport.getJsonFactory());
  return jobMetrics;
}
 

public static Bigquery getNewBigquerryClient(String applicationName) {
  HttpTransport transport = Transport.getTransport();
  JsonFactory jsonFactory = Transport.getJsonFactory();
  Credentials credential = getDefaultCredential();
  return new Bigquery.Builder(transport, jsonFactory, new HttpCredentialsAdapter(credential))
      .setApplicationName(applicationName)
      .build();
}
 

private static Bigquery newBigQueryClient(BigQueryOptions options) {
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace())
      .build();
}
 

@Test
public void testWhenTracingDoesNotMatch() throws Exception {
  String[] args = new String[] {STORAGE_GET_TRACE};
  GcsOptions options = PipelineOptionsFactory.fromArgs(args).as(GcsOptions.class);
  options.setGcpCredential(new TestCredential());

  assertNotNull(options.getGoogleApiTrace());

  Storage.Objects.List request =
      Transport.newStorageClient(options).build().objects().list("testProjectId");
  assertNull(request.get("$trace"));
}