Examples with Source com.google.api.services.dataflow.model.Source used on opensource projects

Example 21 with Source

use of com.google.api.services.dataflow.model.Source in project beam by apache.

the class WorkerCustomSourcesTest method testSplittingProducedInvalidSource.

@Test
public void testSplittingProducedInvalidSource() throws Exception {
    com.google.api.services.dataflow.model.Source cloudSource = translateIOToCloudSource(new SourceProducingInvalidSplits("original", null), options);
    expectedException.expect(IllegalArgumentException.class);
    expectedException.expectMessage(allOf(containsString("Splitting a valid source produced an invalid source"), containsString("original"), containsString("badBundle")));
    expectedException.expectCause(hasMessage(containsString("intentionally invalid")));
    performSplit(cloudSource, options, null, /*desiredBundleSizeBytes*/
    null, /* numBundles limit */
    null);
}

Example 22 with Source

use of com.google.api.services.dataflow.model.Source in project beam by apache.

the class WorkerCustomSourcesTest method testFailureToStartReadingIncludesSourceDetails.

@Test
public void testFailureToStartReadingIncludesSourceDetails() throws Exception {
    com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(new SourceProducingFailingReader(), options);
    // Exception object. So instead we use Throwables.getStackTraceAsString and match on that.
    try {
        readElemsFromSource(options, source);
        fail("Expected to fail");
    } catch (Exception e) {
        assertThat(getStackTraceAsString(e), allOf(containsString("Intentional error"), containsString("Some description")));
    }
}

Example 23 with Source

use of com.google.api.services.dataflow.model.Source in project beam by apache.

the class WorkerCustomSourcesTest method testUnboundedSplits.

@Test
public void testUnboundedSplits() throws Exception {
    com.google.api.services.dataflow.model.Source source = serializeToCloudSource(new TestCountingSource(Integer.MAX_VALUE), options);
    List<String> serializedSplits = getStrings(source.getSpec(), WorkerCustomSources.SERIALIZED_SOURCE_SPLITS, null);
    assertEquals(20, serializedSplits.size());
    for (String serializedSplit : serializedSplits) {
        assertTrue(deserializeFromByteArray(decodeBase64(serializedSplit), "source") instanceof TestCountingSource);
    }
}

Example 24 with Source

use of com.google.api.services.dataflow.model.Source in project beam by apache.

the class WorkerCustomSourcesTest method stepToCloudSource.

private static com.google.api.services.dataflow.model.Source stepToCloudSource(Step step) throws Exception {
    com.google.api.services.dataflow.model.Source res = dictionaryToCloudSource(getDictionary(step.getProperties(), PropertyNames.SOURCE_STEP_INPUT));
    // Encoding is specified in the step, not in the source itself.  This is
    // normal: incoming Dataflow API Source objects in map tasks will have the
    // encoding filled in from the step's output encoding.
    @SuppressWarnings("unchecked") List<Map<String, Object>> outputInfo = (List<Map<String, Object>>) step.getProperties().get(PropertyNames.OUTPUT_INFO);
    CloudObject encoding = CloudObject.fromSpec(getObject(outputInfo.get(0), PropertyNames.ENCODING));
    res.setCodec(encoding);
    return res;
}

Example 25 with Source

use of com.google.api.services.dataflow.model.Source in project beam by apache.

the class LengthPrefixUnknownCoders method forSideInputInfos.

/**
 * Wraps unknown coders on every {@link SideInputInfo} with length prefixes and also replaces the
 * wrapped coder with a byte array coder if requested.
 */
public static List<SideInputInfo> forSideInputInfos(List<SideInputInfo> sideInputInfos, boolean replaceWithByteArrayCoder) {
    ImmutableList.Builder<SideInputInfo> updatedSideInputInfos = ImmutableList.builder();
    for (SideInputInfo sideInputInfo : sideInputInfos) {
        try {
            SideInputInfo updatedSideInputInfo = clone(sideInputInfo, SideInputInfo.class);
            for (Source source : updatedSideInputInfo.getSources()) {
                source.setCodec(forCodec(source.getCodec(), replaceWithByteArrayCoder));
            }
            updatedSideInputInfos.add(updatedSideInputInfo);
        } catch (IOException e) {
            throw new RuntimeException(String.format("Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}", sideInputInfo), e);
        }
    }
    return updatedSideInputInfos.build();
}