From 0f5e9b3f0404b64afc8a3b78013fbd2d973fe8e2 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jan 2020 14:34:30 +0100 Subject: [PATCH 1/2] Add document_count call in collection class --- mongokat/collection.py | 3 +++ tests/test_api.py | 9 +++++++++ 2 files changed, 12 insertions(+) diff --git a/mongokat/collection.py b/mongokat/collection.py index 7133d45..91dbb93 100644 --- a/mongokat/collection.py +++ b/mongokat/collection.py @@ -133,6 +133,9 @@ def exists(self, query, **args): def count(self, *args, **kwargs): return self._collection_with_options(kwargs).count(*args, **kwargs) + def count_documents(self, *args, **kwargs): + return self._collection_with_options(kwargs).count_documents(*args, **kwargs) + def distinct(self, *args, **kwargs): return self._collection_with_options(kwargs).distinct(*args, **kwargs) diff --git a/tests/test_api.py b/tests/test_api.py index 5701ba9..5bfa043 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -89,12 +89,14 @@ def test_document_common_methods(Sample): import collections assert Sample.collection.find().count() == 0 + assert Sample.collection.find().count_documents({}) == 0 # Instanciate new_object = Sample({"name": "XXX", "url": "http://example.com"}) # Should not save to DB yet. assert Sample.collection.find().count() == 0 + assert Sample.collection.find().count_documents({}) == 0 # Now save() new_object.save() @@ -106,6 +108,7 @@ def test_document_common_methods(Sample): assert type(new_object["_id"]) == ObjectId assert Sample.collection.find().count() == 1 + assert Sample.collection.find().count_documents({}) == 1 db_object = Sample.collection.find_one() assert type(db_object) == dict assert db_object["name"] == "XXX" @@ -115,6 +118,7 @@ def test_document_common_methods(Sample): assert type(inserted_object) == ObjectId assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 # Find back with different methods orm_object = Sample.find_by_id(db_object["_id"]) @@ -168,11 +172,13 @@ def test_document_common_methods(Sample): orm_object.save() assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 # FIXME not anymore as we are requesting _id for each query # orm_object.save(force=True) # assert Sample.collection.find().count() == 3 + # assert Sample.collection.find().count_documents({}) == 3 orm_object = Sample.find_by_id(db_object["_id"], fields=["url", "_id"]) assert dict(orm_object) == {"url": "http://example.com", "_id": db_object["_id"]} @@ -186,6 +192,8 @@ def test_document_common_methods(Sample): orm_object.save() assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 + db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" in db_object @@ -193,6 +201,7 @@ def test_document_common_methods(Sample): # Should not add anything new assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" not in db_object From 0ec6e402ea234389ef537391a42df0bbd8c2a844 Mon Sep 17 00:00:00 2001 From: Tristan Date: Thu, 30 Jan 2020 15:30:59 +0100 Subject: [PATCH 2/2] Add estimated_document_count method in collection class --- mongokat/collection.py | 3 +++ tests/test_api.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/mongokat/collection.py b/mongokat/collection.py index 91dbb93..a4fc45a 100644 --- a/mongokat/collection.py +++ b/mongokat/collection.py @@ -136,6 +136,9 @@ def count(self, *args, **kwargs): def count_documents(self, *args, **kwargs): return self._collection_with_options(kwargs).count_documents(*args, **kwargs) + def estimated_document_count(self, **kwargs): + return self._collection_with_options(kwargs).estimated_document_count(**kwargs) + def distinct(self, *args, **kwargs): return self._collection_with_options(kwargs).distinct(*args, **kwargs) diff --git a/tests/test_api.py b/tests/test_api.py index 5bfa043..e2aa8aa 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -90,6 +90,7 @@ def test_document_common_methods(Sample): assert Sample.collection.find().count() == 0 assert Sample.collection.find().count_documents({}) == 0 + assert Sample.collection.find().estimated_document_count() == 0 # Instanciate new_object = Sample({"name": "XXX", "url": "http://example.com"}) @@ -97,6 +98,7 @@ def test_document_common_methods(Sample): # Should not save to DB yet. assert Sample.collection.find().count() == 0 assert Sample.collection.find().count_documents({}) == 0 + assert Sample.collection.find().estimated_document_count() == 0 # Now save() new_object.save() @@ -109,6 +111,7 @@ def test_document_common_methods(Sample): assert Sample.collection.find().count() == 1 assert Sample.collection.find().count_documents({}) == 1 + assert Sample.collection.find().estimated_document_count() == 1 db_object = Sample.collection.find_one() assert type(db_object) == dict assert db_object["name"] == "XXX" @@ -119,6 +122,7 @@ def test_document_common_methods(Sample): assert Sample.collection.find().count() == 2 assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 # Find back with different methods orm_object = Sample.find_by_id(db_object["_id"]) @@ -173,12 +177,14 @@ def test_document_common_methods(Sample): assert Sample.collection.find().count() == 2 assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 # FIXME not anymore as we are requesting _id for each query # orm_object.save(force=True) # assert Sample.collection.find().count() == 3 # assert Sample.collection.find().count_documents({}) == 3 + # assert Sample.collection.find().estimated_document_count() == 3 orm_object = Sample.find_by_id(db_object["_id"], fields=["url", "_id"]) assert dict(orm_object) == {"url": "http://example.com", "_id": db_object["_id"]} @@ -193,6 +199,7 @@ def test_document_common_methods(Sample): assert Sample.collection.find().count() == 2 assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" in db_object @@ -202,6 +209,7 @@ def test_document_common_methods(Sample): # Should not add anything new assert Sample.collection.find().count() == 2 assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" not in db_object