Examples with WorkUnitState org.apache.gobblin.configuration.WorkUnitState used on opensource projects

Example 16 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class CloseOnFlushWriterWrapperTest method testWriteAfterFlush.

@Test
public void testWriteAfterFlush() throws IOException {
    WorkUnitState state = new WorkUnitState();
    state.getJobState().setProp(CloseOnFlushWriterWrapper.WRITER_CLOSE_ON_FLUSH_KEY, "true");
    List<DummyWriter> dummyWriters = new ArrayList<>();
    CloseOnFlushWriterWrapper<byte[]> writer = getCloseOnFlushWriter(dummyWriters, state);
    byte[] record = new byte[] { 'a', 'b', 'c', 'd' };
    writer.writeEnvelope(new RecordEnvelope(record));
    writer.getMessageHandler().handleMessage(FlushControlMessage.builder().build());
    Assert.assertEquals(dummyWriters.size(), 1);
    Assert.assertEquals(dummyWriters.get(0).recordsWritten(), 1);
    Assert.assertEquals(dummyWriters.get(0).flushCount, 1);
    Assert.assertEquals(dummyWriters.get(0).closeCount, 1);
    Assert.assertTrue(dummyWriters.get(0).committed);
    Assert.assertEquals(dummyWriters.get(0).handlerCalled, 1);
    writer.writeEnvelope(new RecordEnvelope(record));
    writer.getMessageHandler().handleMessage(FlushControlMessage.builder().build());
    Assert.assertEquals(dummyWriters.size(), 2);
    Assert.assertEquals(dummyWriters.get(1).recordsWritten(), 1);
    Assert.assertEquals(dummyWriters.get(1).flushCount, 1);
    Assert.assertEquals(dummyWriters.get(1).closeCount, 1);
    Assert.assertTrue(dummyWriters.get(1).committed);
    Assert.assertEquals(dummyWriters.get(1).handlerCalled, 1);
}

Example 17 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class MetadataWriterWrapperTest method testAppendsMetadataWithNormalRecord.

@Test
public void testAppendsMetadataWithNormalRecord() throws IOException {
    state = new WorkUnitState();
    dummyWriter = new MetadataDummyWriter();
    writer = new MetadataWriterWrapper<>(dummyWriter, byte[].class, 1, 0, state.getJobState());
    byte[] recordBytes = new byte[] { 'a', 'b', 'c', 'd' };
    writer.write(recordBytes);
    writer.commit();
    String writerMetadata = state.getProp(ConfigurationKeys.WRITER_METADATA_KEY);
    Assert.assertNotNull(writerMetadata, "Expected there to be metadata");
    Assert.assertNotEquals(writerMetadata.indexOf("\"default-encoding\""), -1, "Expected to find default metadata in metadata");
}

Example 18 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class HiveTask method executePublishQueries.

protected void executePublishQueries(QueryBasedHivePublishEntity publishEntity) {
    Set<String> cleanUpQueries = Sets.newLinkedHashSet();
    Set<String> publishQueries = Sets.newLinkedHashSet();
    List<String> directoriesToDelete = Lists.newArrayList();
    FileSystem fs = null;
    try {
        fs = HiveSource.getSourceFs(workUnitState);
        if (publishEntity.getCleanupQueries() != null) {
            cleanUpQueries.addAll(publishEntity.getCleanupQueries());
        }
        if (publishEntity.getCleanupDirectories() != null) {
            directoriesToDelete.addAll(publishEntity.getCleanupDirectories());
        }
        if (publishEntity.getPublishDirectories() != null) {
            // Publish snapshot / partition directories
            Map<String, String> publishDirectories = publishEntity.getPublishDirectories();
            try {
                for (Map.Entry<String, String> publishDir : publishDirectories.entrySet()) {
                    HadoopUtils.renamePath(fs, new Path(publishDir.getKey()), new Path(publishDir.getValue()), true);
                }
            } catch (Throwable t) {
                throw Throwables.propagate(t);
            }
        }
        if (publishEntity.getPublishQueries() != null) {
            publishQueries.addAll(publishEntity.getPublishQueries());
        }
        WorkUnitState wus = this.workUnitState;
        this.hiveJdbcConnector.executeStatements(publishQueries.toArray(new String[publishQueries.size()]));
        wus.setWorkingState(WorkUnitState.WorkingState.COMMITTED);
        if (wus.getPropAsBoolean(USE_WATERMARKER_KEY, true)) {
            HiveSourceWatermarker watermarker = GobblinConstructorUtils.invokeConstructor(HiveSourceWatermarkerFactory.class, wus.getProp(HiveSource.HIVE_SOURCE_WATERMARKER_FACTORY_CLASS_KEY, HiveSource.DEFAULT_HIVE_SOURCE_WATERMARKER_FACTORY_CLASS)).createFromState(wus);
            watermarker.setActualHighWatermark(wus);
        }
    } catch (RuntimeException re) {
        throw re;
    } catch (Exception e) {
        log.error("Error in HiveMaterializer generate publish queries", e);
    } finally {
        try {
            this.hiveJdbcConnector.executeStatements(cleanUpQueries.toArray(new String[cleanUpQueries.size()]));
            HadoopUtils.deleteDirectories(fs, directoriesToDelete, true, true);
        } catch (RuntimeException re) {
            throw re;
        } catch (Exception e) {
            log.error("Failed to cleanup staging entities.", e);
        }
    }
}

Example 19 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class OldApiHadoopFileInputSourceTest method testGetWorkUnitsAndExtractor.

@Test
public void testGetWorkUnitsAndExtractor() throws IOException, DataRecordException {
    OldApiHadoopFileInputSource<String, Text, LongWritable, Text> fileInputSource = new TestHadoopFileInputSource();
    List<WorkUnit> workUnitList = fileInputSource.getWorkunits(this.sourceState);
    Assert.assertEquals(workUnitList.size(), 1);
    WorkUnitState workUnitState = new WorkUnitState(workUnitList.get(0));
    Closer closer = Closer.create();
    try {
        OldApiHadoopFileInputExtractor<String, Text, LongWritable, Text> extractor = (OldApiHadoopFileInputExtractor<String, Text, LongWritable, Text>) fileInputSource.getExtractor(workUnitState);
        Text text = extractor.readRecord(null);
        Assert.assertEquals(text.toString(), TEXT);
        Assert.assertNull(extractor.readRecord(null));
    } catch (Throwable t) {
        throw closer.rethrow(t);
    } finally {
        closer.close();
    }
}

Example 20 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class ConstructStateTest method test.

@Test
public void test() {
    ConstructState constructState = new ConstructState();
    WorkUnitState workUnitState = new WorkUnitState();
    String overrideKey = "overrideMe";
    String nonOverrideKey = "overrideMe.not";
    String workUnitToken = "workUnit";
    String constructToken = "construct";
    workUnitState.setProp(overrideKey, workUnitToken);
    workUnitState.setProp(nonOverrideKey, workUnitToken);
    constructState.addOverwriteProperties(ImmutableMap.<String, String>builder().put(overrideKey, constructToken).build());
    constructState.setProp(nonOverrideKey, constructToken);
    constructState.mergeIntoWorkUnitState(workUnitState);
    Assert.assertEquals(workUnitState.getProp(overrideKey), constructToken);
    Assert.assertEquals(workUnitState.getProp(nonOverrideKey), workUnitToken);
    Assert.assertEquals(workUnitState.getPropertyNames().size(), 3);
}