Own CBS import commit and rollback in main so the full flow is one transaction.

Author: tkalir

anyway/parsers/cbs/executor.py

@@ -1157,8 +1157,6 @@ def _build_hebrew_tables_and_derived_data():
     with log_duration("Importing safety data tables"):
         sd_utils.load_data(session=db.session, should_commit=False)
     logging.debug("Completed load of safety data tables")
-    db.session.commit()
-    logging.debug("data commited successfully")
 
 
 def main(batch_size, source, load_start_year=None, allow_missing=False):
@@ -1179,7 +1177,10 @@ def main(batch_size, source, load_start_year=None, allow_missing=False):
 
         _log_import_summary(total, started)
         _build_hebrew_tables_and_derived_data()
+        db.session.commit()
+        logging.debug("data committed successfully")
     except Exception as ex:
+        db.session.rollback()
         print("Traceback: {0}".format(traceback.format_exc()))
         raise CBSParsingFailed(message=str(ex))
         # Todo - send an email that an exception occured

tests/parsers/cbs/test_executor.py

@@ -19,6 +19,9 @@ def test_import_streets_is_called_once_when_source_is_s3(monkeypatch, mock_s3_da
     monkeypatch.setattr('anyway.parsers.cbs.executor.delete_cbs_entries', delete_cbs_entries)
     monkeypatch.setattr('anyway.parsers.cbs.executor.fill_db_geo_data', MagicMock())
     monkeypatch.setattr('anyway.parsers.cbs.executor.create_tables', MagicMock())
+    monkeypatch.setattr('anyway.parsers.cbs.executor.recreate_table_for_location_extraction', MagicMock())
+    monkeypatch.setattr('anyway.parsers.cbs.executor.sd_utils.load_data', MagicMock())
+    monkeypatch.setattr('anyway.parsers.cbs.executor.db.session.commit', MagicMock())
 
     # Act
     main(batch_size=MagicMock(), source='s3')
@@ -28,9 +31,14 @@ def test_import_streets_is_called_once_when_source_is_s3(monkeypatch, mock_s3_da
 
 
 def test_cbs_parsing_failed_is_raised_when_something_bad_happens(monkeypatch):
+    monkeypatch.setattr('anyway.parsers.cbs.executor._import_from_s3', MagicMock(return_value=0))
+    monkeypatch.setattr('anyway.parsers.cbs.executor.fill_db_geo_data', MagicMock())
     monkeypatch.setattr('anyway.parsers.cbs.executor.create_tables',
                         MagicMock(side_effect=Exception('something bad')))
+    rollback = MagicMock()
+    monkeypatch.setattr('anyway.parsers.cbs.executor.db.session.rollback', rollback)
 
     with pytest.raises(CBSParsingFailed, match='Exception occurred while loading the cbs data: something bad'):
-        main(batch_size=MagicMock(), source=MagicMock())
+        main(batch_size=MagicMock(), source='s3')
 
+    rollback.assert_called_once()