Remove agents associated with deleted SS, #91. (#214)

AIPscan/Aggregator/database_helpers.py

@@ -98,7 +98,7 @@ def create_event_objects(fs_entry, file_id):
         db.session.commit()
 
 
-def _extract_agent_detail(agent):
+def _extract_agent_detail(agent, storage_service_id):
     """Pull the agent information from the agent record and return an
     agent object ready to insert into the database.
     """
@@ -109,20 +109,22 @@ def _extract_agent_detail(agent):
         linking_type_value=linking_type_value,
         agent_type=agent_type,
         agent_value=agent_value,
+        storage_service_id=storage_service_id,
     )
 
 
-def create_agent_objects(unique_agents):
+def create_agent_objects(unique_agents, storage_service_id):
     """Add our agents to the database. The list is already the
     equivalent of a set by the time it reaches here and so we don't
     need to perform any de-duplication.
     """
     for agent in unique_agents:
-        agent_obj = _extract_agent_detail(agent)
+        agent_obj = _extract_agent_detail(agent, storage_service_id)
         exists = Agent.query.filter_by(
             linking_type_value=agent_obj.linking_type_value,
             agent_type=agent_obj.agent_type,
             agent_value=agent_obj.agent_value,
+            storage_service_id=storage_service_id,
         ).count()
         if exists:
             continue
@@ -392,7 +394,7 @@ def process_aip_data(aip, mets):
     """
     tasks.get_mets.update_state(state="IN PROGRESS")
 
-    create_agent_objects(collect_mets_agents(mets))
+    create_agent_objects(collect_mets_agents(mets), aip.storage_service_id)
 
     all_files = mets.all_files()
 

AIPscan/Aggregator/tasks.py

@@ -21,7 +21,7 @@
 )
 from AIPscan.extensions import celery
 from AIPscan.helpers import file_sha256_hash
-from AIPscan.models import AIP, FetchJob, StorageService, get_mets_tasks
+from AIPscan.models import AIP, Agent, FetchJob, StorageService, get_mets_tasks
 
 logger = get_task_logger(__name__)
 
@@ -395,5 +395,12 @@ def delete_storage_service(storage_service_id):
     for mets_fetch_job in mets_fetch_jobs:
         if os.path.exists(mets_fetch_job.download_directory):
             shutil.rmtree(mets_fetch_job.download_directory)
+
+    # Delete agents associated to the deleted storage service
+    agents = Agent.query.filter_by(storage_service_id=storage_service_id).all()
+    for agent in agents:
+        logger.info("Deleting agent: %s", agent)
+        db.session.delete(agent)
+
     db.session.delete(storage_service)
     db.session.commit()

AIPscan/Aggregator/tests/test_database_helpers.py

@@ -159,6 +159,7 @@ def test_event_creation(
             linking_type_value="some_type_value",
             agent_type="an_agent_type",
             agent_value="an_agent_value",
+            storage_service_id="1",
         ),
     )
     mocked_events = mocker.patch("AIPscan.db.session.add")

AIPscan/Aggregator/tests/test_tasks.py

@@ -25,7 +25,7 @@
     VALID_JSON,
     MockResponse,
 )
-from AIPscan.models import AIP, FetchJob, StorageService
+from AIPscan.models import AIP, Agent, FetchJob, StorageService
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 FIXTURES_DIR = os.path.join(SCRIPT_DIR, "fixtures")
@@ -198,6 +198,9 @@ def test_delete_fetch_job_task(app_instance, tmpdir, mocker):
 def test_delete_storage_service_task(app_instance, tmpdir, mocker):
     """Test that storage service gets deleted by delete storage service job task logic."""
     storage_service = test_helpers.create_test_storage_service()
+    test_helpers.create_test_agent()
+
+    assert Agent.query.filter_by(storage_service_id=storage_service.id).count() == 1
 
     deleted_ss = StorageService.query.filter_by(id=storage_service.id).first()
     assert deleted_ss is not None
@@ -206,6 +209,7 @@ def test_delete_storage_service_task(app_instance, tmpdir, mocker):
 
     deleted_ss = StorageService.query.filter_by(id=storage_service.id).first()
     assert deleted_ss is None
+    assert Agent.query.filter_by(storage_service_id=storage_service.id).count() == 0
 
 
 @pytest.mark.parametrize(

AIPscan/models.py

@@ -470,11 +470,15 @@ class Agent(db.Model):
     linking_type_value = db.Column(db.String(255), index=True)
     agent_type = db.Column(db.String(255), index=True)
     agent_value = db.Column(db.String(255), index=True)
+    storage_service_id = db.Column(
+        db.Integer(), db.ForeignKey("storage_service.id"), nullable=False
+    )
 
-    def __init__(self, linking_type_value, agent_type, agent_value):
+    def __init__(self, linking_type_value, agent_type, agent_value, storage_service_id):
         self.linking_type_value = linking_type_value
         self.agent_type = agent_type
         self.agent_value = agent_value
+        self.storage_service_id = storage_service_id
 
     def __repr__(self):
         return "<Agent '{}: {}'>".format(self.agent_type, self.agent_value)

AIPscan/test_helpers.py

@@ -123,6 +123,7 @@ def create_test_agent(**kwargs):
         agent_value=kwargs.get(
             "agent_value", 'username="user one", first_name="", last_name=""'
         ),
+        storage_service_id=kwargs.get("storage_service_id", 1),
     )
     _add_test_object_to_db(agent)
     return agent