From e9f7e76da0e54b3f831f174d9e37931244088b35 Mon Sep 17 00:00:00 2001
From: David Schultz <david.schultz@icecube.wisc.edu>
Date: Thu, 14 Nov 2024 16:19:23 -0600
Subject: [PATCH] add a flag to not auto-complete filtering datasets

---
 iceprod/rest/handlers/datasets.py             |  3 +++
 iceprod/scheduled_tasks/dataset_completion.py |  3 +++
 tests/rest/datasets_test.py                   | 23 +++++++++++++++++++
 .../dataset_completion_test.py                | 21 +++++++++++++++++
 4 files changed, 50 insertions(+)

diff --git a/iceprod/rest/handlers/datasets.py b/iceprod/rest/handlers/datasets.py
index c6532a4e2..7a929802d 100644
--- a/iceprod/rest/handlers/datasets.py
+++ b/iceprod/rest/handlers/datasets.py
@@ -119,6 +119,7 @@ async def post(self):
             'priority': float,
             'debug': bool,
             'jobs_immutable': bool,
+            'always_active': bool,
             'status': str,
             'auth_groups_read': list,
         }
@@ -168,6 +169,8 @@ async def post(self):
         if 'jobs_immutable' not in data:
             data['jobs_immutable'] = False
         data['truncated'] = False
+        if 'always_active' not in data:
+            data['always_active'] = False
 
         # insert
         ret = await self.db.datasets.insert_one(data)
diff --git a/iceprod/scheduled_tasks/dataset_completion.py b/iceprod/scheduled_tasks/dataset_completion.py
index c38dcd79f..33bda4412 100644
--- a/iceprod/scheduled_tasks/dataset_completion.py
+++ b/iceprod/scheduled_tasks/dataset_completion.py
@@ -45,6 +45,9 @@ async def run(rest_client, debug=False):
         # test if dataset is complete / failed
         try:
             dataset = await rest_client.request('GET', f'/datasets/{dataset_id}')
+            if dataset.get('always_active', False):
+                logger.info('dataset %s always active, skipping', dataset_id)
+                continue
             jobs = await rest_client.request('GET', f'/datasets/{dataset_id}/job_counts/status')
             if dataset.get('truncated', False) or sum(jobs[s] for s in jobs) >= dataset['jobs_submitted']:
                 # jobs are all buffered / materialized
diff --git a/tests/rest/datasets_test.py b/tests/rest/datasets_test.py
index 35a2bd86a..5472a0a7f 100644
--- a/tests/rest/datasets_test.py
+++ b/tests/rest/datasets_test.py
@@ -55,6 +55,29 @@ async def test_rest_datasets_post_bad_role(server):
     assert exc_info.value.response.status_code == 403
 
 
+async def test_rest_datasets_post_always_active(server):
+    client = server(roles=['user'])
+
+    data = {
+        'description': 'blah',
+        'tasks_per_job': 4,
+        'jobs_submitted': 1,
+        'tasks_submitted': 4,
+        'group': 'users',
+        'always_active': True
+    }
+    ret = await client.request('POST', '/datasets', data)
+    dataset_id = ret['result'].split('/')[-1]
+
+    ret = await client.request('GET', '/datasets')
+    assert dataset_id in ret
+
+    ret = await client.request('GET', f'/datasets/{dataset_id}')
+    for k in data:
+        assert k in ret
+        assert data[k] == ret[k]
+
+
 async def test_rest_datasets_err(server):
     client = server(roles=['system'])
     with pytest.raises(requests.exceptions.HTTPError) as exc_info:
diff --git a/tests/scheduled_tasks/dataset_completion_test.py b/tests/scheduled_tasks/dataset_completion_test.py
index 237265fce..2b7e894eb 100644
--- a/tests/scheduled_tasks/dataset_completion_test.py
+++ b/tests/scheduled_tasks/dataset_completion_test.py
@@ -92,3 +92,24 @@ async def client(method, url, args=None):
 
     # check it normally hides the error
     await dataset_completion.run(rc, debug=False)
+
+
+
+async def test_always_active():
+    rc = MagicMock()
+    dataset_summaries = {'processing':['foo']}
+    async def client(method, url, args=None):
+        logger.info('REST: %s, %s', method, url)
+        if url.startswith('/dataset_summaries'):
+            return dataset_summaries
+        elif url == '/config/foo':
+            return {}
+        elif url == '/datasets/foo':
+            return {'jobs_submitted':2, 'tasks_submitted':2, 'always_active': True}
+        else:
+            raise Exception()
+    rc.request = client
+    client.called = False
+
+    await dataset_completion.run(rc, debug=True)
+    assert client.called