invoicing: reduce queysets (#71528)

lingo/invoicing/utils.py

@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import collections
 import datetime
 
 from django.test.client import RequestFactory
@@ -54,7 +55,19 @@ def get_all_subscriptions(agendas, date_start, date_end):
     return all_subscriptions
 
 
-def get_invoice_lines_for_user(agendas, user_external_id, subscriptions, date_start, date_end):
+def get_invoice_lines_for_user(
+    agendas, agendas_pricings, user_external_id, subscriptions, date_start, date_end
+):
+    def get_agenda_pricing(agendas_pricings_for_agenda, date_event):
+        # same logic as AgendaPricing.get_agenda_pricing
+        for agenda_pricing in agendas_pricings_for_agenda:
+            if agenda_pricing.date_start > date_event:
+                continue
+            if agenda_pricing.date_end <= date_event:
+                continue
+            return agenda_pricing
+        raise AgendaPricingNotFound
+
     def get_subscription(subscriptions, event_date):
         # get subscription matching event_date
         for subscription in subscriptions:
@@ -69,6 +82,13 @@ def get_invoice_lines_for_user(agendas, user_external_id, subscriptions, date_st
     if not agendas or not subscriptions:
         return []
 
+    agendas_pricings_by_agendas = collections.defaultdict(list)
+    for agenda_pricing in agendas_pricings:
+        if agenda_pricing.flat_fee_schedule:
+            continue
+        for agenda in agenda_pricing.agendas.all():
+            agendas_pricings_by_agendas[agenda.slug].append(agenda_pricing)
+
     # get check status for user_external_id, on agendas, for the period
     check_status_list = get_check_status(
         agenda_slugs=list(subscriptions.keys()),
@@ -98,9 +118,7 @@ def get_invoice_lines_for_user(agendas, user_external_id, subscriptions, date_st
 
         agenda = agendas_by_slug[serialized_event['agenda']]
         try:
-            agenda_pricing = AgendaPricing.get_agenda_pricing(
-                agenda=agenda, start_date=event_date, flat_fee_schedule=False
-            )  # XXX optimize SQL request
+            agenda_pricing = get_agenda_pricing(agendas_pricings_by_agendas.get(agenda.slug), event_date)
             pricing_data = agenda_pricing.get_pricing_data_for_event(
                 request=request,
                 agenda=agenda,
@@ -155,11 +173,18 @@ def get_invoice_lines_for_user(agendas, user_external_id, subscriptions, date_st
 
 
 def get_all_invoice_lines(agendas, subscriptions, date_start, date_end):
+    agendas_pricings = (
+        AgendaPricing.objects.filter(flat_fee_schedule=False)
+        .extra(where=["(date_start, date_end) OVERLAPS (%s, %s)"], params=[date_start, date_end])
+        .prefetch_related('agendas', 'pricing__criterias', 'pricing__categories')
+    )
+
     lines = []
     for user_external_id, user_subs in subscriptions.items():
         # generate lines for each user
         lines += get_invoice_lines_for_user(
             agendas=agendas,
+            agendas_pricings=agendas_pricings,
             user_external_id=user_external_id,
             subscriptions=user_subs,
             date_start=date_start,

lingo/pricing/models.py

@@ -507,12 +507,16 @@ class AgendaPricing(models.Model):
     def compute_pricing(self, context):
         criterias = {}
         categories = []
-        # for each category (ordered)
-        for category in self.pricing.categories.all().order_by('pricingcriteriacategory__order'):
+        # for each category
+        for category in self.pricing.categories.all():
             criterias[category.slug] = None
             categories.append(category.slug)
             # find the first matching criteria (criterias are ordered)
-            for criteria in self.pricing.criterias.filter(category=category, default=False):
+            for criteria in self.pricing.criterias.all():
+                if criteria.category_id != category.pk:
+                    continue
+                if criteria.default:
+                    continue
                 condition = criteria.compute_condition(context)
                 if condition:
                     criterias[category.slug] = criteria.slug
@@ -520,7 +524,9 @@ class AgendaPricing(models.Model):
             if criterias[category.slug] is not None:
                 continue
             # if no match, take default criteria if only once defined
-            default_criterias = self.pricing.criterias.filter(category=category, default=True)
+            default_criterias = [
+                c for c in self.pricing.criterias.all() if c.default and c.category_id == category.pk
+            ]
             if len(default_criterias) > 1:
                 raise MultipleDefaultCriteriaCondition(details={'category': category.slug})
             if not default_criterias:

tests/invoicing/test_invoice_generation.py

@@ -3,12 +3,14 @@ from unittest import mock
 
 import pytest
 from django.core.management import CommandError, call_command
+from django.db import connection
+from django.test.utils import CaptureQueriesContext
 
 from lingo.agendas.chrono import ChronoError
 from lingo.agendas.models import Agenda
 from lingo.invoicing import utils
 from lingo.invoicing.models import DraftInvoice, DraftInvoiceLine, Regie
-from lingo.pricing.models import AgendaPricing, Pricing, PricingError
+from lingo.pricing.models import AgendaPricing, Criteria, CriteriaCategory, Pricing, PricingError
 
 pytestmark = pytest.mark.django_db
 
@@ -205,6 +207,7 @@ def test_get_invoice_lines_for_user_check_status_error(mock_status):
     with pytest.raises(ChronoError):
         utils.get_invoice_lines_for_user(
             agendas=['foo'],
+            agendas_pricings=[],
             user_external_id='user:1',
             subscriptions={'foo': 'bar'},
             date_start=datetime.date(2022, 9, 1),
@@ -229,6 +232,7 @@ def test_get_invoice_lines_for_user_check_status(mock_pricing_data_event, mock_s
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={
                 'agenda-1': [
@@ -251,6 +255,7 @@ def test_get_invoice_lines_for_user_check_status(mock_pricing_data_event, mock_s
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[agenda1, agenda2],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={},
             date_start=datetime.date(2022, 9, 1),
@@ -266,6 +271,7 @@ def test_get_invoice_lines_for_user_check_status(mock_pricing_data_event, mock_s
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[agenda1, agenda2],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={
                 'agenda-1': [
@@ -306,6 +312,7 @@ def test_get_invoice_lines_for_user_check_status(mock_pricing_data_event, mock_s
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[agenda1, agenda2],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={
                 'unknown': [
@@ -385,6 +392,7 @@ def test_get_invoice_lines_for_user_check_status(mock_pricing_data_event, mock_s
     ]
     lines = utils.get_invoice_lines_for_user(
         agendas=[agenda1, agenda2],
+        agendas_pricings=[agenda_pricing],
         user_external_id='user:1',
         subscriptions={
             # matching dates
@@ -569,6 +577,7 @@ def test_get_invoice_lines_for_user_check_status_subscriptions_dates(mock_pricin
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[agenda],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={
                 'agenda': [
@@ -590,6 +599,7 @@ def test_get_invoice_lines_for_user_check_status_subscriptions_dates(mock_pricin
     assert (
         utils.get_invoice_lines_for_user(
             agendas=[agenda],
+            agendas_pricings=[agenda_pricing],
             user_external_id='user:1',
             subscriptions={
                 'agenda': [
@@ -610,6 +620,7 @@ def test_get_invoice_lines_for_user_check_status_subscriptions_dates(mock_pricin
 
     lines = utils.get_invoice_lines_for_user(
         agendas=[agenda],
+        agendas_pricings=[agenda_pricing],
         user_external_id='user:1',
         subscriptions={
             'agenda': [
@@ -700,6 +711,7 @@ def test_get_invoice_lines_for_user_check_status_agenda_pricing_dates(mock_statu
             mock_pricing_data_event.return_value = {'foo': 'bar', 'pricing': 42}
             lines = utils.get_invoice_lines_for_user(
                 agendas=[agenda],
+                agendas_pricings=AgendaPricing.objects.all(),
                 user_external_id='user:1',
                 subscriptions={
                     'agenda': [
@@ -734,6 +746,7 @@ def test_get_invoice_lines_for_user_check_status_agenda_pricing_dates(mock_statu
             mock_pricing_data_event.return_value = {'foo': 'bar', 'pricing': 42}
             lines = utils.get_invoice_lines_for_user(
                 agendas=[agenda],
+                agendas_pricings=AgendaPricing.objects.all(),
                 user_external_id='user:1',
                 subscriptions={
                     'agenda': [
@@ -767,6 +780,7 @@ def test_get_invoice_lines_for_user_check_status_agenda_pricing_dates(mock_statu
         with pricing_data_event_patch as mock_pricing_data_event:
             lines = utils.get_invoice_lines_for_user(
                 agendas=[agenda],
+                agendas_pricings=AgendaPricing.objects.all(),
                 user_external_id='user:1',
                 subscriptions={
                     'agenda': [
@@ -835,6 +849,7 @@ def test_get_invoice_lines_for_user_check_status_pricing_error(mock_pricing_data
     ]
     lines = utils.get_invoice_lines_for_user(
         agendas=[agenda],
+        agendas_pricings=[agenda_pricing],
         user_external_id='user:1',
         subscriptions={
             'agenda': [
@@ -952,6 +967,7 @@ def test_get_all_invoice_lines(mock_user_lines):
     assert mock_user_lines.call_args_list == [
         mock.call(
             agendas=[agenda1, agenda2],
+            agendas_pricings=mock.ANY,
             user_external_id='user:1',
             subscriptions={'foo': 'bar1'},
             date_start=datetime.date(2022, 9, 1),
@@ -959,6 +975,7 @@ def test_get_all_invoice_lines(mock_user_lines):
         ),
         mock.call(
             agendas=[agenda1, agenda2],
+            agendas_pricings=mock.ANY,
             user_external_id='user:2',
             subscriptions={'foo': 'bar2'},
             date_start=datetime.date(2022, 9, 1),
@@ -967,6 +984,138 @@ def test_get_all_invoice_lines(mock_user_lines):
     ]
 
 
+@mock.patch('lingo.invoicing.utils.get_check_status')
+def test_get_all_invoice_lines_queryset(mock_status):
+    # don't mock get_pricing_data_for_event to check all querysets
+    category1 = CriteriaCategory.objects.create(label='Foo1', slug='foo1')
+    criteria1 = Criteria.objects.create(label='Bar1', slug='bar1', condition='True', category=category1)
+    category2 = CriteriaCategory.objects.create(label='Foo2', slug='foo2')
+    criteria2 = Criteria.objects.create(label='Bar2', slug='bar2', condition='True', category=category2)
+
+    pricing1 = Pricing.objects.create(label='Foo bar 1')
+    pricing1.criterias.add(criteria1, criteria2)
+    pricing1.categories.add(category1, through_defaults={'order': 1})
+    pricing1.categories.add(category2, through_defaults={'order': 2})
+    pricing2 = Pricing.objects.create(label='Foo bar 2')
+    pricing2.criterias.add(criteria1, criteria2)
+    pricing2.categories.add(category1, through_defaults={'order': 1})
+    pricing2.categories.add(category2, through_defaults={'order': 2})
+
+    agenda1 = Agenda.objects.create(label='Agenda 1')
+    agenda_pricing11 = AgendaPricing.objects.create(
+        pricing=pricing1,
+        date_start=datetime.date(year=2022, month=9, day=1),
+        date_end=datetime.date(year=2022, month=10, day=1),
+    )
+    agenda_pricing11.agendas.add(agenda1)
+    agenda_pricing12 = AgendaPricing.objects.create(
+        pricing=pricing1,
+        date_start=datetime.date(year=2022, month=10, day=1),
+        date_end=datetime.date(year=2022, month=11, day=1),
+    )
+    agenda_pricing12.agendas.add(agenda1)
+
+    agenda2 = Agenda.objects.create(label='Agenda 2')
+    agenda_pricing21 = AgendaPricing.objects.create(
+        pricing=pricing2,
+        date_start=datetime.date(year=2022, month=9, day=1),
+        date_end=datetime.date(year=2022, month=10, day=1),
+    )
+    agenda_pricing21.agendas.add(agenda2)
+    agenda_pricing22 = AgendaPricing.objects.create(
+        pricing=pricing2,
+        date_start=datetime.date(year=2022, month=10, day=1),
+        date_end=datetime.date(year=2022, month=11, day=1),
+    )
+    agenda_pricing22.agendas.add(agenda2)
+
+    mock_status.return_value = [
+        {
+            'event': {
+                'agenda': 'agenda-1',
+                'start_datetime': '2022-09-01T12:00:00+02:00',
+                'slug': 'event',
+                'label': 'Event',
+            },
+            'check_status': {'foo': 'bar'},
+            'booking': {'foo': 'baz'},
+        },
+        {
+            'event': {
+                'agenda': 'agenda-1',
+                'start_datetime': '2022-10-01T12:00:00+02:00',
+                'slug': 'event',
+                'label': 'Event',
+            },
+            'check_status': {'foo': 'bar'},
+            'booking': {'foo': 'baz'},
+        },
+        {
+            'event': {
+                'agenda': 'agenda-2',
+                'start_datetime': '2022-09-01T12:00:00+02:00',
+                'slug': 'event',
+                'label': 'Event',
+            },
+            'check_status': {'foo': 'bar'},
+            'booking': {'foo': 'baz'},
+        },
+        {
+            'event': {
+                'agenda': 'agenda-2',
+                'start_datetime': '2022-10-01T12:00:00+02:00',
+                'slug': 'event',
+                'label': 'Event',
+            },
+            'check_status': {'foo': 'bar'},
+            'booking': {'foo': 'baz'},
+        },
+    ]
+
+    with CaptureQueriesContext(connection) as ctx:
+        lines = utils.get_all_invoice_lines(
+            agendas=[agenda1, agenda2],
+            subscriptions={
+                'user:1': {
+                    'agenda-1': [
+                        {
+                            'user_external_id': 'user:1',
+                            'date_start': '2022-09-01',
+                            'date_end': '2022-11-01',
+                        },
+                    ],
+                    'agenda-2': [
+                        {
+                            'user_external_id': 'user:1',
+                            'date_start': '2022-09-01',
+                            'date_end': '2022-11-01',
+                        },
+                    ],
+                },
+                'user:2': {
+                    'agenda-1': [
+                        {
+                            'user_external_id': 'user:2',
+                            'date_start': '2022-09-01',
+                            'date_end': '2022-11-01',
+                        },
+                    ],
+                    'agenda-2': [
+                        {
+                            'user_external_id': 'user:2',
+                            'date_start': '2022-09-01',
+                            'date_end': '2022-11-01',
+                        },
+                    ],
+                },
+            },
+            date_start=datetime.date(2022, 9, 1),
+            date_end=datetime.date(2022, 11, 1),
+        )
+        assert lines
+        assert len(ctx.captured_queries) == 7
+
+
 def test_generate_invoices_from_lines():
     regie1 = Regie.objects.create(label='Regie 1')
     regie2 = Regie.objects.create(label='Regie 2')