Merge pull request #724 from calvinmclean/js_scripts

Add more Jetstream-specific scripts

CHANGELOG.md

@@ -28,8 +28,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 ### Added
   - Jetstream specific scripts
     ([#723](https://github.com/cyverse/atmosphere/pull/723))
+  - More Jetstream specific scripts
+    ([#724](https://github.com/cyverse/atmosphere/pull/724))
+
 ### Fixed
-  - Filter out system-y metadata keys introduced in Rocky which we aren't currently using and make Glance client confused 
+  - Filter out system-y metadata keys introduced in Rocky which we aren't currently using and make Glance client confused
     ([#701](https://github.com/cyverse/atmosphere/pull/701))
 
 ## [v36-4](https://github.com/cyverse/atmosphere/compare/v36-3...v36-4) - 2019-09-09

scripts/jetstream/enforce_special_allocation_quota.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+import argparse
+
+import django
+
+django.setup()
+
+from core.models import Identity, Quota
+
+from service.quota import set_provider_quota, _get_hard_limits
+
+
+def main():
+    """
+    Reset quotas for users on a particular allocation source
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dry-run', action='store_true', help='Print output rather than perform operation')
+    parser.add_argument('--allocation-source', required=True,
+                        help='Allocation source name to reset quotas for, e.g. TG-ASC160018')
+    parser.add_argument('--quota-id', required=True, type=int,
+                        help='Quota ID to set')
+    parser.add_argument('--whitelist-quota-ids', type=lambda s: [int(item.strip()) for item in s.split(',')],
+                        help='Quota IDs that are acceptable and won\'t be overwritten (comma separated)')
+    parser.set_defaults(dry_run=False)
+    args = parser.parse_args()
+    return run_command(args.dry_run,
+                       args.allocation_source,
+                       args.quota_id,
+                       args.whitelist_quota_ids)
+
+
+def run_command(dry_run, allocation_source, quota_id, whitelist_quota_ids):
+    query = '''SELECT identity.*
+FROM
+  (SELECT
+     atmosphere_user.id,
+     atmosphere_user.username,
+     array_agg(allocation_source.name) allocation_sources
+   FROM public.atmosphere_user
+     LEFT JOIN public.user_allocation_source ON atmosphere_user.id = user_allocation_source.user_id
+     LEFT JOIN public.allocation_source ON user_allocation_source.allocation_source_id = allocation_source.id
+   GROUP BY atmosphere_user.id
+   HAVING array_agg(allocation_source.name) = %(allocation_source)s
+   ORDER BY username) AS users
+  LEFT OUTER JOIN public.identity ON identity.created_by_id = users.id
+  LEFT OUTER JOIN public.quota ON identity.quota_id = quota.id
+WHERE
+  identity.quota_id NOT IN %(acceptable_quota_ids)s 
+  AND identity.quota_id IS NOT NULL;'''
+
+    params = {
+        'allocation_source': '{%s}' % allocation_source,
+        'acceptable_quota_ids': tuple(whitelist_quota_ids + [quota_id])
+    }
+    identities = Identity.objects.raw(raw_query=query, params=params)
+    # print('Query to find identities:')
+    # print(identities.query)
+
+    quota_to_set = Quota.objects.get(id=quota_id)
+    # print('\nDesignated quota: {}'.format(quota_to_set))
+
+    for identity in identities:
+        print('\n========================\n')
+        print('Identity: {}'.format(identity))
+        print('Current quota in DB: {}'.format(identity.quota))
+        current_provider_quota = _get_hard_limits(identity)
+        print('Current quota in provider: {}'.format(current_provider_quota))
+        if dry_run:
+            print('DRY-RUN: Not changing quota to {}'.format(quota_id))
+        else:
+            print('Changing provider quota to {}...'.format(quota_id))
+            updated_provider_quota = set_provider_quota(identity.uuid, quota=quota_to_set)
+            print('Updated provider quota: {}'.format(updated_provider_quota))
+            print('Changing DB quota to {}...'.format(quota_id))
+            identity.quota = quota_to_set
+            identity.save()
+            print('Updated DB quota: {}'.format(identity.quota_id))
+
+
+if __name__ == "__main__":
+    main()

scripts/jetstream/enforce_special_allocation_quota.sh

@@ -5,4 +5,4 @@ export DJANGO_SETTINGS_MODULE="atmosphere.settings"
 export PYTHONPATH="$HOME:$PYTHONPATH"
 cd $HOME
 . $VIRTUAL/bin/activate
-python $HOME/scripts/enforce_special_allocation_quota.py --allocation-source TG-ASC160018 --quota-id 33 --whitelist-quota-ids 57
+python $HOME/scripts/jetstream/enforce_special_allocation_quota.py --allocation-source TG-ASC160018 --quota-id 33 --whitelist-quota-ids 57

scripts/jetstream/ezb_users_over_allocation.sh

@@ -6,4 +6,4 @@ echo ""
 WEEK_AGO=$(date --date='a week ago' +%F)
 . /opt/env/atmo/bin/activate
 cd /opt/dev/atmosphere
-cat scripts/ezb_users_over_allocation.sql | sed "s/2017-12-01/${WEEK_AGO}/" | ./manage.py dbshell
+cat scripts/jetstream/ezb_users_over_allocation.sql | sed "s/2017-12-01/${WEEK_AGO}/" | ./manage.py dbshell

scripts/jetstream/ezb_users_over_allocation.sql

@@ -0,0 +1,15 @@
+SELECT
+  atmosphere_user.username,
+  atmosphere_user.date_joined,
+  user_allocation_snapshot.compute_used,
+  user_allocation_snapshot.burn_rate,
+  user_allocation_snapshot.updated
+FROM public.user_allocation_snapshot
+  LEFT JOIN public.allocation_source ON user_allocation_snapshot.allocation_source_id = allocation_source.id
+  LEFT JOIN public.atmosphere_user ON user_allocation_snapshot.user_id = atmosphere_user.id
+WHERE
+  allocation_source.name = 'TG-ASC160018'
+  AND atmosphere_user.is_active
+  AND user_allocation_snapshot.updated > '2017-12-01'
+  AND user_allocation_snapshot.compute_used > 2000
+ORDER BY compute_used DESC;