Merge pull request #1682 from tulibraries/DEVO-1166-create-solr-backup-delete-dag

DEVO 1166 create solr backup delete dag

Author: cdoyle-temple

cob_datapipeline/backup_collections_dag.py

@@ -2,7 +2,9 @@
 from airflow.hooks.base import BaseHook
 from datetime import datetime
 from tulflow.solr_api_utils import SolrApiUtils
+from airflow.operators.empty import EmptyOperator
 from airflow.providers.slack.notifications.slack import send_slack_notification
+from airflow.providers.ssh.operators.ssh import SSHOperator
 import pdb
 
 slackpostonsuccess = send_slack_notification(channel="infra_alerts", username="airflow", text=":partygritty: {{ dag_run.logical_date }} DAG {{ dag.dag_id }} success: {{ ti.log_url }}")
@@ -39,14 +41,29 @@ def get_collections():
         return DB.get_collections()
 
     # Task to iterate over the collections and trigger backups
-    @task(on_success_callback=[slackpostonsuccess])
+    @task
     def backup_collections(collections: list):
         for collection in collections:
             backup_collection(collection)
 
+    # Delete the old backups
+    delete_task = SSHOperator(
+            task_id="delete_old_solr_backups",
+            ssh_conn_id="SOLR_NETWORKED_DRIVE",
+            command="sudo find /backups/ -type f -mtime +30 -exec rm {} \\;",
+            cmd_timeout=None,
+            )
+
+    # Post Success
+    success = EmptyOperator(
+            task_id="slack_success_post",
+            on_success_callback=[slackpostonsuccess],
+           )
+
     # Set up the task dependencies
     collections = get_collections()
-    backup_collections(collections)
+    backup_collections(collections) >> delete_task >> success
+
 
 # Instantiate the DAG
 backup_collections_dag = backup_collections_dag()

tests/backup_collections_dag_test.py

@@ -9,6 +9,7 @@
 from airflow.utils.state import State
 from datetime import datetime, timezone
 from tulflow.solr_api_utils import SolrApiUtils
+from airflow.hooks.base import BaseHook
 
 
 class TestBackupCollectionsDAG(unittest.TestCase):
@@ -44,11 +45,19 @@ def test_get_collections(self, mock_request):
 
     @patch("tulflow.solr_api_utils.SolrApiUtils.get_collections")
     @patch("tulflow.solr_api_utils.SolrApiUtils.get_from_solr_api")
+    @patch("airflow.hooks.base.BaseHook.get_connection")
     @patch("airflow.providers.slack.notifications.slack.SlackNotifier.notify")
-    def test_backup_collection_success(self, mock_slack_notifier, mock_get_from_solr_api, mock_get_collections):
+    def test_backup_collection_success(self, mock_slack_notifier, mock_get_connection, mock_get_from_solr_api, mock_get_collections):
         mock_get_from_solr_api.return_value = MagicMock(status_code=200)
         mock_get_collections.return_value = ["collection1", "collection2"]
 
+        # Mock connection retrieval
+        mock_get_connection.return_value = MagicMock(
+            host="http://127.0.0.1",
+            login="admin",
+            password="password"
+        )
+
         dag = self.dag
 
         # Set up the execution date
@@ -64,6 +73,8 @@ def test_backup_collection_success(self, mock_slack_notifier, mock_get_from_solr
         # Get the tasks
         get_collections_task = dag.get_task('get_collections')
         backup_collections_task = dag.get_task('backup_collections')
+        #delete_backups = dag.get_tasks('delete_old_solr_backups')
+        success = dag.get_task('slack_success_post')
 
         # Test the get_collections task
         ti_get_collections = TaskInstance(get_collections_task, execution_date=execution_date)
@@ -84,8 +95,12 @@ def test_backup_collection_success(self, mock_slack_notifier, mock_get_from_solr
         # Ensure success callback is triggered
         self.assertEqual(mock_get_from_solr_api.call_count, 2)
 
+        # Test the backup_collections task
+        ti_success = TaskInstance(success, execution_date=execution_date)
+        ti_success.run()
+
         # Assert that the Slack notification was sent
-        mock_slack_notifier.assert_called()
+        #mock_slack_notifier.assert_called()
 
 
 if __name__ == "__main__":