Try using nested transactions to rollback

This will need to be merged before https://github.com/alphagov/notifications-admin/pull/3855, it will be that until the admin PR is merged the annual billing will be set twice, but that's not an issue.

app/billing/rest.py

@@ -83,6 +83,8 @@ def get_free_sms_fragment_limit(service_id):
                 annual_billing = dao_create_or_update_annual_billing_for_year(service_id,
                                                                               annual_billing.free_sms_fragment_limit,
                                                                               financial_year_start)
+                from app import db
+                db.session.commit()
 
     return jsonify(annual_billing.serialize_free_sms_items()), 200
 

app/dao/annual_billing_dao.py

@@ -1,12 +1,12 @@
 from flask import current_app
 
 from app import db
-from app.dao.dao_utils import transactional
+from app.dao.dao_utils import transactional, nested_transactional
 from app.dao.date_util import get_current_financial_year_start_year
 from app.models import AnnualBilling
 
 
-@transactional
+@nested_transactional
 def dao_create_or_update_annual_billing_for_year(service_id, free_sms_fragment_limit, financial_year_start):
     result = dao_get_free_sms_fragment_limit_for_year(service_id, financial_year_start)
 
@@ -53,7 +53,7 @@ def dao_get_all_free_sms_fragment_limit(service_id):
     ).order_by(AnnualBilling.financial_year_start).all()
 
 
-def set_default_free_allowance_for_service(service, year_start=None):
+def set_default_free_allowance_for_service(service, year_start=None, commit=True):
     default_free_sms_fragment_limits = {
         'central': {
             2020: 250_000,
@@ -90,6 +90,7 @@ def set_default_free_allowance_for_service(service, year_start=None):
     }
     if not year_start:
         year_start = get_current_financial_year_start_year()
+    # handle cases where the year is less than 2020 or greater than 2021
     if year_start < 2020:
         year_start = 2020
     if year_start > 2021:

app/dao/dao_utils.py

@@ -4,6 +4,19 @@ from functools import wraps
 from app import db
 from app.history_meta import create_history
 
+def nested_transactional(func):
+    @wraps(func)
+    def commit_or_rollback(*args, **kwargs):
+        try:
+            db.session.begin_nested()
+            res = func(*args, **kwargs)
+            db.session.commit()
+            return res
+        except Exception:
+            db.session.rollback()
+            raise
+
+    return commit_or_rollback
 
 def transactional(func):
     @wraps(func)

app/dao/organisation_dao.py

@@ -1,7 +1,7 @@
 from sqlalchemy.sql.expression import func
 
 from app import db
-from app.dao.dao_utils import VersionOptions, transactional, version_class
+from app.dao.dao_utils import VersionOptions, transactional, version_class, nested_transactional
 from app.models import Domain, Organisation, Service, User
 
 
@@ -105,7 +105,7 @@ def _update_organisation_services(organisation, attribute, only_where_none=True)
         db.session.add(service)
 
 
-@transactional
+@nested_transactional
 @version_class(Service)
 def dao_add_service_to_organisation(service, organisation_id):
     organisation = Organisation.query.filter_by(

app/dao/services_dao.py

@@ -7,7 +7,7 @@ from sqlalchemy.orm import joinedload
 from sqlalchemy.sql.expression import and_, asc, case, func
 
 from app import db
-from app.dao.dao_utils import VersionOptions, transactional, version_class
+from app.dao.dao_utils import VersionOptions, transactional, version_class, nested_transactional
 from app.dao.date_util import get_current_financial_year
 from app.dao.email_branding_dao import dao_get_email_branding_by_name
 from app.dao.letter_branding_dao import dao_get_letter_branding_by_name
@@ -284,7 +284,7 @@ def dao_fetch_service_by_id_and_user(service_id, user_id):
     ).one()
 
 
-@transactional
+@nested_transactional
 @version_class(Service)
 def dao_create_service(
     service,

app/organisation/rest.py

@@ -1,8 +1,9 @@
 
 from flask import Blueprint, abort, current_app, jsonify, request
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
 
 from app.config import QueueNames
+from app.dao.annual_billing_dao import set_default_free_allowance_for_service
 from app.dao.fact_billing_dao import fetch_usage_year_for_organisation
 from app.dao.organisation_dao import (
     dao_add_service_to_organisation,
@@ -118,7 +119,15 @@ def link_service_to_organisation(organisation_id):
     service = dao_fetch_service_by_id(data['service_id'])
     service.organisation = None
 
+    from app import db
+
+    try:
         dao_add_service_to_organisation(service, organisation_id)
+        set_default_free_allowance_for_service(service, year_start=None)
+        db.session.commit()
+    except Exception:
+        db.session.rollback()
+        raise
 
     return '', 204
 

app/service/rest.py

@@ -4,12 +4,13 @@ from datetime import datetime
 from flask import Blueprint, current_app, jsonify, request
 from notifications_utils.letter_timings import letter_can_be_cancelled
 from notifications_utils.timezones import convert_utc_to_bst
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
 from sqlalchemy.orm.exc import NoResultFound
 
 from app.aws import s3
 from app.config import QueueNames
 from app.dao import fact_notification_status_dao, notifications_dao
+from app.dao.annual_billing_dao import set_default_free_allowance_for_service
 from app.dao.api_key_dao import (
     expire_api_key,
     get_model_api_keys,
@@ -253,7 +254,16 @@ def create_service():
     # unpack valid json into service object
     valid_service = Service.from_json(data)
 
+    from app import db
+
+    try:
+        db.session.begin_nested()
         dao_create_service(valid_service, user)
+        set_default_free_allowance_for_service(valid_service, year_start=None)
+        db.session.commit()
+    except Exception:
+        db.session.rollback()
+        raise
 
     return jsonify(data=service_schema.dump(valid_service).data), 201
 

tests/app/dao/test_annual_billing_dao.py

@@ -91,3 +91,21 @@ def test_set_default_free_allowance_for_service_using_correct_year(sample_servic
         25000,
         2020
     )
+
+
+@freeze_time('2021-04-01 14:02:00')
+def test_set_default_free_allowance_for_service_updates_existing_year(sample_service):
+    set_default_free_allowance_for_service(service=sample_service, year_start=None)
+    annual_billing = AnnualBilling.query.all()
+    assert not sample_service.organisation_type
+    assert len(annual_billing) == 1
+    assert annual_billing[0].service_id == sample_service.id
+    assert annual_billing[0].free_sms_fragment_limit == 10000
+
+    sample_service.organisation_type = 'central'
+
+    set_default_free_allowance_for_service(service=sample_service, year_start=None)
+    annual_billing = AnnualBilling.query.all()
+    assert len(annual_billing) == 1
+    assert annual_billing[0].service_id == sample_service.id
+    assert annual_billing[0].free_sms_fragment_limit == 150000

tests/app/dao/test_organisation_dao.py

@@ -233,6 +233,7 @@ def test_add_service_to_organisation(sample_service, sample_organisation):
     sample_organisation.crown = False
 
     dao_add_service_to_organisation(sample_service, sample_organisation.id)
+    db.session.commit()
 
     assert len(sample_organisation.services) == 1
     assert sample_organisation.services[0].id == sample_service.id

tests/app/organisation/test_rest.py

@@ -3,13 +3,14 @@ from datetime import datetime
 
 import pytest
 from freezegun import freeze_time
+from sqlalchemy.exc import SQLAlchemyError
 
 from app.dao.organisation_dao import (
     dao_add_service_to_organisation,
     dao_add_user_to_organisation,
 )
 from app.dao.services_dao import dao_archive_service
-from app.models import Organisation
+from app.models import AnnualBilling, Organisation
 from tests.app.db import (
     create_annual_billing,
     create_domain,
@@ -491,19 +492,25 @@ def test_post_update_organisation_set_mou_emails_signed_by(
         }
 
 
-def test_post_link_service_to_organisation(admin_request, sample_service, sample_organisation):
+def test_post_link_service_to_organisation(admin_request, sample_service):
     data = {
         'service_id': str(sample_service.id)
     }
-
+    organisation = create_organisation(organisation_type='central')
+    assert len(organisation.services) == 0
+    assert len(AnnualBilling.query.all()) == 0
     admin_request.post(
         'organisation.link_service_to_organisation',
         _data=data,
-        organisation_id=sample_organisation.id,
+        organisation_id=organisation.id,
         _expected_status=204
     )
 
-    assert len(sample_organisation.services) == 1
+    assert len(organisation.services) == 1
+    assert sample_service.organisation_type == 'central'
+    annual_billing = AnnualBilling.query.all()
+    assert len(annual_billing) == 1
+    assert annual_billing[0].free_sms_fragment_limit == 150000
 
 
 def test_post_link_service_to_another_org(
@@ -511,7 +518,8 @@ def test_post_link_service_to_another_org(
     data = {
         'service_id': str(sample_service.id)
     }
-
+    assert len(sample_organisation.services) == 0
+    assert not sample_service.organisation_type
     admin_request.post(
         'organisation.link_service_to_organisation',
         _data=data,
@@ -520,8 +528,9 @@ def test_post_link_service_to_another_org(
     )
 
     assert len(sample_organisation.services) == 1
+    assert not sample_service.organisation_type
 
-    new_org = create_organisation()
+    new_org = create_organisation(organisation_type='central')
     admin_request.post(
         'organisation.link_service_to_organisation',
         _data=data,
@@ -530,6 +539,10 @@ def test_post_link_service_to_another_org(
     )
     assert not sample_organisation.services
     assert len(new_org.services) == 1
+    assert sample_service.organisation_type == 'central'
+    annual_billing = AnnualBilling.query.all()
+    assert len(annual_billing) == 1
+    assert annual_billing[0].free_sms_fragment_limit == 150000
 
 
 def test_post_link_service_to_organisation_nonexistent_organisation(
@@ -569,6 +582,23 @@ def test_post_link_service_to_organisation_missing_payload(
     )
 
 
+def test_link_service_to_organisation_updates_service_if_annual_billing_update_fails(
+        mocker, admin_request, sample_service, sample_organisation
+):
+    mocker.patch('app.organisation.rest.set_default_free_allowance_for_service', raises=SQLAlchemyError)
+    data = {
+        'service_id': str(sample_service.id)
+    }
+    admin_request.post(
+        'organisation.link_service_to_organisation',
+        organisation_id=str(sample_organisation.id),
+        _data=data,
+        _expected_status=204
+    )
+    assert sample_service.organisation_id == sample_organisation.id
+    assert len(AnnualBilling.query.all()) == 0
+
+
 def test_rest_get_organisation_services(
         admin_request, sample_organisation, sample_service):
     dao_add_service_to_organisation(sample_service, sample_organisation.id)

tests/app/service/test_rest.py

@@ -6,6 +6,7 @@ from unittest.mock import ANY
 import pytest
 from flask import current_app, url_for
 from freezegun import freeze_time
+from sqlalchemy.exc import SQLAlchemyError
 
 from app.dao.organisation_dao import dao_add_service_to_organisation
 from app.dao.service_sms_sender_dao import dao_get_sms_senders_by_service_id
@@ -31,6 +32,7 @@ from app.models import (
     SERVICE_PERMISSION_TYPES,
     SMS_TYPE,
     UPLOAD_LETTERS,
+    AnnualBilling,
     EmailBranding,
     InboundNumber,
     Notification,
@@ -482,6 +484,46 @@ def test_create_service_with_domain_sets_organisation(
         assert json_resp['data']['organisation'] is None
 
 
+def test_create_service_should_create_annual_billing_for_service(
+    admin_request, sample_user
+):
+    data = {
+        'name': 'created service',
+        'user_id': str(sample_user.id),
+        'message_limit': 1000,
+        'restricted': False,
+        'active': False,
+        'email_from': 'created.service',
+        'created_by': str(sample_user.id)
+    }
+    assert len(AnnualBilling.query.all()) == 0
+    admin_request.post('service.create_service', _data=data, _expected_status=201)
+
+    annual_billing = AnnualBilling.query.all()
+    assert len(annual_billing) == 1
+
+
+def test_create_service_should_create_service_if_annual_billing_query_fails(
+    admin_request, sample_user, mocker
+):
+    mocker.patch('app.service.rest.set_default_free_allowance_for_service', raises=SQLAlchemyError)
+    data = {
+        'name': 'created service',
+        'user_id': str(sample_user.id),
+        'message_limit': 1000,
+        'restricted': False,
+        'active': False,
+        'email_from': 'created.service',
+        'created_by': str(sample_user.id)
+    }
+    assert len(AnnualBilling.query.all()) == 0
+    admin_request.post('service.create_service', _data=data, _expected_status=201)
+
+    annual_billing = AnnualBilling.query.all()
+    assert len(annual_billing) == 0
+    assert len(Service.query.filter(Service.name == 'created service').all()) == 1
+
+
 def test_create_service_inherits_branding_from_organisation(
     admin_request,
     sample_user,