org.apache.beam.sdk.options.ValueProvider.StaticValueProvider Java Examples

The following examples show how to use org.apache.beam.sdk.options.ValueProvider.StaticValueProvider. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: WriteFilesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testBuildWrite() {
  SimpleSink<Void> sink = makeSimpleSink();
  WriteFiles<String, ?, String> write = WriteFiles.to(sink).withNumShards(3);
  assertThat((SimpleSink<Void>) write.getSink(), is(sink));
  PTransform<PCollection<String>, PCollectionView<Integer>> originalSharding =
      write.getComputeNumShards();

  assertThat(write.getComputeNumShards(), is(nullValue()));
  assertThat(write.getNumShardsProvider(), instanceOf(StaticValueProvider.class));
  assertThat(write.getNumShardsProvider().get(), equalTo(3));
  assertThat(write.getComputeNumShards(), equalTo(originalSharding));

  WriteFiles<String, ?, ?> write2 = write.withSharding(SHARDING_TRANSFORM);
  assertThat((SimpleSink<Void>) write2.getSink(), is(sink));
  assertThat(write2.getComputeNumShards(), equalTo(SHARDING_TRANSFORM));
  // original unchanged

  WriteFiles<String, ?, ?> writeUnsharded = write2.withRunnerDeterminedSharding();
  assertThat(writeUnsharded.getComputeNumShards(), nullValue());
  assertThat(write.getComputeNumShards(), equalTo(originalSharding));
}
 
Example #2
Source File: WindowedFilenamePolicy.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/**
 * The windowed filename method will construct filenames per window according to the baseFile,
 * suffix, and shardTemplate supplied. Directories with date templates in them will automatically
 * have their values resolved. For example the outputDirectory of /YYYY/MM/DD would resolve to
 * /2017/01/08 on January 8th, 2017.
 */
@Override
public ResourceId windowedFilename(
        int shardNumber,
        int numShards,
        BoundedWindow window,
        PaneInfo paneInfo,
        OutputFileHints outputFileHints) {

    ResourceId outputFile =
            resolveWithDateTemplates(outputDirectory, window)
                    .resolve(outputFilenamePrefix.get(), StandardResolveOptions.RESOLVE_FILE);

    DefaultFilenamePolicy policy =
            DefaultFilenamePolicy.fromStandardParameters(
                    StaticValueProvider.of(outputFile), shardTemplate.get(), suffix.get(), true);
    ResourceId result =
            policy.windowedFilename(shardNumber, numShards, window, paneInfo, outputFileHints);
    LOG.debug("Windowed file name policy created: {}", result.toString());
    return result;
}
 
Example #3
Source File: PubsubIntegrationTest.java    From gcp-ingestion with Mozilla Public License 2.0 6 votes vote down vote up
@Test(timeout = 30000)
public void canSendPubsubOutput() throws Exception {
  final List<String> inputLines = Lines.resources("testdata/pubsub-integration/input.ndjson");

  pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);

  SinkOptions.Parsed sinkOptions = pipeline.getOptions().as(SinkOptions.Parsed.class);
  sinkOptions.setOutput(pipeline.newProvider(topicName.toString()));
  // We would normally use pipeline.newProvider instead of StaticValueProvider in tests,
  // but something about this configuration causes the pipeline to stall when CompressPayload
  // accesses a method on the underlying enum value when defined via pipeline.newProvider.
  sinkOptions.setOutputPubsubCompression(StaticValueProvider.of(Compression.UNCOMPRESSED));

  pipeline.apply(Create.of(inputLines)).apply(InputFileFormat.json.decode())
      .apply(OutputType.pubsub.write(sinkOptions));

  final PipelineResult result = pipeline.run();

  System.err.println("Waiting for subscriber to receive messages published in the pipeline...");
  List<String> expectedLines = Lines.resources("testdata/pubsub-integration/truncated.ndjson");
  List<String> received = receiveLines(expectedLines.size());
  assertThat(received, matchesInAnyOrder(expectedLines));
  result.cancel();
}
 
