sample aggregation stage module - inject

README.md

This is a sample aggregation stage that insert a field into every document.

Steps:

1. Git clone the mongodb source code.
2. Create a directory called 'modules' in src/mongo/db/modules.
3. Git clone this repo as src/mongo/db/modules/inject
4. Build MongoDB
5. Now try the new aggregation stage:

db.foo.aggregate( { $inject : "funField" } );

This gist has been tested with:

• MongoDB master. Currently is at git version 3125749c23f8cf341768ecc37c86e64a2c30a356 or v3.3.1-229-g3125749.
• MongoDB branch v3.2.
• MongoDB tag r3.2.1.

See Blog: Extending the aggregation framework for more information.

build.py

def configure(conf, env):
    return True

inject.cpp

#include "mongo/platform/basic.h"

#include "inject.h"

namespace mongo {

REGISTER_DOCUMENT_SOURCE(inject, DocumentSourceInject::createFromBson);

boost::optional<Document> DocumentSourceInject::getNext() {
    // We check for interrupts so that a user killOp or a $maxTimeMillis can stop the process
    pExpCtx->checkForInterrupt();

    // Pull the next doc from the parent stage.
    // Here we are doing a 1 to 1 transformation, but you could pull N or all docs here.
    // For example if you wanted to compute an average you would just put getNext in a loop.
    boost::optional<Document> input = pSource->getNext();
    if (!input) {
        // input exhausted, we're done here
        return boost::none;
    }

    // This is the meat of our agg stage.
    // You can put whatever you wat here.
    MutableDocument output(*input);
    output.setNestedField(FieldPath(_field), Value(5.17));
    return output.freeze();
}

const char* DocumentSourceInject::getSourceName() const {
    return "$inject";
}

Value DocumentSourceInject::serialize(bool explain) const {
    // TODO: this should output the serialized form.
    invariant(false);
}

DocumentSource::GetDepsReturn DocumentSourceInject::getDependencies(DepsTracker* deps) const {
    // If you can take a subset of fields you can fill this in so the optimizer can do better work.
    return SEE_NEXT;
}

boost::intrusive_ptr<DocumentSource> DocumentSourceInject::createFromBson(
    BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx) {
    // This is where we parse the options to the stage.
    // You could required any BSON you want here.
    boost::intrusive_ptr<DocumentSourceInject> inject(new DocumentSourceInject(pExpCtx));
    inject->_field = elem.String();
    return inject;
}

DocumentSourceInject::DocumentSourceInject(const boost::intrusive_ptr<ExpressionContext>& pExpCtx)
    : DocumentSource(pExpCtx) {}
}

inject.h

#pragma once

#include "mongo/platform/basic.h"

#include "mongo/db/pipeline/document_source.h"

namespace mongo {

class DocumentSourceInject : public DocumentSource {
public:
    // virtuals from DocumentSource
    virtual boost::optional<Document> getNext();
    virtual const char* getSourceName() const;
    virtual Value serialize(bool explain = false) const;

    virtual GetDepsReturn getDependencies(DepsTracker* deps) const;

    /**
      Create a new projection DocumentSource from BSON.

      This is a convenience for directly handling BSON, and relies on the
      above methods.

      @param pBsonElement the BSONElement with an object named $project
      @param pExpCtx the expression context for the pipeline
      @returns the created projection
     */
    static boost::intrusive_ptr<DocumentSource> createFromBson(
        BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx);

private:
    DocumentSourceInject(const boost::intrusive_ptr<ExpressionContext>& pExpCtx);

    std::string _field;
};
}

SConscript

# -*- mode: python -*-
 
Import("env")
 
env.Library(target='inject',
            source=['inject.cpp'],
            LIBDEPS=['$BUILD_DIR/mongo/base'],
            LIBDEPS_DEPENDENTS=['$BUILD_DIR/mongo/db/pipeline/pipeline'])