Refactoring SegmentGenerationAndPushTask for extensibility
Changes the `generateTaskSpec` method to be `protected` so that it can be overridden by a subclass.
@@ -110,10 +110,21 @@ public class SegmentGenerationAndPushTaskExecutor extends BaseTaskExecutor {
SegmentGenerationAndPushResult.Builder resultBuilder = new SegmentGenerationAndPushResult.Builder();
File localTempDir = new File(new File(MinionContext.getInstance().getDataDir(), "SegmentGenerationAndPushResult"),
"tmp-" + UUID.randomUUID());
-
try {
- // Generate Pinot Segment
SegmentGenerationTaskSpec taskSpec = generateTaskSpec(taskConfigs, localTempDir);
+ return generateAndPushSegment(taskSpec, resultBuilder, taskConfigs);
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to execute SegmentGenerationAndPushTask", e);
+ } finally {
+ // Cleanup output dir
+ FileUtils.deleteQuietly(localTempDir);
+ }
+ }
+
+ private SegmentGenerationAndPushResult generateAndPushSegment(SegmentGenerationTaskSpec taskSpec,
+ SegmentGenerationAndPushResult.Builder resultBuilder,
+ Map<String, String> taskConfigs) throws Exception {
+ // Generate Pinot Segment
SegmentGenerationTaskRunner taskRunner = new SegmentGenerationTaskRunner(taskSpec);
String segmentName = taskRunner.run();
@@ -129,12 +140,7 @@ public class SegmentGenerationAndPushTaskExecutor extends BaseTaskExecutor {
// TODO: Make this use SegmentUploader
pushSegment(taskSpec.getTableConfig().getTableName(), taskConfigs, outputSegmentTarURI);
resultBuilder.setSucceed(true);
- } catch (Exception e) {
- throw new RuntimeException("Failed to execute SegmentGenerationAndPushTask", e);
- } finally {
- // Cleanup output dir
- FileUtils.deleteQuietly(localTempDir);
- }
+
return resultBuilder.build();
}
@@ -242,7 +248,7 @@ public class SegmentGenerationAndPushTaskExecutor extends BaseTaskExecutor {
return localSegmentTarFile;
}
- private SegmentGenerationTaskSpec generateTaskSpec(Map<String, String> taskConfigs, File localTempDir)
+ protected SegmentGenerationTaskSpec generateTaskSpec(Map<String, String> taskConfigs, File localTempDir)
throws Exception {
SegmentGenerationTaskSpec taskSpec = new SegmentGenerationTaskSpec();
URI inputFileURI = URI.create(taskConfigs.get(BatchConfigProperties.INPUT_DATA_FILE_URI_KEY));
