document_sequence.go

Documentation: go.mongodb.org/mongo-driver/x/bsonx/bsoncore

     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 bsoncore
     8  
     9  import (
    10  	"errors"
    11  	"io"
    12  
    13  	"go.mongodb.org/mongo-driver/bson/bsontype"
    14  )
    15  
    16  // DocumentSequenceStyle is used to represent how a document sequence is laid out in a slice of
    17  // bytes.
    18  type DocumentSequenceStyle uint32
    19  
    20  // These constants are the valid styles for a DocumentSequence.
    21  const (
    22  	_ DocumentSequenceStyle = iota
    23  	SequenceStyle
    24  	ArrayStyle
    25  )
    26  
    27  // DocumentSequence represents a sequence of documents. The Style field indicates how the documents
    28  // are laid out inside of the Data field.
    29  type DocumentSequence struct {
    30  	Style DocumentSequenceStyle
    31  	Data  []byte
    32  	Pos   int
    33  }
    34  
    35  // ErrCorruptedDocument is returned when a full document couldn't be read from the sequence.
    36  var ErrCorruptedDocument = errors.New("invalid DocumentSequence: corrupted document")
    37  
    38  // ErrNonDocument is returned when a DocumentSequence contains a non-document BSON value.
    39  var ErrNonDocument = errors.New("invalid DocumentSequence: a non-document value was found in sequence")
    40  
    41  // ErrInvalidDocumentSequenceStyle is returned when an unknown DocumentSequenceStyle is set on a
    42  // DocumentSequence.
    43  var ErrInvalidDocumentSequenceStyle = errors.New("invalid DocumentSequenceStyle")
    44  
    45  // DocumentCount returns the number of documents in the sequence.
    46  func (ds *DocumentSequence) DocumentCount() int {
    47  	if ds == nil {
    48  		return 0
    49  	}
    50  	switch ds.Style {
    51  	case SequenceStyle:
    52  		var count int
    53  		var ok bool
    54  		rem := ds.Data
    55  		for len(rem) > 0 {
    56  			_, rem, ok = ReadDocument(rem)
    57  			if !ok {
    58  				return 0
    59  			}
    60  			count++
    61  		}
    62  		return count
    63  	case ArrayStyle:
    64  		_, rem, ok := ReadLength(ds.Data)
    65  		if !ok {
    66  			return 0
    67  		}
    68  
    69  		var count int
    70  		for len(rem) > 1 {
    71  			_, rem, ok = ReadElement(rem)
    72  			if !ok {
    73  				return 0
    74  			}
    75  			count++
    76  		}
    77  		return count
    78  	default:
    79  		return 0
    80  	}
    81  }
    82  
    83  // Empty returns true if the sequence is empty. It always returns true for unknown sequence styles.
    84  func (ds *DocumentSequence) Empty() bool {
    85  	if ds == nil {
    86  		return true
    87  	}
    88  
    89  	switch ds.Style {
    90  	case SequenceStyle:
    91  		return len(ds.Data) == 0
    92  	case ArrayStyle:
    93  		return len(ds.Data) <= 5
    94  	default:
    95  		return true
    96  	}
    97  }
    98  
    99  // ResetIterator resets the iteration point for the Next method to the beginning of the document
   100  // sequence.
   101  func (ds *DocumentSequence) ResetIterator() {
   102  	if ds == nil {
   103  		return
   104  	}
   105  	ds.Pos = 0
   106  }
   107  
   108  // Documents returns a slice of the documents. If nil either the Data field is also nil or could not
   109  // be properly read.
   110  func (ds *DocumentSequence) Documents() ([]Document, error) {
   111  	if ds == nil {
   112  		return nil, nil
   113  	}
   114  	switch ds.Style {
   115  	case SequenceStyle:
   116  		rem := ds.Data
   117  		var docs []Document
   118  		var doc Document
   119  		var ok bool
   120  		for {
   121  			doc, rem, ok = ReadDocument(rem)
   122  			if !ok {
   123  				if len(rem) == 0 {
   124  					break
   125  				}
   126  				return nil, ErrCorruptedDocument
   127  			}
   128  			docs = append(docs, doc)
   129  		}
   130  		return docs, nil
   131  	case ArrayStyle:
   132  		if len(ds.Data) == 0 {
   133  			return nil, nil
   134  		}
   135  		vals, err := Document(ds.Data).Values()
   136  		if err != nil {
   137  			return nil, ErrCorruptedDocument
   138  		}
   139  		docs := make([]Document, 0, len(vals))
   140  		for _, v := range vals {
   141  			if v.Type != bsontype.EmbeddedDocument {
   142  				return nil, ErrNonDocument
   143  			}
   144  			docs = append(docs, v.Data)
   145  		}
   146  		return docs, nil
   147  	default:
   148  		return nil, ErrInvalidDocumentSequenceStyle
   149  	}
   150  }
   151  
   152  // Next retrieves the next document from this sequence and returns it. This method will return
   153  // io.EOF when it has reached the end of the sequence.
   154  func (ds *DocumentSequence) Next() (Document, error) {
   155  	if ds == nil || ds.Pos >= len(ds.Data) {
   156  		return nil, io.EOF
   157  	}
   158  	switch ds.Style {
   159  	case SequenceStyle:
   160  		doc, _, ok := ReadDocument(ds.Data[ds.Pos:])
   161  		if !ok {
   162  			return nil, ErrCorruptedDocument
   163  		}
   164  		ds.Pos += len(doc)
   165  		return doc, nil
   166  	case ArrayStyle:
   167  		if ds.Pos < 4 {
   168  			if len(ds.Data) < 4 {
   169  				return nil, ErrCorruptedDocument
   170  			}
   171  			ds.Pos = 4 // Skip the length of the document
   172  		}
   173  		if len(ds.Data[ds.Pos:]) == 1 && ds.Data[ds.Pos] == 0x00 {
   174  			return nil, io.EOF // At the end of the document
   175  		}
   176  		elem, _, ok := ReadElement(ds.Data[ds.Pos:])
   177  		if !ok {
   178  			return nil, ErrCorruptedDocument
   179  		}
   180  		ds.Pos += len(elem)
   181  		val := elem.Value()
   182  		if val.Type != bsontype.EmbeddedDocument {
   183  			return nil, ErrNonDocument
   184  		}
   185  		return val.Data, nil
   186  	default:
   187  		return nil, ErrInvalidDocumentSequenceStyle
   188  	}
   189  }
   190
View as plain text