Example #4
Source File: BigQueryIOStorageReadTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testTableSourceCreateReader() throws Exception {
  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(
              BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table")),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  thrown.expect(UnsupportedOperationException.class);
  thrown.expectMessage("BigQuery storage source must be split before reading");
  tableSource.createReader(options);
}
 
Example #5
Source File: BigQueryIOStorageReadTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testTableSourceEstimatedSize_WithDefaultProject() throws Exception {
  fakeDatasetService.createDataset("project-id", "dataset", "", "", null);
  TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset.table");
  Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
  fakeDatasetService.createTable(table);

  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}
 
Example #6
Source File: InvoicingUtilsTest.java    From nomulus with Apache License 2.0 6 votes vote down vote up
@Test
public void testDestinationFunction_generatesProperFileParams() {
  SerializableFunction<BillingEvent, Params> destinationFunction =
      InvoicingUtils.makeDestinationFunction("my/directory", StaticValueProvider.of("2017-10"));

  BillingEvent billingEvent = mock(BillingEvent.class);
  // We mock BillingEvent to make the test independent of the implementation of toFilename()
  when(billingEvent.toFilename(any())).thenReturn("invoice_details_2017-10_registrar_tld");

  assertThat(destinationFunction.apply(billingEvent))
      .isEqualTo(
          new Params()
              .withShardTemplate("")
              .withSuffix(".csv")
              .withBaseFilename(
                  FileBasedSink.convertToFileResourceIfPossible(
                      "my/directory/2017-10/invoice_details_2017-10_registrar_tld")));
}
 
Example #7
Source File: BulkCompressorTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/** 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();
}
 
Example #8
Source File: JavascriptTextTransformerTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/**
 * Test {@link TransformTextViaJavascript} allows for script to disregard some data by returning
 * null / undefined.
 */
@Test
@Category(NeedsRunner.class)
public void testDoFnDisregards() {
  List<String> inJson = Arrays.asList("{\"answerToLife\":42}", "{\"answerToLife\":43}");
  List<String> expectedJson = Arrays.asList("{\"answerToLife\":42}");

  PCollection<String> transformedJson =
      pipeline
          .apply("Create", Create.of(inJson))
          .apply(
              TransformTextViaJavascript.newBuilder()
                  .setFileSystemPath(StaticValueProvider.of(TRANSFORM_FILE_PATH))
                  .setFunctionName(StaticValueProvider.of("transformWithFilter"))
                  .build());

  PAssert.that(transformedJson).containsInAnyOrder(expectedJson);

  pipeline.run();
}
 
Example #9
Source File: SplunkEventWriterTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/**
 * Test building {@link SplunkEventWriter} with custom batchcount and certificate validation .
 */
@Test
public void eventWriterCustomBatchCountAndValidation() {

  Integer batchCount = 30;
  Boolean certificateValidation = false;
  SplunkEventWriter writer =
      SplunkEventWriter.newBuilder()
          .withUrl("test-url")
          .withToken("test-token")
          .withInputBatchCount(StaticValueProvider.of(batchCount))
          .withDisableCertificateValidation(StaticValueProvider.of(certificateValidation))
          .build();

  assertThat(writer.inputBatchCount().get()).isEqualTo(batchCount);
  assertThat(writer.disableCertificateValidation().get()).isEqualTo(certificateValidation);
}
 
Example #10
Source File: DualInputNestedvalueProviderTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
@Test
public void testNestedValueProviderStatic() throws Exception {
  ValueProvider<String> xvp = StaticValueProvider.of("foo");
  ValueProvider<Integer> yvp = StaticValueProvider.of(1);
  ValueProvider<String> zvp =
      DualInputNestedValueProvider.of(
          xvp,
          yvp,
          new SerializableFunction<TranslatorInput<String, Integer>, String>() {
            @Override
            public String apply(TranslatorInput<String, Integer> input) {
              return input.getX() + (input.getY() + 1);
            }
          });
  assertTrue(zvp.isAccessible());
  assertEquals("foo2", zvp.get());
}
 
Example #11
Source File: SpannerAccessorTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testCreateOnlyOnce() {
  SpannerConfig config1 =
      SpannerConfig.create()
          .toBuilder()
          .setServiceFactory(serviceFactory)
          .setProjectId(StaticValueProvider.of("project"))
          .setInstanceId(StaticValueProvider.of("test1"))
          .setDatabaseId(StaticValueProvider.of("test1"))
          .build();

  SpannerAccessor acc1 = SpannerAccessor.getOrCreate(config1);
  SpannerAccessor acc2 = SpannerAccessor.getOrCreate(config1);
  SpannerAccessor acc3 = SpannerAccessor.getOrCreate(config1);

  acc1.close();
  acc2.close();
  acc3.close();

  // getDatabaseClient and close() only called once.
  verify(serviceFactory.mockSpanner(), times(1))
      .getDatabaseClient(DatabaseId.of("project", "test1", "test1"));
  verify(serviceFactory.mockSpanner(), times(1)).close();
}
 
Example #12
Source File: RepublishPerNamespace.java    From gcp-ingestion with Mozilla Public License 2.0 6 votes vote down vote up
@Override
public PDone expand(PCollection<PubsubMessage> input) {
  List<Destination> destinations = baseOptions.getPerNamespaceDestinations().entrySet().stream()
      .map(entry -> new Destination(entry.getKey(), entry.getValue()))
      .collect(Collectors.toList());
  int numDestinations = destinations.size();
  int numPartitions = numDestinations + 1;
  PCollectionList<PubsubMessage> partitioned = input.apply("PartitionByNamespace",
      Partition.of(numPartitions, new PartitionFn(destinations)));

  for (int i = 0; i < numDestinations; i++) {
    Destination destination = destinations.get(i);
    RepublisherOptions.Parsed opts = baseOptions.as(RepublisherOptions.Parsed.class);
    opts.setOutput(StaticValueProvider.of(destination.dest));
    String name = String.join("_", "republish", destination.namespace);
    partitioned.get(i).apply(name, opts.getOutputType().write(opts));
  }

  return PDone.in(input.getPipeline());
}
 
Example #13
Source File: BigQueryIOStorageReadTest.java    From beam with Apache License 2.0 6 votes vote down vote up
private void doTableSourceEstimatedSizeTest(boolean useStreamingBuffer) throws Exception {
  fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
  TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");
  Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
  if (useStreamingBuffer) {
    table.setStreamingBuffer(new Streamingbuffer().setEstimatedBytes(BigInteger.TEN));
  }

  fakeDatasetService.createTable(table);

  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(tableRef),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}
 
Example #14
Source File: BigQueryIO.java    From beam with Apache License 2.0 6 votes vote down vote up
/**
 * Returns the table to write, or {@code null} if writing with {@code tableFunction}.
 *
 * <p>If the table's project is not specified, use the executing project.
 */
@Nullable
ValueProvider<TableReference> getTableWithDefaultProject(BigQueryOptions bqOptions) {
  ValueProvider<TableReference> table = getTable();
  if (table == null) {
    return table;
  }

  if (!table.isAccessible()) {
    LOG.info(
        "Using a dynamic value for table input. This must contain a project"
            + " in the table reference: {}",
        table);
    return table;
  }
  if (Strings.isNullOrEmpty(table.get().getProjectId())) {
    // If user does not specify a project we assume the table to be located in
    // the default project.
    TableReference tableRef = table.get();
    tableRef.setProjectId(bqOptions.getProject());
    return NestedValueProvider.of(
        StaticValueProvider.of(BigQueryHelpers.toJsonString(tableRef)),
        new JsonTableRefToTableRef());
  }
  return table;
}
 
Example #15
Source File: PubsubUnboundedSourceTest.java    From beam with Apache License 2.0 6 votes vote down vote up
private void setupOneMessage(Iterable<IncomingMessage> incoming) {
  now = new AtomicLong(REQ_TIME);
  clock = () -> now.get();
  factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming);
  PubsubUnboundedSource source =
      new PubsubUnboundedSource(
          clock,
          factory,
          null,
          null,
          StaticValueProvider.of(SUBSCRIPTION),
          TIMESTAMP_ATTRIBUTE,
          ID_ATTRIBUTE,
          true /* needsAttributes */);
  primSource = new PubsubSource(source);
}
 
