private String readProp(String key, WorkUnitState workUnitState) { String value = workUnitState.getWorkunit().getProp(key); if (StringUtils.isBlank(value)) { value = workUnitState.getProp(key); } if (StringUtils.isBlank(value)) { value = workUnitState.getJobState().getProp(key); } return value; }
/** * Copy WorkUnitState so that work unit also contains job state. FileBasedExtractor needs properties from job state (mostly source.* properties), * where it has been already removed when reached here. * * @param src * @return */ private WorkUnitState copyOf(WorkUnitState src) { WorkUnit copiedWorkUnit = WorkUnit.copyOf(src.getWorkunit()); copiedWorkUnit.addAllIfNotExist(src.getJobState()); WorkUnitState workUnitState = new WorkUnitState(copiedWorkUnit, src.getJobState()); workUnitState.addAll(src); return workUnitState; }
private CloseOnFlushWriterWrapper getCloseOnFlushWriter(List<DummyWriter> dummyWriters, WorkUnitState state) { return new CloseOnFlushWriterWrapper<>(new Supplier<DataWriter<byte[]>>() { @Override public DataWriter<byte[]> get() { DummyWriter writer = new DummyWriter(); dummyWriters.add(writer); return writer; } }, state.getJobState()); }
public ImmutableWorkUnitState(WorkUnitState workUnitState) { super(workUnitState.getWorkunit(), workUnitState.getJobState()); super.addAll(workUnitState.getSpecProperties()); }
public HiveBaseExtractor(WorkUnitState state) throws IOException { if (Boolean.valueOf(state.getPropAsBoolean(PartitionLevelWatermarker.IS_WATERMARK_WORKUNIT_KEY))) { return; } this.hiveWorkUnit = new HiveWorkUnit(state.getWorkunit()); this.hiveDataset = hiveWorkUnit.getHiveDataset(); this.dbName = hiveDataset.getDbAndTable().getDb(); this.tableName = hiveDataset.getDbAndTable().getTable(); this.pool = HiveMetastoreClientPool.get(state.getJobState().getProperties(), Optional.fromNullable(state.getJobState().getProp(HiveDatasetFinder.HIVE_METASTORE_URI_KEY))); }
@Test public void testConverter() { WorkUnitState wuState = new WorkUnitState(); wuState.getJobState().setProp( EncryptionConfigParser.CONVERTER_ENCRYPT_PREFIX + "." + EncryptionConfigParser.ENCRYPTION_ALGORITHM_KEY, "any"); wuState.getJobState().setProp( EncryptionConfigParser.CONVERTER_ENCRYPT_PREFIX + "." + EncryptionConfigParser.ENCRYPTION_KEYSTORE_PATH_KEY, "/tmp/foobar"); wuState.getJobState().setProp( EncryptionConfigParser.CONVERTER_ENCRYPT_PREFIX + "." + EncryptionConfigParser.ENCRYPTION_KEYSTORE_PASSWORD_KEY, "abracadabra"); wuState.getJobState().setProp( EncryptionConfigParser.CONVERTER_ENCRYPT_PREFIX + "." + EncryptionConfigParser.ENCRYPTION_KEY_NAME, "keyname"); wuState.setProp(EncryptionConfigParser.CONVERTER_ENCRYPT_PREFIX + "abc.def", "foobar"); Map<String, Object> parsedProperties = EncryptionConfigParser.getConfigForBranch(EncryptionConfigParser.EntityType.CONVERTER_ENCRYPT, wuState); Assert.assertNotNull(parsedProperties, "Expected parser to only return one record"); Assert.assertEquals(parsedProperties.size(), 4, "Did not expect abc.def to be picked up in config"); Map<String, Object> parsedWriterProperties = EncryptionConfigParser.getConfigForBranch(EncryptionConfigParser.EntityType.WRITER, wuState); Assert.assertNull(parsedWriterProperties, "Did not expect to find writer properties"); }
@Test public void worksNoFork() throws DataConversionException { workUnitState.getJobState() .setProp(DECRYPT_PREFIX + EncryptionConfigParser.ENCRYPTION_ALGORITHM_KEY, "insecure_shift"); converter.init(workUnitState); Iterable<RecordWithMetadata<byte[]>> records = converter.convertRecord("", sampleRecord, workUnitState); Iterator<RecordWithMetadata<byte[]>> recordIterator = records.iterator(); Assert.assertTrue(recordIterator.hasNext()); RecordWithMetadata<byte[]> record = recordIterator.next(); Assert.assertFalse(recordIterator.hasNext()); Assert.assertEquals(record.getMetadata().getGlobalMetadata().getTransferEncoding().get(0), insecureShiftTag); Assert.assertEquals(record.getRecord(), shiftedValue); } }
@BeforeMethod public void setUp() { state = new WorkUnitState(); dummyWriter = new DummyWriter(); writer = new MetadataWriterWrapper<>(dummyWriter, byte[].class, 1, 0, state.getJobState()); }
public TaskState(WorkUnitState workUnitState) { // Since getWorkunit() returns an immutable WorkUnit object, // the WorkUnit object in this object is also immutable. super(workUnitState.getWorkunit(), workUnitState.getJobState(), workUnitState.getTaskBrokerNullable()); addAll(workUnitState); this.jobId = workUnitState.getProp(ConfigurationKeys.JOB_ID_KEY); this.taskId = workUnitState.getProp(ConfigurationKeys.TASK_ID_KEY); this.taskKey = workUnitState.getProp(ConfigurationKeys.TASK_KEY_KEY, "unknown_task_key"); this.taskAttemptId = Optional.fromNullable(workUnitState.getProp(ConfigurationKeys.TASK_ATTEMPT_ID_KEY)); this.setId(this.taskId); }
@Test public void worksNoFork() throws DataConversionException { workUnitState.getJobState() .setProp(ENCRYPT_PREFIX + EncryptionConfigParser.ENCRYPTION_ALGORITHM_KEY, "insecure_shift"); converter.init(workUnitState); Iterable<RecordWithMetadata<byte[]>> records = converter.convertRecord("", sampleRecord, workUnitState); Iterator<RecordWithMetadata<byte[]>> recordIt = records.iterator(); Assert.assertTrue(recordIt.hasNext()); RecordWithMetadata<byte[]> record = recordIt.next(); Assert.assertFalse(recordIt.hasNext()); Assert.assertEquals(record.getMetadata().getGlobalMetadata().getTransferEncoding().get(0), insecureShiftTag); Assert.assertEquals(record.getRecord(), shiftedValue); } }
@Test public void worksWithFork() throws DataConversionException { workUnitState.setProp(ConfigurationKeys.FORK_BRANCH_ID_KEY, 2); workUnitState.getJobState() .setProp(ENCRYPT_PREFIX + EncryptionConfigParser.ENCRYPTION_ALGORITHM_KEY + ".2", "insecure_shift"); converter.init(workUnitState); Iterable<RecordWithMetadata<byte[]>> records = converter.convertRecord("", sampleRecord, workUnitState); Iterator<RecordWithMetadata<byte[]>> recordIt = records.iterator(); Assert.assertTrue(recordIt.hasNext()); RecordWithMetadata<byte[]> record = recordIt.next(); Assert.assertFalse(recordIt.hasNext()); Assert.assertEquals(record.getMetadata().getGlobalMetadata().getTransferEncoding().get(0), insecureShiftTag); Assert.assertEquals(record.getRecord(), shiftedValue); }
@Test public void worksWithFork() throws DataConversionException { workUnitState.setProp(ConfigurationKeys.FORK_BRANCH_ID_KEY, 2); workUnitState.getJobState() .setProp(DECRYPT_PREFIX + EncryptionConfigParser.ENCRYPTION_ALGORITHM_KEY + ".2", "insecure_shift"); converter.init(workUnitState); Iterable<RecordWithMetadata<byte[]>> records = converter.convertRecord("", sampleRecord, workUnitState); Iterator<RecordWithMetadata<byte[]>> recordIterator = records.iterator(); Assert.assertTrue(recordIterator.hasNext()); RecordWithMetadata<byte[]> record = recordIterator.next(); Assert.assertFalse(recordIterator.hasNext()); Assert.assertEquals(record.getMetadata().getGlobalMetadata().getTransferEncoding().get(0), insecureShiftTag); Assert.assertEquals(record.getRecord(), shiftedValue); }
@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"); }
@Test public void testCloseOnFlushEnabled() 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.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); }
@Test public void testConversion() throws DataConversionException, IOException, SchemaConversionException { AvroStringFieldDecryptorConverter converter = new AvroStringFieldDecryptorConverter(); WorkUnitState wuState = new WorkUnitState(); wuState.getJobState().setProp("converter.fieldsToDecrypt", "field1"); wuState.getJobState().setProp("converter.decrypt.AvroStringFieldDecryptorConverter.algorithm", "insecure_shift"); converter.init(wuState); GenericRecord inputRecord = TestUtils.generateRandomAvroRecord(); Schema inputSchema = inputRecord.getSchema(); Schema outputSchema = converter.convertSchema(inputSchema, wuState); String fieldValue = (String) inputRecord.get("field1"); Iterable<GenericRecord> recordIt = converter.convertRecord(outputSchema, inputRecord, wuState); GenericRecord decryptedRecord = recordIt.iterator().next(); Assert.assertEquals(outputSchema, inputSchema); String decryptedValue = (String) decryptedRecord.get("field1"); InsecureShiftCodec codec = new InsecureShiftCodec(Maps.<String, Object>newHashMap()); InputStream in = codec.decodeInputStream(new ByteArrayInputStream(fieldValue.getBytes(StandardCharsets.UTF_8))); byte[] expectedDecryptedValue = new byte[in.available()]; in.read(expectedDecryptedValue); Assert.assertEquals(new String(expectedDecryptedValue, StandardCharsets.UTF_8), decryptedValue); }
@Test public void testCloseAfterFlush() 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.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.close(); // writer should not be closed multiple times Assert.assertEquals(dummyWriters.get(0).closeCount, 1); }
WorkUnitState wuState = new WorkUnitState(); wuState.getJobState().setProp("converter.fieldsToEncrypt", "field1"); wuState.getJobState().setProp("converter.encrypt.algorithm", "insecure_shift");
public HiveMaterializerQueryGenerator(WorkUnitState workUnitState) throws IOException { this.fs = HiveSource.getSourceFs(workUnitState); this.pool = HiveMetastoreClientPool.get(workUnitState.getJobState().getProperties(), Optional.fromNullable(workUnitState.getJobState().getProp(HiveDatasetFinder.HIVE_METASTORE_URI_KEY))); this.workUnitState = workUnitState; this.workUnit = new HiveWorkUnit(workUnitState.getWorkunit()); this.outputTableMetadata = HiveMaterializer.parseStageableTableMetadata(this.workUnit); this.outputDatabaseName = outputTableMetadata.getDestinationDbName(); this.outputTableName = outputTableMetadata.getDestinationTableName(); this.outputDataLocation = HiveConverterUtils.getOutputDataLocation(outputTableMetadata.getDestinationDataPath()); this.destinationTableMeta = HiveConverterUtils.getDestinationTableMeta(this.outputTableMetadata.getDestinationDbName(), this.outputTableMetadata.getDestinationTableName(), workUnitState.getProperties()).getLeft(); this.stagingTableName = HiveConverterUtils.getStagingTableName(this.outputTableMetadata.getDestinationStagingTableName()); this.stagingDataLocation = HiveConverterUtils.getStagingDataLocation(this.outputTableMetadata.getDestinationDataPath(), this.stagingTableName); }
@Test public void testAppendsDefaultMetadata() 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'}; Metadata md = new Metadata(); md.getGlobalMetadata().addTransferEncoding("first"); writer.write(new RecordWithMetadata<>(recordBytes, md)); writer.commit(); String writerMetadata = state.getProp(ConfigurationKeys.WRITER_METADATA_KEY); Assert.assertNotNull(writerMetadata, "Expected there to be metadata"); int firstOccurrence = writerMetadata.indexOf("\"first\""); Assert.assertNotEquals(firstOccurrence, -1, "Expected to find record-level encoding in metadata"); int secondOccurrence = writerMetadata.indexOf("\"default-encoding\""); Assert.assertNotEquals(secondOccurrence, -1, "Expected to find default metadata in metadata"); Assert.assertTrue(firstOccurrence < secondOccurrence, "Expected recordBytes encoding to be present before default encoding"); }
@Test @SuppressWarnings("unchecked") public void testArrayDecryption() throws DataConversionException, IOException, SchemaConversionException { AvroStringFieldDecryptorConverter converter = new AvroStringFieldDecryptorConverter(); WorkUnitState wuState = new WorkUnitState(); wuState.getJobState().setProp("converter.fieldsToDecrypt", "array1"); wuState.getJobState().setProp("converter.decrypt.AvroStringFieldDecryptorConverter.algorithm", "insecure_shift"); converter.init(wuState); GenericRecord inputRecord = generateRecordWithArrays(); Schema inputSchema = inputRecord.getSchema(); Schema outputSchema = converter.convertSchema(inputSchema, wuState); GenericData.Array<String> fieldValue = (GenericData.Array<String>) inputRecord.get("array1"); Iterable<GenericRecord> recordIt = converter.convertRecord(outputSchema, inputRecord, wuState); GenericRecord decryptedRecord = recordIt.iterator().next(); Assert.assertEquals(outputSchema, inputSchema); GenericData.Array<String> decryptedValue = (GenericData.Array<String>) decryptedRecord.get("array1"); for (int i = 0; i < decryptedValue.size(); i++) { assertDecryptedValuesEqual(decryptedValue.get(i), fieldValue.get(i)); } }