Examples with WorkUnit org.apache.gobblin.source.workunit.WorkUnit used on opensource projects

Example 86 with WorkUnit

use of org.apache.gobblin.source.workunit.WorkUnit in project incubator-gobblin by apache.

the class TaskTest method getEmptyTestTaskState.

TaskState getEmptyTestTaskState(String taskId) {
    // Create a TaskState
    WorkUnit workUnit = WorkUnit.create(new Extract(Extract.TableType.SNAPSHOT_ONLY, this.getClass().getName(), this.getClass().getSimpleName()));
    workUnit.setProp(ConfigurationKeys.TASK_KEY_KEY, "taskKey");
    TaskState taskState = new TaskState(new WorkUnitState(workUnit));
    taskState.setProp(ConfigurationKeys.METRICS_ENABLED_KEY, Boolean.toString(false));
    taskState.setTaskId(taskId);
    taskState.setJobId("1234");
    return taskState;
}

Example 87 with WorkUnit

use of org.apache.gobblin.source.workunit.WorkUnit in project incubator-gobblin by apache.

the class TestRecordStream method getEmptyTestTaskState.

TaskState getEmptyTestTaskState(String taskId) {
    // Create a TaskState
    WorkUnit workUnit = WorkUnit.create(new Extract(Extract.TableType.SNAPSHOT_ONLY, this.getClass().getName(), this.getClass().getSimpleName()));
    workUnit.setProp(ConfigurationKeys.TASK_KEY_KEY, "taskKey");
    TaskState taskState = new TaskState(new WorkUnitState(workUnit));
    taskState.setProp(ConfigurationKeys.METRICS_ENABLED_KEY, Boolean.toString(false));
    taskState.setTaskId(taskId);
    taskState.setJobId("1234");
    return taskState;
}

Example 88 with WorkUnit

use of org.apache.gobblin.source.workunit.WorkUnit in project incubator-gobblin by apache.

the class GobblinOutputCommitterTest method setupWorkUnitFiles.

@BeforeClass
public void setupWorkUnitFiles() throws IOException {
    this.conf = new Configuration();
    this.fs = FileSystem.getLocal(this.conf);
    this.stagingDirs = Lists.newArrayList();
    // Create a list of WorkUnits to serialize
    WorkUnit wu1 = createAndSetWorkUnit("wu1");
    WorkUnit wu2 = createAndSetWorkUnit("wu2");
    WorkUnit wu3 = createAndSetWorkUnit("wu3");
    WorkUnit wu4 = createAndSetWorkUnit("wu4");
    // Create a MultiWorkUnit to serialize
    MultiWorkUnit mwu1 = MultiWorkUnit.createEmpty();
    mwu1.setProp(ConfigurationKeys.TASK_ID_KEY, System.nanoTime());
    mwu1.addWorkUnits(Arrays.asList(wu3, wu4));
    Path inputDir = new Path(new Path(OUTPUT_PATH, JOB_NAME), "input");
    // Writer each WorkUnit to a separate file under inputDir
    Closer closer = Closer.create();
    try {
        wu1.write(closer.register(this.fs.create(new Path(inputDir, wu1.getProp(ConfigurationKeys.TASK_ID_KEY) + Path.SEPARATOR + "_").suffix("wu"))));
        wu2.write(closer.register(this.fs.create(new Path(inputDir, wu2.getProp(ConfigurationKeys.TASK_ID_KEY) + Path.SEPARATOR + "_").suffix("wu"))));
        mwu1.write(closer.register(this.fs.create(new Path(inputDir, mwu1.getProp(ConfigurationKeys.TASK_ID_KEY) + Path.SEPARATOR + "_").suffix("mwu"))));
    } finally {
        closer.close();
    }
}

Example 89 with WorkUnit

use of org.apache.gobblin.source.workunit.WorkUnit in project incubator-gobblin by apache.

the class JobStateTest method testSetAndGet.