Example #16
Source File: BigtableServiceImplTest.java    From beam with Apache License 2.0 6 votes vote down vote up
/**
 * This test ensures that protobuf creation and interactions with {@link BigtableDataClient} work
 * as expected.
 *
 * @throws IOException
 * @throws InterruptedException
 */
@Test
public void testRead() throws IOException {
  ByteKey start = ByteKey.copyFrom("a".getBytes(StandardCharsets.UTF_8));
  ByteKey end = ByteKey.copyFrom("b".getBytes(StandardCharsets.UTF_8));
  when(mockBigtableSource.getRanges()).thenReturn(Arrays.asList(ByteKeyRange.of(start, end)));
  when(mockBigtableSource.getTableId()).thenReturn(StaticValueProvider.of("table_name"));
  @SuppressWarnings("unchecked")
  ResultScanner<Row> mockResultScanner = Mockito.mock(ResultScanner.class);
  Row expectedRow = Row.newBuilder().setKey(ByteString.copyFromUtf8("a")).build();
  when(mockResultScanner.next()).thenReturn(expectedRow).thenReturn(null);
  when(mockBigtableDataClient.readRows(any(ReadRowsRequest.class))).thenReturn(mockResultScanner);
  BigtableService.Reader underTest =
      new BigtableServiceImpl.BigtableReaderImpl(mockSession, mockBigtableSource);

  underTest.start();
  Assert.assertEquals(expectedRow, underTest.getCurrentRow());
  Assert.assertFalse(underTest.advance());
  underTest.close();

  verify(mockResultScanner, times(1)).close();
}
 
