11import abc
22import functools
33import json
4+ import logging
45from typing import NamedTuple
56
67from coldfront .core .allocation import models as allocation_models
78from coldfront .core .resource import models as resource_models
89
9- from coldfront_plugin_cloud import attributes
10+ from coldfront_plugin_cloud import attributes , tasks , utils
1011from coldfront_plugin_cloud .models .quota_models import QuotaSpecs
1112
1213
14+ logger = logging .getLogger (__name__ )
15+
16+
1317class ResourceAllocator (abc .ABC ):
1418 resource_type = ""
1519
@@ -45,6 +49,97 @@ def get_or_create_federated_user(self, username):
4549 user = self .create_federated_user (username )
4650 return user
4751
52+ def set_default_quota_on_allocation (self , coldfront_attr ):
53+ resource_quotaspecs = self .resource_quotaspecs
54+ value = resource_quotaspecs .root [coldfront_attr ].quota_by_su_quantity (
55+ self .allocation .quantity
56+ )
57+ utils .set_attribute_on_allocation (self .allocation , coldfront_attr , value )
58+ return value
59+
60+ def set_users (self , project_id , apply ):
61+ coldfront_users = allocation_models .AllocationUser .objects .filter (
62+ allocation = self .allocation , status__name = "Active"
63+ )
64+ cluster_users = self .get_users (project_id )
65+ failed_validation = False
66+
67+ # Create users that exist in coldfront but not in the resource
68+ for coldfront_user in coldfront_users :
69+ if coldfront_user .user .username not in cluster_users :
70+ failed_validation = True
71+ logger .info (
72+ f"{ coldfront_user .user .username } is not part of { project_id } "
73+ )
74+ if apply :
75+ tasks .add_user_to_allocation (coldfront_user .pk )
76+
77+ # remove users that are in the resource but not in coldfront
78+ users = set (
79+ [coldfront_user .user .username for coldfront_user in coldfront_users ]
80+ )
81+ for allocation_user in cluster_users :
82+ if allocation_user not in users :
83+ failed_validation = True
84+ logger .info (
85+ f"{ allocation_user } exists in the resource { project_id } but not in coldfront"
86+ )
87+ if apply :
88+ self .remove_role_from_user (allocation_user , project_id )
89+
90+ return failed_validation
91+
92+ def check_and_apply_quota_attr (
93+ self ,
94+ project_id ,
95+ attr : str ,
96+ expected_quota : int ,
97+ current_quota : int ,
98+ apply : bool ,
99+ ):
100+ if current_quota is None and expected_quota is None :
101+ msg = (
102+ f"Value for quota for { attr } is not set anywhere"
103+ f" on { self .allocation_str } "
104+ )
105+
106+ if apply :
107+ expected_quota = self .set_default_quota_on_allocation (attr )
108+ msg = f"Added default quota for { attr } to { self .allocation_str } to { expected_quota } "
109+ logger .info (msg )
110+ elif current_quota is not None and expected_quota is None :
111+ msg = (
112+ f'Attribute "{ attr } " expected on { self .allocation_str } but not set.'
113+ f" Current quota is { current_quota } ."
114+ )
115+
116+ if apply :
117+ utils .set_attribute_on_allocation (self .allocation , attr , current_quota )
118+ expected_quota = (
119+ current_quota # To pass `current_quota != expected_quota` check
120+ )
121+ msg = f"{ msg } Attribute set to match current quota."
122+ logger .info (msg )
123+
124+ if current_quota != expected_quota :
125+ msg = (
126+ f"Value for quota for { attr } = { current_quota } does not match expected"
127+ f" value of { expected_quota } on { self .allocation_str } "
128+ )
129+ logger .info (msg )
130+
131+ if apply :
132+ try :
133+ self .set_quota (project_id )
134+ logger .info (f"Quota for { project_id } was out of date. Reapplied!" )
135+ except Exception as e :
136+ logger .info (f"setting cluster-side quota failed: { e } " )
137+ return
138+
139+ @functools .cached_property
140+ def allocation_str (self ):
141+ return f'allocation { self .allocation .pk } of project "{ self .allocation .project .title } "'
142+
48143 @functools .cached_property
49144 def auth_url (self ):
50145 return self .resource .get_attribute (attributes .RESOURCE_AUTH_URL ).rstrip ("/" )
@@ -54,7 +149,11 @@ def member_role_name(self):
54149 return self .resource .get_attribute (attributes .RESOURCE_ROLE ) or "member"
55150
56151 @abc .abstractmethod
57- def set_project_configuration (self , project_id , dry_run = False ):
152+ def set_project_configuration (self , project_id , apply = True ):
153+ pass
154+
155+ @abc .abstractmethod
156+ def get_project (self , project_id ):
58157 pass
59158
60159 @abc .abstractmethod
@@ -85,6 +184,10 @@ def get_quota(self, project_id):
85184 def create_federated_user (self , unique_id ):
86185 pass
87186
187+ @abc .abstractmethod
188+ def get_users (self , unique_id ):
189+ pass
190+
88191 @abc .abstractmethod
89192 def get_federated_user (self , unique_id ):
90193 pass
0 commit comments