From b1d37fea8fd7904c2b2d11c91393cf1b989314b6 Mon Sep 17 00:00:00 2001
From: Vojtech Bocek <vojtech.bocek@avast.com>
Date: Mon, 29 Jul 2019 16:40:27 +0200
Subject: [PATCH] Fixed #28107 -- Added
 DatabaseFeatures.allows_group_by_selected_pks_on_model() to allow enabling
 optimization for unmanaged models.

---
 django/db/backends/base/features.py |  5 +++++
 django/db/models/sql/compiler.py    |  6 +++++-
 docs/releases/3.0.txt               | 11 +++++++++++
 tests/aggregation_regress/tests.py  | 27 +++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/django/db/backends/base/features.py b/django/db/backends/base/features.py
index 425abf85404a3..89b417c7dee97 100644
--- a/django/db/backends/base/features.py
+++ b/django/db/backends/base/features.py
@@ -313,3 +313,8 @@ def supports_transactions(self):
             count, = cursor.fetchone()
             cursor.execute('DROP TABLE ROLLBACK_TEST')
         return count == 0
+
+    def allows_group_by_selected_pks_on_model(self, model):
+        if not self.allows_group_by_selected_pks:
+            return False
+        return model._meta.managed
diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py
index 5193362a92cdb..5da862f716e5a 100644
--- a/django/db/models/sql/compiler.py
+++ b/django/db/models/sql/compiler.py
@@ -171,7 +171,11 @@ def collapse_group_by(self, expressions, having):
             # database views on which the optimization might not be allowed.
             pks = {
                 expr for expr in expressions
-                if hasattr(expr, 'target') and expr.target.primary_key and expr.target.model._meta.managed
+                if (
+                    hasattr(expr, 'target') and
+                    expr.target.primary_key and
+                    self.connection.features.allows_group_by_selected_pks_on_model(expr.target.model)
+                )
             }
             aliases = {expr.alias for expr in pks}
             expressions = [
diff --git a/docs/releases/3.0.txt b/docs/releases/3.0.txt
index 9891119e66e1a..5d7ab98851797 100644
--- a/docs/releases/3.0.txt
+++ b/docs/releases/3.0.txt
@@ -355,6 +355,17 @@ Models
 
 * :class:`~django.db.models.CheckConstraint` is now supported on MySQL 8.0.16+.
 
+* The new ``allows_group_by_selected_pks_on_model()`` method of
+  ``django.db.backends.base.BaseDatabaseFeatures`` allows optimization of
+  ``GROUP BY`` clauses to require only the selected models' primary keys. By
+  default, it's supported only for managed models on PostgreSQL.
+
+  To enable the ``GROUP BY`` primary key-only optimization for unmanaged
+  models, you have to subclass the PostgreSQL database engine, overriding the
+  features class ``allows_group_by_selected_pks_on_model()`` method as you
+  require. See :ref:`Subclassing the built-in database backends
+  <subclassing-database-backends>` for an example.
+
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/aggregation_regress/tests.py b/tests/aggregation_regress/tests.py
index 4142643020a0f..6234d3c590599 100644
--- a/tests/aggregation_regress/tests.py
+++ b/tests/aggregation_regress/tests.py
@@ -1333,6 +1333,33 @@ def assertQuerysetResults(queryset):
                 self.assertIn(field.name, grouping[index + 1][0])
             assertQuerysetResults(queryset)
 
+    @skipUnlessDBFeature('allows_group_by_selected_pks')
+    def test_aggregate_unmanaged_model_as_tables(self):
+        qs = Book.objects.select_related('contact').annotate(num_authors=Count('authors'))
+        # Force treating unmanaged models as tables.
+        with mock.patch(
+            'django.db.connection.features.allows_group_by_selected_pks_on_model',
+            return_value=True,
+        ):
+            with mock.patch.object(Book._meta, 'managed', False), \
+                    mock.patch.object(Author._meta, 'managed', False):
+                _, _, grouping = qs.query.get_compiler(using='default').pre_sql_setup()
+                self.assertEqual(len(grouping), 2)
+                self.assertIn('id', grouping[0][0])
+                self.assertIn('id', grouping[1][0])
+                self.assertQuerysetEqual(
+                    qs.order_by('name'),
+                    [
+                        ('Artificial Intelligence: A Modern Approach', 2),
+                        ('Paradigms of Artificial Intelligence Programming: Case Studies in Common Lisp', 1),
+                        ('Practical Django Projects', 1),
+                        ('Python Web Development with Django', 3),
+                        ('Sams Teach Yourself Django in 24 Hours', 1),
+                        ('The Definitive Guide to Django: Web Development Done Right', 2),
+                    ],
+                    attrgetter('name', 'num_authors'),
+                )
+
     def test_reverse_join_trimming(self):
         qs = Author.objects.annotate(Count('book_contact_set__contact'))
         self.assertIn(' JOIN ', str(qs.query))