Example #17
Source File: JdbcAvroIO.java    From dbeam with Apache License 2.0 6 votes vote down vote up
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);
}
 
Example #18
Source File: WindowedFilenamePolicyTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/**
 * 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")));
}
 
Example #19
Source File: WindowedFilenamePolicyTest.java    From DataflowTemplates with Apache License 2.0 6 votes vote down vote up
/**
 * Tests that windowedFilename() constructs the filename correctly according to the parameters
 * when using ValueProviders.
 */
@Test
public void testWindowedFilenameFormatValueProvider() 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(".txt"));

  // 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.txt")));
}
 
Example #20
Source File: DatastoreV1Test.java    From beam with Apache License 2.0 6 votes vote down vote up
private void datastoreWriterFnTest(int numMutations) throws Exception {
  // Create the requested number of mutations.
  List<Mutation> mutations = new ArrayList<>(numMutations);
  for (int i = 0; i < numMutations; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  DatastoreWriterFn datastoreWriter =
      new DatastoreWriterFn(
          StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
  DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(mutations);

  int start = 0;
  while (start < numMutations) {
    int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    commitRequest.addAllMutations(mutations.subList(start, end));
    // Verify all the batch requests were made with the expected mutations.
    verify(mockDatastore, times(1)).commit(commitRequest.build());
    start = end;
  }
}
 
Example #21
Source File: TestPipeline.java    From beam with Apache License 2.0 5 votes vote down vote up
/** Like {@link #run} but with the given potentially modified options. */
@Override
public PipelineResult run(PipelineOptions options) {
  checkState(
      enforcement.isPresent(),
      "Is your TestPipeline declaration missing a @Rule annotation? Usage: "
          + "@Rule public final transient TestPipeline pipeline = TestPipeline.create();");

  final PipelineResult pipelineResult;
  try {
    enforcement.get().beforePipelineExecution();
    PipelineOptions updatedOptions =
        MAPPER.convertValue(MAPPER.valueToTree(options), PipelineOptions.class);
    updatedOptions
        .as(TestValueProviderOptions.class)
        .setProviderRuntimeValues(StaticValueProvider.of(providerRuntimeValues));
    pipelineResult = super.run(updatedOptions);
    verifyPAssertsSucceeded(this, pipelineResult);
  } catch (RuntimeException exc) {
    Throwable cause = exc.getCause();
    if (cause instanceof AssertionError) {
      throw (AssertionError) cause;
    } else {
      throw exc;
    }
  }

  // If we reach this point, the pipeline has been run and no exceptions have been thrown during
  // its execution.
  enforcement.get().afterPipelineExecution();
  return pipelineResult;
}
 
Example #22
Source File: FhirIO.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * 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;
  }
}
 
