Try a more magical way of doing nested transactions

db update/insert.

Using a savepoint for the multiple transactions allows us to rollback if
there is an error when executing the second db transaction.
However, this does add a bit of complexity. Developers need to manage
the db session when calling multiple nested tranactions.

Unit tests have been added to test this functionality and some end to
end tests have been done to make sure all transactions are rollback if
there is an exception while executing the transaction.

app/dao/annual_billing_dao.py

@@ -6,7 +6,6 @@ from app.dao.date_util import get_current_financial_year_start_year
 from app.models import AnnualBilling
 
 
-@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 +52,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, commit=True):
+def set_default_free_allowance_for_service(service, year_start=None):
     default_free_sms_fragment_limits = {
         'central': {
             2020: 250_000,

app/dao/dao_utils.py

@@ -1,16 +1,34 @@
 import itertools
+from contextlib import contextmanager
 from functools import wraps
 
 from app import db
 from app.history_meta import create_history
 
 
+@contextmanager
+def nested_transaction():
+    try:
+        db.session.begin_nested()
+        yield
+        db.session.commit()
+
+        if not db.session.registry().transaction.nested:
+            db.session.commit()
+    except Exception:
+        db.session.rollback()
+        raise
+
+
 def transactional(func):
     @wraps(func)
     def commit_or_rollback(*args, **kwargs):
         try:
             res = func(*args, **kwargs)
-            db.session.commit()
+
+            if not db.session.registry().transaction.nested:
+                db.session.commit()
+
             return res
         except Exception:
             db.session.rollback()

app/organisation/rest.py

@@ -2,6 +2,8 @@
 from flask import Blueprint, abort, current_app, jsonify, request
 from sqlalchemy.exc import IntegrityError, SQLAlchemyError
 
+from app.dao.dao_utils import nested_transaction
+
 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
@@ -119,17 +121,9 @@ def link_service_to_organisation(organisation_id):
     service = dao_fetch_service_by_id(data['service_id'])
     service.organisation = None
 
-    dao_add_service_to_organisation(service, organisation_id)
-    # Need to do the annual billing update in a separate transaction because the both the
-    # dao_add_service_to_organisation and set_default_free_allowance_for_service are wrapped in a transaction.
-    # Catch and report an error if the annual billing doesn't happen - but don't rollback the service update.
-    try:
+    with nested_transaction():
+        dao_add_service_to_organisation(service, organisation_id)
         set_default_free_allowance_for_service(service, year_start=None)
-    except SQLAlchemyError:
-        # No need to worry about key errors because service.organisation_type has a foreign key to organisation_types
-        current_app.logger.exception(
-            f"Exception caught when trying to update annual billing when the organisation "
-            f"changed for service: {service.id} to organisation: {organisation_id}")
 
     return '', 204
 

app/service/rest.py

@@ -18,7 +18,7 @@ from app.dao.api_key_dao import (
     save_model_api_key,
 )
 from app.dao.broadcast_service_dao import set_broadcast_service_type
-from app.dao.dao_utils import dao_rollback
+from app.dao.dao_utils import dao_rollback, nested_transaction
 from app.dao.date_util import get_financial_year
 from app.dao.fact_notification_status_dao import (
     fetch_monthly_template_usage_for_service,
@@ -254,18 +254,9 @@ def create_service():
     # unpack valid json into service object
     valid_service = Service.from_json(data)
 
-    dao_create_service(valid_service, user)
-
-    # Need to do the annual billing update in a separate transaction because the both the
-    # dao_add_service_to_organisation and set_default_free_allowance_for_service are wrapped in a transaction.
-    # Catch and report an error if the annual billing doesn't happen - but don't rollback the service update.
-    try:
+    with nested_transaction():
+        dao_create_service(valid_service, user)
         set_default_free_allowance_for_service(valid_service, year_start=None)
-    except SQLAlchemyError:
-        # No need to worry about key errors because service.organisation_type has a foreign key to organisation_types
-        current_app.logger.exception(
-            f"Exception caught when trying to insert annual billing creating a service {valid_service.id} "
-            f"for organisation_type {valid_service.organisation_type}")
 
     return jsonify(data=service_schema.dump(valid_service).data), 201
 

tests/app/organisation/test_rest.py

@@ -505,14 +505,54 @@ def test_post_link_service_to_organisation(admin_request, sample_service):
         organisation_id=organisation.id,
         _expected_status=204
     )
-
     assert len(organisation.services) == 1
     assert sample_service.organisation_type == 'central'
+
+
+def test_post_link_service_to_organisation_inserts_annual_billing(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=organisation.id,
+        _expected_status=204
+    )
+
     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_rollback_service_if_annual_billing_update_fails(
+        admin_request, sample_service, mocker
+):
+    mocker.patch('app.dao.annual_billing_dao.dao_create_or_update_annual_billing_for_year',
+                 side_effect=SQLAlchemyError)
+    data = {
+        'service_id': str(sample_service.id)
+    }
+    assert not sample_service.organisation_type
+
+    organisation = create_organisation(organisation_type='central')
+    assert len(organisation.services) == 0
+    assert len(AnnualBilling.query.all()) == 0
+    with pytest.raises(expected_exception=SQLAlchemyError):
+        admin_request.post(
+                'organisation.link_service_to_organisation',
+                _data=data,
+                organisation_id=organisation.id,
+                _expected_status=404
+            )
+    assert not sample_service.organisation_type
+    assert len(organisation.services) == 0
+    assert len(AnnualBilling.query.all()) == 0
+
+
 def test_post_link_service_to_another_org(
         admin_request, sample_service, sample_organisation):
     data = {
@@ -582,20 +622,20 @@ def test_post_link_service_to_organisation_missing_payload(
     )
 
 
-def test_link_service_to_organisation_updates_service_if_annual_billing_update_fails(
+def test_link_service_to_organisation_does_not_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)
+    mocker.patch('app.organisation.rest.set_default_free_allowance_for_service', side_effect=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
+    with pytest.raises(expected_exception=SQLAlchemyError):
+        admin_request.post(
+            'organisation.link_service_to_organisation',
+            organisation_id=str(sample_organisation.id),
+            _data=data,
+        )
+    assert not sample_service.organisation_id
     assert len(AnnualBilling.query.all()) == 0
 
 

tests/app/service/test_rest.py

@@ -506,7 +506,7 @@ def test_create_service_should_create_annual_billing_for_service(
 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)
+    mocker.patch('app.service.rest.set_default_free_allowance_for_service', side_effect=SQLAlchemyError)
     data = {
         'name': 'created service',
         'user_id': str(sample_user.id),