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

Example 86 with WorkUnitState

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

the class EnvelopePayloadConverterTest method testConverter.

@Test
public void testConverter() throws IOException, DataConversionException, SchemaRegistryException {
    Schema inputSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/envelope.avsc"));
    GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(inputSchema);
    File tmp = File.createTempFile(getClass().getSimpleName(), null);
    FileUtils.copyInputStreamToFile(getClass().getResourceAsStream("/converter/envelope.avro"), tmp);
    DataFileReader<GenericRecord> dataFileReader = new DataFileReader<>(tmp, datumReader);
    GenericRecord inputRecord = dataFileReader.next();
    Schema latestPayloadSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/record.avsc"));
    when(mockRegistry.getLatestSchemaByTopic(any())).thenReturn(latestPayloadSchema);
    when(mockRegistry.getSchemaByKey(any())).thenReturn(inputSchema.getField("nestedRecord").schema());
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.setProp(BaseEnvelopeSchemaConverter.PAYLOAD_SCHEMA_TOPIC, "test");
    workUnitState.setProp(BaseEnvelopeSchemaConverter.PAYLOAD_SCHEMA_ID_FIELD, "metadata.payloadSchemaId");
    workUnitState.setProp(BaseEnvelopeSchemaConverter.KAFKA_REGISTRY_FACTORY, MockKafkaAvroSchemaRegistryFactory.class.getName());
    EnvelopePayloadConverter converter = new EnvelopePayloadConverter();
    converter.init(workUnitState);
    Schema outputSchema = converter.convertSchema(inputSchema, workUnitState);
    List<GenericRecord> outputRecords = new ArrayList<>();
    Iterables.addAll(outputRecords, converter.convertRecord(outputSchema, inputRecord, workUnitState));
    Assert.assertTrue(outputRecords.size() == 1);
    GenericRecord outputRecord = outputRecords.get(0);
    GenericRecord payload = (GenericRecord) outputRecord.get("payload");
    // While making the test envelope avro record, its nestedRecord was intentionally set to the deserialized payload
    GenericRecord expectedPayload = (GenericRecord) outputRecord.get("nestedRecord");
    Schema payloadSchema = payload.getSchema();
    Schema expectedPayloadSchema = expectedPayload.getSchema();
    // The expected payload schema has the same number of fields as payload schema but in different order
    Assert.assertTrue(expectedPayloadSchema.getName().equals(payloadSchema.getName()));
    Assert.assertTrue(expectedPayloadSchema.getNamespace().equals(payloadSchema.getNamespace()));
    Assert.assertTrue(expectedPayloadSchema.getFields().size() == payloadSchema.getFields().size());
    for (Schema.Field field : payload.getSchema().getFields()) {
        Assert.assertTrue(expectedPayload.get(field.name()).equals(payload.get(field.name())));
    }
}

Example 87 with WorkUnitState

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

the class EnvelopePayloadExtractingConverterTest method testConverter.

@Test
public void testConverter() throws Exception {
    Schema inputSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/envelope.avsc"));
    GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(inputSchema);
    File tmp = File.createTempFile(getClass().getSimpleName(), null);
    FileUtils.copyInputStreamToFile(getClass().getResourceAsStream("/converter/envelope.avro"), tmp);
    DataFileReader<GenericRecord> dataFileReader = new DataFileReader<>(tmp, datumReader);
    GenericRecord inputRecord = dataFileReader.next();
    Schema latestPayloadSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/record.avsc"));
    when(mockRegistry.getLatestSchemaByTopic(any())).thenReturn(latestPayloadSchema);
    when(mockRegistry.getSchemaByKey(any())).thenReturn(inputSchema.getField("nestedRecord").schema());
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.setProp(BaseEnvelopeSchemaConverter.PAYLOAD_SCHEMA_TOPIC, "test");
    workUnitState.setProp(BaseEnvelopeSchemaConverter.PAYLOAD_SCHEMA_ID_FIELD, "metadata.payloadSchemaId");
    workUnitState.setProp(BaseEnvelopeSchemaConverter.KAFKA_REGISTRY_FACTORY, EnvelopePayloadExtractingConverterTest.MockKafkaAvroSchemaRegistryFactory.class.getName());
    EnvelopePayloadExtractingConverter converter = new EnvelopePayloadExtractingConverter();
    converter.init(workUnitState);
    Schema outputSchema = converter.convertSchema(inputSchema, workUnitState);
    Assert.assertTrue(outputSchema.equals(latestPayloadSchema));
    List<GenericRecord> outputRecords = new ArrayList<>();
    Iterables.addAll(outputRecords, converter.convertRecord(outputSchema, inputRecord, workUnitState));
    Assert.assertTrue(outputRecords.size() == 1);
    GenericRecord payload = outputRecords.get(0);
    // While making the test envelope avro input record, its nestedRecord was intentionally set to the deserialized payload
    GenericRecord expectedPayload = (GenericRecord) inputRecord.get("nestedRecord");
    Schema payloadSchema = payload.getSchema();
    Schema expectedPayloadSchema = expectedPayload.getSchema();
    // The expected payload schema has the same number of fields as payload schema but in different order
    Assert.assertTrue(expectedPayloadSchema.getName().equals(payloadSchema.getName()));
    Assert.assertTrue(expectedPayloadSchema.getNamespace().equals(payloadSchema.getNamespace()));
    Assert.assertTrue(expectedPayloadSchema.getFields().size() == payloadSchema.getFields().size());
    for (Schema.Field field : payload.getSchema().getFields()) {
        Assert.assertTrue(expectedPayload.get(field.name()).equals(payload.get(field.name())));
    }
}

Example 88 with WorkUnitState

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

the class MetadataConverterWrapperTest method testConvertsMetadataMultiOutput.

@Test
public void testConvertsMetadataMultiOutput() throws DataConversionException {
    final int numRecordsToReturn = 2;
    DummyConverter dummyConverter = new DummyConverter(numRecordsToReturn);
    MetadataConverterWrapper<String, String, String, String> wrapper = new MetadataConverterWrapper<>(dummyConverter);
    Iterable<RecordWithMetadata<String>> records1 = wrapper.convertRecord("foo", new RecordWithMetadata<String>("bar", buildMetadata(1)), new WorkUnitState());
    Iterable<RecordWithMetadata<String>> records2 = wrapper.convertRecord("foo", new RecordWithMetadata<String>("baz", buildMetadata(2)), new WorkUnitState());
    Iterator<RecordWithMetadata<String>> record1It = records1.iterator();
    Iterator<RecordWithMetadata<String>> record2It = records2.iterator();
    for (int i = 0; i < numRecordsToReturn; i++) {
        RecordWithMetadata<String> record1 = record1It.next();
        Assert.assertEquals(record1.getRecord(), "converted" + String.valueOf(i));
        Assert.assertEquals(record1.getMetadata().getGlobalMetadata().getDatasetUrn(), "dataset-id:1");
        RecordWithMetadata<String> record2 = record2It.next();
        Assert.assertEquals(record2.getRecord(), "converted" + String.valueOf(i));
        Assert.assertEquals(record2.getMetadata().getGlobalMetadata().getDatasetUrn(), "dataset-id:2");
    }
}

Example 89 with WorkUnitState

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

the class AvroToJdbcEntryConverterTest method testFlattening.

@Test
public void testFlattening() throws IOException, SchemaConversionException, SQLException, URISyntaxException, DataConversionException {
    final String db = "db";
    final String table = "users";
    Map<String, JdbcType> dateColums = new HashMap<>();
    dateColums.put("date_of_birth", JdbcType.DATE);
    dateColums.put("last_modified", JdbcType.TIME);
    dateColums.put("created", JdbcType.TIMESTAMP);
    JdbcWriterCommands mockWriterCommands = mock(JdbcWriterCommands.class);
    when(mockWriterCommands.retrieveDateColumns(db, table)).thenReturn(dateColums);
    JdbcWriterCommandsFactory factory = mock(JdbcWriterCommandsFactory.class);
    when(factory.newInstance(any(State.class), any(Connection.class))).thenReturn(mockWriterCommands);
    List<JdbcEntryMetaDatum> jdbcEntryMetaData = new ArrayList<>();
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("name", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("favorite_number", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("favorite_color", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("date_of_birth", JdbcType.DATE));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("last_modified", JdbcType.TIME));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("created", JdbcType.TIMESTAMP));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("nested1_nested1_string", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("nested1_nested1_int", JdbcType.INTEGER));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("nested1_nested2_union_nested2_string", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("nested1_nested2_union_nested2_int", JdbcType.INTEGER));
    JdbcEntrySchema expected = new JdbcEntrySchema(jdbcEntryMetaData);
    Schema inputSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/pickfields_nested_with_union.avsc"));
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.appendToListProp(JdbcPublisher.JDBC_PUBLISHER_FINAL_TABLE_NAME, table);
    AvroToJdbcEntryConverter converter = new AvroToJdbcEntryConverter(workUnitState);
    Map<String, JdbcType> dateColumnMapping = Maps.newHashMap();
    dateColumnMapping.put("date_of_birth", JdbcType.DATE);
    dateColumnMapping.put("last_modified", JdbcType.TIME);
    dateColumnMapping.put("created", JdbcType.TIMESTAMP);
    workUnitState.appendToListProp(AvroToJdbcEntryConverter.CONVERTER_AVRO_JDBC_DATE_FIELDS, new Gson().toJson(dateColumnMapping));
    JdbcEntrySchema actualSchema = converter.convertSchema(inputSchema, workUnitState);
    Assert.assertEquals(expected, actualSchema);
    try (DataFileReader<GenericRecord> srcDataFileReader = new DataFileReader<GenericRecord>(new File(getClass().getResource("/converter/pickfields_nested_with_union.avro").toURI()), new GenericDatumReader<GenericRecord>(inputSchema))) {
        List<JdbcEntryData> entries = new ArrayList<>();
        while (srcDataFileReader.hasNext()) {
            JdbcEntryData actualData = converter.convertRecord(actualSchema, srcDataFileReader.next(), workUnitState).iterator().next();
            entries.add(actualData);
        }
        final JsonSerializer<JdbcEntryDatum> datumSer = new JsonSerializer<JdbcEntryDatum>() {

            @Override
            public JsonElement serialize(JdbcEntryDatum datum, Type typeOfSrc, JsonSerializationContext context) {
                JsonObject jso = new JsonObject();
                if (datum.getVal() == null) {
                    jso.add(datum.getColumnName(), null);
                    return jso;
                }
                if (datum.getVal() instanceof Date) {
                    jso.addProperty(datum.getColumnName(), ((Date) datum.getVal()).getTime());
                } else if (datum.getVal() instanceof Timestamp) {
                    jso.addProperty(datum.getColumnName(), ((Timestamp) datum.getVal()).getTime());
                } else if (datum.getVal() instanceof Time) {
                    jso.addProperty(datum.getColumnName(), ((Time) datum.getVal()).getTime());
                } else {
                    jso.addProperty(datum.getColumnName(), datum.getVal().toString());
                }
                return jso;
            }
        };
        JsonSerializer<JdbcEntryData> serializer = new JsonSerializer<JdbcEntryData>() {

            @Override
            public JsonElement serialize(JdbcEntryData src, Type typeOfSrc, JsonSerializationContext context) {
                JsonArray arr = new JsonArray();
                for (JdbcEntryDatum datum : src) {
                    arr.add(datumSer.serialize(datum, datum.getClass(), context));
                }
                return arr;
            }
        };
        Gson gson = new GsonBuilder().registerTypeAdapter(JdbcEntryData.class, serializer).serializeNulls().create();
        JsonElement actualSerialized = gson.toJsonTree(entries);
        JsonElement expectedSerialized = new JsonParser().parse(new InputStreamReader(getClass().getResourceAsStream("/converter/pickfields_nested_with_union.json")));
        Assert.assertEquals(actualSerialized, expectedSerialized);
    }
    converter.close();
}

Example 90 with WorkUnitState

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

the class AvroToJdbcEntryConverterTest method testDateConversion.

@Test
public void testDateConversion() throws IOException, SchemaConversionException, SQLException {
    final String db = "db";
    final String table = "users";
    Map<String, JdbcType> dateColums = new HashMap<>();
    dateColums.put("date_of_birth", JdbcType.DATE);
    dateColums.put("last_modified", JdbcType.TIME);
    dateColums.put("created", JdbcType.TIMESTAMP);
    JdbcWriterCommands mockWriterCommands = mock(JdbcWriterCommands.class);
    when(mockWriterCommands.retrieveDateColumns(db, table)).thenReturn(dateColums);
    JdbcWriterCommandsFactory factory = mock(JdbcWriterCommandsFactory.class);
    when(factory.newInstance(any(State.class), any(Connection.class))).thenReturn(mockWriterCommands);
    List<JdbcEntryMetaDatum> jdbcEntryMetaData = new ArrayList<>();
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("name", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("favorite_number", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("favorite_color", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("date_of_birth", JdbcType.DATE));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("last_modified", JdbcType.TIME));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("created", JdbcType.TIMESTAMP));
    JdbcEntrySchema expected = new JdbcEntrySchema(jdbcEntryMetaData);
    Schema inputSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/fieldPickInput.avsc"));
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.appendToListProp(JdbcPublisher.JDBC_PUBLISHER_FINAL_TABLE_NAME, table);
    AvroToJdbcEntryConverter converter = new AvroToJdbcEntryConverter(workUnitState);
    Map<String, JdbcType> dateColumnMapping = Maps.newHashMap();
    dateColumnMapping.put("date_of_birth", JdbcType.DATE);
    dateColumnMapping.put("last_modified", JdbcType.TIME);
    dateColumnMapping.put("created", JdbcType.TIMESTAMP);
    workUnitState.appendToListProp(AvroToJdbcEntryConverter.CONVERTER_AVRO_JDBC_DATE_FIELDS, new Gson().toJson(dateColumnMapping));
    JdbcEntrySchema actual = converter.convertSchema(inputSchema, workUnitState);
    Assert.assertEquals(expected, actual);
}