update starfish-related commands

fasrc · Jan 13, 2024 · 69d8fc6 · 69d8fc6
1 parent c870a9a
commit 69d8fc6
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/coldfront/plugins/fasrc/management/commands/id_import_new_allocations.py b/coldfront/plugins/fasrc/management/commands/id_import_new_allocations.py
@@ -10,6 +10,7 @@
 from django.core.exceptions import ValidationError
 from django.core.management.base import BaseCommand
 
+from coldfront.core.settings import ENV
 from coldfront.core.project.models import ProjectStatusChoice
 from coldfront.core.allocation.models import (
     AllocationUser,
@@ -23,8 +24,7 @@
 if ENV.bool('PLUGIN_SFTOCF', default=False):
     from coldfront.plugins.sftocf.utils import (
         StarFishRedash,
-        STARFISH_SERVER,
-        pull_sf_push_cf_redash
+        RedashDataPipeline,
     )
 from coldfront.plugins.fasrc.utils import (
     AllTheThingsConn,
@@ -58,7 +58,7 @@ def handle(self, *args, **options):
         # Remove allocations for labs not in Coldfront, add those labs to a list
         result_json_cleaned, proj_models = match_entries_with_projects(resp_json_by_lab)
 
-        redash_api = StarFishRedash(STARFISH_SERVER)
+        redash_api = StarFishRedash()
         allocation_usages = redash_api.return_query_results(query='subdirectory')
         subdir_type = AllocationAttributeType.objects.get(name='Subdirectory')
 
@@ -137,5 +137,7 @@ def handle(self, *args, **options):
             update_csv(added_allocations_df.to_dict(orient='records'),
                 './local_data/', f'added_allocations_{datetime.today().date()}.csv')
         push_quota_data(result_file)
-        pull_sf_push_cf_redash()
+        data_pull = RedashDataPipeline()
+        allocationquerymatch_objs, user_models = data_pull.clean_collected_data()
+        data_pull.update_coldfront_objects(allocationquerymatch_objs, user_models)
         return json.dumps(command_report, indent=2)
diff --git a/coldfront/plugins/sftocf/management/commands/id_new_storage_allocations.py b/coldfront/plugins/sftocf/management/commands/id_new_storage_allocations.py
@@ -8,7 +8,7 @@
 from django.core.management.base import BaseCommand
 
 from coldfront.core.allocation.models import Allocation
-from coldfront.plugins.sftocf.utils import StarFishRedash, STARFISH_SERVER
+from coldfront.plugins.sftocf.utils import StarFishRedash
 
 
 logger = logging.getLogger(__name__)
@@ -20,7 +20,7 @@ class Command(BaseCommand):
 
 
     def handle(self, *args, **kwargs):
-        redash = StarFishRedash(STARFISH_SERVER)
+        redash = StarFishRedash()
         subdir_results = redash.submit_query('subdirectory')
         data = subdir_results['query_result']['data']['rows']
         # remove entries with no group name or with group names that are not projects

diff --git a/coldfront/plugins/sftocf/management/commands/import_allocation_filepaths.py b/coldfront/plugins/sftocf/management/commands/import_allocation_filepaths.py
@@ -5,7 +5,7 @@
 
 from coldfront.core.resource.models import Resource
 from coldfront.core.allocation.models import Allocation
-from coldfront.plugins.sftocf.utils import StarFishRedash, compare_cf_sf_volumes, STARFISH_SERVER
+from coldfront.plugins.sftocf.utils import StarFishRedash, StarFishServer
 
 
 logger = logging.getLogger(__name__)
@@ -27,12 +27,12 @@ class Command(BaseCommand):
 
     def handle(self, *args, **kwargs):
         errors = []
-        redash = StarFishRedash(STARFISH_SERVER)
+        redash = StarFishRedash()
         subdir_results = redash.submit_query("subdirectory")
         data = subdir_results['query_result']['data']['rows']
         data = [result for result in data if result['group_name']]
-        vols_to_collect = compare_cf_sf_volumes()
-        searched_resources = [Resource.objects.get(name__contains=vol) for vol in vols_to_collect]
+        starfishserver = StarFishServer()
+        searched_resources = starfishserver.get_volumes_in_coldfront()
         allocations = Allocation.objects.filter(resources__in=searched_resources)
         for allocation in allocations:
             lab = allocation.project.title