@Test
public void testSetAndGet() {
    this.jobState.setId(this.jobState.getJobId());
    this.startTime = System.currentTimeMillis();
    this.jobState.setStartTime(this.startTime);
    this.jobState.setEndTime(this.startTime + 1000);
    this.jobState.setDuration(1000);
    this.jobState.setState(JobState.RunningState.COMMITTED);
    this.jobState.setTaskCount(3);
    this.jobState.setProp("foo", "bar");
    for (int i = 0; i < 3; i++) {
        WorkUnit workUnit = WorkUnit.createEmpty();
        Properties initialProps = new Properties();
        initialProps.setProperty("common1", "1");
        initialProps.setProperty("common2", "2");
        initialProps.setProperty("spec1", String.valueOf(i));
        initialProps.setProperty("spec2", "spec" + i);
        workUnit.setProps(new Properties(), initialProps);
        WorkUnitState workUnitState = new WorkUnitState(workUnit);
        workUnitState.setProp(ConfigurationKeys.JOB_ID_KEY, "TestJob-1");
        workUnitState.setProp(ConfigurationKeys.TASK_ID_KEY, "TestTask-" + i);
        workUnitState.setProp(ConfigurationKeys.DATASET_URN_KEY, "TestDataset" + i);
        TaskState taskState = new TaskState(workUnitState);
        taskState.setTaskId("TestTask-" + i);
        taskState.setId(taskState.getTaskId());
        taskState.setStartTime(this.startTime);
        taskState.setEndTime(this.startTime + 1000);
        taskState.setTaskDuration(1000);
        taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED);
        taskState.setProp("foo", "bar");
        this.jobState.addTaskState(taskState);
    }
    doAsserts(this.jobState, true, true);
}

Example 90 with WorkUnit

use of org.apache.gobblin.source.workunit.WorkUnit in project incubator-gobblin by apache.

the class LimiterStopEventTest method testGetLimiterStopMetadataCase1.

@Test
public void testGetLimiterStopMetadataCase1() throws InterruptedException {
    Properties properties = new Properties();
    String key1 = "topic";
    String key2 = "partition.id";
    String keyList = Joiner.on(',').join(key1, key2);
    String subKey1 = key2 + ".0";
    String subKey2 = key2 + ".1";
    String subKey3 = key2 + ".2";
    String subKey4 = key2 + ".3";
    String subKey5 = "partition";
    properties.setProperty(LimiterConfigurationKeys.LIMITER_REPORT_KEY_LIST, keyList);
    properties.setProperty(subKey1, "1111");
    properties.setProperty(subKey2, "1111");
    properties.setProperty(subKey3, "1111");
    properties.setProperty(subKey4, "1111");
    Extractor extractor = mock(Extractor.class);
    Limiter limiter = mock(Limiter.class);
    TaskState taskState = mock(TaskState.class);
    WorkUnit workUnit = mock(WorkUnit.class);
    Mockito.when(taskState.getWorkunit()).thenReturn(workUnit);
    Mockito.when(taskState.getJobId()).thenReturn("123");
    Mockito.when(taskState.getTaskAttemptId()).thenReturn(Optional.of("555"));
    Mockito.when(taskState.getTaskId()).thenReturn("888");
    Mockito.when(limiter.acquirePermits(1)).thenReturn(null);
    Mockito.when(taskState.getProp(ConfigurationKeys.DATASET_URN_KEY, ConfigurationKeys.DEFAULT_DATASET_URN)).thenReturn("file://xyz");
    Mockito.when(workUnit.getProperties()).thenReturn(properties);
    LimitingExtractorDecorator<String, String> decorator = new LimitingExtractorDecorator<>(extractor, limiter, taskState);
    try {
        Method method = LimitingExtractorDecorator.class.getDeclaredMethod("getLimiterStopMetadata");
        method.setAccessible(true);
        ImmutableMap<String, String> metaData = (ImmutableMap<String, String>) method.invoke(decorator);
        Assert.assertEquals(metaData.containsKey(subKey1), true);
        Assert.assertEquals(metaData.containsKey(subKey2), true);
        Assert.assertEquals(metaData.containsKey(subKey3), true);
        Assert.assertEquals(metaData.containsKey(subKey4), true);
        Assert.assertEquals(metaData.containsKey(subKey5), false);
    } catch (Exception e) {
        Assert.fail();
    }
}