Skip to content

Commit

Permalink
Fix daily limit counting when scheduling jobs (#2112)
Browse files Browse the repository at this point in the history
* Fix daily limit counting when scheduling jobs

- When scheduling a job, the daily limits in Redis will only be
  incremented if the job was scheduled for the current day
- When refreshing daily count keys from the DB the underlying query will
  now include counts from jobs scheduled for the current day
- Experimenting with a new decorator for feature flags
- Fixed an issue with the dev container where the poetry installation will
  sometimes not be detected by adding the installation dir to the $PATH

* Bump utils version and update lock file

* fix tests?

* Fix tests

* Bump waffles commit sha in test action
  • Loading branch information
whabanks authored Feb 22, 2024
1 parent f4c7a95 commit 9d1d445
Show file tree
Hide file tree
Showing 10 changed files with 164 additions and 40 deletions.
6 changes: 4 additions & 2 deletions .devcontainer/scripts/notify-dev-entrypoint.sh
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,10 @@ cd /workspace
echo -e "fpath+=/.zfunc" >> ~/.zshrc
echo -e "autoload -Uz compinit && compinit"

pip install poetry==${POETRY_VERSION} \
&& poetry --version
pip install poetry==${POETRY_VERSION}
export PATH=$PATH:/home/vscode/.local/bin/
which poetry
poetry --version

# Initialize poetry autocompletions
mkdir ~/.zfunc
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/test.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ jobs:
run: |
cp -f .env.example .env
- name: Checks for new endpoints against AWS WAF rules
uses: cds-snc/notification-utils/.github/actions/waffles@2da74685e0ffb220f0403e1f2584e783be99bbad # 52.1.0
uses: cds-snc/notification-utils/.github/actions/waffles@06a40db6286f525fe3551e029418458d33342592 # 52.1.0
with:
app-loc: '/github/workspace'
app-libs: '/github/workspace/env/site-packages'
Expand Down
43 changes: 28 additions & 15 deletions app/dao/services_dao.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
CROWN_ORGANISATION_TYPES,
EMAIL_TYPE,
INTERNATIONAL_SMS_TYPE,
JOB_STATUS_SCHEDULED,
KEY_TYPE_TEST,
NHS_ORGANISATION_TYPES,
NON_CROWN_ORGANISATION_TYPES,
Expand Down Expand Up @@ -424,20 +425,24 @@ def dao_fetch_todays_stats_for_service(service_id):

def fetch_todays_total_message_count(service_id):
midnight = get_midnight(datetime.now(tz=pytz.utc))
scheduled = (
db.session.query(func.coalesce(func.sum(Job.notification_count), 0).label("count")).filter(
Job.service_id == service_id,
Job.job_status == JOB_STATUS_SCHEDULED,
Job.scheduled_for >= midnight,
Job.scheduled_for < midnight + timedelta(days=1),
)
).first()

result = (
db.session.query(func.count(Notification.id).label("count"))
.filter(
db.session.query(func.coalesce(func.count(Notification.id), 0).label("count")).filter(
Notification.service_id == service_id,
Notification.key_type != KEY_TYPE_TEST,
Notification.created_at > midnight,
Notification.created_at >= midnight,
)
.group_by(
Notification.notification_type,
Notification.status,
)
.first()
)
return 0 if result is None else result.count
).first()

return result.count + scheduled.count


def fetch_todays_total_sms_count(service_id):
Expand All @@ -461,17 +466,25 @@ def fetch_service_email_limit(service_id: uuid.UUID) -> int:

def fetch_todays_total_email_count(service_id: uuid.UUID) -> int:
midnight = get_midnight(datetime.now(tz=pytz.utc))
scheduled = (
db.session.query(func.coalesce(func.sum(Job.notification_count), 0).label("total_scheduled_notifications")).filter(
Job.service_id == service_id,
Job.job_status == JOB_STATUS_SCHEDULED,
Job.scheduled_for > midnight,
Job.scheduled_for < midnight + timedelta(hours=23, minutes=59, seconds=59),
)
).first()

result = (
db.session.query(func.count(Notification.id).label("total_email_notifications"))
.filter(
db.session.query(func.coalesce(func.count(Notification.id), 0).label("total_email_notifications")).filter(
Notification.service_id == service_id,
Notification.key_type != KEY_TYPE_TEST,
Notification.created_at > midnight,
Notification.notification_type == "email",
)
.first()
)
return 0 if result is None else result.total_email_notifications
).first()

return result.total_email_notifications + scheduled.total_scheduled_notifications


def _stats_for_service_query(service_id):
Expand Down
11 changes: 9 additions & 2 deletions app/email_limit_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from flask import current_app
from notifications_utils.clients.redis import email_daily_count_cache_key
from notifications_utils.decorators import requires_feature

from app import redis_store
from app.dao.services_dao import fetch_todays_total_email_count
Expand All @@ -20,9 +21,15 @@ def fetch_todays_email_count(service_id: UUID) -> int:
return int(total_email_count)


@requires_feature("REDIS_ENABLED")
def increment_todays_email_count(service_id: UUID, increment_by: int) -> None:
if not current_app.config["REDIS_ENABLED"]:
return
fetch_todays_email_count(service_id) # to make sure it's set in redis
cache_key = email_daily_count_cache_key(service_id)
redis_store.incrby(cache_key, increment_by)


@requires_feature("REDIS_ENABLED")
def decrement_todays_email_count(service_id: UUID, decrement_by: int) -> None:
fetch_todays_email_count(service_id)
cache_key = email_daily_count_cache_key(service_id)
redis_store.decrby(cache_key, decrement_by)
35 changes: 20 additions & 15 deletions app/job/rest.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from datetime import datetime

import dateutil
from flask import Blueprint, current_app, jsonify, request
from notifications_utils.recipients import RecipientCSV
Expand All @@ -20,6 +22,7 @@
from app.dao.notifications_dao import get_notifications_for_job
from app.dao.services_dao import dao_fetch_service_by_id
from app.dao.templates_dao import dao_get_template_by_id
from app.email_limit_utils import decrement_todays_email_count
from app.errors import InvalidRequest, register_errors
from app.models import (
EMAIL_TYPE,
Expand Down Expand Up @@ -67,7 +70,7 @@ def cancel_job(service_id, job_id):
job = dao_get_future_scheduled_job_by_id_and_service_id(job_id, service_id)
job.job_status = JOB_STATUS_CANCELLED
dao_update_job(job)

decrement_todays_email_count(service_id, job.notification_count)
return get_job_by_service_and_job_id(service_id, job_id)


Expand Down Expand Up @@ -137,15 +140,23 @@ def create_job(service_id):
raise InvalidRequest("Create job is not allowed: service is inactive ", 403)

data = request.get_json()

data.update({"service": service_id})

try:
data.update(**get_job_metadata_from_s3(service_id, data["id"]))
except KeyError:
raise InvalidRequest({"id": ["Missing data for required field."]}, status_code=400)

if data.get("valid") != "True":
raise InvalidRequest("File is not valid, can't create job", 400)

data["template"] = data.pop("template_id")

template = dao_get_template_by_id(data["template"])
template_errors = unarchived_template_schema.validate({"archived": template.archived})

if template_errors:
raise InvalidRequest(template_errors, status_code=400)

job = get_job_from_s3(service_id, data["id"])
recipient_csv = RecipientCSV(
Expand All @@ -170,22 +181,16 @@ def create_job(service_id):

if not is_test_notification:
check_sms_daily_limit(service, len(recipient_csv))
increment_sms_daily_count_send_warnings_if_needed(service, len(recipient_csv))

if template.template_type == EMAIL_TYPE:
elif template.template_type == EMAIL_TYPE:
check_email_daily_limit(service, len(list(recipient_csv.get_rows())))
scheduled_for = datetime.fromisoformat(form.get("scheduled_for")) if form.get("scheduled_for") else None # noqa: F821

if data.get("valid") != "True":
raise InvalidRequest("File is not valid, can't create job", 400)

errors = unarchived_template_schema.validate({"archived": template.archived})

if errors:
raise InvalidRequest(errors, status_code=400)

if template.template_type == SMS_TYPE and not is_test_notification:
increment_sms_daily_count_send_warnings_if_needed(service, len(recipient_csv))
elif template.template_type == EMAIL_TYPE:
increment_email_daily_count_send_warnings_if_needed(service, len(list(recipient_csv.get_rows())))
if scheduled_for is None or not scheduled_for.date() > datetime.today().date():
increment_email_daily_count_send_warnings_if_needed(
authenticated_service, len(list(recipient_csv.get_rows())) # noqa: F821
)

data.update({"template_version": template.version})

Expand Down
2 changes: 1 addition & 1 deletion app/notifications/rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ def send_notification(notification_type: NotificationType):
)

simulated = simulated_recipient(notification_form["to"], notification_type)
if not simulated != api_user.key_type == KEY_TYPE_TEST:
if not simulated != api_user.key_type == KEY_TYPE_TEST and notification_type == EMAIL_TYPE:
check_email_daily_limit(authenticated_service, 1)

check_template_is_for_notification_type(notification_type, template.template_type)
Expand Down
5 changes: 4 additions & 1 deletion app/v2/notifications/post_notifications.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,10 @@ def post_bulk():

if template.template_type == EMAIL_TYPE and api_user.key_type != KEY_TYPE_TEST:
check_email_daily_limit(authenticated_service, len(list(recipient_csv.get_rows())))
increment_email_daily_count_send_warnings_if_needed(authenticated_service, len(list(recipient_csv.get_rows())))
scheduled_for = datetime.fromisoformat(form.get("scheduled_for")) if form.get("scheduled_for") else None

if scheduled_for is None or not scheduled_for.date() > datetime.today().date():
increment_email_daily_count_send_warnings_if_needed(authenticated_service, len(list(recipient_csv.get_rows())))

if template.template_type == SMS_TYPE:
# calculate the number of simulated recipients
Expand Down
97 changes: 95 additions & 2 deletions tests/app/dao/test_services_dao.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from app.models import (
EMAIL_TYPE,
INTERNATIONAL_SMS_TYPE,
JOB_STATUS_SCHEDULED,
KEY_TYPE_NORMAL,
KEY_TYPE_TEAM,
KEY_TYPE_TEST,
Expand All @@ -68,6 +69,7 @@
user_folder_permissions,
)
from app.schemas import service_schema
from tests.app.conftest import create_sample_job
from tests.app.db import (
create_annual_billing,
create_api_key,
Expand Down Expand Up @@ -1396,7 +1398,7 @@ def create_email_sms_letter_template():


class TestServiceEmailLimits:
def test_get_email_count_for_service(self, notify_db_session):
def test_get_email_count_for_service(self):
active_user_1 = create_user(email="[email protected]", state="active")
service = Service(
name="service_name",
Expand Down Expand Up @@ -1432,7 +1434,98 @@ def test_dao_fetch_todays_total_message_count_returns_0_with_yesterday_messages(
notification = save_notification(
create_notification(
created_at=yesterday,
template=create_template(service=create_service(service_name="tester"), template_type="email"),
template=create_template(service=create_service(service_name="tester123"), template_type="email"),
)
)
assert fetch_todays_total_message_count(notification.service.id) == 0

def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today(
self, notify_db, notify_db_session
):
service = create_service(service_name="tester12")
template = create_template(service=service, template_type="email")
today = datetime.utcnow().date()

create_sample_job(
notify_db,
notify_db_session,
service=service,
template=template,
scheduled_for=today,
job_status=JOB_STATUS_SCHEDULED,
notification_count=10,
)
save_notification(
create_notification(
created_at=today,
template=template,
)
)
assert fetch_todays_total_message_count(service.id) == 11

def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today_but_not_after_today(
self, notify_db, notify_db_session
):
service = create_service()
template = create_template(service=service, template_type="email")
today = datetime.utcnow().date()

create_sample_job(
notify_db,
notify_db_session,
service=service,
template=template,
scheduled_for=today,
job_status=JOB_STATUS_SCHEDULED,
notification_count=10,
)
save_notification(
create_notification(
created_at=today,
template=template,
)
)
create_sample_job(
notify_db,
notify_db_session,
service=service,
template=template,
scheduled_for=today + timedelta(days=1),
job_status=JOB_STATUS_SCHEDULED,
notification_count=10,
)

assert fetch_todays_total_message_count(service.id) == 11

def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today_but_not_before_today(
self, notify_db, notify_db_session
):
service = create_service()
template = create_template(service=service, template_type="email")
today = datetime.utcnow().date()

create_sample_job(
notify_db,
notify_db_session,
service=service,
template=template,
scheduled_for=today,
job_status=JOB_STATUS_SCHEDULED,
notification_count=10,
)
create_sample_job(
notify_db,
notify_db_session,
service=service,
template=template,
scheduled_for=today - timedelta(days=1),
job_status=JOB_STATUS_SCHEDULED,
notification_count=10,
)
save_notification(
create_notification(
created_at=today,
template=template,
)
)
assert fetch_todays_total_message_count(service.id) == 11
1 change: 1 addition & 0 deletions tests/app/job/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -451,6 +451,7 @@ def test_create_job_returns_404_if_template_does_not_exist(client, sample_servic
"app.job.rest.get_job_metadata_from_s3",
return_value={
"template_id": str(sample_service.id),
"valid": "True",
},
)
data = {
Expand Down
2 changes: 1 addition & 1 deletion tests/app/notifications/rest/test_send_notification.py
Original file line number Diff line number Diff line change
Expand Up @@ -466,7 +466,7 @@ def test_should_allow_api_call_if_under_day_limit_regardless_of_type(
sms_template = create_sample_template(notify_db, notify_db_session, service=service)
create_sample_notification(notify_db, notify_db_session, template=email_template, service=service)

data = {"to": sample_user.mobile_number, "template": str(sms_template.id)}
data = {"to": sample_user.mobile_number, "template": str(sms_template.id), "valid": "True"}

auth_header = create_authorization_header(service_id=service.id)

Expand Down

0 comments on commit 9d1d445

Please sign in to comment.