Java Code Examples for org.apache.beam.sdk.options.ValueProvider.StaticValueProvider#of()

public static PTransform<PCollection<String>, WriteFilesResult<Void>> createWrite(
    String filenamePrefix, String filenameSuffix, Schema schema, JdbcAvroArgs jdbcAvroArgs) {
  filenamePrefix = filenamePrefix.replaceAll("/+$", "") + "/part";
  ValueProvider<ResourceId> prefixProvider =
      StaticValueProvider.of(FileBasedSink.convertToFileResourceIfPossible(filenamePrefix));
  FileBasedSink.FilenamePolicy filenamePolicy =
      DefaultFilenamePolicy.fromStandardParameters(
          prefixProvider, DEFAULT_SHARD_TEMPLATE, filenameSuffix, false);

  final DynamicAvroDestinations<String, Void, String> destinations =
      AvroIO.constantDestinations(
          filenamePolicy,
          schema,
          ImmutableMap.of(),
          // since Beam does not support zstandard
          CodecFactory.nullCodec(),
          SerializableFunctions.identity());
  final FileBasedSink<String, Void, String> sink =
      new JdbcAvroSink<>(prefixProvider, destinations, jdbcAvroArgs);
  return WriteFiles.to(sink);
}
 

/**
 * Tests that windowedFilename() constructs the filename correctly according to the parameters
 * when the filename suffix is a null value.
 */
@Test
public void testWindowedFilenameNullSuffix() throws IOException {
  // Arrange
  //
  ResourceId outputDirectory = getBaseTempDirectory();
  WindowedContext context = mock(WindowedContext.class);
  BoundedWindow window = mock(BoundedWindow.class);
  PaneInfo paneInfo = PaneInfo.createPane(false, true, Timing.ON_TIME, 0, 0);
  WindowedFilenamePolicy policy =
      new WindowedFilenamePolicy(
          StaticValueProvider.of(outputDirectory.toString()),
          StaticValueProvider.of("output"),
          StaticValueProvider.of("-SSS-of-NNN"),
          StaticValueProvider.of(null));

  // Act
  //
  ResourceId filename =
      policy.windowedFilename(1, 1, window, paneInfo, new TestOutputFileHints());
  // Assert
  //
  assertThat(filename, is(notNullValue()));
  assertThat(filename.getFilename(), is(equalTo("output-001-of-001")));
}
 

@Override
public List<PubsubSource> split(int desiredNumSplits, PipelineOptions options)
    throws Exception {
  List<PubsubSource> result = new ArrayList<>(desiredNumSplits);
  PubsubSource splitSource = this;
  if (subscriptionPath == null) {
    splitSource =
        new PubsubSource(
            outer, StaticValueProvider.of(outer.createRandomSubscription(options)));
  }
  for (int i = 0; i < desiredNumSplits * SCALE_OUT; i++) {
    // Since the source is immutable and Pubsub automatically shards we simply
    // replicate ourselves the requested number of times
    result.add(splitSource);
  }
  return result;
}
 

/** Tests the {@link BulkCompressor.Compressor} performs compression properly. */
@Test
public void testCompressFile() throws Exception {
  // Setup test
  final Compression compression = Compression.GZIP;

  final ValueProvider<String> outputDirectoryProvider =
      pipeline.newProvider(tempFolderCompressedPath.toString());

  final ValueProvider<Compression> compressionProvider = StaticValueProvider.of(compression);

  final Metadata metadata = FileSystems.matchSingleFileSpec(textFile.toString());

  // Execute the compressor
  PCollection<String> lines = pipeline
      .apply("Create File Input", Create.of(metadata))
      .apply("Compress", ParDo.of(new Compressor(outputDirectoryProvider, compressionProvider)))
      .apply("Read the Files", TextIO.readAll().withCompression(Compression.AUTO));

  // Test the result
  PAssert.that(lines).containsInAnyOrder(FILE_CONTENT);
  pipeline.run();
}
 

/**
 * Instantiates a new Import.
 *
 * @param fhirStore the fhir store
 * @param tempGcsPath the temp gcs path
 * @param deadLetterGcsPath the dead letter gcs path
 * @param contentStructure the content structure
 */
Import(
    String fhirStore,
    String tempGcsPath,
    String deadLetterGcsPath,
    @Nullable ContentStructure contentStructure) {
  this.fhirStore = StaticValueProvider.of(fhirStore);
  this.tempGcsPath = StaticValueProvider.of(tempGcsPath);
  this.deadLetterGcsPath = StaticValueProvider.of(deadLetterGcsPath);
  if (contentStructure == null) {
    this.contentStructure = ContentStructure.CONTENT_STRUCTURE_UNSPECIFIED;
  } else {
    this.contentStructure = contentStructure;
  }
}
 

DatastoreWriterFn(String projectId, @Nullable String localhost) {
  this(
      StaticValueProvider.of(projectId),
      localhost,
      new V1DatastoreFactory(),
      new WriteBatcherImpl());
}
 

