org.apache.beam.sdk.extensions.gcp.options.GcpOptions Java Examples

@Test
public void testGetBigtableServiceWithConfigurator() {
  SerializableFunction<BigtableOptions.Builder, BigtableOptions.Builder> configurator =
      (SerializableFunction<BigtableOptions.Builder, BigtableOptions.Builder>)
          input ->
              input
                  .setInstanceId(INSTANCE_ID.get() + INSTANCE_ID.get())
                  .setProjectId(PROJECT_ID.get() + PROJECT_ID.get())
                  .setBulkOptions(new BulkOptions.Builder().setUseBulkApi(true).build());

  BigtableService service =
      config
          .withProjectId(PROJECT_ID)
          .withInstanceId(INSTANCE_ID)
          .withBigtableOptionsConfigurator(configurator)
          .getBigtableService(PipelineOptionsFactory.as(GcpOptions.class));

  assertEquals(PROJECT_ID.get(), service.getBigtableOptions().getProjectId());
  assertEquals(INSTANCE_ID.get(), service.getBigtableOptions().getInstanceId());
  assertTrue(service.getBigtableOptions().getBulkOptions().useBulkApi());
}
 

private GrpcWindmillServer(String name, boolean enableStreamingEngine) {
  this.options = PipelineOptionsFactory.create().as(StreamingDataflowWorkerOptions.class);
  this.streamingRpcBatchLimit = Integer.MAX_VALUE;
  options.setProject("project");
  options.setJobId("job");
  options.setWorkerId("worker");
  if (enableStreamingEngine) {
    List<String> experiments = this.options.getExperiments();
    if (experiments == null) {
      experiments = new ArrayList<>();
    }
    experiments.add(GcpOptions.STREAMING_ENGINE_EXPERIMENT);
    options.setExperiments(experiments);
  }
  this.stubList.add(CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel(name)));
}
 

/**
 * Helper function that either returns the mock Bigtable service supplied by {@link
 * #withBigtableService} or creates and returns an implementation that talks to {@code Cloud
 * Bigtable}.
 *
 * <p>Also populate the credentials option from {@link GcpOptions#getGcpCredential()} if the
 * default credentials are being used on {@link BigtableOptions}.
 */
@VisibleForTesting
BigtableService getBigtableService(PipelineOptions pipelineOptions) {
  if (getBigtableService() != null) {
    return getBigtableService();
  }

  BigtableOptions.Builder bigtableOptions = effectiveUserProvidedBigtableOptions();

  bigtableOptions.setUserAgent(pipelineOptions.getUserAgent());

  if (bigtableOptions.build().getCredentialOptions().getCredentialType()
      == CredentialOptions.CredentialType.DefaultCredentials) {
    bigtableOptions.setCredentialOptions(
        CredentialOptions.credential(pipelineOptions.as(GcpOptions.class).getGcpCredential()));
  }

  // Default option that should be forced
  bigtableOptions.setUseCachedDataPool(true);

  return new BigtableServiceImpl(bigtableOptions.build());
}
 

public FhirIOReadIT(String version) {
  this.version = version;
  long testTime = System.currentTimeMillis();
  this.fhirStoreName =
      "FHIR_store_" + version + "_write_it_" + testTime + "_" + (new SecureRandom().nextInt(32));
  this.project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  this.pubsubTopic =
      "projects/"
          + project
          + "/topics/FhirIO-IT-"
          + version
          + "-notifications-"
          + testTime
          + "-"
          + (new SecureRandom().nextInt(32));
  this.pubsubSubscription = pubsubTopic.replaceAll("topic", "subscription");
  pipelineOptions = TestPipeline.testingPipelineOptions().as(TestPubsubOptions.class);
}
 

@BeforeClass
public static void setupTestEnvironment() throws Exception {
  PipelineOptionsFactory.register(BigQueryToTableOptions.class);
  project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  // Create one BQ dataset for all test cases.
  BQ_CLIENT.createNewDataset(project, BIG_QUERY_DATASET_ID);

  // Create table and insert data for new type query test cases.
  BQ_CLIENT.createNewTable(
      project,
      BIG_QUERY_DATASET_ID,
      new Table()
          .setSchema(BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_SCHEMA)
          .setTableReference(
              new TableReference()
                  .setTableId(BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_NAME)
                  .setDatasetId(BIG_QUERY_DATASET_ID)
                  .setProjectId(project)));
  BQ_CLIENT.insertDataToTable(
      project,
      BIG_QUERY_DATASET_ID,
      BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_NAME,
      BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_DATA);
}
 