Example #23
Source File: SplunkIO.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Disables ssl certificate validation.
 *
 * @param disableCertificateValidation for disabling certificate validation
 */
public Write withDisableCertificateValidation(Boolean disableCertificateValidation) {
  checkArgument(
      disableCertificateValidation != null,
      "withDisableCertificateValidation(disableCertificateValidation) called with null input.");
  return toBuilder()
      .setDisableCertificateValidation(StaticValueProvider.of((disableCertificateValidation)))
      .build();
}
 
Example #24
Source File: PubsubUnboundedSourceTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@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));
}
 
Example #25
Source File: PubsubIOTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testReadSubscriptionDisplayData() {
  String subscription = "projects/project/subscriptions/subscription";
  PubsubIO.Read<String> read =
      PubsubIO.readStrings()
          .fromSubscription(StaticValueProvider.of(subscription))
          .withTimestampAttribute("myTimestamp")
          .withIdAttribute("myId");

  DisplayData displayData = DisplayData.from(read);

  assertThat(displayData, hasDisplayItem("subscription", subscription));
  assertThat(displayData, hasDisplayItem("timestampAttribute", "myTimestamp"));
  assertThat(displayData, hasDisplayItem("idAttribute", "myId"));
}
 
Example #26
Source File: PubsubIntegrationTest.java    From gcp-ingestion with Mozilla Public License 2.0 5 votes vote down vote up
@Test(timeout = 30000)
public void canSendPubsubErrorOutput() throws Exception {
  final List<String> inputLines = Lines
      .resources("testdata/pubsub-integration/error-input.ndjson");

  pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);

  SinkOptions.Parsed sinkOptions = pipeline.getOptions().as(SinkOptions.Parsed.class);
  sinkOptions.setInput(pipeline.newProvider("test input"));
  sinkOptions.setJobName("test job name");
  sinkOptions.setErrorOutput(pipeline.newProvider(topicName.toString()));
  // We would normally use pipeline.newProvider instead of StaticValueProvider in tests,
  // but something about this configuration causes the pipeline to stall when CompressPayload
  // accesses a method on the underlying enum value when defined via pipeline.newProvider.
  sinkOptions.setErrorOutputPubsubCompression(StaticValueProvider.of(Compression.UNCOMPRESSED));

  pipeline.apply(Create.of(inputLines)).apply(InputFileFormat.json.decode())
      .apply(ErrorOutputType.pubsub.write(sinkOptions));

  final PipelineResult result = pipeline.run();

  System.err.println("Waiting for subscriber to receive messages published in the pipeline...");
  List<String> expectedLines = Lines.resources("testdata/pubsub-integration/error-output.ndjson");
  List<String> received = receiveLines(expectedLines.size());
  assertThat(received, matchesInAnyOrder(expectedLines));
  result.cancel();
}
 
