8
8
from ansible_base .authentication .models import Authenticator , AuthenticatorUser
9
9
from ansible_base .oauth2_provider .models import OAuth2Application
10
10
from ansible_base .rbac import permission_registry
11
- from ansible_base .rbac .models import DABContentType , RoleDefinition
11
+ from ansible_base .rbac .models import DABContentType , DABPermission , RoleDefinition
12
12
from test_app .models import EncryptionModel , InstanceGroup , Inventory , Organization , Team , User
13
13
14
14
@@ -20,6 +20,9 @@ def create_large(self, data_counts):
20
20
start = time .time ()
21
21
self .stdout .write ('' )
22
22
self .stdout .write ('About to create large demo data set. This will take a while.' )
23
+
24
+ # Create standard models first
25
+ created_org_ids = []
23
26
for cls in (Organization , Team , User ):
24
27
count = data_counts [cls ._meta .model_name ]
25
28
for i in range (count ):
@@ -28,9 +31,74 @@ def create_large(self, data_counts):
28
31
if cls is User :
29
32
data = {'username' : name }
30
33
elif cls is Team :
31
- data ['organization_id' ] = i + 1 # fudged, teams fewer than orgs
32
- cls .objects .create (** data )
34
+ # Use actual created organization IDs, cycling through them
35
+ if created_org_ids :
36
+ data ['organization_id' ] = created_org_ids [i % len (created_org_ids )]
37
+ else :
38
+ raise ValueError ("Teams cannot be created before organizations" )
39
+ obj = cls .objects .create (** data )
40
+ # Collect organization IDs for team creation
41
+ if cls is Organization :
42
+ created_org_ids .append (obj .id )
33
43
self .stdout .write (f'Created { count } { cls ._meta .model_name } ' )
44
+
45
+ # Create RoleDefinitions with permissions
46
+ if 'roledefinition' in data_counts :
47
+ rd_count = data_counts ['roledefinition' ]
48
+ org_ct = DABContentType .objects .get_for_model (Organization )
49
+
50
+ for i in range (rd_count ):
51
+ # Create some sample permissions for each role definition
52
+ perm1 = DABPermission .objects .create (name = f'Can view large role { i } ' , codename = f'view_large_role_{ i } ' , content_type = org_ct )
53
+ perm2 = DABPermission .objects .create (name = f'Can edit large role { i } ' , codename = f'edit_large_role_{ i } ' , content_type = org_ct )
54
+
55
+ # Create role definition with Organization content type
56
+ rd = RoleDefinition .objects .create (name = f'Large Role Definition { i } ' , description = f'Large demo role definition { i } ' , content_type = org_ct )
57
+
58
+ # Add permissions to the role definition
59
+ rd .permissions .add (perm1 , perm2 )
60
+
61
+ self .stdout .write (f'Created { rd_count } role definitions with permissions' )
62
+
63
+ # Create permission assignments for users and teams
64
+ if created_org_ids and 'user' in data_counts and 'team' in data_counts :
65
+ # Get created users and teams
66
+ large_users = list (User .objects .filter (username__startswith = 'large_user_' ))
67
+ large_teams = list (Team .objects .filter (name__startswith = 'large_team_' ))
68
+ large_orgs = list (Organization .objects .filter (name__startswith = 'large_organization_' ))
69
+ large_rds = list (RoleDefinition .objects .filter (name__startswith = 'Large Role Definition' ))
70
+
71
+ # Give over 25 permissions to users
72
+ user_permissions_given = 0
73
+ for user in large_users :
74
+ for rd in large_rds :
75
+ for org in large_orgs :
76
+ rd .give_permission (user , org )
77
+ user_permissions_given += 1
78
+ if user_permissions_given >= 25 :
79
+ break
80
+ if user_permissions_given >= 25 :
81
+ break
82
+ if user_permissions_given >= 25 :
83
+ break
84
+
85
+ # Give over 25 permissions to teams
86
+ team_permissions_given = 0
87
+ for team in large_teams :
88
+ for rd in large_rds :
89
+ for org in large_orgs :
90
+ rd .give_permission (team , org )
91
+ team_permissions_given += 1
92
+ if team_permissions_given >= 25 :
93
+ break
94
+ if team_permissions_given >= 25 :
95
+ break
96
+ if team_permissions_given >= 25 :
97
+ break
98
+
99
+ self .stdout .write (f'Assigned { user_permissions_given } permissions to users' )
100
+ self .stdout .write (f'Assigned { team_permissions_given } permissions to teams' )
101
+
34
102
self .stdout .write (f'Finished creating large demo data in { time .time () - start :.2f} seconds' )
35
103
36
104
def handle (self , * args , ** kwargs ):
0 commit comments