private ProjectTopicName getPubSubTopic() {
  if (options.getInputSubscriptions() != null && !options.getInputSubscriptions().isEmpty()) {
    try {
      return PubsubUtils.getPubSubTopicFromSubscription(
          options.as(GcpOptions.class).getProject(), options.getInputSubscriptions());
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  } else {
    Preconditions.checkArgument(
        options.getInputTopics() != null && !options.getInputTopics().isEmpty(),
        "Must provide an inputSubscriptions or inputTopics parameter.");
    return ProjectTopicName.of(
        options.as(GcpOptions.class).getProject(), options.getInputTopics());
  }
}
 

@BeforeClass
public static void setUpClass() {
  project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();

  setProject = String.format("SET project = '%s';", project);

  createPubsubTableStatement =
      "CREATE EXTERNAL TABLE taxi_rides (\n"
          + "         event_timestamp TIMESTAMP,\n"
          + "         attributes MAP<VARCHAR, VARCHAR>,\n"
          + "         payload ROW<\n"
          + "           ride_id VARCHAR,\n"
          + "           point_idx INT,\n"
          + "           latitude DOUBLE,\n"
          + "           longitude DOUBLE,\n"
          + "           meter_reading DOUBLE,\n"
          + "           meter_increment DOUBLE,\n"
          + "           ride_status VARCHAR,\n"
          + "           passenger_count TINYINT>)\n"
          + "       TYPE pubsub \n"
          + "       LOCATION '%s'\n"
          + "       TBLPROPERTIES '{\"timestampAttributeKey\": \"ts\"}';";

  dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
}
 

@Before
public void setup() throws Exception {
  PipelineOptionsFactory.register(BigtableTestOptions.class);
  options = TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class);
  project = options.as(GcpOptions.class).getProject();

  bigtableOptions =
      new Builder()
          .setProjectId(project)
          .setInstanceId(options.getInstanceId())
          .setUserAgent("apache-beam-test")
          .build();

  session =
      new BigtableSession(
          bigtableOptions
              .toBuilder()
              .setCredentialOptions(
                  CredentialOptions.credential(options.as(GcpOptions.class).getGcpCredential()))
              .build());
  tableAdminClient = session.getTableAdminClient();
}
 

@Test
public void deidentifiesText() {
  String projectId = testPipeline.getOptions().as(GcpOptions.class).getProject();

  PCollection<KV<String, DeidentifyContentResponse>> deidentificationResult =
      testPipeline
          .apply(Create.of(KV.of("", IDENTIFYING_TEXT)))
          .apply(
              DLPDeidentifyText.newBuilder()
                  .setBatchSizeBytes(524000)
                  .setProjectId(projectId)
                  .setDeidentifyConfig(getDeidentifyConfig())
                  .build());
  PAssert.that(deidentificationResult)
      .satisfies(new VerifyDeidentificationResult("####################"));
  testPipeline.run().waitUntilFinish();
}
 

@Test
public void testE2EBigtableRead() throws Exception {
  PipelineOptionsFactory.register(BigtableTestOptions.class);
  BigtableTestOptions options =
      TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class);

  String project = options.getBigtableProject();
  if (project.equals("")) {
    project = options.as(GcpOptions.class).getProject();
  }

  BigtableOptions.Builder bigtableOptionsBuilder =
      new BigtableOptions.Builder().setProjectId(project).setInstanceId(options.getInstanceId());

  final String tableId = "BigtableReadTest";
  final long numRows = 1000L;

  Pipeline p = Pipeline.create(options);
  PCollection<Long> count =
      p.apply(BigtableIO.read().withBigtableOptions(bigtableOptionsBuilder).withTableId(tableId))
          .apply(Count.globally());
  PAssert.thatSingleton(count).isEqualTo(numRows);
  p.run();
}
 

@Before
public void setup() {
  PipelineOptionsFactory.register(V1TestOptions.class);
  options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class);
  project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  ancestor = UUID.randomUUID().toString();
}
 

private void setupTestEnvironment(String recordSize) {
  PipelineOptionsFactory.register(BigQueryIOReadOptions.class);
  options = TestPipeline.testingPipelineOptions().as(BigQueryIOReadOptions.class);
  options.setNumRecords(numOfRecords.get(recordSize));
  options.setTempLocation(options.getTempRoot() + "/temp-it/");
  project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  options.setInputTable(project + ":" + datasetId + "." + tablePrefix + recordSize);
}
 

private void setUpTestEnvironment(String tableSize) {
  PipelineOptionsFactory.register(BigQueryIOStorageReadOptions.class);
  options = TestPipeline.testingPipelineOptions().as(BigQueryIOStorageReadOptions.class);
  options.setNumRecords(EXPECTED_NUM_RECORDS.get(tableSize));
  String project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  options.setInputTable(project + ":" + DATASET_ID + "." + TABLE_PREFIX + tableSize);
}
 

private void setUpTestEnvironment(String tableSize) {
  PipelineOptionsFactory.register(BigQueryIOStorageQueryOptions.class);
  options = TestPipeline.testingPipelineOptions().as(BigQueryIOStorageQueryOptions.class);
  options.setNumRecords(EXPECTED_NUM_RECORDS.get(tableSize));
  String project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  options.setInputTable(project + '.' + DATASET_ID + '.' + TABLE_PREFIX + tableSize);
}
 

/** Tests that credentials are used from PipelineOptions if not supplied by BigtableOptions. */
@Test
public void testUsePipelineOptionsCredentialsIfNotSpecifiedInBigtableOptions() throws Exception {
  BigtableOptions options =
      BIGTABLE_OPTIONS
          .toBuilder()
          .setCredentialOptions(CredentialOptions.defaultCredentials())
          .build();
  GcpOptions pipelineOptions = PipelineOptionsFactory.as(GcpOptions.class);
  pipelineOptions.setGcpCredential(new TestCredential());
  BigtableService readService =
      BigtableIO.read()
          .withBigtableOptions(options)
          .withTableId("TEST-TABLE")
          .getBigtableConfig()
          .getBigtableService(pipelineOptions);
  BigtableService writeService =
      BigtableIO.write()
          .withBigtableOptions(options)
          .withTableId("TEST-TABLE")
          .getBigtableConfig()
          .getBigtableService(pipelineOptions);
  assertEquals(
      CredentialType.SuppliedCredentials,
      readService.getBigtableOptions().getCredentialOptions().getCredentialType());
  assertEquals(
      CredentialType.SuppliedCredentials,
      writeService.getBigtableOptions().getCredentialOptions().getCredentialType());
}
 

/** Tests that credentials are not used from PipelineOptions if supplied by BigtableOptions. */
@Test
public void testDontUsePipelineOptionsCredentialsIfSpecifiedInBigtableOptions() throws Exception {
  BigtableOptions options =
      BIGTABLE_OPTIONS
          .toBuilder()
          .setCredentialOptions(CredentialOptions.nullCredential())
          .build();
  GcpOptions pipelineOptions = PipelineOptionsFactory.as(GcpOptions.class);
  pipelineOptions.setGcpCredential(new TestCredential());
  BigtableService readService =
      BigtableIO.read()
          .withBigtableOptions(options)
          .withTableId("TEST-TABLE")
          .getBigtableConfig()
          .getBigtableService(pipelineOptions);
  BigtableService writeService =
      BigtableIO.write()
          .withBigtableOptions(options)
          .withTableId("TEST-TABLE")
          .getBigtableConfig()
          .getBigtableService(pipelineOptions);
  assertEquals(
      CredentialType.None,
      readService.getBigtableOptions().getCredentialOptions().getCredentialType());
  assertEquals(
      CredentialType.None,
      writeService.getBigtableOptions().getCredentialOptions().getCredentialType());
}
 

private SubscriptionPath createRandomSubscription(PipelineOptions options) {
  TopicPath topicPath = topic.get();

  ProjectPath projectPath;
  if (project != null) {
    projectPath = project.get();
  } else {
    String projectId = options.as(GcpOptions.class).getProject();
    checkState(
        projectId != null,
        "Cannot create subscription to topic %s because pipeline option 'project' not specified",
        topicPath);
    projectPath = PubsubClient.projectPathFromId(options.as(GcpOptions.class).getProject());
  }

  try {
    try (PubsubClient pubsubClient =
        pubsubFactory.newClient(
            timestampAttribute, idAttribute, options.as(PubsubOptions.class))) {
      SubscriptionPath subscriptionPath =
          pubsubClient.createRandomSubscription(projectPath, topicPath, DEAULT_ACK_TIMEOUT_SEC);
      LOG.warn(
          "Created subscription {} to topic {}."
              + " Note this subscription WILL NOT be deleted when the pipeline terminates",
          subscriptionPath,
          topic);
      return subscriptionPath;
    }
  } catch (Exception e) {
    throw new RuntimeException(
        String.format(
            "Failed to create subscription to topic %s on project %s: %s",
            topicPath, projectPath, e.getMessage()),
        e);
  }
}
 

@Before
public void setup() throws Exception {
  PipelineOptionsFactory.register(V1TestOptions.class);
  options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class);
  project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  ancestor = UUID.randomUUID().toString();
  // Create entities and write them to datastore
  writeEntitiesToDatastore(options, project, ancestor, numEntities);
}
 

@BeforeClass
public static void createHL7v2tore() throws IOException {
  String project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  healthcareDataset = String.format(HEALTHCARE_DATASET_TEMPLATE, project);
  HealthcareApiClient client = new HttpHealthcareApiClient();
  Hl7V2Store store = client.createHL7v2Store(healthcareDataset, HL7V2_STORE_NAME);
  store.getParserConfig();
}
 

@BeforeClass
public static void createHL7v2tores() throws IOException {
  String project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  healthcareDataset = String.format(HEALTHCARE_DATASET_TEMPLATE, project);
  HealthcareApiClient client = new HttpHealthcareApiClient();
  client.createHL7v2Store(healthcareDataset, INPUT_HL7V2_STORE_NAME);
  client.createHL7v2Store(healthcareDataset, OUTPUT_HL7V2_STORE_NAME);
}
 

@BeforeClass
public static void createHL7v2tore() throws IOException {
  String project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  healthcareDataset = String.format(HEALTHCARE_DATASET_TEMPLATE, project);
  HealthcareApiClient client = new HttpHealthcareApiClient();
  client.createHL7v2Store(healthcareDataset, HL7V2_STORE_NAME);
}
 

@Before
public void setUp() throws Exception {
  PipelineOptionsFactory.register(SpannerTestPipelineOptions.class);
  options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class);

  project = options.getInstanceProjectId();
  if (project == null) {
    project = options.as(GcpOptions.class).getProject();
  }

  spanner = SpannerOptions.newBuilder().setProjectId(project).build().getService();

  databaseName = generateDatabaseName();

  databaseAdminClient = spanner.getDatabaseAdminClient();

  // Delete database if exists.
  databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName);

  OperationFuture<Database, CreateDatabaseMetadata> op =
      databaseAdminClient.createDatabase(
          options.getInstanceId(),
          databaseName,
          Collections.singleton(
              "CREATE TABLE "
                  + options.getTable()
                  + " ("
                  + "  Key           INT64,"
                  + "  Value         STRING(MAX),"
                  + ") PRIMARY KEY (Key)"));
  op.get();
  makeTestData();
}
 

@Before
public void setUp() throws Exception {
  PipelineOptionsFactory.register(SpannerTestPipelineOptions.class);
  options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class);

  project = options.getInstanceProjectId();
  if (project == null) {
    project = options.as(GcpOptions.class).getProject();
  }

  spanner = SpannerOptions.newBuilder().setProjectId(project).build().getService();

  databaseName = generateDatabaseName();

  databaseAdminClient = spanner.getDatabaseAdminClient();

  // Delete database if exists.
  databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName);

  OperationFuture<Database, CreateDatabaseMetadata> op =
      databaseAdminClient.createDatabase(
          options.getInstanceId(),
          databaseName,
          Collections.singleton(
              "CREATE TABLE "
                  + options.getTable()
                  + " ("
                  + "  Key           INT64,"
                  + "  Value         STRING(MAX) NOT NULL,"
                  + ") PRIMARY KEY (Key)"));
  op.get();
}
 

@VisibleForTesting
static void validateWorkerSettings(GcpOptions gcpOptions) {
  Preconditions.checkArgument(
      gcpOptions.getZone() == null || gcpOptions.getWorkerRegion() == null,
      "Cannot use option zone with workerRegion. Prefer either workerZone or workerRegion.");
  Preconditions.checkArgument(
      gcpOptions.getZone() == null || gcpOptions.getWorkerZone() == null,
      "Cannot use option zone with workerZone. Prefer workerZone.");
  Preconditions.checkArgument(
      gcpOptions.getWorkerRegion() == null || gcpOptions.getWorkerZone() == null,
      "workerRegion and workerZone options are mutually exclusive.");

  DataflowPipelineOptions dataflowOptions = gcpOptions.as(DataflowPipelineOptions.class);
  boolean hasExperimentWorkerRegion = false;
  if (dataflowOptions.getExperiments() != null) {
    for (String experiment : dataflowOptions.getExperiments()) {
      if (experiment.startsWith("worker_region")) {
        hasExperimentWorkerRegion = true;
        break;
      }
    }
  }
  Preconditions.checkArgument(
      !hasExperimentWorkerRegion || gcpOptions.getWorkerRegion() == null,
      "Experiment worker_region and option workerRegion are mutually exclusive.");
  Preconditions.checkArgument(
      !hasExperimentWorkerRegion || gcpOptions.getWorkerZone() == null,
      "Experiment worker_region and option workerZone are mutually exclusive.");

  if (gcpOptions.getZone() != null) {
    LOG.warn("Option --zone is deprecated. Please use --workerZone instead.");
    gcpOptions.setWorkerZone(gcpOptions.getZone());
    gcpOptions.setZone(null);
  }
}
 

@Test
public void testZoneAndWorkerRegionMutuallyExclusive() {
  GcpOptions options = PipelineOptionsFactory.as(GcpOptions.class);
  options.setZone("us-east1-b");
  options.setWorkerRegion("us-east1");
  assertThrows(
      IllegalArgumentException.class, () -> DataflowRunner.validateWorkerSettings(options));
}
 

@Override
public String create(PipelineOptions options) {
  return "projects/"
      + options.as(GcpOptions.class).getProject()
      + "/topics/"
      + options.getJobName();
}
 

public static void main(String[] args) throws IOException, GeneralSecurityException {

    TokenizePipelineOptions options =
        PipelineOptionsFactory.fromArgs(args).withValidation().as(TokenizePipelineOptions.class);

    Pipeline p = Pipeline.create(options);
    p.apply(
            FileIO.match()
                .filepattern(options.getInputFile())
                .continuously(
                    Duration.standardSeconds(options.getPollingInterval()), Watch.Growth.never()))
        .apply(FileIO.readMatches().withCompression(Compression.UNCOMPRESSED))
        .apply(
            "Text File Reader",
            ParDo.of(
                new TextFileReader(
                    options.as(GcpOptions.class).getProject(),
                    options.getFileDecryptKeyName(),
                    options.getFileDecryptKey(),
                    options.getBatchSize(),
                    options.getCsek(),
                    options.getCsekhash())))
        .apply(
            "Tokenize Data",
            ParDo.of(
                new TokenizeData(
                    options.as(GcpOptions.class).getProject(),
                    options.getDeidentifyTemplateName(),
                    options.getInspectTemplateName())))
        .apply(
            Window.<String>into(FixedWindows.of(Duration.standardMinutes(options.getInterval()))))
        .apply(new WriteOneFilePerWindow(options.getOutputFile(), 1));

    p.run();
  }
 

public static void main(String[] args) throws Exception {

    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
    // Enforce that this pipeline is always run in streaming mode.
    options.setStreaming(true);
    ExampleUtils exampleUtils = new ExampleUtils(options);
    Pipeline pipeline = Pipeline.create(options);

    pipeline
        // Read game events from Pub/Sub using custom timestamps, which are extracted from the
        // pubsub data elements, and parse the data.
        .apply(
            PubsubIO.readStrings()
                .withTimestampAttribute(GameConstants.TIMESTAMP_ATTRIBUTE)
                .fromTopic(options.getTopic()))
        .apply("ParseGameEvent", ParDo.of(new ParseEventFn()))
        // Create <team, GameActionInfo> mapping. UpdateTeamScore uses team name as key.
        .apply(
            "MapTeamAsKey",
            MapElements.into(
                    TypeDescriptors.kvs(
                        TypeDescriptors.strings(), TypeDescriptor.of(GameActionInfo.class)))
                .via((GameActionInfo gInfo) -> KV.of(gInfo.team, gInfo)))
        // Outputs a team's score every time it passes a new multiple of the threshold.
        .apply("UpdateTeamScore", ParDo.of(new UpdateTeamScoreFn(options.getThresholdScore())))
        // Write the results to BigQuery.
        .apply(
            "WriteTeamLeaders",
            new WriteWindowedToBigQuery<>(
                options.as(GcpOptions.class).getProject(),
                options.getDataset(),
                options.getLeaderBoardTableName() + "_team_leader",
                configureCompleteWindowedTableWrite()));

    // Run the pipeline and wait for the pipeline to finish; capture cancellation requests from the
    // command line.
    PipelineResult result = pipeline.run();
    exampleUtils.waitToFinish(result);
  }
 

@Override
public String create(PipelineOptions options) {
  return "projects/"
      + options.as(GcpOptions.class).getProject()
      + "/topics/"
      + options.getJobName();
}
 

@Override
public String create(PipelineOptions options) {
  return "projects/"
      + options.as(GcpOptions.class).getProject()
      + "/subscriptions/"
      + options.getJobName();
}