From 8f8d159a4209f4c294aba16bce3b1ca5789dc2d2 Mon Sep 17 00:00:00 2001
From: Grzegorz Janik <3079032+glujan@users.noreply.github.com>
Date: Thu, 6 Jan 2022 11:09:30 +0100
Subject: [PATCH] Use correct db in a multi db setup

SQLCompiler has a .connection instance which is already aware of what
database it should use. As such, when running EXPLAIN query let's use
that connection instead of connection to a default database.

Fixes #522.
---
 project/tests/test_execute_sql.py | 34 ++++++++++++++++++++-----------
 silk/sql.py                       |  5 ++---
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/project/tests/test_execute_sql.py b/project/tests/test_execute_sql.py
index fbee3f0f..7e9c5b20 100644
--- a/project/tests/test_execute_sql.py
+++ b/project/tests/test_execute_sql.py
@@ -10,12 +10,26 @@
 
 
 def mock_sql():
-    mock_sql_query = Mock(spec_set=['_execute_sql', 'query', 'as_sql'])
+    mock_sql_query = Mock(spec_set=['_execute_sql', 'query', 'as_sql', 'connection'])
     mock_sql_query._execute_sql = Mock()
     mock_sql_query.query = NonCallableMock(spec_set=['model'])
     mock_sql_query.query.model = Mock()
     query_string = 'SELECT * from table_name'
     mock_sql_query.as_sql = Mock(return_value=(query_string, ()))
+
+    mock_sql_query.connection = NonCallableMock(
+        spec_set=['cursor', 'features', 'ops'],
+        cursor=Mock(
+            spec_set=['__call__'],
+            return_value=NonCallableMagicMock(spec_set=['__enter__', '__exit__', 'execute'])
+        ),
+        features=NonCallableMock(
+            spec_set=['supports_explaining_query_execution'],
+            supports_explaining_query_execution=True
+        ),
+        ops=NonCallableMock(spec_set=['explain_query_prefix']),
+    )
+
     return mock_sql_query, query_string
 
 
@@ -29,8 +43,7 @@ def call_execute_sql(cls, request):
     }
     cls.args = [1, 2]
     cls.kwargs = kwargs
-    with patch('silk.sql.connection'):
-        execute_sql(cls.mock_sql, *cls.args, **cls.kwargs)
+    execute_sql(cls.mock_sql, *cls.args, **cls.kwargs)
 
 
 class TestCallNoRequest(TestCase):
@@ -86,16 +99,14 @@ def _query(self):
     def test_request(self):
         DataCollector().configure(request=Request.objects.create(path='/path/to/somewhere'))
         sql, _ = mock_sql()
-        with patch('silk.sql.connection'):
-            execute_sql(sql)
+        execute_sql(sql)
         query = self._query()
         self.assertEqual(query['request'], DataCollector().request)
 
     def test_registration(self):
         DataCollector().configure(request=Request.objects.create(path='/path/to/somewhere'))
         sql, _ = mock_sql()
-        with patch('silk.sql.connection'):
-            execute_sql(sql)
+        execute_sql(sql)
         query = self._query()
         self.assertIn(query, DataCollector().queries.values())
 
@@ -103,8 +114,7 @@ def test_explain(self):
         DataCollector().configure(request=Request.objects.create(path='/path/to/somewhere'))
         sql, qs = mock_sql()
         prefix = "EXPLAIN"
-        with patch('silk.sql.connection') as m:
-            mock_cursor = m.cursor.return_value.__enter__.return_value
-            m.ops.explain_query_prefix.return_value = prefix
-            execute_sql(sql)
-            mock_cursor.execute.assert_called_once_with(f"{prefix} {qs}", ())
+        mock_cursor = sql.connection.cursor.return_value.__enter__.return_value
+        sql.connection.ops.explain_query_prefix.return_value = prefix
+        execute_sql(sql)
+        mock_cursor.execute.assert_called_once_with(f"{prefix} {qs}", ())
diff --git a/silk/sql.py b/silk/sql.py
index fa8c1dd4..aa4a71be 100644
--- a/silk/sql.py
+++ b/silk/sql.py
@@ -2,7 +2,6 @@
 import traceback
 
 from django.core.exceptions import EmptyResultSet
-from django.db import connection
 from django.utils import timezone
 
 from silk.collector import DataCollector
@@ -29,7 +28,7 @@ def _unpack_explanation(result):
             yield row
 
 
-def _explain_query(q, params):
+def _explain_query(connection, q, params):
     if connection.features.supports_explaining_query_execution:
         if SilkyConfig().SILKY_ANALYZE_QUERIES:
             # Work around some DB engines not supporting analyze option
@@ -91,7 +90,7 @@ def execute_sql(self, *args, **kwargs):
             if request:
                 query_dict['request'] = request
             if self.query.model.__module__ != 'silk.models':
-                query_dict['analysis'] = _explain_query(q, params)
+                query_dict['analysis'] = _explain_query(self.connection, q, params)
                 DataCollector().register_query(query_dict)
             else:
                 DataCollector().register_silk_query(query_dict)