diff --git a/AIPscan/Aggregator/database_helpers.py b/AIPscan/Aggregator/database_helpers.py index 8bcf8c16..b3503a7f 100644 --- a/AIPscan/Aggregator/database_helpers.py +++ b/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() diff --git a/AIPscan/Aggregator/tasks.py b/AIPscan/Aggregator/tasks.py index 866226db..abfd316e 100644 --- a/AIPscan/Aggregator/tasks.py +++ b/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() diff --git a/AIPscan/Aggregator/tests/test_database_helpers.py b/AIPscan/Aggregator/tests/test_database_helpers.py index 159a85e7..06de425b 100644 --- a/AIPscan/Aggregator/tests/test_database_helpers.py +++ b/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") diff --git a/AIPscan/Aggregator/tests/test_tasks.py b/AIPscan/Aggregator/tests/test_tasks.py index 66791b5b..8d0695f5 100644 --- a/AIPscan/Aggregator/tests/test_tasks.py +++ b/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( diff --git a/AIPscan/models.py b/AIPscan/models.py index 21b4c82f..8a77c2dd 100644 --- a/AIPscan/models.py +++ b/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 "".format(self.agent_type, self.agent_value) diff --git a/AIPscan/test_helpers.py b/AIPscan/test_helpers.py index 2f126e2f..a7634b71 100644 --- a/AIPscan/test_helpers.py +++ b/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