import typing from django.conf import settings from django.core.validators import MinLengthValidator from django.db import models, transaction from apps.alerts.models import AlertReceiveChannel, ChannelFilter from apps.metrics_exporter.helpers import metrics_add_integration_to_cache, metrics_bulk_update_team_label_cache from apps.metrics_exporter.metrics_cache_manager import MetricsCacheManager from common.public_primary_keys import generate_public_primary_key, increase_public_primary_key_length if typing.TYPE_CHECKING: from django.db.models.manager import RelatedManager from apps.alerts.models import AlertGroupLogRecord from apps.grafana_plugin.helpers.client import GrafanaAPIClient from apps.schedules.models import CustomOnCallShift from apps.user_management.models import Organization, User def generate_public_primary_key_for_team() -> str: prefix = "T" new_public_primary_key = generate_public_primary_key(prefix) failure_counter = 0 while Team.objects.filter(public_primary_key=new_public_primary_key).exists(): new_public_primary_key = increase_public_primary_key_length( failure_counter=failure_counter, prefix=prefix, model_name="Team" ) failure_counter += 1 return new_public_primary_key class TeamManager(models.Manager["Team"]): @staticmethod def sync_for_organization( organization: "Organization", api_teams: typing.List["GrafanaAPIClient.Types.GrafanaTeam"] ) -> None: grafana_teams = {team["id"]: team for team in api_teams} existing_team_ids: typing.Set[int] = set(organization.teams.all().values_list("team_id", flat=True)) # create missing teams teams_to_create = tuple( Team( organization_id=organization.pk, team_id=team["id"], name=team["name"], email=team["email"], avatar_url=team["avatarUrl"], ) for team in grafana_teams.values() if team["id"] not in existing_team_ids ) with transaction.atomic(): organization.teams.bulk_create(teams_to_create, batch_size=5000) # Retrieve primary keys for the newly created users # # If the model’s primary key is an AutoField, the primary key attribute can only be retrieved # on certain databases (currently PostgreSQL, MariaDB 10.5+, and SQLite 3.35+). # On other databases, it will not be set. # https://docs.djangoproject.com/en/4.1/ref/models/querysets/#django.db.models.query.QuerySet.bulk_create created_teams = organization.teams.exclude(team_id__in=existing_team_ids) direct_paging_integrations_to_create = [] for team in created_teams: alert_receive_channel = AlertReceiveChannel( organization=organization, team=team, integration=AlertReceiveChannel.INTEGRATION_DIRECT_PAGING, verbal_name=f"Direct paging ({team.name if team else 'No'} team)", ) direct_paging_integrations_to_create.append(alert_receive_channel) AlertReceiveChannel.objects.bulk_create(direct_paging_integrations_to_create, batch_size=5000) created_direct_paging_integrations = AlertReceiveChannel.objects.filter( organization=organization, integration=AlertReceiveChannel.INTEGRATION_DIRECT_PAGING, ).exclude(team__team_id__in=existing_team_ids) default_channel_filters_to_create = [] for integration in created_direct_paging_integrations: channel_filter = ChannelFilter( alert_receive_channel=integration, filtering_term=None, is_default=True, order=0, ) default_channel_filters_to_create.append(channel_filter) ChannelFilter.objects.bulk_create(default_channel_filters_to_create, batch_size=5000) # Add direct paging integrations to metrics cache for integration in direct_paging_integrations_to_create: metrics_add_integration_to_cache(integration) # delete excess teams team_ids_to_delete = existing_team_ids - grafana_teams.keys() organization.teams.filter(team_id__in=team_ids_to_delete).delete() # collect teams diffs to update metrics cache metrics_teams_to_update: MetricsCacheManager.TeamsDiffMap = {} for team_id in team_ids_to_delete: metrics_teams_to_update = MetricsCacheManager.update_team_diff( metrics_teams_to_update, team_id, deleted=True ) # update existing teams if any fields have changed teams_to_update = [] for team in organization.teams.filter(team_id__in=existing_team_ids): grafana_team = grafana_teams[team.team_id] if ( team.name != grafana_team["name"] or team.email != grafana_team["email"] or team.avatar_url != grafana_team["avatarUrl"] ): if team.name != grafana_team["name"]: # collect teams diffs to update metrics cache metrics_teams_to_update = MetricsCacheManager.update_team_diff( metrics_teams_to_update, team.id, new_name=grafana_team["name"] ) team.name = grafana_team["name"] team.email = grafana_team["email"] team.avatar_url = grafana_team["avatarUrl"] teams_to_update.append(team) organization.teams.bulk_update(teams_to_update, ["name", "email", "avatar_url"], batch_size=5000) metrics_bulk_update_team_label_cache(metrics_teams_to_update, organization.id) class Team(models.Model): current_team_users: "RelatedManager['User']" custom_on_call_shifts: "RelatedManager['CustomOnCallShift']" oncall_schedules: "RelatedManager['AlertGroupLogRecord']" users: "RelatedManager['User']" public_primary_key = models.CharField( max_length=20, validators=[MinLengthValidator(settings.PUBLIC_PRIMARY_KEY_MIN_LENGTH + 1)], unique=True, default=generate_public_primary_key_for_team, ) objects = TeamManager() team_id = models.PositiveIntegerField() organization = models.ForeignKey( to="user_management.Organization", related_name="teams", on_delete=models.deletion.CASCADE, ) users = models.ManyToManyField(to="user_management.User", related_name="teams") name = models.CharField(max_length=300) email = models.CharField(max_length=300, null=True, blank=True, default=None) avatar_url = models.URLField() # If is_sharing_resources_to_all is False only team members and admins can access it and it's resources # if it's True every oncall organization user can access it is_sharing_resources_to_all = models.BooleanField(default=False)