count.go

Documentation: go.mongodb.org/mongo-driver/x/mongo/driver/operation

     1  // Copyright (C) MongoDB, Inc. 2019-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 operation
     8  
     9  import (
    10  	"context"
    11  	"errors"
    12  	"fmt"
    13  	"time"
    14  
    15  	"go.mongodb.org/mongo-driver/bson/bsontype"
    16  	"go.mongodb.org/mongo-driver/event"
    17  	"go.mongodb.org/mongo-driver/internal/driverutil"
    18  	"go.mongodb.org/mongo-driver/mongo/description"
    19  	"go.mongodb.org/mongo-driver/mongo/readconcern"
    20  	"go.mongodb.org/mongo-driver/mongo/readpref"
    21  	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
    22  	"go.mongodb.org/mongo-driver/x/mongo/driver"
    23  	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
    24  )
    25  
    26  // Count represents a count operation.
    27  type Count struct {
    28  	maxTime        *time.Duration
    29  	query          bsoncore.Document
    30  	session        *session.Client
    31  	clock          *session.ClusterClock
    32  	collection     string
    33  	comment        bsoncore.Value
    34  	monitor        *event.CommandMonitor
    35  	crypt          driver.Crypt
    36  	database       string
    37  	deployment     driver.Deployment
    38  	readConcern    *readconcern.ReadConcern
    39  	readPreference *readpref.ReadPref
    40  	selector       description.ServerSelector
    41  	retry          *driver.RetryMode
    42  	result         CountResult
    43  	serverAPI      *driver.ServerAPIOptions
    44  	timeout        *time.Duration
    45  }
    46  
    47  // CountResult represents a count result returned by the server.
    48  type CountResult struct {
    49  	// The number of documents found
    50  	N int64
    51  }
    52  
    53  func buildCountResult(response bsoncore.Document) (CountResult, error) {
    54  	elements, err := response.Elements()
    55  	if err != nil {
    56  		return CountResult{}, err
    57  	}
    58  	cr := CountResult{}
    59  	for _, element := range elements {
    60  		switch element.Key() {
    61  		case "n": // for count using original command
    62  			var ok bool
    63  			cr.N, ok = element.Value().AsInt64OK()
    64  			if !ok {
    65  				return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s",
    66  					element.Value().Type)
    67  			}
    68  		case "cursor": // for count using aggregate with $collStats
    69  			firstBatch, err := element.Value().Document().LookupErr("firstBatch")
    70  			if err != nil {
    71  				return cr, err
    72  			}
    73  
    74  			// get count value from first batch
    75  			val := firstBatch.Array().Index(0)
    76  			count, err := val.Document().LookupErr("n")
    77  			if err != nil {
    78  				return cr, err
    79  			}
    80  
    81  			// use count as Int64 for result
    82  			var ok bool
    83  			cr.N, ok = count.AsInt64OK()
    84  			if !ok {
    85  				return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s",
    86  					element.Value().Type)
    87  			}
    88  		}
    89  	}
    90  	return cr, nil
    91  }
    92  
    93  // NewCount constructs and returns a new Count.
    94  func NewCount() *Count {
    95  	return &Count{}
    96  }
    97  
    98  // Result returns the result of executing this operation.
    99  func (c *Count) Result() CountResult { return c.result }
   100  
   101  func (c *Count) processResponse(info driver.ResponseInfo) error {
   102  	var err error
   103  	c.result, err = buildCountResult(info.ServerResponse)
   104  	return err
   105  }
   106  
   107  // Execute runs this operations and returns an error if the operation did not execute successfully.
   108  func (c *Count) Execute(ctx context.Context) error {
   109  	if c.deployment == nil {
   110  		return errors.New("the Count operation must have a Deployment set before Execute can be called")
   111  	}
   112  
   113  	err := driver.Operation{
   114  		CommandFn:         c.command,
   115  		ProcessResponseFn: c.processResponse,
   116  		RetryMode:         c.retry,
   117  		Type:              driver.Read,
   118  		Client:            c.session,
   119  		Clock:             c.clock,
   120  		CommandMonitor:    c.monitor,
   121  		Crypt:             c.crypt,
   122  		Database:          c.database,
   123  		Deployment:        c.deployment,
   124  		MaxTime:           c.maxTime,
   125  		ReadConcern:       c.readConcern,
   126  		ReadPreference:    c.readPreference,
   127  		Selector:          c.selector,
   128  		ServerAPI:         c.serverAPI,
   129  		Timeout:           c.timeout,
   130  		Name:              driverutil.CountOp,
   131  	}.Execute(ctx)
   132  
   133  	// Swallow error if NamespaceNotFound(26) is returned from aggregate on non-existent namespace
   134  	if err != nil {
   135  		dErr, ok := err.(driver.Error)
   136  		if ok && dErr.Code == 26 {
   137  			err = nil
   138  		}
   139  	}
   140  	return err
   141  }
   142  
   143  func (c *Count) command(dst []byte, _ description.SelectedServer) ([]byte, error) {
   144  	dst = bsoncore.AppendStringElement(dst, "count", c.collection)
   145  	if c.query != nil {
   146  		dst = bsoncore.AppendDocumentElement(dst, "query", c.query)
   147  	}
   148  	if c.comment.Type != bsontype.Type(0) {
   149  		dst = bsoncore.AppendValueElement(dst, "comment", c.comment)
   150  	}
   151  	return dst, nil
   152  }
   153  
   154  // MaxTime specifies the maximum amount of time to allow the query to run on the server.
   155  func (c *Count) MaxTime(maxTime *time.Duration) *Count {
   156  	if c == nil {
   157  		c = new(Count)
   158  	}
   159  
   160  	c.maxTime = maxTime
   161  	return c
   162  }
   163  
   164  // Query determines what results are returned from find.
   165  func (c *Count) Query(query bsoncore.Document) *Count {
   166  	if c == nil {
   167  		c = new(Count)
   168  	}
   169  
   170  	c.query = query
   171  	return c
   172  }
   173  
   174  // Session sets the session for this operation.
   175  func (c *Count) Session(session *session.Client) *Count {
   176  	if c == nil {
   177  		c = new(Count)
   178  	}
   179  
   180  	c.session = session
   181  	return c
   182  }
   183  
   184  // ClusterClock sets the cluster clock for this operation.
   185  func (c *Count) ClusterClock(clock *session.ClusterClock) *Count {
   186  	if c == nil {
   187  		c = new(Count)
   188  	}
   189  
   190  	c.clock = clock
   191  	return c
   192  }
   193  
   194  // Collection sets the collection that this command will run against.
   195  func (c *Count) Collection(collection string) *Count {
   196  	if c == nil {
   197  		c = new(Count)
   198  	}
   199  
   200  	c.collection = collection
   201  	return c
   202  }
   203  
   204  // Comment sets a value to help trace an operation.
   205  func (c *Count) Comment(comment bsoncore.Value) *Count {
   206  	if c == nil {
   207  		c = new(Count)
   208  	}
   209  
   210  	c.comment = comment
   211  	return c
   212  }
   213  
   214  // CommandMonitor sets the monitor to use for APM events.
   215  func (c *Count) CommandMonitor(monitor *event.CommandMonitor) *Count {
   216  	if c == nil {
   217  		c = new(Count)
   218  	}
   219  
   220  	c.monitor = monitor
   221  	return c
   222  }
   223  
   224  // Crypt sets the Crypt object to use for automatic encryption and decryption.
   225  func (c *Count) Crypt(crypt driver.Crypt) *Count {
   226  	if c == nil {
   227  		c = new(Count)
   228  	}
   229  
   230  	c.crypt = crypt
   231  	return c
   232  }
   233  
   234  // Database sets the database to run this operation against.
   235  func (c *Count) Database(database string) *Count {
   236  	if c == nil {
   237  		c = new(Count)
   238  	}
   239  
   240  	c.database = database
   241  	return c
   242  }
   243  
   244  // Deployment sets the deployment to use for this operation.
   245  func (c *Count) Deployment(deployment driver.Deployment) *Count {
   246  	if c == nil {
   247  		c = new(Count)
   248  	}
   249  
   250  	c.deployment = deployment
   251  	return c
   252  }
   253  
   254  // ReadConcern specifies the read concern for this operation.
   255  func (c *Count) ReadConcern(readConcern *readconcern.ReadConcern) *Count {
   256  	if c == nil {
   257  		c = new(Count)
   258  	}
   259  
   260  	c.readConcern = readConcern
   261  	return c
   262  }
   263  
   264  // ReadPreference set the read preference used with this operation.
   265  func (c *Count) ReadPreference(readPreference *readpref.ReadPref) *Count {
   266  	if c == nil {
   267  		c = new(Count)
   268  	}
   269  
   270  	c.readPreference = readPreference
   271  	return c
   272  }
   273  
   274  // ServerSelector sets the selector used to retrieve a server.
   275  func (c *Count) ServerSelector(selector description.ServerSelector) *Count {
   276  	if c == nil {
   277  		c = new(Count)
   278  	}
   279  
   280  	c.selector = selector
   281  	return c
   282  }
   283  
   284  // Retry enables retryable mode for this operation. Retries are handled automatically in driver.Operation.Execute based
   285  // on how the operation is set.
   286  func (c *Count) Retry(retry driver.RetryMode) *Count {
   287  	if c == nil {
   288  		c = new(Count)
   289  	}
   290  
   291  	c.retry = &retry
   292  	return c
   293  }
   294  
   295  // ServerAPI sets the server API version for this operation.
   296  func (c *Count) ServerAPI(serverAPI *driver.ServerAPIOptions) *Count {
   297  	if c == nil {
   298  		c = new(Count)
   299  	}
   300  
   301  	c.serverAPI = serverAPI
   302  	return c
   303  }
   304  
   305  // Timeout sets the timeout for this operation.
   306  func (c *Count) Timeout(timeout *time.Duration) *Count {
   307  	if c == nil {
   308  		c = new(Count)
   309  	}
   310  
   311  	c.timeout = timeout
   312  	return c
   313  }
   314
View as plain text