zilencer-models: Create helper function for remote user counts.

Pulls out shared code in get_remote_realm_guest_and_non_guest_count
and get_remote_server_guest_and_non_guest_count for generating the
RemoteCustomerUserCount so that it can be used in logic for getting
these counts for all remote servers and remote realms on an
installation.
This commit is contained in:
Lauryn Menard
2024-01-19 17:03:13 +01:00
committed by Tim Abbott
parent 536aef854c
commit c08e266956

View File

@@ -463,6 +463,26 @@ class RemoteCustomerUserCount:
non_guest_user_count: int non_guest_user_count: int
def get_remote_customer_user_count(
audit_logs: List[RemoteRealmAuditLog],
) -> RemoteCustomerUserCount:
guest_count = 0
non_guest_count = 0
for log in audit_logs:
humans_count_dict = log.extra_data[RemoteRealmAuditLog.ROLE_COUNT][
RemoteRealmAuditLog.ROLE_COUNT_HUMANS
]
for role_type in UserProfile.ROLE_TYPES:
if role_type == UserProfile.ROLE_GUEST:
guest_count += humans_count_dict.get(str(role_type), 0)
else:
non_guest_count += humans_count_dict.get(str(role_type), 0)
return RemoteCustomerUserCount(
non_guest_user_count=non_guest_count, guest_user_count=guest_count
)
def get_remote_server_guest_and_non_guest_count( def get_remote_server_guest_and_non_guest_count(
server_id: int, event_time: datetime = timezone_now() server_id: int, event_time: datetime = timezone_now()
) -> RemoteCustomerUserCount: ) -> RemoteCustomerUserCount:
@@ -486,26 +506,11 @@ def get_remote_server_guest_and_non_guest_count(
.values_list("max_id", flat=True) .values_list("max_id", flat=True)
) )
extra_data_list = RemoteRealmAuditLog.objects.filter( realm_audit_logs = RemoteRealmAuditLog.objects.filter(id__in=list(realm_last_audit_log_ids))
id__in=list(realm_last_audit_log_ids)
).values_list("extra_data", flat=True)
# Now we add up the user counts from the different realms. # Now we add up the user counts from the different realms.
guest_count = 0 user_count = get_remote_customer_user_count(list(realm_audit_logs))
non_guest_count = 0 return user_count
for extra_data in extra_data_list:
humans_count_dict = extra_data[RemoteRealmAuditLog.ROLE_COUNT][
RemoteRealmAuditLog.ROLE_COUNT_HUMANS
]
for role_type in UserProfile.ROLE_TYPES:
if role_type == UserProfile.ROLE_GUEST:
guest_count += humans_count_dict.get(str(role_type), 0)
else:
non_guest_count += humans_count_dict.get(str(role_type), 0)
return RemoteCustomerUserCount(
non_guest_user_count=non_guest_count, guest_user_count=guest_count
)
def get_remote_realm_guest_and_non_guest_count( def get_remote_realm_guest_and_non_guest_count(
@@ -526,21 +531,12 @@ def get_remote_realm_guest_and_non_guest_count(
.exclude(extra_data={}).last() .exclude(extra_data={}).last()
) )
guest_count = 0
non_guest_count = 0
if latest_audit_log is not None: if latest_audit_log is not None:
humans_count_dict = latest_audit_log.extra_data[RemoteRealmAuditLog.ROLE_COUNT][ assert latest_audit_log is not None
RemoteRealmAuditLog.ROLE_COUNT_HUMANS user_count = get_remote_customer_user_count([latest_audit_log])
]
for role_type in UserProfile.ROLE_TYPES:
if role_type == UserProfile.ROLE_GUEST:
guest_count += humans_count_dict.get(str(role_type), 0)
else: else:
non_guest_count += humans_count_dict.get(str(role_type), 0) user_count = RemoteCustomerUserCount(guest_user_count=0, non_guest_user_count=0)
return user_count
return RemoteCustomerUserCount(
non_guest_user_count=non_guest_count, guest_user_count=guest_count
)
def has_stale_audit_log(server: RemoteZulipServer) -> bool: def has_stale_audit_log(server: RemoteZulipServer) -> bool: