[GitHub] nifi pull request #2619: NIFI-5059 Updated MongoDBLookupService to be able t...

[MikeThomsen <gi...@git.apache.org>]

Github user MikeThomsen commented on a diff in the pull request:

    https://github.com/apache/nifi/pull/2619#discussion_r189434124
  
    --- Diff: nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-services/src/main/java/org/apache/nifi/mongodb/MongoDBLookupService.java ---
    @@ -52,68 +54,125 @@
         "The query is limited to the first result (findOne in the Mongo documentation). If no \"Lookup Value Field\" is specified " +
         "then the entire MongoDB result document minus the _id field will be returned as a record."
     )
    -public class MongoDBLookupService extends MongoDBControllerService implements LookupService<Object> {
    +public class MongoDBLookupService extends SchemaRegistryService implements LookupService<Object> {
    +    public static final PropertyDescriptor CONTROLLER_SERVICE = new PropertyDescriptor.Builder()
    +        .name("mongo-lookup-client-service")
    +        .displayName("Client Service")
    +        .description("A MongoDB controller service to use with this lookup service.")
    +        .required(true)
    +        .identifiesControllerService(MongoDBControllerService.class)
    +        .build();
     
         public static final PropertyDescriptor LOOKUP_VALUE_FIELD = new PropertyDescriptor.Builder()
    -            .name("mongo-lookup-value-field")
    -            .displayName("Lookup Value Field")
    -            .description("The field whose value will be returned when the lookup key(s) match a record. If not specified then the entire " +
    -                    "MongoDB result document minus the _id field will be returned as a record.")
    -            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
    -            .required(false)
    -            .build();
    +        .name("mongo-lookup-value-field")
    +        .displayName("Lookup Value Field")
    +        .description("The field whose value will be returned when the lookup key(s) match a record. If not specified then the entire " +
    +                "MongoDB result document minus the _id field will be returned as a record.")
    +        .addValidator(Validator.VALID)
    +        .required(false)
    +        .build();
    +    public static final PropertyDescriptor PROJECTION = new PropertyDescriptor.Builder()
    +        .name("mongo-lookup-projection")
    +        .displayName("Projection")
    +        .description("Specifies a projection for limiting which fields will be returned.")
    +        .required(false)
    +        .build();
     
         private String lookupValueField;
     
    -    private static final List<PropertyDescriptor> lookupDescriptors;
    -
    -    static {
    -        lookupDescriptors = new ArrayList<>();
    -        lookupDescriptors.addAll(descriptors);
    -        lookupDescriptors.add(LOOKUP_VALUE_FIELD);
    -    }
    -
         @Override
         public Optional<Object> lookup(Map<String, Object> coordinates) throws LookupFailureException {
    -        Map<String, Object> clean = new HashMap<>();
    -        clean.putAll(coordinates);
    +        Map<String, Object> clean = coordinates.entrySet().stream()
    +            .filter(e -> !schemaNameProperty.equals(String.format("${%s}", e.getKey())))
    +            .collect(Collectors.toMap(
    +                e -> e.getKey(),
    +                e -> e.getValue()
    +            ));
             Document query = new Document(clean);
     
             if (coordinates.size() == 0) {
                 throw new LookupFailureException("No keys were configured. Mongo query would return random documents.");
             }
     
             try {
    -            Document result = this.findOne(query);
    +            Document result = projection != null ? controllerService.findOne(query, projection) : controllerService.findOne(query);
     
                 if(result == null) {
                     return Optional.empty();
                 } else if (!StringUtils.isEmpty(lookupValueField)) {
                     return Optional.ofNullable(result.get(lookupValueField));
                 } else {
    -                final List<RecordField> fields = new ArrayList<>();
    +                RecordSchema schema = loadSchema(coordinates);
     
    -                for (String key : result.keySet()) {
    -                    if (key.equals("_id")) {
    -                        continue;
    -                    }
    -                    fields.add(new RecordField(key, RecordFieldType.STRING.getDataType()));
    -                }
    -
    -                final RecordSchema schema = new SimpleRecordSchema(fields);
    -                return Optional.ofNullable(new MapRecord(schema, result));
    +                RecordSchema toUse = schema != null ? schema : convertSchema(result);
    +                return Optional.ofNullable(new MapRecord(toUse, result));
                 }
             } catch (Exception ex) {
                 getLogger().error("Error during lookup {}", new Object[]{ query.toJson() }, ex);
                 throw new LookupFailureException(ex);
             }
         }
     
    -    @Override
    +    private RecordSchema loadSchema(Map<String, Object> coordinates) {
    +        Map<String, String> variables = coordinates.entrySet().stream()
    +            .collect(Collectors.toMap(
    +                e -> e.getKey(),
    +                e -> e.getValue().toString()
    +            ));
    +        try {
    +            return getSchema(variables, null);
    +        } catch (Exception ex) {
    +            return null;
    +        }
    +    }
    +
    +    private RecordSchema convertSchema(Map<String, Object> result) {
    +        List<RecordField> fields = new ArrayList<>();
    +        for (Map.Entry<String, Object> entry : result.entrySet()) {
    +
    +            RecordField field;
    +            if (entry.getValue() instanceof Integer) {
    --- End diff --
    
    At some point once this and the ES Lookup Service are merged, I'll refactor this into a helper method.


---