Python pymongo.collection.Collection() Examples

def test_mongo_date_range_query():
    self = create_autospec(TickStore)
    self._collection = create_autospec(Collection)
    self._symbol_query.return_value = {"sy": {"$in" : ["s1" , "s2"]}}
    self._collection.aggregate.return_value = iter([{"_id": "s1", "start": dt(2014, 1, 1, 0, 0, tzinfo=mktz())},
                                                    {"_id": "s2", "start": dt(2014, 1, 1, 12, 0, tzinfo=mktz())}])

    self._collection.find_one.side_effect = [
        {'e': dt(2014, 1, 1, 15, 0, tzinfo=mktz())},
        {'e': dt(2014, 1, 2, 12, 0, tzinfo=mktz())}]

    query = TickStore._mongo_date_range_query(self, 'sym', DateRange(dt(2014, 1, 2, 0, 0, tzinfo=mktz()),
                                                                     dt(2014, 1, 3, 0, 0, tzinfo=mktz())))

    assert self._collection.aggregate.call_args_list == [call([
     {"$match": {"s": {"$lte": dt(2014, 1, 2, 0, 0, tzinfo=mktz())}, "sy": {"$in" : ["s1" , "s2"]}}},
     {"$project": {"_id": 0, "s": 1, "sy": 1}},
     {"$group": {"_id": "$sy", "start": {"$max": "$s"}}},
     {"$sort": {"start": 1}}])]

    assert self._collection.find_one.call_args_list == [
        call({'sy': 's1', 's': dt(2014, 1, 1, 0, 0, tzinfo=mktz())}, {'e': 1}),
        call({'sy': 's2', 's': dt(2014, 1, 1, 12, 0, tzinfo=mktz())}, {'e': 1})]

    assert query == {'s': {'$gte': dt(2014, 1, 1, 12, 0, tzinfo=mktz()), '$lte': dt(2014, 1, 3, 0, 0, tzinfo=mktz())}} 

def drop_collection(self, name_or_collection):
        """Drop a collection.

        :Parameters:
          - `name_or_collection`: the name of a collection to drop or the
            collection object itself
        """
        name = name_or_collection
        if isinstance(name, Collection):
            name = name.name

        if not isinstance(name, str):
            raise TypeError("name_or_collection must be an instance of "
                            "%s or Collection" % (str.__name__,))

        self.__connection._purge_index(self.__name, name)

        self.command("drop", str(name), allowable_errors=["ns not found"],
                     read_preference=ReadPreference.PRIMARY) 

def dereference(self, dbref, **kwargs):
        """Dereference a :class:`~bson.dbref.DBRef`, getting the
        document it points to.

        Raises :class:`TypeError` if `dbref` is not an instance of
        :class:`~bson.dbref.DBRef`. Returns a document, or ``None`` if
        the reference does not point to a valid document.  Raises
        :class:`ValueError` if `dbref` has a database specified that
        is different from the current database.

        :Parameters:
          - `dbref`: the reference
          - `**kwargs` (optional): any additional keyword arguments
            are the same as the arguments to
            :meth:`~pymongo.collection.Collection.find`.
        """
        if not isinstance(dbref, DBRef):
            raise TypeError("cannot dereference a %s" % type(dbref))
        if dbref.database is not None and dbref.database != self.__name:
            raise ValueError("trying to dereference a DBRef that points to "
                             "another database (%r not %r)" % (dbref.database,
                                                               self.__name))
        return self[dbref.collection].find_one({"_id": dbref.id}, **kwargs) 

def load_mongo_col(col: Collection) -> dict:
    """Load the pymongo collection to a dictionary.

    In the dictionary. The key will be the '_id' and in each value which is a dictionary there will also be a
    key '_id' so that the structure will be the same as the filesystem collection.

    Parameters
    ----------
    col : Collection
        The mongodb collection.

    Returns
    -------
    dct : dict
        A dictionary with all the info in the collection.
    """
    return {
        doc['_id']: doc for doc in col.find({})
    } 

