Skip to content

Instantly share code, notes, and snippets.

@erh
Last active January 23, 2018 11:33
Show Gist options
  • Save erh/957e3193b9bd79b16cb1 to your computer and use it in GitHub Desktop.
Save erh/957e3193b9bd79b16cb1 to your computer and use it in GitHub Desktop.
sample aggregation stage module - inject

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.

def configure(conf, env):
return True
#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) {}
}
#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;
};
}
# -*- mode: python -*-
Import("env")
env.Library(target='inject',
source=['inject.cpp'],
LIBDEPS=['$BUILD_DIR/mongo/base'],
LIBDEPS_DEPENDENTS=['$BUILD_DIR/mongo/db/pipeline/pipeline'])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment