Python context_managers.query_counter函数代码示例

    def test_list_item_dereference(self):
        """Ensure that DBRef items in ListFields are dereferenced.
        """

        class User(Document):
            name = StringField()

        class Group(Document):
            members = ListField(ReferenceField(User))

        User.drop_collection()
        Group.drop_collection()

        for i in xrange(1, 51):
            user = User(name="user %s" % i)
            user.save()

        group = Group(members=User.objects)
        group.save()

        group = Group(members=User.objects)
        group.save()

        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first()
            self.assertEqual(q, 1)

            len(group_obj._data["members"])
            self.assertEqual(q, 1)

            len(group_obj.members)
            self.assertEqual(q, 2)

            [m for m in group_obj.members]
            self.assertEqual(q, 2)

        # Document select_related
        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first().select_related()
            self.assertEqual(q, 2)
            [m for m in group_obj.members]
            self.assertEqual(q, 2)

        # Queryset select_related
        with query_counter() as q:
            self.assertEqual(q, 0)
            group_objs = Group.objects.select_related()
            self.assertEqual(q, 2)
            for group_obj in group_objs:
                [m for m in group_obj.members]
                self.assertEqual(q, 2)

        User.drop_collection()
        Group.drop_collection()

    def test_list_item_dereference_dref_false(self):
        """Ensure that DBRef items in ListFields are dereferenced.
        """
        class User(Document):
            name = StringField()

        class Group(Document):
            members = ListField(ReferenceField(User, dbref=False))

        User.drop_collection()
        Group.drop_collection()

        for i in range(1, 51):
            user = User(name='user %s' % i)
            user.save()

        group = Group(members=User.objects)
        group.save()
        group.reload()  # Confirm reload works

        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first()
            self.assertEqual(q, 1)

            [m for m in group_obj.members]
            self.assertEqual(q, 2)
            self.assertTrue(group_obj._data['members']._dereferenced)

            # verifies that no additional queries gets executed
            # if we re-iterate over the ListField once it is
            # dereferenced
            [m for m in group_obj.members]
            self.assertEqual(q, 2)
            self.assertTrue(group_obj._data['members']._dereferenced)

        # Document select_related
        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first().select_related()

            self.assertEqual(q, 2)
            [m for m in group_obj.members]
            self.assertEqual(q, 2)

        # Queryset select_related
        with query_counter() as q:
            self.assertEqual(q, 0)
            group_objs = Group.objects.select_related()
            self.assertEqual(q, 2)
            for group_obj in group_objs:
                [m for m in group_obj.members]
                self.assertEqual(q, 2)

    def test_query_counter(self):
        connect('mongoenginetest')
        db = get_db()

        collection = db.query_counter
        collection.drop()

        def issue_1_count_query():
            collection.find({}).count()

        def issue_1_insert_query():
            collection.insert_one({'test': 'garbage'})

        def issue_1_find_query():
            collection.find_one()

        counter = 0
        with query_counter() as q:
            self.assertEqual(q, counter)
            self.assertEqual(q, counter)    # Ensures previous count query did not get counted

            for _ in range(10):
                issue_1_insert_query()
                counter += 1
            self.assertEqual(q, counter)

            for _ in range(4):
                issue_1_find_query()
                counter += 1
            self.assertEqual(q, counter)

            for _ in range(3):
                issue_1_count_query()
                counter += 1
            self.assertEqual(q, counter)

    def test_select_related_follows_embedded_referencefields(self):

        class Song(Document):
            title = StringField()

        class PlaylistItem(EmbeddedDocument):
            song = ReferenceField("Song")

        class Playlist(Document):
            items = ListField(EmbeddedDocumentField("PlaylistItem"))

        Playlist.drop_collection()
        Song.drop_collection()

        songs = [Song.objects.create(title="song %d" % i) for i in range(3)]
        items = [PlaylistItem(song=song) for song in songs]
        playlist = Playlist.objects.create(items=items)

        with query_counter() as q:
            self.assertEqual(q, 0)

            playlist = Playlist.objects.first().select_related()
            songs = [item.song for item in playlist.items]

            self.assertEqual(q, 2)

    def test_list_item_dereference_dref_true_save_doesnt_cause_extra_queries(self):
        """Ensure that DBRef items in ListFields are dereferenced.
        """
        class User(Document):
            name = StringField()

        class Group(Document):
            name = StringField()
            members = ListField(ReferenceField(User, dbref=True))

        User.drop_collection()
        Group.drop_collection()

        for i in xrange(1, 51):
            User(name='user %s' % i).save()

        Group(name="Test", members=User.objects).save()

        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first()
            self.assertEqual(q, 1)

            group_obj.name = "new test"
            group_obj.save()

            self.assertEqual(q, 2)

    def test_collision(self):
        """
        Make sure that the right object is returned, even in case of a hash
        collision.
        """
        faux_obj = self.BlogPost.objects.get(title='Second Title')
        with query_counter() as q:
            obj1 = self.mongo_cache.get(self.BlogPost, 'title', 'First Title')

            # manually override the data in First Post's cache_key to point to
            # the Second Title's object, simulating a collision
            cache_key, flag_key = self.mongo_cache.get_keys(
                self.BlogPost,'title', 'First Title'
            )
            self.mongo_cache.redis.set(
                cache_key,
                self.mongo_cache.serialize(faux_obj)
            )

            obj2 = self.mongo_cache.get(self.BlogPost, 'title', 'First Title') # this should return the valid obj and invalidate the cache
            obj3 = self.mongo_cache.get(self.BlogPost, 'title', 'First Title') # this should cache the right obj
            obj4 = self.mongo_cache.get(self.BlogPost, 'title', 'First Title') # this should already have the valid obj cached

            self.assertEqual(obj1, obj2)
            self.assertEqual(obj1, obj3)
            self.assertEqual(obj1, obj4)

            self.assertEqual(q, 3) # 1st query + query after a cache collision + query that cached the right obj

    def test_generic_reference_orphan_dbref(self):
        """Ensure that generic orphan DBRef items in ListFields are dereferenced.
        """

        class UserA(Document):
            name = StringField()

        class UserB(Document):
            name = StringField()

        class UserC(Document):
            name = StringField()

        class Group(Document):
            members = ListField(GenericReferenceField())

        UserA.drop_collection()
        UserB.drop_collection()
        UserC.drop_collection()
        Group.drop_collection()

        members = []
        for i in range(1, 51):
            a = UserA(name='User A %s' % i)
            a.save()

            b = UserB(name='User B %s' % i)
            b.save()

            c = UserC(name='User C %s' % i)
            c.save()

            members += [a, b, c]

        group = Group(members=members)
        group.save()

        # Delete one UserA instance so that there is
        # an orphan DBRef in the GenericReference ListField
        UserA.objects[0].delete()
        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first()
            self.assertEqual(q, 1)

            [m for m in group_obj.members]
            self.assertEqual(q, 4)
            self.assertTrue(group_obj._data['members']._dereferenced)

            [m for m in group_obj.members]
            self.assertEqual(q, 4)
            self.assertTrue(group_obj._data['members']._dereferenced)

        UserA.drop_collection()
        UserB.drop_collection()
        UserC.drop_collection()
        Group.drop_collection()

    def test_invalidate(self):
        """Make sure simple cache invalidation works."""

        with query_counter() as q:
            self.mongo_cache.get(self.BlogPost, 'title', 'First Title')
            self.mongo_cache.invalidate(self.BlogPost, 'title', 'First Title')
            bp = self.mongo_cache.get(self.BlogPost, 'title', 'First Title')
            self.mongo_cache.invalidate(bp.__class__, 'title', bp.title)
            self.mongo_cache.get(self.BlogPost, 'title', 'First Title')
            self.assertEqual(q, 3)

    def test_datetime(self):
        resp = self.app.post('/datetime/', data=json.dumps({
            'datetime': '2010-01-01T00:00:00',
        }))
        response_success(resp)
        datetime = json.loads(resp.data)
        self.assertEqual(datetime['datetime'], '2010-01-01T00:00:00+00:00')

        with query_counter() as c:
            resp = self.app.put('/datetime/%s/' % datetime['id'], data=json.dumps({
                'datetime': '2010-01-02T00:00:00',
            }))
            response_success(resp)
            datetime = json.loads(resp.data)
            self.assertEqual(datetime['datetime'], '2010-01-02T00:00:00+00:00')

            self.assertEqual(c, 3) # query, update, query (reload)

        with query_counter() as c:
            resp = self.app.put('/datetime/%s/' % datetime['id'], data=json.dumps({
                'datetime': '2010-01-02T00:00:00',
            }))
            response_success(resp)
            datetime = json.loads(resp.data)
            self.assertEqual(datetime['datetime'], '2010-01-02T00:00:00+00:00')

            # Ideally this would be one query since we're not modifying, but
            # in the generic case the save method may have other side effects
            # and we don't know if the object was modified, so we currently
            # always reload.
            self.assertEqual(c, 2) # 2x query (with reload)

        # Same as above, with no body
        with query_counter() as c:
            resp = self.app.put('/datetime/%s/' % datetime['id'], data=json.dumps({
            }))
            response_success(resp)
            datetime = json.loads(resp.data)
            self.assertEqual(datetime['datetime'], '2010-01-02T00:00:00+00:00')

            self.assertEqual(c, 2) # 2x query (with reload)

    def test_query_counter(self):
        connect('mongoenginetest')
        db = get_db()
        db.test.find({})

        with query_counter() as q:
            self.assertEqual(0, q)

            for i in xrange(1, 51):
                db.test.find({}).count()

            self.assertEqual(50, q)

    def test_query_counter_counts_getmore_queries(self):
        connect('mongoenginetest')
        db = get_db()

        collection = db.query_counter
        collection.drop()

        many_docs = [{'test': 'garbage %s' % i} for i in range(150)]
        collection.insert_many(many_docs)   # first batch of documents contains 101 documents

        with query_counter() as q:
            self.assertEqual(q, 0)
            list(collection.find())
            self.assertEqual(q, 2)  # 1st select + 1 getmore

    def test_get_db_queries(self):
        """
        Make sure consecutive calls to mongo_cache.get don't hit the database
        for the same doc.
        """
        with query_counter() as q:
            for i in range(3):
                self.mongo_cache.get(self.BlogPost, 'title', 'First Title')

            self.mongo_cache.get(self.BlogPost, 'title', 'Second Title')

            # Only 2 queries to Mongo should be performed - one for caching
            # First Title and one for Second Title
            self.assertEqual(q, 2)

    def test_query_counter_temporarily_modifies_profiling_level(self):
        connect('mongoenginetest')
        db = get_db()

        initial_profiling_level = db.profiling_level()

        try:
            NEW_LEVEL = 1
            db.set_profiling_level(NEW_LEVEL)
            self.assertEqual(db.profiling_level(), NEW_LEVEL)
            with query_counter() as q:
                self.assertEqual(db.profiling_level(), 2)
            self.assertEqual(db.profiling_level(), NEW_LEVEL)
        except Exception:
            db.set_profiling_level(initial_profiling_level)    # Ensures it gets reseted no matter the outcome of the test
            raise

    def test_query_counter_ignores_particular_queries(self):
        connect('mongoenginetest')
        db = get_db()

        collection = db.query_counter
        collection.insert_many([{'test': 'garbage %s' % i} for i in range(10)])

        with query_counter() as q:
            self.assertEqual(q, 0)
            cursor = collection.find()
            self.assertEqual(q, 0)      # cursor wasn't opened yet
            _ = next(cursor)            # opens the cursor and fires the find query
            self.assertEqual(q, 1)

            cursor.close()              # issues a `killcursors` query that is ignored by the context
            self.assertEqual(q, 1)
            _ = db.system.indexes.find_one()    # queries on db.system.indexes are ignored as well
            self.assertEqual(q, 1)

    def test_list_item_dereference_orphan_dbref(self):
        """Ensure that orphan DBRef items in ListFields are dereferenced.
        """
        class User(Document):
            name = StringField()

        class Group(Document):
            members = ListField(ReferenceField(User, dbref=False))

        User.drop_collection()
        Group.drop_collection()

        for i in range(1, 51):
            user = User(name='user %s' % i)
            user.save()

        group = Group(members=User.objects)
        group.save()
        group.reload()  # Confirm reload works

        # Delete one User so one of the references in the
        # Group.members list is an orphan DBRef
        User.objects[0].delete()
        with query_counter() as q:
            self.assertEqual(q, 0)

            group_obj = Group.objects.first()
            self.assertEqual(q, 1)

            [m for m in group_obj.members]
            self.assertEqual(q, 2)
            self.assertTrue(group_obj._data['members']._dereferenced)

            # verifies that no additional queries gets executed
            # if we re-iterate over the ListField once it is
            # dereferenced
            [m for m in group_obj.members]
            self.assertEqual(q, 2)
            self.assertTrue(group_obj._data['members']._dereferenced)

        User.drop_collection()
        Group.drop_collection()