@Test
public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
  TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
  factory = PubsubTestClient.createFactoryForCreateSubscription();
  PubsubUnboundedSource source =
      new PubsubUnboundedSource(
          factory,
          StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
          StaticValueProvider.of(topicPath),
          null /* subscription */,
          null /* timestampLabel */,
          null /* idLabel */,
          false /* needsAttributes */);
  assertThat(source.getSubscription(), nullValue());

  assertThat(source.getSubscription(), nullValue());

  PipelineOptions options = PipelineOptionsFactory.create();
  PubsubSource actualSource = new PubsubSource(source);
  PubsubReader reader = actualSource.createReader(options, null);
  SubscriptionPath createdSubscription = reader.subscription;
  assertThat(createdSubscription, not(nullValue()));

  PubsubCheckpoint checkpoint = reader.getCheckpointMark();
  assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));

  checkpoint.finalizeCheckpoint();
  PubsubCheckpoint deserCheckpoint =
      CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
  assertThat(checkpoint.subscriptionPath, not(nullValue()));
  assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));

  PubsubReader readerFromOriginal = actualSource.createReader(options, checkpoint);
  PubsubReader readerFromDeser = actualSource.createReader(options, deserCheckpoint);

  assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
  assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
}
 

@Test
public void testNestedValueProviderStatic() throws Exception {
  SerializableFunction<String, String> function = from -> from + "bar";
  ValueProvider<String> svp = StaticValueProvider.of("foo");
  ValueProvider<String> nvp = NestedValueProvider.of(svp, function);
  assertTrue(nvp.isAccessible());
  assertEquals("foobar", nvp.get());
  assertEquals("foobar", nvp.toString());
  assertEquals(nvp, NestedValueProvider.of(svp, function));
}
 

public TestFileBasedSource(
    String fileOrPattern,
    EmptyMatchTreatment emptyMatchTreatment,
    long minBundleSize,
    String splitHeader) {
  super(StaticValueProvider.of(fileOrPattern), emptyMatchTreatment, minBundleSize);
  this.splitHeader = splitHeader;
}
 

public static BeamRowToBigtableFn create(
    ValueProvider<String> defaultRowKeySeparator, ValueProvider<String> defaultColumnFamily) {
  return new BeamRowToBigtableFn(
      defaultRowKeySeparator,
      defaultColumnFamily,
      StaticValueProvider.of(DEFAULT_SPLIT_LARGE_ROWS),
      MAX_MUTATION_PER_REQUEST);
}
 

public BigQueryMapper(String projectId) {
  this.projectIdProvider = StaticValueProvider.of(projectId);
}
 

/** Read all HL7v2 Messages from a single store. */
public static ListHL7v2Messages read(String hl7v2Store) {
  return new ListHL7v2Messages(
      StaticValueProvider.of(Collections.singletonList(hl7v2Store)),
      StaticValueProvider.of(null));
}
 

/**
 * Tests the end-to-end Spec11 pipeline with mocked out API calls.
 *
 * <p>We suppress the (Serializable & Supplier) dual-casted lambda warnings because the supplier
 * produces an explicitly serializable mock, which is safe to cast.
 */
