github.com/mdaxf/iac@v0.0.0-20240519030858-58a061660378/vendor_skip/go.mongodb.org/mongo-driver/x/mongo/driver/batches.go (about)

     1  // Copyright (C) MongoDB, Inc. 2022-present.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License"); you may
     4  // not use this file except in compliance with the License. You may obtain
     5  // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
     6  
     7  package driver
     8  
     9  import (
    10  	"errors"
    11  
    12  	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
    13  )
    14  
    15  // ErrDocumentTooLarge occurs when a document that is larger than the maximum size accepted by a
    16  // server is passed to an insert command.
    17  var ErrDocumentTooLarge = errors.New("an inserted document is too large")
    18  
    19  // Batches contains the necessary information to batch split an operation. This is only used for write
    20  // oeprations.
    21  type Batches struct {
    22  	Identifier string
    23  	Documents  []bsoncore.Document
    24  	Current    []bsoncore.Document
    25  	Ordered    *bool
    26  }
    27  
    28  // Valid returns true if Batches contains both an identifier and the length of Documents is greater
    29  // than zero.
    30  func (b *Batches) Valid() bool { return b != nil && b.Identifier != "" && len(b.Documents) > 0 }
    31  
    32  // ClearBatch clears the Current batch. This must be called before AdvanceBatch will advance to the
    33  // next batch.
    34  func (b *Batches) ClearBatch() { b.Current = b.Current[:0] }
    35  
    36  // AdvanceBatch splits the next batch using maxCount and targetBatchSize. This method will do nothing if
    37  // the current batch has not been cleared. We do this so that when this is called during execute we
    38  // can call it without first needing to check if we already have a batch, which makes the code
    39  // simpler and makes retrying easier.
    40  // The maxDocSize parameter is used to check that any one document is not too large. If the first document is bigger
    41  // than targetBatchSize but smaller than maxDocSize, a batch of size 1 containing that document will be created.
    42  func (b *Batches) AdvanceBatch(maxCount, targetBatchSize, maxDocSize int) error {
    43  	if len(b.Current) > 0 {
    44  		return nil
    45  	}
    46  
    47  	if maxCount <= 0 {
    48  		maxCount = 1
    49  	}
    50  
    51  	splitAfter := 0
    52  	size := 0
    53  	for i, doc := range b.Documents {
    54  		if i == maxCount {
    55  			break
    56  		}
    57  		if len(doc) > maxDocSize {
    58  			return ErrDocumentTooLarge
    59  		}
    60  		if size+len(doc) > targetBatchSize {
    61  			break
    62  		}
    63  
    64  		size += len(doc)
    65  		splitAfter++
    66  	}
    67  
    68  	// if there are no documents, take the first one.
    69  	// this can happen if there is a document that is smaller than maxDocSize but greater than targetBatchSize.
    70  	if splitAfter == 0 {
    71  		splitAfter = 1
    72  	}
    73  
    74  	b.Current, b.Documents = b.Documents[:splitAfter], b.Documents[splitAfter:]
    75  	return nil
    76  }