Merge pull request #788 from alphagov/fix-billable-units-query

Fix performance of billable units query

app/dao/notifications_dao.py

@@ -1,14 +1,13 @@
+
 import pytz
 from datetime import (
     datetime,
     timedelta,
-    date
-)
-from itertools import groupby
+    date)
 
 from flask import current_app
 from werkzeug.datastructures import MultiDict
-from sqlalchemy import (desc, func, or_, and_, asc)
+from sqlalchemy import (desc, func, or_, and_, asc, extract)
 from sqlalchemy.orm import joinedload
 
 from app import db, create_uuid
@@ -226,24 +225,31 @@ def get_notifications_for_job(service_id, job_id, filter_dict=None, page=1, page
 def get_notification_billable_unit_count_per_month(service_id, year):
     start, end = get_financial_year(year)
 
+    """
+     The query needs to sum the billable_units per month, but this needs to be the month in BST (British Standard Time).
+     The database stores all timestamps as UTC without the timezone.
+      - First set the timezone on created_at to UTC
+      - then convert the timezone to BST (or Europe/London)
+      - lastly truncate the datetime to month to group the sum of the billable_units
+    """
+    month = func.date_trunc("month",
+                            func.timezone("Europe/London", func.timezone("UTC",
+                                                                         NotificationHistory.created_at)))
     notifications = db.session.query(
-        NotificationHistory.created_at,
-        NotificationHistory.billable_units
-    ).order_by(
-        NotificationHistory.created_at
+        month,
+        func.sum(NotificationHistory.billable_units)
     ).filter(
         NotificationHistory.billable_units != 0,
         NotificationHistory.service_id == service_id,
         NotificationHistory.created_at >= start,
         NotificationHistory.created_at < end
-    )
+    ).group_by(
+        month
+    ).order_by(
+        month
+    ).all()
 
-    return [
-        (month, sum(count for _, count in row))
-        for month, row in groupby(
-            notifications, lambda row: get_bst_month(row[0])
-        )
-    ]
+    return [(datetime.strftime(x[0], "%B"), x[1]) for x in notifications]
 
 
 @statsd(namespace="dao")
@@ -394,13 +400,11 @@ def get_financial_year(year):
 
 
 def get_april_fools(year):
-    return datetime(
-        year, 4, 1, 0, 0, 0, 0,
-        pytz.timezone("Europe/London")
-    ).astimezone(pytz.utc)
-
-
-def get_bst_month(datetime):
-    return pytz.utc.localize(datetime).astimezone(
-        pytz.timezone("Europe/London")
-    ).strftime('%B')
+    """
+     This function converts the start of the financial year April 1, 00:00 as BST (British Standard Time) to UTC,
+     the tzinfo is lastly removed from the datetime becasue the database stores the timestamps without timezone.
+     :param year: the year to calculate the April 1, 00:00 BST for
+     :return: the datetime of April 1 for the given year, for example 2016 = 2016-03-31 23:00:00
+    """
+    return pytz.timezone('Europe/London').localize(datetime(year, 4, 1, 0, 0, 0)).astimezone(pytz.UTC).replace(
+        tzinfo=None)

tests/app/dao/test_notification_dao.py

@@ -39,7 +39,8 @@ from app.dao.notifications_dao import (
     update_notification_status_by_reference,
     dao_delete_notifications_and_history_by_id,
     dao_timeout_notifications,
-    get_financial_year)
+    get_financial_year,
+    get_april_fools)
 
 from app.dao.services_dao import dao_update_service
 
@@ -150,11 +151,11 @@ def test_should_by_able_to_get_template_count_from_notifications_history(notify_
 
 
 def test_template_history_should_ignore_test_keys(
-        notify_db,
-        notify_db_session,
-        sample_team_api_key,
-        sample_test_api_key,
-        sample_api_key
+    notify_db,
+    notify_db_session,
+    sample_team_api_key,
+    sample_test_api_key,
+    sample_api_key
 ):
     sms = sample_template(notify_db, notify_db_session)
 
@@ -764,28 +765,31 @@ def test_get_all_notifications_for_job_by_status(notify_db, notify_db_session, s
 
 
 def test_get_notification_billable_unit_count_per_month(notify_db, notify_db_session, sample_service):
-    for year, month, day in (
-            (2017, 1, 15),  # ↓ 2016 financial year
-            (2016, 8, 1),
-            (2016, 7, 15),
-            (2016, 4, 15),
-            (2016, 4, 15),
-            (2016, 4, 1),  # ↓ 2015 financial year
-            (2016, 3, 31),
-            (2016, 1, 15)
+
+    for year, month, day, hour, minute, second in (
+        (2017, 1, 15, 23, 59, 59),  # ↓ 2016 financial year
+        (2016, 9, 30, 23, 59, 59),  # counts in October with BST conversion
+        (2016, 6, 30, 23, 50, 20),
+        (2016, 7, 15, 9, 20, 25),
+        (2016, 4, 1, 1, 1, 00),
+        (2016, 4, 1, 0, 0, 00),
+        (2016, 3, 31, 23, 00, 1),  # counts in April with BST conversion
+        (2015, 4, 1, 13, 8, 59),  # ↓ 2015 financial year
+        (2015, 11, 20, 22, 40, 45),
+        (2016, 1, 31, 23, 30, 40)  # counts in January no BST conversion in winter
     ):
         sample_notification(
             notify_db, notify_db_session, service=sample_service,
             created_at=datetime(
-                year, month, day, 0, 0, 0, 0
-            ) - timedelta(hours=1, seconds=1)  # one second before midnight
+                year, month, day, hour, minute, second, 0
+            )
         )
 
     for financial_year, months in (
-            (2017, []),
-            (2016, [('April', 2), ('July', 2), ('January', 1)]),
-            (2015, [('January', 1), ('March', 2)]),
-            (2014, [])
+        (2017, []),
+        (2016, [('April', 3), ('July', 2), ('October', 1), ('January', 1)]),
+        (2015, [('April', 1), ('November', 1), ('January', 1)]),
+        (2014, [])
     ):
         assert get_notification_billable_unit_count_per_month(
             sample_service.id, financial_year
@@ -1294,7 +1298,11 @@ def test_should_exclude_test_key_notifications_by_default(
 
 def test_get_financial_year():
     start, end = get_financial_year(2000)
-    assert start.tzinfo == pytz.utc
-    assert start.isoformat() == '2000-04-01T00:01:00+00:00'
-    assert end.tzinfo == pytz.utc
-    assert end.isoformat() == '2001-04-01T00:01:00+00:00'
+    assert str(start) == '2000-03-31 23:00:00'
+    assert str(end) == '2001-03-31 23:00:00'
+
+
+def test_get_april_fools():
+    april_fools = get_april_fools(2016)
+    assert str(april_fools) == '2016-03-31 23:00:00'
+    assert april_fools.tzinfo is None