@Test public void testGetPreviousWatermarksForAllTables() { { State prevJobState = new SourceState(); prevJobState.setProp(ConfigurationKeys.JOB_COMMIT_POLICY_KEY, "full"); Extract[] extracts = new Extract[3]; SourceEntity[] sourceEntities = new SourceEntity[extracts.length]; List<WorkUnitState> prevWuStates = new ArrayList<>(); for (int i = 0; i < extracts.length; ++i) { String sourceEntityName = "Table$" + i; SourceEntity sourceEntity = SourceEntity.fromSourceEntityName(sourceEntityName); sourceEntities[i] = sourceEntity; extracts[i] = new Extract(TableType.APPEND_ONLY, "", sourceEntity.getDestTableName()); for (int j = 0; j < 3; ++j) { WorkUnit wu = new WorkUnit(extracts[i]); wu.setProp(ConfigurationKeys.SOURCE_ENTITY, sourceEntity.getSourceEntityName()); wu.setProp(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY, 10 * i); wu.setProp(ConfigurationKeys.EXTRACT_TABLE_NAME_KEY, sourceEntity.getDestTableName()); WorkUnitState wuState = new WorkUnitState(wu, prevJobState); wuState.setProp(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK, 20 * i); wuState.setProp(ConfigurationKeys.WORK_UNIT_WORKING_STATE_KEY, i == 1 ? WorkingState.FAILED.toString() : WorkingState.SUCCESSFUL.toString() ); wuState.setProp(ConfigurationKeys.EXTRACTOR_ROWS_EXPECTED, (i + j) * 5); prevWuStates.add(wuState); } } SourceState prevState = new SourceState(prevJobState, prevWuStates); Map<SourceEntity, Long> previousWM = QueryBasedSource.getPreviousWatermarksForAllTables(prevState); Assert.assertEquals(previousWM.size(), 3); Assert.assertEquals(previousWM.get(sourceEntities[0]), Long.valueOf(0L)); Assert.assertEquals(previousWM.get(sourceEntities[1]), Long.valueOf(10L)); Assert.assertEquals(previousWM.get(sourceEntities[2]), Long.valueOf(40L)); } }
static Map<SourceEntity, Long> getPreviousWatermarksForAllTables(SourceState state) { Map<SourceEntity, Long> result = Maps.newHashMap(); Map<SourceEntity, Long> prevLowWatermarksByTable = Maps.newHashMap(); Map<SourceEntity, Long> prevActualHighWatermarksByTable = Maps.newHashMap(); Set<SourceEntity> tablesWithFailedTasks = Sets.newHashSet(); Set<SourceEntity> tablesWithNoUpdatesOnPreviousRun = Sets.newHashSet(); boolean commitOnFullSuccess = JobCommitPolicy.getCommitPolicy(state) == JobCommitPolicy.COMMIT_ON_FULL_SUCCESS; for (WorkUnitState previousWus : state.getPreviousWorkUnitStates()) { Optional<SourceEntity> sourceEntity = SourceEntity.fromState(previousWus); if (!sourceEntity.isPresent()) { log.warn("Missing source entity for WorkUnit state: " + previousWus); continue; } SourceEntity table = sourceEntity.get(); long lowWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; LongWatermark waterMarkObj = previousWus.getWorkunit().getLowWatermark(LongWatermark.class); if(waterMarkObj != null){ lowWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)){ lowWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)); log.warn("can not find low water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY + " low water mark " + lowWm); } if (!prevLowWatermarksByTable.containsKey(table)) { prevLowWatermarksByTable.put(table, lowWm); } else { prevLowWatermarksByTable.put(table, Math.min(prevLowWatermarksByTable.get(table), lowWm)); } long highWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; waterMarkObj = previousWus.getActualHighWatermark(LongWatermark.class); if(waterMarkObj != null){ highWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)){ highWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)); log.warn("can not find high water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK + " high water mark " + highWm); } if (!prevActualHighWatermarksByTable.containsKey(table)) { prevActualHighWatermarksByTable.put(table, highWm); } else { prevActualHighWatermarksByTable.put(table, Math.max(prevActualHighWatermarksByTable.get(table), highWm)); } if (commitOnFullSuccess && !isSuccessfulOrCommited(previousWus)) { tablesWithFailedTasks.add(table); } if (!isAnyDataProcessed(previousWus)) { tablesWithNoUpdatesOnPreviousRun.add(table); } } for (Map.Entry<SourceEntity, Long> entry : prevLowWatermarksByTable.entrySet()) { if (tablesWithFailedTasks.contains(entry.getKey())) { log.info("Resetting low watermark to {} because previous run failed.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else if (tablesWithNoUpdatesOnPreviousRun.contains(entry.getKey()) && state.getPropAsBoolean(ConfigurationKeys.SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK, ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK)) { log.info("Resetting low watermakr to {} because previous run processed no data.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else { result.put(entry.getKey(), prevActualHighWatermarksByTable.get(entry.getKey())); } } return result; }
QueryBasedSource extends AbstractSource<S, D> { static Map<SourceEntity, Long> getPreviousWatermarksForAllTables(SourceState state) { Map<SourceEntity, Long> result = Maps.newHashMap(); Map<SourceEntity, Long> prevLowWatermarksByTable = Maps.newHashMap(); Map<SourceEntity, Long> prevActualHighWatermarksByTable = Maps.newHashMap(); Set<SourceEntity> tablesWithFailedTasks = Sets.newHashSet(); Set<SourceEntity> tablesWithNoUpdatesOnPreviousRun = Sets.newHashSet(); boolean commitOnFullSuccess = JobCommitPolicy.getCommitPolicy(state) == JobCommitPolicy.COMMIT_ON_FULL_SUCCESS; for (WorkUnitState previousWus : state.getPreviousWorkUnitStates()) { Optional<SourceEntity> sourceEntity = SourceEntity.fromState(previousWus); if (!sourceEntity.isPresent()) { log.warn("Missing source entity for WorkUnit state: " + previousWus); continue; } SourceEntity table = sourceEntity.get(); long lowWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; LongWatermark waterMarkObj = previousWus.getWorkunit().getLowWatermark(LongWatermark.class); if(waterMarkObj != null){ lowWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)){ lowWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)); log.warn("can not find low water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY + " low water mark " + lowWm); } if (!prevLowWatermarksByTable.containsKey(table)) { prevLowWatermarksByTable.put(table, lowWm); } else { prevLowWatermarksByTable.put(table, Math.min(prevLowWatermarksByTable.get(table), lowWm)); } long highWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; waterMarkObj = previousWus.getActualHighWatermark(LongWatermark.class); if(waterMarkObj != null){ highWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)){ highWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)); log.warn("can not find high water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK + " high water mark " + highWm); } if (!prevActualHighWatermarksByTable.containsKey(table)) { prevActualHighWatermarksByTable.put(table, highWm); } else { prevActualHighWatermarksByTable.put(table, Math.max(prevActualHighWatermarksByTable.get(table), highWm)); } if (commitOnFullSuccess && !isSuccessfulOrCommited(previousWus)) { tablesWithFailedTasks.add(table); } if (!isAnyDataProcessed(previousWus)) { tablesWithNoUpdatesOnPreviousRun.add(table); } } for (Map.Entry<SourceEntity, Long> entry : prevLowWatermarksByTable.entrySet()) { if (tablesWithFailedTasks.contains(entry.getKey())) { log.info("Resetting low watermark to {} because previous run failed.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else if (tablesWithNoUpdatesOnPreviousRun.contains(entry.getKey()) && state.getPropAsBoolean(ConfigurationKeys.SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK, ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK)) { log.info("Resetting low watermakr to {} because previous run processed no data.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else { result.put(entry.getKey(), prevActualHighWatermarksByTable.get(entry.getKey())); } } return result; } }
QueryBasedSource extends AbstractSource<S, D> { static Map<SourceEntity, Long> getPreviousWatermarksForAllTables(SourceState state) { Map<SourceEntity, Long> result = Maps.newHashMap(); Map<SourceEntity, Long> prevLowWatermarksByTable = Maps.newHashMap(); Map<SourceEntity, Long> prevActualHighWatermarksByTable = Maps.newHashMap(); Set<SourceEntity> tablesWithFailedTasks = Sets.newHashSet(); Set<SourceEntity> tablesWithNoUpdatesOnPreviousRun = Sets.newHashSet(); boolean commitOnFullSuccess = JobCommitPolicy.getCommitPolicy(state) == JobCommitPolicy.COMMIT_ON_FULL_SUCCESS; for (WorkUnitState previousWus : state.getPreviousWorkUnitStates()) { Optional<SourceEntity> sourceEntity = SourceEntity.fromState(previousWus); if (!sourceEntity.isPresent()) { log.warn("Missing source entity for WorkUnit state: " + previousWus); continue; } SourceEntity table = sourceEntity.get(); long lowWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; LongWatermark waterMarkObj = previousWus.getWorkunit().getLowWatermark(LongWatermark.class); if(waterMarkObj != null){ lowWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)){ lowWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)); log.warn("can not find low water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY + " low water mark " + lowWm); } if (!prevLowWatermarksByTable.containsKey(table)) { prevLowWatermarksByTable.put(table, lowWm); } else { prevLowWatermarksByTable.put(table, Math.min(prevLowWatermarksByTable.get(table), lowWm)); } long highWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; waterMarkObj = previousWus.getActualHighWatermark(LongWatermark.class); if(waterMarkObj != null){ highWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)){ highWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)); log.warn("can not find high water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK + " high water mark " + highWm); } if (!prevActualHighWatermarksByTable.containsKey(table)) { prevActualHighWatermarksByTable.put(table, highWm); } else { prevActualHighWatermarksByTable.put(table, Math.max(prevActualHighWatermarksByTable.get(table), highWm)); } if (commitOnFullSuccess && !isSuccessfulOrCommited(previousWus)) { tablesWithFailedTasks.add(table); } if (!isAnyDataProcessed(previousWus)) { tablesWithNoUpdatesOnPreviousRun.add(table); } } for (Map.Entry<SourceEntity, Long> entry : prevLowWatermarksByTable.entrySet()) { if (tablesWithFailedTasks.contains(entry.getKey())) { log.info("Resetting low watermark to {} because previous run failed.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else if (tablesWithNoUpdatesOnPreviousRun.contains(entry.getKey()) && state.getPropAsBoolean(ConfigurationKeys.SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK, ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK)) { log.info("Resetting low watermakr to {} because previous run processed no data.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else { result.put(entry.getKey(), prevActualHighWatermarksByTable.get(entry.getKey())); } } return result; }  }
QueryBasedSource extends AbstractSource<S, D> { static Map<SourceEntity, Long> getPreviousWatermarksForAllTables(SourceState state) { Map<SourceEntity, Long> result = Maps.newHashMap(); Map<SourceEntity, Long> prevLowWatermarksByTable = Maps.newHashMap(); Map<SourceEntity, Long> prevActualHighWatermarksByTable = Maps.newHashMap(); Set<SourceEntity> tablesWithFailedTasks = Sets.newHashSet(); Set<SourceEntity> tablesWithNoUpdatesOnPreviousRun = Sets.newHashSet(); boolean commitOnFullSuccess = JobCommitPolicy.getCommitPolicy(state) == JobCommitPolicy.COMMIT_ON_FULL_SUCCESS; for (WorkUnitState previousWus : state.getPreviousWorkUnitStates()) { Optional<SourceEntity> sourceEntity = SourceEntity.fromState(previousWus); if (!sourceEntity.isPresent()) { log.warn("Missing source entity for WorkUnit state: " + previousWus); continue; } SourceEntity table = sourceEntity.get(); long lowWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; LongWatermark waterMarkObj = previousWus.getWorkunit().getLowWatermark(LongWatermark.class); if(waterMarkObj != null){ lowWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)){ lowWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)); log.warn("can not find low water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY + " low water mark " + lowWm); } if (!prevLowWatermarksByTable.containsKey(table)) { prevLowWatermarksByTable.put(table, lowWm); } else { prevLowWatermarksByTable.put(table, Math.min(prevLowWatermarksByTable.get(table), lowWm)); } long highWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; waterMarkObj = previousWus.getActualHighWatermark(LongWatermark.class); if(waterMarkObj != null){ highWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)){ highWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)); log.warn("can not find high water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK + " high water mark " + highWm); } if (!prevActualHighWatermarksByTable.containsKey(table)) { prevActualHighWatermarksByTable.put(table, highWm); } else { prevActualHighWatermarksByTable.put(table, Math.max(prevActualHighWatermarksByTable.get(table), highWm)); } if (commitOnFullSuccess && !isSuccessfulOrCommited(previousWus)) { tablesWithFailedTasks.add(table); } if (!isAnyDataProcessed(previousWus)) { tablesWithNoUpdatesOnPreviousRun.add(table); } } for (Map.Entry<SourceEntity, Long> entry : prevLowWatermarksByTable.entrySet()) { if (tablesWithFailedTasks.contains(entry.getKey())) { log.info("Resetting low watermark to {} because previous run failed.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else if (tablesWithNoUpdatesOnPreviousRun.contains(entry.getKey()) && state.getPropAsBoolean(ConfigurationKeys.SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK, ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK)) { log.info("Resetting low watermakr to {} because previous run processed no data.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else { result.put(entry.getKey(), prevActualHighWatermarksByTable.get(entry.getKey())); } } return result; }  @Override List<WorkUnit> getWorkunits(SourceState state); static Map<SourceEntity, State> getTableSpecificPropsFromState(
          Iterable<SourceEntity> entities,
          SourceState state); @Override void shutdown(SourceState state); }