def restore_pymongo():
    """Restores pymongo"""
    try:
        from pymongo.collection import Collection
    except ImportError:  # pragma: no cover
        pass
    else:
        for class_method in (
            "bulk_write",
            "delete_many",
            "delete_one",
            "insert_many",
            "insert_one",
            "replace_one",
            "update_many",
            "update_one",
        ):
            if hasattr(getattr(Collection, class_method), "__wrapped__"):
                setattr(
                    Collection,
                    class_method,
                    getattr(Collection, class_method).__wrapped__,
                ) 

def dereference(self, dbref, **kwargs):
        """Dereference a :class:`~bson.dbref.DBRef`, getting the
        document it points to.

        Raises :class:`TypeError` if `dbref` is not an instance of
        :class:`~bson.dbref.DBRef`. Returns a document, or ``None`` if
        the reference does not point to a valid document.  Raises
        :class:`ValueError` if `dbref` has a database specified that
        is different from the current database.

        :Parameters:
          - `dbref`: the reference
          - `**kwargs` (optional): any additional keyword arguments
            are the same as the arguments to
            :meth:`~pymongo.collection.Collection.find`.
        """
        if not isinstance(dbref, DBRef):
            raise TypeError("cannot dereference a %s" % type(dbref))
        if dbref.database is not None and dbref.database != self.__name:
            raise ValueError("trying to dereference a DBRef that points to "
                             "another database (%r not %r)" % (dbref.database,
                                                               self.__name))
        return self[dbref.collection].find_one({"_id": dbref.id}, **kwargs) 

def ensure_installed():
    global have_patched_collection

    logger.debug("Instrumenting pymongo.")

    if Collection is None:
        logger.debug("Couldn't import pymongo.Collection - probably not installed.")
    elif not have_patched_collection:
        for name in COLLECTION_METHODS:
            try:
                setattr(
                    Collection, name, wrap_collection_method(getattr(Collection, name))
                )
            except Exception as exc:
                logger.warning(
                    "Failed to instrument pymongo.Collection.%s: %r",
                    name,
                    exc,
                    exc_info=exc,
                )
        have_patched_collection = True 

def dereference(self, dbref, **kwargs):
        """Dereference a :class:`~bson.dbref.DBRef`, getting the
        document it points to.

        Raises :class:`TypeError` if `dbref` is not an instance of
        :class:`~bson.dbref.DBRef`. Returns a document, or ``None`` if
        the reference does not point to a valid document.  Raises
        :class:`ValueError` if `dbref` has a database specified that
        is different from the current database.

        :Parameters:
          - `dbref`: the reference
          - `**kwargs` (optional): any additional keyword arguments
            are the same as the arguments to
            :meth:`~pymongo.collection.Collection.find`.
        """
        if not isinstance(dbref, DBRef):
            raise TypeError("cannot dereference a %s" % type(dbref))
        if dbref.database is not None and dbref.database != self.__name:
            raise ValueError("trying to dereference a DBRef that points to "
                             "another database (%r not %r)" % (dbref.database,
                                                               self.__name))
        return self[dbref.collection].find_one({"_id": dbref.id}, **kwargs) 

def test_read_handles_operation_failure():
    self = Mock(spec=VersionStore)
    self._read_preference.return_value = sentinel.read_preference
    self._collection = create_autospec(Collection)
    self._read_metadata.side_effect = [sentinel.meta1, sentinel.meta2]
    self._read_metadata.__name__ = 'name'
    self._do_read.__name__ = 'name'  # feh: mongo_retry decorator cares about this
    self._do_read.side_effect = [OperationFailure('error'), sentinel.read]
    VersionStore.read(self, sentinel.symbol, sentinel.as_of,
                      from_version=sentinel.from_version,
                      date_range=sentinel.date_range,
                      other_kwarg=sentinel.other_kwarg)
    # Assert that, for the two read calls, the second uses the new metadata
    assert self._do_read.call_args_list == [call(sentinel.symbol, sentinel.meta1, 
                                                 sentinel.from_version,
                                                 date_range=sentinel.date_range,
                                                 other_kwarg=sentinel.other_kwarg,
                                                 read_preference=sentinel.read_preference)]
    assert self._do_read_retry.call_args_list == [call(sentinel.symbol, sentinel.meta2,
                                                       sentinel.from_version,
                                                       date_range=sentinel.date_range,
                                                       other_kwarg=sentinel.other_kwarg,
                                                       read_preference=ReadPreference.PRIMARY)] 

