Examples with Source org.apache.flink.api.connector.source.Source used on opensource projects

Example 1 with Source

use of org.apache.flink.api.connector.source.Source in project flink by apache.

the class HybridSourceTest method testBuilderWithEnumeratorSuperclass.

@Test
public void testBuilderWithEnumeratorSuperclass() {
    HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator> sourceFactory = (HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator>) context -> {
        MockSplitEnumerator enumerator = context.getPreviousEnumerator();
        return new MockBaseSource(1, 1, Boundedness.BOUNDED);
    };
    HybridSource<Integer> source = new HybridSource.HybridSourceBuilder<Integer, MockSplitEnumerator>().<ExtendedMockSplitEnumerator, Source<Integer, ?, ?>>addSource(new MockBaseSource(1, 1, Boundedness.BOUNDED)).addSource(sourceFactory, Boundedness.BOUNDED).build();
    assertNotNull(source);
}

Example 2 with Source

use of org.apache.flink.api.connector.source.Source in project flink by apache.

the class HybridSourceReaderTest method testReader.

@Test
public void testReader() throws Exception {
    TestingReaderContext readerContext = new TestingReaderContext();
    TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
    MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED);
    // 2 underlying readers to exercise switch
    SourceReader<Integer, MockSourceSplit> mockSplitReader1 = source.createReader(readerContext);
    SourceReader<Integer, MockSourceSplit> mockSplitReader2 = source.createReader(readerContext);
    HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
    Assert.assertThat(readerContext.getSentEvents(), Matchers.emptyIterable());
    reader.start();
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    Assert.assertNull(currentReader(reader));
    Assert.assertEquals(InputStatus.NOTHING_AVAILABLE, reader.pollNext(readerOutput));
    Source source1 = new MockSource(null, 0) {

        @Override
        public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
            return mockSplitReader1;
        }
    };
    reader.handleSourceEvents(new SwitchSourceEvent(0, source1, false));
    MockSourceSplit mockSplit = new MockSourceSplit(0, 0, 1);
    mockSplit.addRecord(0);
    SwitchedSources switchedSources = new SwitchedSources();
    switchedSources.put(0, source);
    HybridSourceSplit hybridSplit = HybridSourceSplit.wrapSplit(mockSplit, 0, switchedSources);
    reader.addSplits(Collections.singletonList(hybridSplit));
    // drain splits
    InputStatus status = reader.pollNext(readerOutput);
    while (readerOutput.getEmittedRecords().isEmpty() || status == InputStatus.MORE_AVAILABLE) {
        status = reader.pollNext(readerOutput);
        Thread.sleep(10);
    }
    Assert.assertThat(readerOutput.getEmittedRecords(), Matchers.contains(0));
    reader.pollNext(readerOutput);
    Assert.assertEquals("before notifyNoMoreSplits", InputStatus.NOTHING_AVAILABLE, reader.pollNext(readerOutput));
    reader.notifyNoMoreSplits();
    reader.pollNext(readerOutput);
    assertAndClearSourceReaderFinishedEvent(readerContext, 0);
    Assert.assertEquals("reader before switch source event", mockSplitReader1, currentReader(reader));
    Source source2 = new MockSource(null, 0) {

        @Override
        public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
            return mockSplitReader2;
        }
    };
    reader.handleSourceEvents(new SwitchSourceEvent(1, source2, true));
    Assert.assertEquals("reader after switch source event", mockSplitReader2, currentReader(reader));
    reader.notifyNoMoreSplits();
    Assert.assertEquals("reader 1 after notifyNoMoreSplits", InputStatus.END_OF_INPUT, reader.pollNext(readerOutput));
    reader.close();
}

Example 3 with Source

use of org.apache.flink.api.connector.source.Source in project flink by apache.

the class HybridSourceSplitSerializerTest method testSerialization.

@Test
public void testSerialization() throws Exception {
    Map<Integer, Source> switchedSources = new HashMap<>();
    switchedSources.put(0, new MockSource(null, 0));
    byte[] splitBytes = { 1, 2, 3 };
    HybridSourceSplitSerializer serializer = new HybridSourceSplitSerializer();
    HybridSourceSplit split = new HybridSourceSplit(0, splitBytes, 0, "splitId");
    byte[] serialized = serializer.serialize(split);
    HybridSourceSplit clonedSplit = serializer.deserialize(0, serialized);
    Assert.assertEquals(split, clonedSplit);
    try {
        serializer.deserialize(1, serialized);
        Assert.fail();
    } catch (IOException e) {
    // expected invalid version
    }
}

Example 4 with Source

use of org.apache.flink.api.connector.source.Source in project flink by apache.

the class HybridSourceTest method testBuilderWithSourceFactory.

@Test
public void testBuilderWithSourceFactory() {
    HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator> sourceFactory = new HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator>() {

        @Override
        public Source<Integer, ?, ?> create(HybridSource.SourceSwitchContext<MockSplitEnumerator> context) {
            MockSplitEnumerator enumerator = context.getPreviousEnumerator();
            return new MockBaseSource(1, 1, Boundedness.BOUNDED);
        }
    };
    HybridSource<Integer> source = new HybridSource.HybridSourceBuilder<Integer, MockSplitEnumerator>().<MockSplitEnumerator, Source<Integer, ?, ?>>addSource(new MockBaseSource(1, 1, Boundedness.BOUNDED)).addSource(sourceFactory, Boundedness.BOUNDED).build();
    assertNotNull(source);
}

Example 5 with Source

use of org.apache.flink.api.connector.source.Source in project flink by apache.

the class HybridSourceReader method setCurrentReader.

private void setCurrentReader(int index) {
    Preconditions.checkArgument(index != currentSourceIndex);
    if (currentReader != null) {
        try {
            currentReader.close();
        } catch (Exception e) {
            throw new RuntimeException("Failed to close current reader", e);
        }
        LOG.debug("Reader closed: subtask={} sourceIndex={} currentReader={}", readerContext.getIndexOfSubtask(), currentSourceIndex, currentReader);
    }
    // TODO: track previous readers splits till checkpoint
    Source source = switchedSources.sourceOf(index);
    SourceReader<T, ?> reader;
    try {
        reader = source.createReader(readerContext);
    } catch (Exception e) {
        throw new RuntimeException("Failed tp create reader", e);
    }
    reader.start();
    currentSourceIndex = index;
    currentReader = reader;
    currentReader.isAvailable().whenComplete((result, ex) -> {
        if (ex == null) {
            availabilityFuture.complete(result);
        } else {
            availabilityFuture.completeExceptionally(ex);
        }
    });
    LOG.debug("Reader started: subtask={} sourceIndex={} {}", readerContext.getIndexOfSubtask(), currentSourceIndex, reader);
    // add restored splits
    if (!restoredSplits.isEmpty()) {
        List<HybridSourceSplit> splits = new ArrayList<>(restoredSplits.size());
        Iterator<HybridSourceSplit> it = restoredSplits.iterator();
        while (it.hasNext()) {
            HybridSourceSplit hybridSplit = it.next();
            if (hybridSplit.sourceIndex() == index) {
                splits.add(hybridSplit);
                it.remove();
            }
        }
        addSplits(splits);
    }
}