Initial commit
This commit is contained in:
235
venv/lib/python3.8/site-packages/pymongo/aggregation.py
Normal file
235
venv/lib/python3.8/site-packages/pymongo/aggregation.py
Normal file
@@ -0,0 +1,235 @@
|
||||
# Copyright 2019-present MongoDB, Inc.
|
||||
#
|
||||
# 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
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
# implied. See the License for the specific language governing
|
||||
# permissions and limitations under the License.
|
||||
|
||||
"""Perform aggregation operations on a collection or database."""
|
||||
|
||||
from bson.son import SON
|
||||
|
||||
from pymongo import common
|
||||
from pymongo.collation import validate_collation_or_none
|
||||
from pymongo.errors import ConfigurationError
|
||||
from pymongo.read_preferences import ReadPreference
|
||||
|
||||
|
||||
class _AggregationCommand(object):
|
||||
"""The internal abstract base class for aggregation cursors.
|
||||
|
||||
Should not be called directly by application developers. Use
|
||||
:meth:`pymongo.collection.Collection.aggregate`, or
|
||||
:meth:`pymongo.database.Database.aggregate` instead.
|
||||
"""
|
||||
def __init__(self, target, cursor_class, pipeline, options,
|
||||
explicit_session, user_fields=None, result_processor=None):
|
||||
if "explain" in options:
|
||||
raise ConfigurationError("The explain option is not supported. "
|
||||
"Use Database.command instead.")
|
||||
|
||||
self._target = target
|
||||
|
||||
common.validate_list('pipeline', pipeline)
|
||||
self._pipeline = pipeline
|
||||
self._performs_write = False
|
||||
if pipeline and ("$out" in pipeline[-1] or "$merge" in pipeline[-1]):
|
||||
self._performs_write = True
|
||||
|
||||
common.validate_is_mapping('options', options)
|
||||
self._options = options
|
||||
|
||||
# This is the batchSize that will be used for setting the initial
|
||||
# batchSize for the cursor, as well as the subsequent getMores.
|
||||
self._batch_size = common.validate_non_negative_integer_or_none(
|
||||
"batchSize", self._options.pop("batchSize", None))
|
||||
|
||||
# If the cursor option is already specified, avoid overriding it.
|
||||
self._options.setdefault("cursor", {})
|
||||
# If the pipeline performs a write, we ignore the initial batchSize
|
||||
# since the server doesn't return results in this case.
|
||||
if self._batch_size is not None and not self._performs_write:
|
||||
self._options["cursor"]["batchSize"] = self._batch_size
|
||||
|
||||
self._cursor_class = cursor_class
|
||||
self._explicit_session = explicit_session
|
||||
self._user_fields = user_fields
|
||||
self._result_processor = result_processor
|
||||
|
||||
self._collation = validate_collation_or_none(
|
||||
options.pop('collation', None))
|
||||
|
||||
self._max_await_time_ms = options.pop('maxAwaitTimeMS', None)
|
||||
|
||||
@property
|
||||
def _aggregation_target(self):
|
||||
"""The argument to pass to the aggregate command."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def _cursor_namespace(self):
|
||||
"""The namespace in which the aggregate command is run."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def _cursor_collection(self, cursor_doc):
|
||||
"""The Collection used for the aggregate command cursor."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def _database(self):
|
||||
"""The database against which the aggregation command is run."""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _check_compat(sock_info):
|
||||
"""Check whether the server version in-use supports aggregation."""
|
||||
pass
|
||||
|
||||
def _process_result(self, result, session, server, sock_info, slave_ok):
|
||||
if self._result_processor:
|
||||
self._result_processor(
|
||||
result, session, server, sock_info, slave_ok)
|
||||
|
||||
def get_read_preference(self, session):
|
||||
if self._performs_write:
|
||||
return ReadPreference.PRIMARY
|
||||
return self._target._read_preference_for(session)
|
||||
|
||||
def get_cursor(self, session, server, sock_info, slave_ok):
|
||||
# Ensure command compatibility.
|
||||
self._check_compat(sock_info)
|
||||
|
||||
# Serialize command.
|
||||
cmd = SON([("aggregate", self._aggregation_target),
|
||||
("pipeline", self._pipeline)])
|
||||
cmd.update(self._options)
|
||||
|
||||
# Apply this target's read concern if:
|
||||
# readConcern has not been specified as a kwarg and either
|
||||
# - server version is >= 4.2 or
|
||||
# - server version is >= 3.2 and pipeline doesn't use $out
|
||||
if (('readConcern' not in cmd) and
|
||||
((sock_info.max_wire_version >= 4 and
|
||||
not self._performs_write) or
|
||||
(sock_info.max_wire_version >= 8))):
|
||||
read_concern = self._target.read_concern
|
||||
else:
|
||||
read_concern = None
|
||||
|
||||
# Apply this target's write concern if:
|
||||
# writeConcern has not been specified as a kwarg and pipeline doesn't
|
||||
# perform a write operation
|
||||
if 'writeConcern' not in cmd and self._performs_write:
|
||||
write_concern = self._target._write_concern_for(session)
|
||||
else:
|
||||
write_concern = None
|
||||
|
||||
# Run command.
|
||||
result = sock_info.command(
|
||||
self._database.name,
|
||||
cmd,
|
||||
slave_ok,
|
||||
self.get_read_preference(session),
|
||||
self._target.codec_options,
|
||||
parse_write_concern_error=True,
|
||||
read_concern=read_concern,
|
||||
write_concern=write_concern,
|
||||
collation=self._collation,
|
||||
session=session,
|
||||
client=self._database.client,
|
||||
user_fields=self._user_fields)
|
||||
|
||||
self._process_result(result, session, server, sock_info, slave_ok)
|
||||
|
||||
# Extract cursor from result or mock/fake one if necessary.
|
||||
if 'cursor' in result:
|
||||
cursor = result['cursor']
|
||||
else:
|
||||
# Pre-MongoDB 2.6 or unacknowledged write. Fake a cursor.
|
||||
cursor = {
|
||||
"id": 0,
|
||||
"firstBatch": result.get("result", []),
|
||||
"ns": self._cursor_namespace,
|
||||
}
|
||||
|
||||
# Create and return cursor instance.
|
||||
return self._cursor_class(
|
||||
self._cursor_collection(cursor), cursor, sock_info.address,
|
||||
batch_size=self._batch_size or 0,
|
||||
max_await_time_ms=self._max_await_time_ms,
|
||||
session=session, explicit_session=self._explicit_session)
|
||||
|
||||
|
||||
class _CollectionAggregationCommand(_AggregationCommand):
|
||||
def __init__(self, *args, **kwargs):
|
||||
# Pop additional option and initialize parent class.
|
||||
use_cursor = kwargs.pop("use_cursor", True)
|
||||
super(_CollectionAggregationCommand, self).__init__(*args, **kwargs)
|
||||
|
||||
# Remove the cursor document if the user has set use_cursor to False.
|
||||
self._use_cursor = use_cursor
|
||||
if not self._use_cursor:
|
||||
self._options.pop("cursor", None)
|
||||
|
||||
@property
|
||||
def _aggregation_target(self):
|
||||
return self._target.name
|
||||
|
||||
@property
|
||||
def _cursor_namespace(self):
|
||||
return self._target.full_name
|
||||
|
||||
def _cursor_collection(self, cursor):
|
||||
"""The Collection used for the aggregate command cursor."""
|
||||
return self._target
|
||||
|
||||
@property
|
||||
def _database(self):
|
||||
return self._target.database
|
||||
|
||||
|
||||
class _CollectionRawAggregationCommand(_CollectionAggregationCommand):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(_CollectionRawAggregationCommand, self).__init__(*args, **kwargs)
|
||||
|
||||
# For raw-batches, we set the initial batchSize for the cursor to 0.
|
||||
if self._use_cursor and not self._performs_write:
|
||||
self._options["cursor"]["batchSize"] = 0
|
||||
|
||||
|
||||
class _DatabaseAggregationCommand(_AggregationCommand):
|
||||
@property
|
||||
def _aggregation_target(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def _cursor_namespace(self):
|
||||
return "%s.$cmd.aggregate" % (self._target.name,)
|
||||
|
||||
@property
|
||||
def _database(self):
|
||||
return self._target
|
||||
|
||||
def _cursor_collection(self, cursor):
|
||||
"""The Collection used for the aggregate command cursor."""
|
||||
# Collection level aggregate may not always return the "ns" field
|
||||
# according to our MockupDB tests. Let's handle that case for db level
|
||||
# aggregate too by defaulting to the <db>.$cmd.aggregate namespace.
|
||||
_, collname = cursor.get("ns", self._cursor_namespace).split(".", 1)
|
||||
return self._database[collname]
|
||||
|
||||
@staticmethod
|
||||
def _check_compat(sock_info):
|
||||
# Older server version don't raise a descriptive error, so we raise
|
||||
# one instead.
|
||||
if not sock_info.max_wire_version >= 6:
|
||||
err_msg = "Database.aggregate() is only supported on MongoDB 3.6+."
|
||||
raise ConfigurationError(err_msg)
|
||||
Reference in New Issue
Block a user