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 diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 173e02f2..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."""