From d3420918cd9804243900c8566b7e155044e668cb Mon Sep 17 00:00:00 2001 From: Eloi Zalczer Date: Mon, 18 Nov 2019 17:16:06 +0100 Subject: [PATCH 1/2] Added alias parameter in query_counter --- mongoengine/context_managers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index d8dfeaac..5920b724 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -182,10 +182,10 @@ class query_counter(object): - Some queries are ignored by default by the counter (killcursors, db.system.indexes) """ - def __init__(self): + def __init__(self, alias=DEFAULT_CONNECTION_NAME): """Construct the query_counter """ - self.db = get_db() + self.db = get_db(alias=alias) self.initial_profiling_level = None self._ctx_query_counter = 0 # number of queries issued by the context From 0bf08db7b943eba85d7e0dd85d161df4e615a371 Mon Sep 17 00:00:00 2001 From: Eloi Zalczer Date: Mon, 2 Dec 2019 10:07:33 +0100 Subject: [PATCH 2/2] Added test case for query_counter alias --- tests/document/test_instance.py | 40 ++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 1d3e18d0..c8ad2ff3 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -2825,6 +2825,44 @@ class TestInstance(MongoDBTestCase): assert "testdb-1" == B._meta.get("db_alias") + def test_query_counter_alias(self): + """query_counter works properly with db aliases?""" + # Register a connection with db_alias testdb-1 + register_connection("testdb-1", "mongoenginetest2") + + class A(Document): + """Uses default db_alias + """ + + name = StringField() + + class B(Document): + """Uses testdb-1 db_alias + """ + + name = StringField() + meta = {"db_alias": "testdb-1"} + + with query_counter() as q: + assert q == 0 + a = A.objects.create(name="A") + assert q == 1 + a = A.objects.first() + assert q == 2 + a.name = "Test A" + a.save() + assert q == 3 + + with query_counter(alias="testdb-1") as q: + assert q == 0 + b = B.objects.create(name="B") + assert q == 1 + b = B.objects.first() + assert q == 2 + b.name = "Test B" + b.save() + assert q == 3 + def test_db_ref_usage(self): """DB Ref usage in dict_fields.""" @@ -3644,7 +3682,7 @@ class TestInstance(MongoDBTestCase): User.objects().select_related() def test_embedded_document_failed_while_loading_instance_when_it_is_not_a_dict( - self + self, ): class LightSaber(EmbeddedDocument): color = StringField()