MongoDB Java Driver for Polymorphism

When using MongoDB with a strong OOP language like Java, it’s a no brainer to want to hack the MongoDB Java Driver to serialize data to different sub-classes of model classes given different shapes of data, because MongoDB doesn’t have a strict schema restriction for a collection, which makes it a perfect match to store data models with polymorphism or inheritance into the same collection. This post will explore into how can we adapt polymorphism to MongoDB Java Driver.

Example

The post will use invoice line items as an example to demonstrate how to achieve polymorphism in MongoDB.

For line item, it always comes with a SKU (Stock Keeping Unit) representing the type of products. When we store line items into a collection, line items with different SKUs might have different shapes. For different SKUs, we might want to store different extra information about the product. For example, if it’s a T-shirt, we’d like to include the size info for the specific item, and we also want to include the expiration date if the item is food. To represent them in Java code, we can either use inheritance or composition. We can achieve both with MongoDB Java Driver, but before we jumping into those 2 directions, I’ll first introduce some common setups for both directions.

Setup

We first need to have a MongoCollection object to access the collection.

Main.java

public class Main {
  public static void main() {
    MongoCollection<LineItem> collection =   
      MongoClients.create("mongodb://localhost:27017")
        .getDatabase("testdb")
        .getCollection("lineItems", LineItem.class)
        .withCodecRegistry(getCodecRegistry());
  }
}

Note that we need to provide a CodecRegistry object, the getCodecRegistry will look like the following:

Main.java

protected static CodecRegistry getCodecRegistry() {

    return fromRegistries(MongoClientSettings.getDefaultCodecRegistry(),
            fromProviders(PojoCodecProvider.builder()
                    .conventions(DEFAULT_CONVENTIONS)
                    .register(getClassesToRegister())
                    .automatic(true).build()));
}

The code snippet does the following things:

1. Set the conventions to use DEFAULT_CONVENTIONS . The DEFAULT_CONVENTIONS has 3 different conventions — we will need 2 of them here: CLASS_AND_PROPERTY_CONVENTION and ANNOTATION_CONVENTION . More documentation for each convention can be found on their respective documentation page: CLASS_AND_PROPERTY_CONVENTION and ANNOTATION_CONVENTION
2. Register other classes that could be used while encoding a LineItem object. It has different implication for inheritance and composition style, so getClassesToRegister will be elaborated for both of them respectively later.

Polymorphism from Inheritance

For inheritance style, the concrete class will contain detailed info for different types respectively. I’ll make the assumption that all different line items will have their special info, thus the following code:

LineItem.java

@BsonDiscriminator(key = "sku")
abstract public class BaseLineItem<T> {
    @BsonId
    protected ObjectId _id;

    @BsonProperty("sku")
    protected String sku;
    @BsonProperty("quantity")
    protected double quantity;
    @BsonProperty
    protected String unit;

    @BsonProperty("info")
    protected T info;

    public BaseLineItem() {

    }

    public double getQuantity() {
        return quantity;
    }

    public ObjectId getId() {
        return _id;
    }

    public String getSku() {
        return sku;
    }

    public String getUnit() {
        return unit;
    }

    public T getInfo() {
        return info;
    }

    public void setInfo(T info) {
        this.info = info;
    }

    public void setId(ObjectId _id) {
        this._id = _id;
    }

    public void setQuantity(double quantity) {
        this.quantity = quantity;
    }

    public void setSku(String sku) {
        this.sku = sku;
    }

    public void setUnit(String unit) {
        this.unit = unit;
    }
}

A few things to note here:

1. As the base class for line item, @BsonDiscriminator(key = “sku”) needs to be set to let MongoDB Java Driver to know that the sku field will be used to differentiate which sub-class to use.
2. All the fields needs to be saved in the database should be annotated with BsonProperty , and the id field needs to be annotated with BsonId . We also need to define the setters and getters for those fields.
3. An empty constructor needs to be provided;
4. Note that we can’t use BaseLineItem directly for encoding even it’s not abstract . The reason is that the driver can’t use a class with generic for codec. more detail, see the official doc. This actually makes sense, because for a class with generic, it’s easy to insert it into database, but it will have trouble loading the data from the database because the driver doesn’t know which class to use. To address this, we can make sub-class of BaseLineItem without generic, which the driver will happily accept.

MilkLineItem.java

@BsonDiscriminator(key = "sku", value = "MILK")
public class MilkLineItem extends BaseLineItem<MilkInfo> {
    public MilkLineItem() {}

    public MilkLineItem(ObjectId _id,
                        String sku,
                        double quantity,
                        String unit,
                        MilkInfo info) {
        super(_id, sku, quantity, unit, info);
    }
}

Things to note:

1. @BsonDiscriminator(key = “sku”, value = “MILK”) is provided here, letting the driver know that only when the sku equals MILK this sub-class will be used.
2. An empty constructor is also needed here.

MilkInfo.java

public class MilkInfo {
    @BsonProperty
    Date expirationDate;

    public MilkInfo() {}

    public MilkInfo(Date date) {this.expirationDate = date;}

    public Date getExpirationDate() {
        return expirationDate;
    }

    public void setExpirationDate(Date expirationDate) {
        this.expirationDate = expirationDate;
    }
}

We mentioned earlier that we will need to register other classes needs to be used by the driver. Here, we need to register MilkLineItem because the driver needs to know what classes to use when the discriminator key — sku, has its value equals MILK .

we will add the following:

Main.java

protected static Class<?>[] getSubClasses() {
    return new Class<?>[]{MilkLineItem.class};
}

We can test the code by adding the following lines to the main method:

Main.java

collection.insertOne(new MilkLineItem(new ObjectId(), "MILK", 11, "HOUR", new MilkInfo(new Date())));
BaseLineItem first = document.find().first();
System.out.println("line item quantity: " + first.getQuantity());

In mongoDB, we will see the following document:

{
    "_id" : ObjectId("5fcdc455aef85552c266f37d"),
    "sku" : "MILK",
    "info" : {
        "expirationDate" : ISODate("2020-12-07T05:57:41.610Z")
    },
    "quantity" : 11,
    "unit" : "LITER"
}

Polymorphism from Composition

With composition style, we also assume that all the line items will have their specialized info. And since we use composition, we will use LineItem class directly to represent all the variations.

@BsonDiscriminator
public class LineItem {
    @BsonId
    protected ObjectId _id;

    @BsonProperty("sku")
    protected String sku;
    @BsonProperty("quantity")
    protected double quantity;
    @BsonProperty
    protected String unit;

    @BsonProperty(value = "info", useDiscriminator = true)
    protected Info info;

    public LineItem() {

    }

    public LineItem(ObjectId _id, String sku,
                    double quantity, String unit, Info info) {
        this._id = _id;
        this.sku = sku;
        this.quantity = quantity;
        this.unit = unit;
        this.info = info;
    }

    public double getQuantity() {
        return quantity;
    }

    public ObjectId getId() {
        return _id;
    }

    public String getSku() {
        return sku;
    }

    public String getUnit() {
        return unit;
    }

    public Info getInfo() {
        return info;
    }

    public void setInfo(Info info) {
        this.info = info;
    }

    public void setId(ObjectId _id) {
        this._id = _id;
    }

    public void setQuantity(double quantity) {
        this.quantity = quantity;
    }

    public void setSku(String sku) {
        this.sku = sku;
    }

    public void setUnit(String unit) {
        this.unit = unit;
    }
}

Note:

1. We still have to note 2 and 3 listed for BaseLineItem in previous section
2. Instead of @BsonDiscriminator(key=”sku”) , we simply put @BsonDiscriminator , which means that we will store a special field _t in the database, and the value will the full class name with it’s package path.

Info.java

@BsonDiscriminator
public interface Info {
}

MilkInfo.java

@BsonDiscriminator
public class MilkInfo implements Info {
    @BsonProperty
    Date expirationDate;
    public MilkInfo() {}

    public MilkInfo(Date date) {this.expirationDate = date;}

    public Date getExpirationDate() {
        return expirationDate;
    }

    public void setExpirationDate(Date expirationDate) {
        this.expirationDate = expirationDate;
    }
}

With composition style, as Info sub-classes are the ones holding BsonDiscriminator , the concrete Info classes needs to be registered, so we need to add the following to the main class:

Main.java

protected static Class<?>[] getSubClasses() {
    return new Class<?>[]{MilkInfo.class, TShirtInfo.class};
}

Again we can add the following code to the main method to test it works:

Main.java

document.insertOne(new LineItem(new ObjectId(), "MILK", 11, "LITER", new MilkInfo(new Date())));

LineItem first = document.find().first();
System.out.println("line item quantity: " + first.getQuantity());

In the database, the document looks a bit different from that when using inheritance:

> db.lineItems.find().pretty()
{
    "_id" : ObjectId("5fcdc76699b6323dcfc1882d"),
    "_t" : "src.model.billing.LineItem",
    "info" : {
        "_t" : "src.model.metrics.MilkInfo",
        "expirationDate" : ISODate("2020-12-07T06:10:46.108Z")
    },
    "quantity" : 11,
    "sku" : "MILK",
    "unit" : "LITER"
}

Conclusion

That’s all you have to know when using MongoDB with Java to support Polymorphism. I’ll update later with a repo with sample code.