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,102 @@ 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+ attr : str ,
95+ expected_quota : int ,
96+ current_quota : int ,
97+ apply : bool ,
98+ ) -> bool :
99+ failed_validation = False
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+ failed_validation = True
106+
107+ if apply :
108+ expected_quota = self .set_default_quota_on_allocation (attr )
109+ msg = f"Added default quota for { attr } to { self .allocation_str } to { expected_quota } "
110+ logger .info (msg )
111+ elif current_quota is not None and expected_quota is None :
112+ msg = (
113+ f'Attribute "{ attr } " expected on { self .allocation_str } but not set.'
114+ f" Current quota is { current_quota } ."
115+ )
116+ failed_validation = True
117+
118+ if apply :
119+ utils .set_attribute_on_allocation (self .allocation , attr , current_quota )
120+ expected_quota = (
121+ current_quota # To pass `current_quota != expected_quota` check
122+ )
123+ msg = f"{ msg } Attribute set to match current quota."
124+ logger .info (msg )
125+
126+ if current_quota != expected_quota :
127+ msg = (
128+ f"Value for quota for { attr } = { current_quota } does not match expected"
129+ f" value of { expected_quota } on { self .allocation_str } "
130+ )
131+ logger .info (msg )
132+ failed_validation = True
133+
134+ return failed_validation
135+
136+ # if apply:
137+ # try:
138+ # self.set_quota(project_id)
139+ # logger.info(f"Quota for {project_id} was out of date. Reapplied!")
140+ # except Exception as e:
141+ # logger.info(f"setting cluster-side quota failed: {e}")
142+ # return
143+
144+ @functools .cached_property
145+ def allocation_str (self ):
146+ return f'allocation { self .allocation .pk } of project "{ self .allocation .project .title } "'
147+
48148 @functools .cached_property
49149 def auth_url (self ):
50150 return self .resource .get_attribute (attributes .RESOURCE_AUTH_URL ).rstrip ("/" )
@@ -54,7 +154,11 @@ def member_role_name(self):
54154 return self .resource .get_attribute (attributes .RESOURCE_ROLE ) or "member"
55155
56156 @abc .abstractmethod
57- def set_project_configuration (self , project_id , dry_run = False ):
157+ def set_project_configuration (self , project_id , apply = True ):
158+ pass
159+
160+ @abc .abstractmethod
161+ def get_project (self , project_id ):
58162 pass
59163
60164 @abc .abstractmethod
@@ -85,6 +189,10 @@ def get_quota(self, project_id):
85189 def create_federated_user (self , unique_id ):
86190 pass
87191
192+ @abc .abstractmethod
193+ def get_users (self , unique_id ):
194+ pass
195+
88196 @abc .abstractmethod
89197 def get_federated_user (self , unique_id ):
90198 pass
0 commit comments