def test_prune_previous_versions_0_timeout():
    self = create_autospec(VersionStore, _versions=Mock())
    self.name = sentinel.name
    self._versions = create_autospec(Collection)
    self._versions.with_options.return_value.find.__name__ = 'find'
    self._versions.with_options.return_value.find.return_value = []
    with patch('arctic.store.version_store.dt') as dt:
        dt.utcnow.return_value = datetime.datetime(2013, 10, 1)
        VersionStore._find_prunable_version_ids(self, sentinel.symbol, keep_mins=0)
    assert self._versions.with_options.call_args_list == [call(read_preference=ReadPreference.PRIMARY)]
    assert self._versions.with_options.return_value.find.call_args_list == [
                                                  call({'$or': [{'parent': {'$exists': False}},
                                                                {'parent': []}],
                                                        'symbol': sentinel.symbol,
                                                        '_id': {'$lt': bson.ObjectId('524a10810000000000000000')}},
                                                       sort=[('version', -1)],
                                                       skip=1,
                                                       projection={'FW_POINTERS_CONFIG': 1, '_id': 1, 'SEGMENT_SHAS': 1}
                                                       )] 

def test_concat_and_rewrite_checks_written():
    self = create_autospec(NdarrayStore)
    collection = create_autospec(Collection)
    version = {'_id': sentinel.version_id,
               'segment_count': 1}
    previous_version = {'_id': sentinel.id,
                       'up_to': sentinel.up_to,
                        'base_version_id': sentinel.base_version_id,
                        'version': sentinel.version,
                        'segment_count' : 5,
                        'append_count' : 3}
    symbol = sentinel.symbol
    item = []

    collection.find.return_value = [{'_id': sentinel.id, 'segment': 47, 'compressed': True, 'sha': 'abc0'},
                                    {'_id': sentinel.id_2, 'segment': 48, 'compressed': True, 'sha': 'abc1'},
                                    # 3 appended items
                                    {'_id': sentinel.id_3, 'segment': 49, 'compressed': False, 'sha': 'abc2'},
                                    {'_id': sentinel.id_4, 'segment': 50, 'compressed': False, 'sha': 'abc3'},
                                    {'_id': sentinel.id_5, 'segment': 51, 'compressed': False, 'sha': 'abc4'}]
    collection.update_many.return_value = create_autospec(UpdateResult, matched_count=1)
    NdarrayStore._concat_and_rewrite(self, collection, version, symbol, item, previous_version)
    assert self.check_written.call_count == 1 

def test_concat_and_rewrite_checks_chunk_count():
    self = create_autospec(NdarrayStore)
    collection = create_autospec(Collection)
    version = {}
    previous_version = {'_id': sentinel.id,
                        'base_version_id': sentinel.base_version_id,
                        'version': sentinel.version,
                        'segment_count' : 3,
                        'append_count' : 1,
                        'up_to': sentinel.up_to}
    symbol = sentinel.symbol
    item = sentinel.item

    collection.find.return_value = [{'compressed': True, 'segment': 1},
                                    {'compressed': False, 'segment': 2}]
    with pytest.raises(DataIntegrityException) as e:
        NdarrayStore._concat_and_rewrite(self, collection, version, symbol, item, previous_version)
    assert str(e.value) == 'Symbol: sentinel.symbol:sentinel.version expected 1 segments but found 0' 

def delete_expired_chains(db: Collection, valid_hours: int) -> int:
    """Delete any Duplicate Chains that were created before some point in time.

    Returns the number of Duplicate Chains deleted.
    """

    last_valid_creation_time = datetime.utcnow() - timedelta(hours=valid_hours)

    result = db.delete_many({"created": {"$lt": last_valid_creation_time}})

    return result.deleted_count 

def setUp(self):
        self._key_pair_serializer = MagicMock()
        self._collection = MagicMock()

        self._wallet_storage = MongoWalletStorageImpl(
            serializer=cast(KeyPairSerializer, self._key_pair_serializer),
            collection=cast(Collection, self._collection)) 

def __init__(self, db_name: str):
        self._coll: Collection = get_db(db_name).hussars 

def __init__(self, db_name: str):
        self._coll: Collection = get_db(db_name).members 

def test_get_collection(self):
        ioc = GreaseContainer()
        self.assertTrue(isinstance(ioc.getMongo(), Mongo))
        coll = ioc.getCollection('TestCollection')
        self.assertTrue(isinstance(coll, Collection))
        self.assertEqual(coll.name, "TestCollection") 

def __init__(self, db_name: str):
        self._coll: Collection = get_db(db_name).quarantine 

def __init__(self, config: ConfigWrapper, db: Collection):
        self.db = db

        self.engineQueueDB = EngineQueueDB(db[config['queue coll engine']])
        self.irwinQueueDB = IrwinQueueDB(db[config['queue coll irwin']]) 

def _legacy_add_user(self, name, password, read_only, **kwargs):
        """Uses v1 system to add users, i.e. saving to system.users.
        """
        # Use a Collection with the default codec_options.
        system_users = self._collection_default_options('system.users')
        user = system_users.find_one({"user": name}) or {"user": name}
        if password is not None:
            user["pwd"] = auth._password_digest(name, password)
        if read_only is not None:
            user["readOnly"] = read_only
        user.update(kwargs)

        # We don't care what the _id is, only that it has one
        # for the replace_one call below.
        user.setdefault("_id", ObjectId())
        try:
            system_users.replace_one({"_id": user["_id"]}, user, True)
        except OperationFailure as exc:
            # First admin user add fails gle in MongoDB >= 2.1.2
            # See SERVER-4225 for more information.
            if 'login' in str(exc):
                pass
            # First admin user add fails gle from mongos 2.0.x
            # and 2.2.x.
            elif (exc.details and
                  'getlasterror' in exc.details.get('note', '')):
                pass
            else:
                raise 

def __init__(self, base_object: Union[collection.Collection, redis.Redis], threshold: int = 100, max_time_interval=10, is_print_log: bool = True):
        if str(base_object) not in self.bulk_helper_map:
            self._custom_init(base_object, threshold, max_time_interval, is_print_log)
            self.bulk_helper_map[str(base_object)] = self 

def get_collection(self, collection_name: str) -> Collection:
        return self.mongodb.get_collection(
            collection_name, codec_options=self.codec_options) 

def test_mongodb_connection(db):
    from pymongo.collection import Collection

    try:
        test_collection = Collection(db, "auth_check", create=True)
        test_collection.drop()
        return True
    except:
        return False 

def test_mongodb_connection(db):
    try:
        test_collection = Collection(db, "auth_check", create=True)
        test_collection.drop()
        return True
    except:
        return False 

def test_mongo_date_range_query_asserts():
    self = create_autospec(TickStore)
    self._collection = create_autospec(Collection)
    self._collection.find_one.return_value = {'s': sentinel.start}
    with pytest.raises(AssertionError):
        TickStore._mongo_date_range_query(self, 'sym', DateRange(None, None, CLOSED_OPEN))

    with pytest.raises(AssertionError):
        TickStore._mongo_date_range_query(self, 'sym', DateRange(dt(2014, 1, 1), None))

    with pytest.raises(AssertionError):
        TickStore._mongo_date_range_query(self, 'sym', DateRange(None, dt(2014, 1, 1))) 