QueryBasedSource extends AbstractSource<S, D> { static Map<SourceEntity, Long> getPreviousWatermarksForAllTables(SourceState state) { Map<SourceEntity, Long> result = Maps.newHashMap(); Map<SourceEntity, Long> prevLowWatermarksByTable = Maps.newHashMap(); Map<SourceEntity, Long> prevActualHighWatermarksByTable = Maps.newHashMap(); Set<SourceEntity> tablesWithFailedTasks = Sets.newHashSet(); Set<SourceEntity> tablesWithNoUpdatesOnPreviousRun = Sets.newHashSet(); boolean commitOnFullSuccess = JobCommitPolicy.getCommitPolicy(state) == JobCommitPolicy.COMMIT_ON_FULL_SUCCESS; for (WorkUnitState previousWus : state.getPreviousWorkUnitStates()) { Optional<SourceEntity> sourceEntity = SourceEntity.fromState(previousWus); if (!sourceEntity.isPresent()) { log.warn("Missing source entity for WorkUnit state: " + previousWus); continue; } SourceEntity table = sourceEntity.get(); long lowWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; LongWatermark waterMarkObj = previousWus.getWorkunit().getLowWatermark(LongWatermark.class); if(waterMarkObj != null){ lowWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)){ lowWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY)); log.warn("can not find low water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY + " low water mark " + lowWm); } if (!prevLowWatermarksByTable.containsKey(table)) { prevLowWatermarksByTable.put(table, lowWm); } else { prevLowWatermarksByTable.put(table, Math.min(prevLowWatermarksByTable.get(table), lowWm)); } long highWm = ConfigurationKeys.DEFAULT_WATERMARK_VALUE; waterMarkObj = previousWus.getActualHighWatermark(LongWatermark.class); if(waterMarkObj != null){ highWm = waterMarkObj.getValue(); } else if(previousWus.getProperties().containsKey(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)){ highWm = Long.parseLong(previousWus.getProperties().getProperty(ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK)); log.warn("can not find high water mark in json format, getting value from " + ConfigurationKeys.WORK_UNIT_STATE_RUNTIME_HIGH_WATER_MARK + " high water mark " + highWm); } if (!prevActualHighWatermarksByTable.containsKey(table)) { prevActualHighWatermarksByTable.put(table, highWm); } else { prevActualHighWatermarksByTable.put(table, Math.max(prevActualHighWatermarksByTable.get(table), highWm)); } if (commitOnFullSuccess && !isSuccessfulOrCommited(previousWus)) { tablesWithFailedTasks.add(table); } if (!isAnyDataProcessed(previousWus)) { tablesWithNoUpdatesOnPreviousRun.add(table); } } for (Map.Entry<SourceEntity, Long> entry : prevLowWatermarksByTable.entrySet()) { if (tablesWithFailedTasks.contains(entry.getKey())) { log.info("Resetting low watermark to {} because previous run failed.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else if (tablesWithNoUpdatesOnPreviousRun.contains(entry.getKey()) && state.getPropAsBoolean(ConfigurationKeys.SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK, ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_RESET_EMPTY_PARTITION_WATERMARK)) { log.info("Resetting low watermakr to {} because previous run processed no data.", entry.getValue()); result.put(entry.getKey(), entry.getValue()); } else { result.put(entry.getKey(), prevActualHighWatermarksByTable.get(entry.getKey())); } } return result; }  @Override List<WorkUnit> getWorkunits(SourceState state); static Map<SourceEntity, State> getTableSpecificPropsFromState(
          Iterable<SourceEntity> entities,
          SourceState state); @Override void shutdown(SourceState state); static final String ENTITY_BLACKLIST; static final String ENTITY_WHITELIST; static final String SOURCE_OBTAIN_TABLE_PROPS_FROM_CONFIG_STORE; static final boolean DEFAULT_SOURCE_OBTAIN_TABLE_PROPS_FROM_CONFIG_STORE; static final String WORK_UNIT_STATE_VERSION_KEY; static final Integer CURRENT_WORK_UNIT_STATE_VERSION; }