private List<WorkUnit> initializeWorkUnits() { List<WorkUnit> workUnits = Lists.newArrayList(); for (int i = 0; i < NUM_WORK_UNITS; i++) { WorkUnit workUnit = WorkUnit.create(createExtract(Extract.TableType.SNAPSHOT_ONLY, NAMESPACE, TABLE)); workUnit.setLowWaterMark(i * NUM_RECORDS_TO_EXTRACT_PER_EXTRACTOR + 1); workUnit.setHighWaterMark((i + 1) * NUM_RECORDS_TO_EXTRACT_PER_EXTRACTOR); workUnit.setProp(WORK_UNIT_INDEX_KEY, i); workUnits.add(workUnit); } return workUnits; } }
@BeforeClass public void setUp() { this.multiWorkUnit = new MultiWorkUnit(); WorkUnit workUnit1 = WorkUnit.createEmpty(); workUnit1.setHighWaterMark(1000); workUnit1.setLowWaterMark(0); workUnit1.setProp("k1", "v1"); this.multiWorkUnit.addWorkUnit(workUnit1); WorkUnit workUnit2 = WorkUnit.createEmpty(); workUnit2.setHighWaterMark(2000); workUnit2.setLowWaterMark(1001); workUnit2.setProp("k2", "v2"); this.multiWorkUnit.addWorkUnit(workUnit2); }
@Override public List<WorkUnit> getWorkunits(SourceState sourceState) { sourceState.setProp(FOO, BAR); if (Iterables.isEmpty(sourceState.getPreviousWorkUnitStates())) { return initializeWorkUnits(); } List<WorkUnit> workUnits = Lists.newArrayList(); for (WorkUnitState workUnitState : sourceState.getPreviousWorkUnitStates()) { WorkUnit workUnit = WorkUnit.create(createExtract(Extract.TableType.SNAPSHOT_ONLY, NAMESPACE, TABLE)); workUnit.setLowWaterMark(workUnitState.getPropAsInt(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY) + NUM_WORK_UNITS * NUM_RECORDS_TO_EXTRACT_PER_EXTRACTOR); workUnit.setHighWaterMark(workUnitState.getPropAsInt(ConfigurationKeys.WORK_UNIT_HIGH_WATER_MARK_KEY) + NUM_WORK_UNITS * NUM_RECORDS_TO_EXTRACT_PER_EXTRACTOR); workUnit.setProp(WORK_UNIT_INDEX_KEY, workUnitState.getPropAsInt(WORK_UNIT_INDEX_KEY)); workUnits.add(workUnit); } return workUnits; }