@Test
@SuppressWarnings("unchecked")
public void testEndToEndPipeline_generatesExpectedFiles() throws Exception {
  // Establish mocks for testing
  ImmutableList<Subdomain> inputRows = getInputDomains();
  CloseableHttpClient httpClient = mock(CloseableHttpClient.class, withSettings().serializable());

  // Return a mock HttpResponse that returns a JSON response based on the request.
  when(httpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder());

  EvaluateSafeBrowsingFn evalFn =
      new EvaluateSafeBrowsingFn(
          StaticValueProvider.of("apikey"),
          new Retrier(new FakeSleeper(new FakeClock()), 3),
          (Serializable & Supplier) () -> httpClient);

  // Apply input and evaluation transforms
  PCollection<Subdomain> input = p.apply(Create.of(inputRows));
  spec11Pipeline.evaluateUrlHealth(input, evalFn, StaticValueProvider.of("2018-06-01"));
  p.run();

  // Verify header and 4 threat matches for 3 registrars are found
  ImmutableList<String> generatedReport = resultFileContents();
  assertThat(generatedReport).hasSize(4);
  assertThat(generatedReport.get(0))
      .isEqualTo("Map from registrar email / name to detected subdomain threats:");

  // The output file can put the registrar emails and bad URLs in any order.
  // We cannot rely on the JSON toString to sort because the keys are not always in the same
  // order, so we must rely on length even though that's not ideal.
  ImmutableList<String> sortedLines =
      ImmutableList.sortedCopyOf(
          Comparator.comparingInt(String::length), generatedReport.subList(1, 4));

  JSONObject noEmailRegistrarJSON = new JSONObject(sortedLines.get(0));
  assertThat(noEmailRegistrarJSON.get("registrarEmailAddress")).isEqualTo("");
  assertThat(noEmailRegistrarJSON.get("registrarClientId")).isEqualTo("noEmailRegistrar");
  assertThat(noEmailRegistrarJSON.has("threatMatches")).isTrue();
  JSONArray noEmailThreatMatch = noEmailRegistrarJSON.getJSONArray("threatMatches");
  assertThat(noEmailThreatMatch.length()).isEqualTo(1);
  assertThat(noEmailThreatMatch.getJSONObject(0).get("fullyQualifiedDomainName"))
      .isEqualTo("no-email.com");
  assertThat(noEmailThreatMatch.getJSONObject(0).get("threatType"))
      .isEqualTo("MALWARE");

  JSONObject someRegistrarJSON = new JSONObject(sortedLines.get(1));
  assertThat(someRegistrarJSON.get("registrarEmailAddress")).isEqualTo("[email protected]");
  assertThat(someRegistrarJSON.get("registrarClientId")).isEqualTo("someRegistrar");
  assertThat(someRegistrarJSON.has("threatMatches")).isTrue();
  JSONArray someThreatMatch = someRegistrarJSON.getJSONArray("threatMatches");
  assertThat(someThreatMatch.length()).isEqualTo(1);
  assertThat(someThreatMatch.getJSONObject(0).get("fullyQualifiedDomainName"))
      .isEqualTo("444.com");
  assertThat(someThreatMatch.getJSONObject(0).get("threatType"))
      .isEqualTo("MALWARE");

  // theRegistrar has two ThreatMatches, we have to parse it explicitly
  JSONObject theRegistrarJSON = new JSONObject(sortedLines.get(2));
  assertThat(theRegistrarJSON.get("registrarEmailAddress")).isEqualTo("[email protected]");
  assertThat(theRegistrarJSON.get("registrarClientId")).isEqualTo("theRegistrar");
  assertThat(theRegistrarJSON.has("threatMatches")).isTrue();
  JSONArray theThreatMatches = theRegistrarJSON.getJSONArray("threatMatches");
  assertThat(theThreatMatches.length()).isEqualTo(2);
  ImmutableList<String> threatMatchStrings =
      ImmutableList.of(
          theThreatMatches.getJSONObject(0).toString(),
          theThreatMatches.getJSONObject(1).toString());
  assertThat(threatMatchStrings)
      .containsExactly(
          new JSONObject()
              .put("fullyQualifiedDomainName", "111.com")
              .put("threatType", "MALWARE")
              .toString(),
          new JSONObject()
              .put("fullyQualifiedDomainName", "222.com")
              .put("threatType", "MALWARE")
              .toString());
}
 

@Override
public FileBasedSource<byte[]> apply(String input) {
  return new TFRecordSource(StaticValueProvider.of(input));
}
 

@Test
public void testParseDurationValueProvider() {
  final ValueProvider<String> provider = StaticValueProvider.of("13s");
  assertEquals(Time.parseDuration(provider).get().getMillis(), SECONDS.toMillis(13));
}
 

@Override
public XmlSource<T> apply(String input) {
  return new XmlSource<>(StaticValueProvider.of(input), configuration, 1L);
}
 

@Override
public FileBasedSource<String> apply(String input) {
  return new TextSource(
      StaticValueProvider.of(input), EmptyMatchTreatment.DISALLOW, delimiter);
}
 

/**
 * Constructs a new {@link WindowedFilenamePolicy} with the supplied baseFile used for output
 * files.
 *
 * @param outputDirectory The output directory for all files.
 * @param outputFilenamePrefix The common prefix for output files.
 * @param shardTemplate The template used to create uniquely named sharded files.
 * @param suffix The suffix to append to all files output by the policy.
 */
public WindowedFilenamePolicy(
    String outputDirectory, String outputFilenamePrefix, String shardTemplate, String suffix) {
  this(
      StaticValueProvider.of(outputDirectory),
      StaticValueProvider.of(outputFilenamePrefix),
      StaticValueProvider.of(shardTemplate),
      StaticValueProvider.of(suffix));
}
 

/**
 * Instantiates a new Import fn.
 *
 * @param fhirStore the fhir store
 * @param tempGcsPath the temp gcs path
 * @param deadLetterGcsPath the dead letter gcs path
 */
WriteBundlesToFilesFn(String fhirStore, String tempGcsPath, String deadLetterGcsPath) {
  this.fhirStore = StaticValueProvider.of(fhirStore);
  this.tempGcsPath = StaticValueProvider.of(tempGcsPath);
  this.deadLetterGcsPath = StaticValueProvider.of(deadLetterGcsPath);
}
 

/**
 * Like {@link #BlockBasedSource(String, EmptyMatchTreatment, long)} but with a default {@link
 * EmptyMatchTreatment} of {@link EmptyMatchTreatment#DISALLOW}.
 */
public BlockBasedSource(String fileOrPatternSpec, long minBundleSize) {
  this(StaticValueProvider.of(fileOrPatternSpec), minBundleSize);
}