diff --git a/coldfront/plugins/fasrc/management/commands/id_import_new_allocations.py b/coldfront/plugins/fasrc/management/commands/id_import_new_allocations.py index 86864b7e7..e5ba48177 100644 --- 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 index d4c8b05ec..218314d46 100644 --- 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 index e7cf24e22..9100ad458 100644 --- 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