Example #27
Source File: FhirIO.java    From beam with Apache License 2.0 5 votes vote down vote up
@Override
public Result expand(PCollection<String> input) {
  PCollection<HealthcareIOError<String>> failedBundles;
  PCollection<HealthcareIOError<String>> failedImports;
  switch (this.getWriteMethod()) {
    case IMPORT:
      LOG.warn(
          "Make sure the Cloud Healthcare Service Agent has permissions when using import:"
              + " https://cloud.google.com/healthcare/docs/how-tos/permissions-healthcare-api-gcp-products#fhir_store_cloud_storage_permissions");
      ValueProvider<String> deadPath =
          getImportGcsDeadLetterPath().orElseThrow(IllegalArgumentException::new);
      FhirIO.Import.ContentStructure contentStructure =
          getContentStructure().orElseThrow(IllegalArgumentException::new);
      ValueProvider<String> tempPath =
          getImportGcsTempPath()
              .orElse(
                  StaticValueProvider.of(input.getPipeline().getOptions().getTempLocation()));

      return input.apply(new Import(getFhirStore(), tempPath, deadPath, contentStructure));
    case EXECUTE_BUNDLE:
    default:
      failedBundles =
          input
              .apply(
                  "Execute FHIR Bundles",
                  ParDo.of(new ExecuteBundles.ExecuteBundlesFn(this.getFhirStore())))
              .setCoder(HealthcareIOErrorCoder.of(StringUtf8Coder.of()));
  }
  return Result.in(input.getPipeline(), failedBundles);
}
 
Example #28
Source File: PubsubUnboundedSinkTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void sendOneMessage() throws IOException {
  List<OutgoingMessage> outgoing =
      ImmutableList.of(
          OutgoingMessage.of(
              com.google.pubsub.v1.PubsubMessage.newBuilder()
                  .setData(ByteString.copyFromUtf8(DATA))
                  .putAllAttributes(ATTRIBUTES)
                  .build(),
              TIMESTAMP,
              getRecordId(DATA)));
  int batchSize = 1;
  int batchBytes = 1;
  try (PubsubTestClientFactory factory =
      PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.of())) {
    PubsubUnboundedSink sink =
        new PubsubUnboundedSink(
            factory,
            StaticValueProvider.of(TOPIC),
            TIMESTAMP_ATTRIBUTE,
            ID_ATTRIBUTE,
            NUM_SHARDS,
            batchSize,
            batchBytes,
            Duration.standardSeconds(2),
            RecordIdMethod.DETERMINISTIC);
    p.apply(Create.of(ImmutableList.of(DATA))).apply(ParDo.of(new Stamp(ATTRIBUTES))).apply(sink);
    p.run();
  }
  // The PubsubTestClientFactory will assert fail on close if the actual published
  // message does not match the expected publish message.
}
 
Example #29
Source File: PubsubIOTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testValueProviderTopic() {
  StaticValueProvider<String> provider = StaticValueProvider.of("projects/project/topics/topic");
  Read<String> pubsubRead = PubsubIO.readStrings().fromTopic(provider);
  Pipeline.create().apply(pubsubRead);
  assertThat(pubsubRead.getTopicProvider(), not(nullValue()));
  assertThat(pubsubRead.getTopicProvider().isAccessible(), is(true));
  assertThat(pubsubRead.getTopicProvider().get().asPath(), equalTo(provider.get()));
}
 
Example #30
Source File: WindowedFilenamePolicyTest.java    From DataflowTemplates with Apache License 2.0 5 votes vote down vote up
/**
 * Tests that windowedFilename() produces the correct directory when the directory has dynamic
 * date variables.
 */
@Test
public void testWithDynamicDirectory() throws IOException {
  // Arrange
  //
  ResourceId outputDirectory = getBaseTempDirectory()
      .resolve("YYYY/MM/DD/HH:mm", StandardResolveOptions.RESOLVE_DIRECTORY);
  WindowedContext context = mock(WindowedContext.class);
  IntervalWindow window = mock(IntervalWindow.class);
  PaneInfo paneInfo = PaneInfo.createPane(false, true, Timing.ON_TIME, 0, 0);
  Instant windowBegin = new DateTime(2017, 1, 8, 10, 55, 0).toInstant();
  Instant windowEnd = new DateTime(2017, 1, 8, 10, 56, 0).toInstant();
  when(window.maxTimestamp()).thenReturn(windowEnd);
  when(window.start()).thenReturn(windowBegin);
  when(window.end()).thenReturn(windowEnd);

  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.getCurrentDirectory().toString().endsWith("2017/01/08/10:56/"), is(equalTo(true)));
  assertThat(filename.getFilename(), is(equalTo("output-001-of-001")));
}