Merge pull request #205 from boxine/fix_connection_confusion

Fix issue where only the last database's connection was used in QueryAsserter

Author: jedie

README.md

@@ -152,7 +152,7 @@ ModelField, FormField and validators for GTIN/UPC/EAN numbers
 
 #### bx_django_utils.dbperf.query_recorder
 
-* [`SQLQueryRecorder()`](https://github.com/boxine/bx_django_utils/blob/master/bx_django_utils/dbperf/query_recorder.py#L95-L176) - A context manager that allows recording SQL queries executed during its lifetime.
+* [`SQLQueryRecorder()`](https://github.com/boxine/bx_django_utils/blob/master/bx_django_utils/dbperf/query_recorder.py#L106-L185) - A context manager that allows recording SQL queries executed during its lifetime.
 
 ### bx_django_utils.feature_flags
 
@@ -261,7 +261,7 @@ Utilities / helper for writing tests.
 
 #### bx_django_utils.test_utils.assert_queries
 
-* [`AssertQueries()`](https://github.com/boxine/bx_django_utils/blob/master/bx_django_utils/test_utils/assert_queries.py#L34-L287) - Assert executed database queries: Check table names, duplicate/similar Queries.
+* [`AssertQueries()`](https://github.com/boxine/bx_django_utils/blob/master/bx_django_utils/test_utils/assert_queries.py#L34-L295) - Assert executed database queries: Check table names, duplicate/similar Queries.
 
 #### bx_django_utils.test_utils.cache
 
@@ -383,7 +383,15 @@ apt-get install pipx
 pipx install uv
 ```
 
-Clone the project and just use our `Makefile` e.g.:
+you should be able to then do
+```bash
+make install
+make playwright-install
+```
+
+and validate everything works with `make test`
+
+For other options, you can check out our makefile:
 
 ```bash
 ~$ git clone https://github.com/boxine/bx_django_utils.git

bx_django_utils/admin_utils/tests/test_log_entry.py

@@ -16,6 +16,7 @@
 class LogEntryTestCase(TestCase):
     """"""  # noqa - don't add in README
 
+    databases = ['default', 'second']
     maxDiff = None
 
     def test_basic(self):

bx_django_utils/dbperf/query_recorder.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from collections.abc import Callable, Iterable
+from functools import partial
 from pprint import saferepr
 
 from django.db import connections
@@ -92,6 +93,16 @@ def dump(self, aggregate_queries=True):
         return results
 
 
+def _get_cursor_wrapper(*, cursor, connection, logger, collect_stacktrace, query_explain):
+    return RecordingCursorWrapper(
+        cursor(),
+        connection,
+        logger,
+        collect_stacktrace=collect_stacktrace,
+        query_explain=query_explain,
+    )
+
+
 class SQLQueryRecorder:
     """
     A context manager that allows recording SQL queries executed during its lifetime.
@@ -135,26 +146,23 @@ def __enter__(self):
             connection._recording_cursor = connection.cursor
             connection._recording_chunked_cursor = connection.chunked_cursor
 
-            def cursor():
-                return RecordingCursorWrapper(
-                    connection._recording_cursor(),  # noqa:B023
-                    connection,  # noqa:B023
-                    self.logger,
-                    collect_stacktrace=self.collect_stacktrace,
-                    query_explain=self.query_explain,
-                )
-
-            def chunked_cursor():
-                return RecordingCursorWrapper(
-                    connection._recording_chunked_cursor(),  # noqa:B023
-                    connection,  # noqa:B023
-                    self.logger,
-                    collect_stacktrace=self.collect_stacktrace,
-                    query_explain=self.query_explain,
-                )
-
-            connection.cursor = cursor
-            connection.chunked_cursor = chunked_cursor
+            common_kwargs = {
+                'connection': connection,
+                'logger': self.logger,
+                'collect_stacktrace': self.collect_stacktrace,
+                'query_explain': self.query_explain
+            }
+
+            connection.cursor = partial(
+                _get_cursor_wrapper,
+                cursor=connection._recording_cursor,
+                **common_kwargs
+            )
+            connection.chunked_cursor = partial(
+                _get_cursor_wrapper,
+                cursor=connection._recording_chunked_cursor,
+                **common_kwargs
+            )
 
         self.running = True
         return self
@@ -166,6 +174,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
             # undo the cursor wrapping so the connection is 'clean' again
             del connection._recording_cursor
+            del connection._recording_chunked_cursor
             del connection.cursor
             del connection.chunked_cursor
 

bx_django_utils/test_utils/assert_queries.py

@@ -175,6 +175,14 @@ def assert_table_names(self, *expected_table_names):
             ))
             raise AssertionError(self.build_error_message(f'Table names does not match:\n{diff}'))
 
+    def assert_databases_touched(self, *expected_database_aliases):
+        touched_keys = self.logger._databases.keys()
+        if set(expected_database_aliases) != touched_keys:
+            raise AssertionError(
+                self.build_error_message(f'Not all expected tables were touched, '
+                                         f'expected: {expected_database_aliases}, touched: {touched_keys}')
+            )
+
     def assert_table_counts(self, table_counts: Counter | dict, exclude: tuple[str, ...] | None = None):
         if not isinstance(table_counts, Counter):
             table_counts = Counter(table_counts)

bx_django_utils_tests/test_app/migrations/0001_initial.py

@@ -40,6 +40,17 @@ class Migration(migrations.Migration):
                 ),
             ],
         ),
+        migrations.CreateModel(
+            name='ColorFieldTestModelSecondary',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('required_color', bx_django_utils.models.color_field.ColorModelField(max_length=7)),
+                (
+                    'optional_color',
+                    bx_django_utils.models.color_field.ColorModelField(blank=True, max_length=7, null=True),
+                ),
+            ],
+        ),
         migrations.CreateModel(
             name='ConnectedUniqueSlugModel1',
             fields=[

bx_django_utils_tests/test_app/models.py

@@ -46,6 +46,11 @@ class ColorFieldTestModel(models.Model):
     optional_color = ColorModelField(blank=True, null=True)
 
 
+class ColorFieldTestModelSecondary(models.Model):
+    required_color = ColorModelField()
+    optional_color = ColorModelField(blank=True, null=True)
+
+
 class StoreSaveModel(models.Model):
     name = models.CharField(max_length=64)
     _save_calls = threading.local()

bx_django_utils_tests/test_project/routers.py

@@ -0,0 +1,21 @@
+
+
+class MultiDBRouter:
+
+    def _model_name_to_db(self, model_name: str):
+        if 'secondary' in model_name.lower():
+            return 'second'
+        else:
+            return 'default'
+
+    def db_for_read(self, model, **hints):
+        return self._model_name_to_db(model.__name__)
+
+    def db_for_write(self, model, **hints):
+        return self._model_name_to_db(model.__name__)
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        if model_name:
+            return self._model_name_to_db(model_name) == db
+        else:
+            return 'default'

bx_django_utils_tests/test_project/settings.py

@@ -98,8 +98,13 @@
     'default': {
         'ENGINE': 'django.db.backends.sqlite3',
         'NAME': str(BASE_DIR / 'db.sqlite3'),
+    },
+    'second': {
+        'ENGINE': 'django.db.backends.sqlite3',
+        'NAME': str(BASE_DIR / 'db_second.sqlite3'),
     }
 }
+DATABASE_ROUTERS = ['bx_django_utils_tests.test_project.routers.MultiDBRouter']
 DEFAULT_AUTO_FIELD = 'django.db.models.AutoField'
 
 # Password validation

bx_django_utils_tests/tests/test_assert_queries.py

@@ -10,6 +10,7 @@
 from django.test import TestCase
 
 from bx_django_utils.test_utils.assert_queries import AssertQueries
+from bx_django_utils_tests.test_app.models import ColorFieldTestModelSecondary
 
 
 def make_database_queries(count=1):
@@ -22,6 +23,7 @@ def make_database_queries2(count=1):
 
 
 class AssertQueriesTestCase(TestCase):
+    databases = ['default', 'second']
 
     def get_instance(self):
         with AssertQueries() as queries:
@@ -89,6 +91,7 @@ def test_assert_table_counts(self):
         queries = self.get_instance()
         queries.assert_table_counts(Counter(auth_permission=1))
         queries.assert_table_counts({'auth_permission': 1})
+        queries.assert_databases_touched('default')
 
         with self.assertRaises(AssertionError) as err:
             queries.assert_table_counts(Counter(auth_permission=2, foo=1, bar=3))
@@ -101,6 +104,18 @@ def test_assert_table_counts(self):
         assert 'Captured queries were:\n' in msg
         assert '1. SELECT "auth_permission"' in msg
 
+    def test_assert_two_touched_dbs(self):
+        with AssertQueries(databases=['default', 'second']) as queries:
+            Permission.objects.all().first()
+        with self.assertRaises(AssertionError):
+            # only 1 touched.
+            queries.assert_databases_touched('default', 'second')
+
+        with AssertQueries(databases=['default', 'second']) as queries:
+            Permission.objects.all().first()
+            ColorFieldTestModelSecondary.objects.count()
+        queries.assert_databases_touched('default', 'second')
+
     def test_assert_table_counts_exclude(self):
         with AssertQueries() as queries:
             Permission.objects.all().first()

bx_django_utils_tests/tests/test_routers.py

@@ -0,0 +1,31 @@
+from django.db import OperationalError
+from django.test import TestCase
+from model_bakery import baker
+
+from bx_django_utils_tests.test_app.models import ColorFieldTestModel, ColorFieldTestModelSecondary
+
+
+class DatabaseRoutersTestCase(TestCase):
+    databases = ['default', 'second']
+
+    def test_database_router(self):
+        baker.make(
+            ColorFieldTestModel,
+            pk=1,
+            required_color='#000002',
+            optional_color='#000003',
+        )
+        with self.assertRaises(OperationalError):
+            ColorFieldTestModel.objects.using("second").count()
+        self.assertEqual(ColorFieldTestModel.objects.using("default").count(), 1)
+
+    def test_database_router_secondary(self):
+        baker.make(
+            ColorFieldTestModelSecondary,
+            pk=1,
+            required_color='#000002',
+            optional_color='#000003',
+        )
+        with self.assertRaises(OperationalError):
+            ColorFieldTestModelSecondary.objects.using("default").count()
+        self.assertEqual(ColorFieldTestModelSecondary.objects.using("second").count(), 1)