summaryrefslogtreecommitdiff
path: root/src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go
diff options
context:
space:
mode:
Diffstat (limited to 'src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go')
-rw-r--r--src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go122
1 files changed, 0 insertions, 122 deletions
diff --git a/src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go b/src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go
deleted file mode 100644
index 4abfdadcce0..00000000000
--- a/src/mongo/gotools/src/github.com/mongodb/mongo-tools/vendor/github.com/mongodb/mongo-tools-common/db/buffered_bulk.go
+++ /dev/null
@@ -1,122 +0,0 @@
-// Copyright (C) MongoDB, Inc. 2014-present.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"); you may
-// not use this file except in compliance with the License. You may obtain
-// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
-
-package db
-
-import (
- "context"
- "fmt"
-
- "go.mongodb.org/mongo-driver/bson"
- "go.mongodb.org/mongo-driver/mongo"
- "go.mongodb.org/mongo-driver/mongo/options"
-)
-
-// BufferedBulkInserter implements a bufio.Writer-like design for queuing up
-// documents and inserting them in bulk when the given doc limit (or max
-// message size) is reached. Must be flushed at the end to ensure that all
-// documents are written.
-type BufferedBulkInserter struct {
- collection *mongo.Collection
- writeModels []mongo.WriteModel
- docLimit int
- docCount int
- bulkWriteOpts *options.BulkWriteOptions
- upsert bool
-}
-
-func newBufferedBulkInserter(collection *mongo.Collection, docLimit int, ordered bool) *BufferedBulkInserter {
- bb := &BufferedBulkInserter{
- collection: collection,
- bulkWriteOpts: options.BulkWrite().SetOrdered(ordered),
- docLimit: docLimit,
- writeModels: make([]mongo.WriteModel, 0, docLimit),
- }
- return bb
-}
-
-// NewOrderedBufferedBulkInserter returns an initialized BufferedBulkInserter for performing ordered bulk writes.
-func NewOrderedBufferedBulkInserter(collection *mongo.Collection, docLimit int) *BufferedBulkInserter {
- return newBufferedBulkInserter(collection, docLimit, true)
-}
-
-// NewOrderedBufferedBulkInserter returns an initialized BufferedBulkInserter for performing unordered bulk writes.
-func NewUnorderedBufferedBulkInserter(collection *mongo.Collection, docLimit int) *BufferedBulkInserter {
- return newBufferedBulkInserter(collection, docLimit, false)
-}
-
-func (bb *BufferedBulkInserter) SetOrdered(ordered bool) *BufferedBulkInserter {
- bb.bulkWriteOpts.SetOrdered(ordered)
- return bb
-}
-
-func (bb *BufferedBulkInserter) SetBypassDocumentValidation(bypass bool) *BufferedBulkInserter {
- bb.bulkWriteOpts.SetBypassDocumentValidation(bypass)
- return bb
-}
-
-func (bb *BufferedBulkInserter) SetUpsert(upsert bool) *BufferedBulkInserter {
- bb.upsert = upsert
- return bb
-}
-
-// throw away the old bulk and init a new one
-func (bb *BufferedBulkInserter) resetBulk() {
- bb.writeModels = bb.writeModels[:0]
- bb.docCount = 0
-}
-
-// Insert adds a document to the buffer for bulk insertion. If the buffer becomes full, the bulk write is performed, returning
-// any error that occurs.
-func (bb *BufferedBulkInserter) Insert(doc interface{}) (*mongo.BulkWriteResult, error) {
- rawBytes, err := bson.Marshal(doc)
- if err != nil {
- return nil, fmt.Errorf("bson encoding error: %v", err)
- }
-
- return bb.InsertRaw(rawBytes)
-}
-
-// Update adds a document to the buffer for bulk update. If the buffer becomes full, the bulk write is performed, returning
-// any error that occurs.
-func (bb *BufferedBulkInserter) Update(selector, update bson.D) (*mongo.BulkWriteResult, error) {
- return bb.addModel(mongo.NewUpdateOneModel().SetFilter(selector).SetUpdate(update).SetUpsert(bb.upsert))
-}
-
-// Replace adds a document to the buffer for bulk replacement. If the buffer becomes full, the bulk write is performed, returning
-// any error that occurs.
-func (bb *BufferedBulkInserter) Replace(selector, replacement bson.D) (*mongo.BulkWriteResult, error) {
- return bb.addModel(mongo.NewReplaceOneModel().SetFilter(selector).SetReplacement(replacement).SetUpsert(bb.upsert))
-}
-
-// InsertRaw adds a document, represented as raw bson bytes, to the buffer for bulk insertion. If the buffer becomes full,
-// the bulk write is performed, returning any error that occurs.
-func (bb *BufferedBulkInserter) InsertRaw(rawBytes []byte) (*mongo.BulkWriteResult, error) {
- return bb.addModel(mongo.NewInsertOneModel().SetDocument(rawBytes))
-}
-
-// addModel adds a WriteModel to the buffer. If the buffer becomes full, the bulk write is performed, returning any error
-// that occurs.
-func (bb *BufferedBulkInserter) addModel(model mongo.WriteModel) (*mongo.BulkWriteResult, error) {
- bb.docCount++
- bb.writeModels = append(bb.writeModels, model)
-
- if bb.docCount >= bb.docLimit {
- return bb.Flush()
- }
-
- return nil, nil
-}
-
-// Flush writes all buffered documents in one bulk write and then resets the buffer.
-func (bb *BufferedBulkInserter) Flush() (*mongo.BulkWriteResult, error) {
- if bb.docCount == 0 {
- return nil, nil
- }
-
- defer bb.resetBulk()
- return bb.collection.BulkWrite(context.Background(), bb.writeModels, bb.bulkWriteOpts)
-}