Database Sequence Lookup with ScriptedLookupService

While addressing this question, I realized it would make a good example of how to use ScriptedLookupService to provide a sequence number from a database to a field in a record. This can be useful when using PutDatabaseRecord downstream, when you want some field to contain a sequence number from a database.

In my example I'm using a CSV file with a few dummy values in it, adding an "id" field containing a sequence number, and writing out the records as JSON objects:

The LookupRecord configuration is as follows, note that the processor requires the value of any user-defined property to be a RecordPath that evaluates to a non-null value. I chose "/Age" as a dummy RecordPath here, since I know the record contains a non-null Age field:

The Groovy script for ScriptedLookupService is as follows, I am connecting to a PostgreSQL database so I generated the corresponding SQL using nextval(), you can change this to the SQL appropriate for your database type:

import org.apache.nifi.controller.ControllerServiceInitializationContext
import org.apache.nifi.reporting.InitializationException
import org.apache.nifi.dbcp.DBCPService
import java.sql.*

class SequenceLookupService implements LookupService<String> {

final String ID = UUID.randomUUID().toString()

public static final PropertyDescriptor DBCP_SERVICE = new PropertyDescriptor.Builder()
            .name("Database Connection Pooling Service")
            .description("The Controller Service that is used to obtain connection to database")
            .required(true)
            .identifiesControllerService(DBCPService)
            .build()

Connection conn = null
Statement stmt = null
ComponentLog log = null

    @Override
    Optional<String> lookup(Map<String, String> coordinates) {
        final String key = coordinates.keySet().first()
        ResultSet rs = stmt?.executeQuery("select nextval('${key}')".toString())
        return rs.next() ? Optional.ofNullable(rs.getLong(1)) : null
    }
    
    Set<String> getRequiredKeys() {
        return java.util.Collections.emptySet();
    }
    
    @Override
    Class<?> getValueType() {
        return String
    }

    @Override
    void initialize(ControllerServiceInitializationContext context) throws InitializationException {
    }

    @Override
    Collection<ValidationResult> validate(ValidationContext context) {
       null
    }

    @Override
    PropertyDescriptor getPropertyDescriptor(String name) {
       name.equals(DBCP_SERVICE.name) ? DBCP_SERVICE : null
    }

    @Override
    void onPropertyModified(PropertyDescriptor descriptor, String oldValue, String newValue) {
    }

    @Override
    List<PropertyDescriptor> getPropertyDescriptors() {
        [DBCP_SERVICE] as List;
    }

    @Override
    String getIdentifier() {
       ID
    }

    def onEnabled(context) {
        def db = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService)
        conn = db.connection
        stmt = conn.createStatement()
    }

    def onDisabled() {
      conn?.close()
    }

    def setLogger(logger) {
        log = logger
    }
}

lookupService = new SequenceLookupService()

One caveat I should mention is that currently, any properties defined by the script do not get added to the ScriptedLookupService dialog until the service is enabled, I have written an improvement Jira (NIFI-5547) to fix that.

The HCC question also refers to the nextInt() feature of NiFi Expression Language, this is MUCH faster than retrieving a sequence from a database. So is the UpdateAttribute approach, to let NiFi handle the "sequence" rather than an external database. But if you do need to use an external database sequence, this script should allow you to do that.

Note in my lookup function I only grabbed the first entry in the map, this can easily be extended to do multiple lookups of multiple sequences, but note that user-defined properties are unique by name, so you can't fill in multiple fields from the same sequence.

The above template (with script) is available as a Gist here. As always, please let me know if you try this whether it works for you, and I welcome all comments, questions, and suggestions. Cheers!