def patch_pymongo(context):
    """
    Monkey patches pymongo client, if available. Overloads the
    query methods to add tracing and metrics collection.
    """

    def wrapper(wrapped, instance, args, kwargs):
        if not hasattr(context, "iopipe") or not hasattr(
            context.iopipe, "mark"
        ):  # pragma: no cover
            return wrapped(*args, **kwargs)

        id = ensure_utf8(str(uuid.uuid4()))
        with context.iopipe.mark(id):
            response = wrapped(*args, **kwargs)
        trace = context.iopipe.mark.measure(id)
        context.iopipe.mark.delete(id)
        collect_pymongo_metrics(context, trace, instance, response)
        return response

    for module, attr, _wrapper in [
        ("pymongo.collection", "Collection.find", wrapper),
        ("pymongo.collection", "Collection.bulk_write", wrapper),
        ("pymongo.collection", "Collection.delete_many", wrapper),
        ("pymongo.collection", "Collection.delete_one", wrapper),
        ("pymongo.collection", "Collection.insert_many", wrapper),
        ("pymongo.collection", "Collection.insert_one", wrapper),
        ("pymongo.collection", "Collection.replace_one", wrapper),
        ("pymongo.collection", "Collection.update_many", wrapper),
        ("pymongo.collection", "Collection.update_one", wrapper),
    ]:
        try:
            wrapt.wrap_function_wrapper(module, attr, _wrapper)
        except Exception:  # pragma: no cover
            pass 

def test_patch_db_requests_no_iopipe(mock_context,):
    """Asserts that monkey patching does not occur if IOpipe not present"""
    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__")

    delattr(mock_context, "iopipe")

    patch_db_requests(mock_context)

    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__") 

def test_patch_db_requests(mock_context_wrapper,):
    """Asserts that monkey patching occurs if iopipe present"""
    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__")

    patch_db_requests(mock_context_wrapper)

    assert hasattr(MySQLdb.connect, "__wrapped__")
    assert hasattr(pymysql.connect, "__wrapped__")

    assert isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert hasattr(Redis.execute_command, "__wrapped__")
    assert hasattr(Pipeline.execute, "__wrapped__")
    assert hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert hasattr(getattr(Collection, class_method), "__wrapped__")

    restore_db_requests()

    assert not hasattr(MySQLdb.connect, "__wrapped__")
    assert not hasattr(pymysql.connect, "__wrapped__")

    assert not isinstance(psycopg2.connect, wrapt.ObjectProxy)

    assert not hasattr(Redis.execute_command, "__wrapped__")
    assert not hasattr(Pipeline.execute, "__wrapped__")
    assert not hasattr(Pipeline.immediate_execute_command, "__wrapped__")

    for class_method in pymongo_collection_class_methods:
        assert not hasattr(getattr(Collection, class_method), "__wrapped__") 

def __getattr__(self, k):
        attr = getattr(self.__collect, k)
        if isinstance(attr, collection.Collection):
            raise AttributeError(
                "model object '%s' has not attribute '%s'" % (self.name, k))
        return attr 

def __init__(self, root_collection, file_id=None, file_document=None):
        """Read a file from GridFS

        Application developers should generally not need to
        instantiate this class directly - instead see the methods
        provided by :class:`~gridfs.GridFS`.

        Either `file_id` or `file_document` must be specified,
        `file_document` will be given priority if present. Raises
        :class:`TypeError` if `root_collection` is not an instance of
        :class:`~pymongo.collection.Collection`.

        :Parameters:
          - `root_collection`: root collection to read from
          - `file_id` (optional): value of ``"_id"`` for the file to read
          - `file_document` (optional): file document from
            `root_collection.files`

        .. versionchanged:: 3.0
           Creating a GridOut does not immediately retrieve the file metadata
           from the server. Metadata is fetched when first needed.
        """
        if not isinstance(root_collection, Collection):
            raise TypeError("root_collection must be an "
                            "instance of Collection")

        self.__chunks = root_collection.chunks
        self.__files = root_collection.files
        self.__file_id = file_id
        self.__buffer = EMPTY
        self.__position = 0
        self._file = file_document