private static void migrateStateForJob(DatasetStateStore srcDatasetStateStore, DatasetStateStore dstDatasetStateStore, String jobName, boolean deleteFromSource) throws IOException { Map<String, JobState.DatasetState> map = srcDatasetStateStore.getLatestDatasetStatesByUrns(jobName); for (Map.Entry<String, JobState.DatasetState> entry : map.entrySet()) { dstDatasetStateStore.persistDatasetState(entry.getKey(), entry.getValue()); } if (deleteFromSource) { try { srcDatasetStateStore.delete(jobName); } catch (IOException ioe) { log.warn("The source state store has been deleted", ioe); } } }
@Override public CombinedWorkUnitAndDatasetState getCombinedWorkUnitAndDatasetState(String datasetUrn) throws Exception { Map<String, JobState.DatasetState> datasetStateMap = ImmutableMap.of(); List<WorkUnitState> workUnitStates = new ArrayList<>(); if (Strings.isNullOrEmpty(datasetUrn)) { datasetStateMap = this.datasetStateStore.getLatestDatasetStatesByUrns(this.jobName); workUnitStates = JobState.workUnitStatesFromDatasetStates(datasetStateMap.values()); } else { JobState.DatasetState datasetState = (JobState.DatasetState) this.datasetStateStore.getLatestDatasetState(this.jobName, datasetUrn); if (datasetState != null) { datasetStateMap = ImmutableMap.of(datasetUrn, datasetState); workUnitStates = JobState.workUnitStatesFromDatasetStates(Arrays.asList(datasetState)); } } return new CombinedWorkUnitAndDatasetState(workUnitStates, datasetStateMap); } }
@Override public void run(String[] args) throws Exception { CliObjectFactory<Command> factory = new ConstructorAndPublicMethodsCliObjectFactory<>(Command.class); Command command = factory.buildObject(args, 1, true, args[0]); FileSystem fs = FileSystem.get(new Configuration()); FSDataInputStream inputStream = fs.open(command.path); Config config = ConfigFactory.parseReader(new InputStreamReader(inputStream, Charset.defaultCharset())); Preconditions.checkNotNull(config.getObject(SOURCE_KEY)); Preconditions.checkNotNull(config.getObject(DESTINATION_KEY)); DatasetStateStore dstDatasetStateStore = DatasetStateStore.buildDatasetStateStore(config.getConfig(DESTINATION_KEY)); DatasetStateStore srcDatasetStateStore = DatasetStateStore.buildDatasetStateStore(config.getConfig(SOURCE_KEY)); Map<String, JobState.DatasetState> map; // if migrating state for all jobs then list the store names (job names) and copy the current jst files if (ConfigUtils.getBoolean(config, MIGRATE_ALL_JOBS, Boolean.valueOf(DEFAULT_MIGRATE_ALL_JOBS))) { List<String> jobNames = srcDatasetStateStore.getStoreNames(Predicates.alwaysTrue()); for (String jobName : jobNames) { migrateStateForJob(srcDatasetStateStore, dstDatasetStateStore, jobName, command.deleteSourceStateStore); } } else { Preconditions.checkNotNull(config.getString(JOB_NAME_KEY)); migrateStateForJob(srcDatasetStateStore, dstDatasetStateStore, config.getString(JOB_NAME_KEY), command.deleteSourceStateStore); } }
@Test(dependsOnMethods = "testGetPreviousDatasetStatesByUrns") public void testDeleteDatasetJobState() throws IOException { JobState.DatasetState datasetState = zkDatasetStateStore.get(TEST_JOB_NAME, TEST_DATASET_URN + "-" + zkDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + zkDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, TEST_DATASET_URN); Assert.assertNotNull(datasetState); Assert.assertEquals(datasetState.getJobId(), TEST_JOB_ID); zkDatasetStateStore.delete(TEST_JOB_NAME); datasetState = zkDatasetStateStore.get(TEST_JOB_NAME, TEST_DATASET_URN + "-" + zkDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + zkDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, TEST_DATASET_URN); Assert.assertNull(datasetState); }
@AfterClass public void tearDown() throws IOException { dbJobStateStore.delete(TEST_JOB_NAME); dbDatasetStateStore.delete(TEST_JOB_NAME); } }
@Override public void execute() throws IOException { log.info("Persisting dataset state for dataset " + this.datasetUrn); getDatasetStateStore().persistDatasetState(this.datasetUrn, this.datasetState); }
@Override public boolean isCompleted() throws IOException { Preconditions.checkNotNull(this.datasetState); return this.datasetState .equals(getDatasetStateStore().getLatestDatasetState(this.datasetState.getJobName(), this.datasetUrn)); }
public KafkaJobMonitor(String topic, MutableJobCatalog catalog, Config config) { super(topic, ConfigUtils.getConfigOrEmpty(config, KAFKA_JOB_MONITOR_PREFIX), 1); this.jobCatalog = catalog; try { this.datasetStateStore = DatasetStateStore.buildDatasetStateStore(config); } catch (Exception e) { log.warn("DatasetStateStore could not be created.", e); } }
this.dbDatasetStateStore.persistDatasetState(TEST_DATASET_URN1, datasetState); this.dbDatasetStateStore.persistDatasetState(TEST_DATASET_URN1, datasetState); this.dbDatasetStateStore.persistDatasetState(TEST_DATASET_URN1, datasetState); this.dbDatasetStateStore.persistDatasetState(TEST_DATASET_URN2, datasetState); datasetState.setJobName(TEST_JOB_NAME2); this.dbDatasetStateStore.persistDatasetState(TEST_DATASET_URN2, datasetState); Assert.assertTrue(this.dbDatasetStateStore.exists(TEST_JOB_NAME1, TEST_DATASET_URN1 + "-" + MysqlDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + MysqlDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX)); Assert.assertTrue(this.dbDatasetStateStore.exists(TEST_JOB_NAME1, TEST_DATASET_URN1 + "-" + getJobId(TEST_JOB_ID, 1) + MysqlDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX)); Assert.assertTrue(this.dbDatasetStateStore.exists(TEST_JOB_NAME1, TEST_DATASET_URN1 + "-" + getJobId(TEST_JOB_ID, 2) + MysqlDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX)); Assert.assertTrue(this.dbDatasetStateStore.exists(TEST_JOB_NAME1, TEST_DATASET_URN1 + "-" + getJobId(TEST_JOB_ID, 3) + MysqlDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX)); Assert.assertTrue(this.dbDatasetStateStore.exists(TEST_JOB_NAME1, TEST_DATASET_URN2 + "-" + MysqlDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + MysqlDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX));
@Test(dependsOnMethods = { "testGetDatasetState" }) public void testGetStoreNames() throws IOException { List<String> storeNames = this.zkJobStateStore.getStoreNames(Predicates.alwaysTrue()); Collections.sort(storeNames); Assert.assertTrue(storeNames.size() == 2); Assert.assertEquals(storeNames.get(0), TEST_JOB_NAME); Assert.assertEquals(storeNames.get(1), TEST_JOB_NAME2); storeNames = this.zkDatasetStateStore.getStoreNames(Predicates.alwaysTrue()); Collections.sort(storeNames); Assert.assertTrue(storeNames.size() == 2); Assert.assertEquals(storeNames.get(0), TEST_JOB_NAME); Assert.assertEquals(storeNames.get(1), TEST_JOB_NAME2); }
@Test(dependsOnMethods = "testGetStoreNames") public void testGetPreviousDatasetStatesByUrns() throws IOException { Map<String, JobState.DatasetState> datasetStatesByUrns = zkDatasetStateStore.getLatestDatasetStatesByUrns(TEST_JOB_NAME); Assert.assertEquals(datasetStatesByUrns.size(), 2); JobState.DatasetState datasetState = datasetStatesByUrns.get(TEST_DATASET_URN); Assert.assertEquals(datasetState.getDatasetUrn(), TEST_DATASET_URN); Assert.assertEquals(datasetState.getJobName(), TEST_JOB_NAME); Assert.assertEquals(datasetState.getJobId(), TEST_JOB_ID); Assert.assertEquals(datasetState.getState(), JobState.RunningState.COMMITTED); Assert.assertEquals(datasetState.getStartTime(), this.startTime); Assert.assertEquals(datasetState.getEndTime(), this.startTime + 1000); Assert.assertEquals(datasetState.getDuration(), 1000); datasetState = datasetStatesByUrns.get(TEST_DATASET_URN2); Assert.assertEquals(datasetState.getDatasetUrn(), TEST_DATASET_URN2); Assert.assertEquals(datasetState.getJobName(), TEST_JOB_NAME); Assert.assertEquals(datasetState.getJobId(), TEST_JOB_ID); Assert.assertEquals(datasetState.getState(), JobState.RunningState.COMMITTED); Assert.assertEquals(datasetState.getStartTime(), this.startTime); Assert.assertEquals(datasetState.getEndTime(), this.startTime + 1000); Assert.assertEquals(datasetState.getDuration(), 2000); }
static String buildTableName(DatasetStateStore store, String storeName, String stateId, String datasetUrn) throws IOException { return Strings.isNullOrEmpty(datasetUrn) ? stateId + DATASET_STATE_STORE_TABLE_SUFFIX : store.sanitizeDatasetStatestoreNameFromDatasetURN(storeName,datasetUrn) + "-" + stateId + DATASET_STATE_STORE_TABLE_SUFFIX; }
@Override public List<DatasetStoreDataset> findDatasets() throws IOException { List<DatasetStateStoreEntryManager> entries = this.store.getMetadataForTables(this.predicate); Map<DatasetStoreDataset.Key, List<DatasetStateStoreEntryManager>> entriesGroupedByDataset = entries.stream().collect(Collectors.groupingBy(DatasetStoreDataset.Key::new)); return entriesGroupedByDataset.entrySet().stream(). map(entry -> new DatasetStoreDataset(entry.getKey(), entry.getValue())).collect(Collectors.toList()); }
if (this.datasetStatesByUrns.isPresent()) { this.logger.info("Persisting dataset urns."); this.datasetStateStore.persistDatasetURNs(this.jobName, this.datasetStatesByUrns.get().keySet());
/** * Persist dataset state of a given dataset identified by the dataset URN. */ private void persistDatasetState(String datasetUrn, JobState.DatasetState datasetState) throws IOException { log.info("Persisting dataset state for dataset " + datasetUrn); this.jobContext.getDatasetStateStore().persistDatasetState(datasetUrn, datasetState); }
@Test(dependsOnMethods = "testPersistDatasetState") public void testGetDatasetState() throws IOException { JobState.DatasetState datasetState = zkDatasetStateStore.getLatestDatasetState(TEST_JOB_NAME, TEST_DATASET_URN); Assert.assertEquals(datasetState.getDatasetUrn(), TEST_DATASET_URN); Assert.assertEquals(datasetState.getJobName(), TEST_JOB_NAME); Assert.assertEquals(datasetState.getJobId(), TEST_JOB_ID); Assert.assertEquals(datasetState.getState(), JobState.RunningState.COMMITTED); Assert.assertEquals(datasetState.getStartTime(), this.startTime); Assert.assertEquals(datasetState.getEndTime(), this.startTime + 1000); Assert.assertEquals(datasetState.getDuration(), 1000); Assert.assertEquals(datasetState.getCompletedTasks(), 3); for (int i = 0; i < datasetState.getCompletedTasks(); i++) { TaskState taskState = datasetState.getTaskStates().get(i); Assert.assertEquals(taskState.getJobId(), TEST_JOB_ID); Assert.assertEquals(taskState.getTaskId(), TEST_TASK_ID_PREFIX + i); Assert.assertEquals(taskState.getId(), TEST_TASK_ID_PREFIX + i); Assert.assertEquals(taskState.getWorkingState(), WorkUnitState.WorkingState.COMMITTED); } }
@Test(dependsOnMethods = "testGetPreviousDatasetStatesByUrns") public void testDeleteDatasetJobState() throws IOException { JobState.DatasetState datasetState = dbDatasetStateStore.get(TEST_JOB_NAME, TEST_DATASET_URN + "-" + dbDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + dbDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, TEST_DATASET_URN); Assert.assertNotNull(datasetState); Assert.assertEquals(datasetState.getJobId(), TEST_JOB_ID); dbDatasetStateStore.delete(TEST_JOB_NAME); datasetState = dbDatasetStateStore.get(TEST_JOB_NAME, TEST_DATASET_URN + "-" + dbDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + dbDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, TEST_DATASET_URN); Assert.assertNull(datasetState); }
public JobStateToJsonConverter(Properties props, String storeUrl, boolean keepConfig) throws IOException { Configuration conf = new Configuration(); JobConfigurationUtils.putPropertiesIntoConfiguration(props, conf); if (StringUtils.isNotBlank(storeUrl)) { props.setProperty(ConfigurationKeys.STATE_STORE_ROOT_DIR_KEY, storeUrl); } this.keepConfig = keepConfig; this.jobStateStore = (StateStore) DatasetStateStore.buildDatasetStateStore(ConfigUtils.propertiesToConfig(props)); }
@Test(dependsOnMethods = { "testGetDatasetState" }) public void testGetStoreNames() throws IOException { List<String> storeNames = this.dbJobStateStore.getStoreNames(Predicates.alwaysTrue()); Collections.sort(storeNames); Assert.assertTrue(storeNames.size() == 3); Assert.assertEquals(storeNames.get(0), TEST_JOB_NAME); Assert.assertEquals(storeNames.get(1), TEST_JOB_NAME2); Assert.assertEquals(storeNames.get(2), TEST_JOB_NAME_LOWER); storeNames = this.dbDatasetStateStore.getStoreNames(Predicates.alwaysTrue()); Collections.sort(storeNames); Assert.assertTrue(storeNames.size() == 2); Assert.assertEquals(storeNames.get(0), TEST_JOB_NAME); Assert.assertEquals(storeNames.get(1), TEST_JOB_NAME2); }