From 4871b3a781a3d6b73eacd2be8ef83d4cf3a4d42e Mon Sep 17 00:00:00 2001 From: Vadim Stepanov Date: Tue, 5 Nov 2024 09:57:49 +0000 Subject: [PATCH 1/4] PD migrator: populate important OnCall notification rules (#5226) # What this PR does Makes so that user notification rules from PD are migrated into both "default" and "important" OnCall user notification rules instead of just "default". Before this PR: Screenshot 2024-11-04 at 16 54 28 After this PR: Screenshot 2024-11-04 at 16 54 22 ## Which issue(s) this PR closes Related to [Slack thread](https://raintank-corp.slack.com/archives/C07HMCM59TK/p1730306579122409?thread_ts=1730303532.031559&cid=C07HMCM59TK) ## Checklist - [x] Unit, integration, and e2e (if applicable) tests updated - [x] Documentation added (or `pr:no public docs` PR label added if not required) - [x] Added the relevant release notes label (see labels prefixed w/ `release:`). These labels dictate how your PR will show up in the autogenerated release notes. --- tools/migrators/README.md | 5 +- .../pagerduty/resources/notification_rules.py | 34 +++++--- .../test_migrate_notification_rules.py | 85 +++++++++++++++++++ 3 files changed, 108 insertions(+), 16 deletions(-) create mode 100644 tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py diff --git a/tools/migrators/README.md b/tools/migrators/README.md index 460873ff..61881ebf 100644 --- a/tools/migrators/README.md +++ b/tools/migrators/README.md @@ -172,8 +172,9 @@ Configuration is done via environment variables passed to the docker container. The tool is capable of migrating user notification rules from PagerDuty to Grafana OnCall. Notification rules from the `"When a high-urgency incident is assigned to me..."` section in PagerDuty settings are -taken into account and will be migrated to default notification rules in Grafana OnCall for each user. Note that delays -between notification rules may be slightly different in Grafana OnCall, see [Limitations](#limitations) for more info. +taken into account and will be migrated to both default and important notification rules in Grafana OnCall +for each user. Note that delays between notification rules may be slightly different in Grafana OnCall, +see [Limitations](#limitations) for more info. When running the migration, existing notification rules in Grafana OnCall will be deleted for every affected user. diff --git a/tools/migrators/lib/pagerduty/resources/notification_rules.py b/tools/migrators/lib/pagerduty/resources/notification_rules.py index 58a49dd0..7f712d7c 100644 --- a/tools/migrators/lib/pagerduty/resources/notification_rules.py +++ b/tools/migrators/lib/pagerduty/resources/notification_rules.py @@ -27,21 +27,25 @@ def migrate_notification_rules(user: dict) -> None: rule for rule in user["notification_rules"] if rule["urgency"] == "high" ] - oncall_rules = transform_notification_rules( - notification_rules, user["oncall_user"]["id"] - ) + for important in (False, True): + oncall_rules = transform_notification_rules( + notification_rules, user["oncall_user"]["id"], important + ) - for rule in oncall_rules: - OnCallAPIClient.create("personal_notification_rules", rule) + for rule in oncall_rules: + OnCallAPIClient.create("personal_notification_rules", rule) - if oncall_rules: - # delete old notification rules if any new rules were created - for rule in user["oncall_user"]["notification_rules"]: - OnCallAPIClient.delete("personal_notification_rules/{}".format(rule["id"])) + if oncall_rules: + # delete old notification rules if any new rules were created + for rule in user["oncall_user"]["notification_rules"]: + if rule["important"] == important: + OnCallAPIClient.delete( + "personal_notification_rules/{}".format(rule["id"]) + ) def transform_notification_rules( - notification_rules: list[dict], user_id: str + notification_rules: list[dict], user_id: str, important: bool ) -> list[dict]: """ Transform PagerDuty user notification rules to Grafana OnCall personal notification rules. @@ -58,7 +62,9 @@ def transform_notification_rules( previous_delay = notification_rules[idx - 1]["start_delay_in_minutes"] delay -= previous_delay - oncall_notification_rules += transform_notification_rule(rule, delay, user_id) + oncall_notification_rules += transform_notification_rule( + rule, delay, user_id, important + ) oncall_notification_rules = remove_duplicate_rules_between_waits( oncall_notification_rules @@ -68,12 +74,12 @@ def transform_notification_rules( def transform_notification_rule( - notification_rule: dict, delay: int, user_id: str + notification_rule: dict, delay: int, user_id: str, important: bool ) -> list[dict]: contact_method_type = notification_rule["contact_method"]["type"] oncall_type = PAGERDUTY_TO_ONCALL_CONTACT_METHOD_MAP[contact_method_type] - notify_rule = {"user_id": user_id, "type": oncall_type, "important": False} + notify_rule = {"user_id": user_id, "type": oncall_type, "important": important} if not delay: return [notify_rule] @@ -82,6 +88,6 @@ def transform_notification_rule( "user_id": user_id, "type": "wait", "duration": transform_wait_delay(delay), - "important": "False", + "important": important, } return [wait_rule, notify_rule] diff --git a/tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py b/tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py new file mode 100644 index 00000000..566d1614 --- /dev/null +++ b/tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py @@ -0,0 +1,85 @@ +from unittest.mock import call, patch + +from lib.oncall.api_client import OnCallAPIClient +from lib.pagerduty.resources.notification_rules import migrate_notification_rules + + +@patch.object(OnCallAPIClient, "delete") +@patch.object(OnCallAPIClient, "create") +def test_migrate_notification_rules(api_client_create_mock, api_client_delete_mock): + migrate_notification_rules( + { + "notification_rules": [ + { + "contact_method": {"type": "sms_contact_method"}, + "start_delay_in_minutes": 0, + "urgency": "high", + }, + { + "contact_method": {"type": "push_notification_contact_method"}, + "start_delay_in_minutes": 5, + "urgency": "high", + }, + ], + "oncall_user": { + "id": "EXISTING_USER_ID", + "notification_rules": [ + {"id": "EXISTING_RULE_ID_1", "important": False}, + {"id": "EXISTING_RULE_ID_2", "important": True}, + ], + }, + } + ) + + assert api_client_create_mock.call_args_list == [ + call( + "personal_notification_rules", + { + "user_id": "EXISTING_USER_ID", + "type": "notify_by_sms", + "important": False, + }, + ), + call( + "personal_notification_rules", + { + "user_id": "EXISTING_USER_ID", + "type": "wait", + "duration": 300, + "important": False, + }, + ), + call( + "personal_notification_rules", + { + "user_id": "EXISTING_USER_ID", + "type": "notify_by_mobile_app", + "important": False, + }, + ), + call( + "personal_notification_rules", + {"user_id": "EXISTING_USER_ID", "type": "notify_by_sms", "important": True}, + ), + call( + "personal_notification_rules", + { + "user_id": "EXISTING_USER_ID", + "type": "wait", + "duration": 300, + "important": True, + }, + ), + call( + "personal_notification_rules", + { + "user_id": "EXISTING_USER_ID", + "type": "notify_by_mobile_app", + "important": True, + }, + ), + ] + assert api_client_delete_mock.call_args_list == [ + call("personal_notification_rules/EXISTING_RULE_ID_1"), + call("personal_notification_rules/EXISTING_RULE_ID_2"), + ] From 7501304e963fd4396bc51ee8ed135aac3c91b3e0 Mon Sep 17 00:00:00 2001 From: Joey Orlando Date: Tue, 5 Nov 2024 05:54:38 -0500 Subject: [PATCH 2/4] feat: add `grafana_irm_enabled` to `GET /organization` endpoint response (#5230) # What this PR does Backend portion of https://github.com/grafana/oncall-mobile-app/issues/1021 ## Checklist - [x] Unit, integration, and e2e (if applicable) tests updated - [x] Documentation added (or `pr:no public docs` PR label added if not required) - [x] Added the relevant release notes label (see labels prefixed w/ `release:`). These labels dictate how your PR will show up in the autogenerated release notes. --- engine/apps/api/serializers/organization.py | 3 ++ engine/apps/api/tests/test_organization.py | 33 ++++++++++++++++++--- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/engine/apps/api/serializers/organization.py b/engine/apps/api/serializers/organization.py index 163c80df..e502e8a3 100644 --- a/engine/apps/api/serializers/organization.py +++ b/engine/apps/api/serializers/organization.py @@ -26,6 +26,7 @@ class OrganizationSerializer(EagerLoadingMixin, serializers.ModelSerializer): rbac_enabled = serializers.BooleanField(read_only=True, source="is_rbac_permissions_enabled") grafana_incident_enabled = serializers.BooleanField(read_only=True, source="is_grafana_incident_enabled") + grafana_irm_enabled = serializers.BooleanField(read_only=True, source="is_grafana_irm_enabled") SELECT_RELATED = ["slack_team_identity", "slack_channel"] @@ -39,6 +40,7 @@ class OrganizationSerializer(EagerLoadingMixin, serializers.ModelSerializer): "slack_channel", "rbac_enabled", "grafana_incident_enabled", + "grafana_irm_enabled", "direct_paging_prefer_important_policy", ] read_only_fields = [ @@ -46,6 +48,7 @@ class OrganizationSerializer(EagerLoadingMixin, serializers.ModelSerializer): "slack_team_identity", "rbac_enabled", "grafana_incident_enabled", + "grafana_irm_enabled", ] diff --git a/engine/apps/api/tests/test_organization.py b/engine/apps/api/tests/test_organization.py index f98408a0..b00ef956 100644 --- a/engine/apps/api/tests/test_organization.py +++ b/engine/apps/api/tests/test_organization.py @@ -36,7 +36,10 @@ def test_get_organization( client = APIClient() url = reverse("api-internal:api-organization") - expected_result = { + response = client.get(url, format="json", **make_user_auth_headers(user, token)) + + assert response.status_code == status.HTTP_200_OK + assert response.json() == { "pk": organization.public_primary_key, "name": organization.org_title, "stack_slug": organization.stack_slug, @@ -44,14 +47,12 @@ def test_get_organization( "slack_channel": None, "rbac_enabled": organization.is_rbac_permissions_enabled, "grafana_incident_enabled": organization.is_grafana_incident_enabled, + "grafana_irm_enabled": organization.is_grafana_irm_enabled, "direct_paging_prefer_important_policy": organization.direct_paging_prefer_important_policy, "is_resolution_note_required": False, "env_status": mock_env_status, "banner": mock_banner, } - response = client.get(url, format="json", **make_user_auth_headers(user, token)) - assert response.status_code == status.HTTP_200_OK - assert response.json() == expected_result @pytest.mark.django_db @@ -70,6 +71,30 @@ def test_get_organization_rbac_enabled(make_organization_and_user_with_plugin_to assert response.json()["rbac_enabled"] == organization.is_rbac_permissions_enabled +# NOTE: we need to patch the following because when is_grafana_irm_enabled is True, it alters how +# API authz works. For the purpose of this test, we don't care about testing that behaviour (it's already tested), +# just want to test the serializer essentially. +@patch("apps.api.permissions.user_is_authorized", return_value=True) +@pytest.mark.django_db +@pytest.mark.parametrize("is_grafana_irm_enabled", [True, False]) +def test_get_organization_grafana_irm_enabled( + _mock_user_is_authorized, + make_organization_and_user_with_plugin_token, + make_user_auth_headers, + is_grafana_irm_enabled, +): + organization, user, token = make_organization_and_user_with_plugin_token() + organization.is_grafana_irm_enabled = is_grafana_irm_enabled + organization.save() + + client = APIClient() + url = reverse("api-internal:api-organization") + + response = client.get(url, format="json", **make_user_auth_headers(user, token)) + assert response.status_code == status.HTTP_200_OK + assert response.json()["grafana_irm_enabled"] is is_grafana_irm_enabled + + @pytest.mark.django_db def test_update_organization_settings(make_organization_and_user_with_plugin_token, make_user_auth_headers): organization, user, token = make_organization_and_user_with_plugin_token() From effaa0a33069565f51f6d6b41d9b263bc0beacbf Mon Sep 17 00:00:00 2001 From: Innokentii Konstantinov Date: Wed, 6 Nov 2024 11:37:23 +0800 Subject: [PATCH 3/4] Always emit insight logs to simplify debugging (#4988) --- engine/common/insight_log/insight_logs_enabled_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/engine/common/insight_log/insight_logs_enabled_check.py b/engine/common/insight_log/insight_logs_enabled_check.py index be24250b..771b066e 100644 --- a/engine/common/insight_log/insight_logs_enabled_check.py +++ b/engine/common/insight_log/insight_logs_enabled_check.py @@ -21,4 +21,4 @@ def is_insight_logs_enabled(organization: "Organization") -> bool: f"ONCALL_BACKEND_REGION={settings.ONCALL_BACKEND_REGION} " f"cluster_slug={organization.cluster_slug}" ) - return not settings.IS_OPEN_SOURCE and settings.ONCALL_BACKEND_REGION == organization.cluster_slug + return not settings.IS_OPEN_SOURCE From 53ac2bcc12fe3f197f9fe920a65b50e39db59de2 Mon Sep 17 00:00:00 2001 From: Joey Orlando Date: Wed, 6 Nov 2024 06:02:21 -0500 Subject: [PATCH 4/4] fix: improve performance of recent `SlackChannel` related migrations (#5233) # What this PR does After deploying [`r439-v1.12.0`](https://github.com/grafana/oncall-private/releases/tag/r439-v1.12.0) to staging, I noticed that the migrations were taking a long time, and caused some wonkiness (see https://raintank-corp.slack.com/archives/C08063QES5N). ```bash Apply all migrations: [redacted secret grafana-admin-creds:admin-user], alerts, auth, auth_token, base, contenttypes, email, exotel, fcm_django, google, heartbeat, labels, mobile_app, oss_installation, phone_notifications, schedules, sessions, slack, social_django, telegram, twilioapp, user_management, webhooks, zvonok Running migrations: source=engine:app google_trace_id=none logger=apps.alerts.migrations.0063_migrate_channelfilter_slack_channel_id Starting migration to populate slack_channel field. source=engine:app google_trace_id=none logger=apps.alerts.migrations.0063_migrate_channelfilter_slack_channel_id Bulk updated 1 ChannelFilters with their Slack channel. source=engine:app google_trace_id=none logger=apps.alerts.migrations.0063_migrate_channelfilter_slack_channel_id Finished migration to populate slack_channel field. Applying alerts.0063_migrate_channelfilter_slack_channel_id... OK source=engine:app google_trace_id=none logger=apps.alerts.migrations.0064_migrate_resolutionnoteslackmessage_slack_channel_id Starting migration to populate slack_channel field. source=engine:app google_trace_id=none logger=apps.alerts.migrations.0064_migrate_resolutionnoteslackmessage_slack_channel_id Bulk updated 1 ResolutionNoteSlackMessage records with their Slack channel. source=engine:app google_trace_id=none logger=apps.alerts.migrations.0064_migrate_resolutionnoteslackmessage_slack_channel_id Finished migration to populate slack_channel field. Applying alerts.0064_migrate_resolutionnoteslackmessage_slack_channel_id... OK source=engine:app google_trace_id=none logger=apps.schedules.migrations.0019_auto_20241021_1735 Starting migration to populate slack_channel field. source=engine:app google_trace_id=none logger=apps.schedules.migrations.0019_auto_20241021_1735 Bulk updated 6 OnCallSchedules with their Slack channel. source=engine:app google_trace_id=none logger=apps.schedules.migrations.0019_auto_20241021_1735 Finished migration to populate slack_channel field. Applying schedules.0019_auto_20241021_1735... OK source=engine:app google_trace_id=none logger=apps.user_management.migrations.0026_auto_20241017_1919 Starting migration to populate default_slack_channel field. source=engine:app google_trace_id=none logger=apps.user_management.migrations.0026_auto_20241017_1919 Bulk updated 1 organizations with their default Slack channel. source=engine:app google_trace_id=none logger=apps.user_management.migrations.0026_auto_20241017_1919 Finished migration to populate default_slack_channel field. Applying user_management.0026_auto_20241017_1919... OK ``` **NOTE**: wrt these migrations already being run for certain OSS stacks; it shouldn't have much of an impact on OSS deployments, as it's really only an issue for _very large_ versions of these tables (particularly the `ResolutionNoteSlackMessage` table, which by its nature, has a tendency to generate a lot of data). ## Checklist - [ ] Unit, integration, and e2e (if applicable) tests updated - [ ] Documentation added (or `pr:no public docs` PR label added if not required) - [ ] Added the relevant release notes label (see labels prefixed w/ `release:`). These labels dictate how your PR will show up in the autogenerated release notes. --- ..._migrate_channelfilter_slack_channel_id.py | 55 ++++++----------- ...lutionnoteslackmessage_slack_channel_id.py | 61 ++++++------------- .../migrations/0019_auto_20241021_1735.py | 51 +++++----------- .../migrations/0026_auto_20241017_1919.py | 48 ++++----------- 4 files changed, 64 insertions(+), 151 deletions(-) diff --git a/engine/apps/alerts/migrations/0063_migrate_channelfilter_slack_channel_id.py b/engine/apps/alerts/migrations/0063_migrate_channelfilter_slack_channel_id.py index d7a46b95..dab5a459 100644 --- a/engine/apps/alerts/migrations/0063_migrate_channelfilter_slack_channel_id.py +++ b/engine/apps/alerts/migrations/0063_migrate_channelfilter_slack_channel_id.py @@ -9,49 +9,28 @@ logger = logging.getLogger(__name__) def populate_slack_channel(apps, schema_editor): ChannelFilter = apps.get_model("alerts", "ChannelFilter") SlackChannel = apps.get_model("slack", "SlackChannel") + AlertReceiveChannel = apps.get_model("alerts", "AlertReceiveChannel") + Organization = apps.get_model("user_management", "Organization") logger.info("Starting migration to populate slack_channel field.") - queryset = ChannelFilter.objects.filter( - _slack_channel_id__isnull=False, - alert_receive_channel__organization__slack_team_identity__isnull=False, - ) - total_channel_filters = queryset.count() - updated_channel_filters = 0 - missing_channel_filters = 0 - channel_filters_to_update = [] + sql = f""" + UPDATE {ChannelFilter._meta.db_table} AS cf + JOIN {AlertReceiveChannel._meta.db_table} AS arc ON arc.id = cf.alert_receive_channel_id + JOIN {Organization._meta.db_table} AS org ON org.id = arc.organization_id + JOIN {SlackChannel._meta.db_table} AS sc ON sc.slack_id = cf._slack_channel_id + AND sc.slack_team_identity_id = org.slack_team_identity_id + SET cf.slack_channel_id = sc.id + WHERE cf._slack_channel_id IS NOT NULL + AND org.slack_team_identity_id IS NOT NULL; + """ - logger.info(f"Total channel filters to process: {total_channel_filters}") + with schema_editor.connection.cursor() as cursor: + cursor.execute(sql) + updated_rows = cursor.rowcount # Number of rows updated - for channel_filter in queryset: - slack_id = channel_filter._slack_channel_id - slack_team_identity = channel_filter.alert_receive_channel.organization.slack_team_identity - - try: - slack_channel = SlackChannel.objects.get(slack_id=slack_id, slack_team_identity=slack_team_identity) - channel_filter.slack_channel = slack_channel - channel_filters_to_update.append(channel_filter) - - updated_channel_filters += 1 - logger.info( - f"ChannelFilter {channel_filter.id} updated with SlackChannel {slack_channel.id} " - f"(slack_id: {slack_id})." - ) - except SlackChannel.DoesNotExist: - missing_channel_filters += 1 - logger.warning( - f"SlackChannel with slack_id {slack_id} and slack_team_identity {slack_team_identity} " - f"does not exist for ChannelFilter {channel_filter.id}." - ) - - if channel_filters_to_update: - ChannelFilter.objects.bulk_update(channel_filters_to_update, ["slack_channel"]) - logger.info(f"Bulk updated {len(channel_filters_to_update)} ChannelFilters with their Slack channel.") - - logger.info( - f"Finished migration. Total channel filters processed: {total_channel_filters}. " - f"Channel filters updated: {updated_channel_filters}. Missing SlackChannels: {missing_channel_filters}." - ) + logger.info(f"Bulk updated {updated_rows} ChannelFilters with their Slack channel.") + logger.info("Finished migration to populate slack_channel field.") class Migration(migrations.Migration): diff --git a/engine/apps/alerts/migrations/0064_migrate_resolutionnoteslackmessage_slack_channel_id.py b/engine/apps/alerts/migrations/0064_migrate_resolutionnoteslackmessage_slack_channel_id.py index ccb304ed..4f492e31 100644 --- a/engine/apps/alerts/migrations/0064_migrate_resolutionnoteslackmessage_slack_channel_id.py +++ b/engine/apps/alerts/migrations/0064_migrate_resolutionnoteslackmessage_slack_channel_id.py @@ -10,53 +10,30 @@ logger = logging.getLogger(__name__) def populate_slack_channel(apps, schema_editor): ResolutionNoteSlackMessage = apps.get_model("alerts", "ResolutionNoteSlackMessage") SlackChannel = apps.get_model("slack", "SlackChannel") + AlertGroup = apps.get_model("alerts", "AlertGroup") + AlertReceiveChannel = apps.get_model("alerts", "AlertReceiveChannel") + Organization = apps.get_model("user_management", "Organization") logger.info("Starting migration to populate slack_channel field.") - queryset = ResolutionNoteSlackMessage.objects.filter( - _slack_channel_id__isnull=False, - alert_group__channel__organization__slack_team_identity__isnull=False, - ) - total_resolution_notes = queryset.count() - updated_resolution_notes = 0 - missing_resolution_notes = 0 - resolution_notes_to_update = [] + sql = f""" + UPDATE {ResolutionNoteSlackMessage._meta.db_table} AS rsm + JOIN {AlertGroup._meta.db_table} AS ag ON ag.id = rsm.alert_group_id + JOIN {AlertReceiveChannel._meta.db_table} AS arc ON arc.id = ag.channel_id + JOIN {Organization._meta.db_table} AS org ON org.id = arc.organization_id + JOIN {SlackChannel._meta.db_table} AS sc ON sc.slack_id = rsm._slack_channel_id + AND sc.slack_team_identity_id = org.slack_team_identity_id + SET rsm.slack_channel_id = sc.id + WHERE rsm._slack_channel_id IS NOT NULL + AND org.slack_team_identity_id IS NOT NULL; + """ - logger.info(f"Total resolution note slack messages to process: {total_resolution_notes}") - - for resolution_note in queryset: - slack_id = resolution_note._slack_channel_id - slack_team_identity = resolution_note.alert_group.channel.organization.slack_team_identity - - try: - slack_channel = SlackChannel.objects.get(slack_id=slack_id, slack_team_identity=slack_team_identity) - resolution_note.slack_channel = slack_channel - resolution_notes_to_update.append(resolution_note) - - updated_resolution_notes += 1 - logger.info( - f"ResolutionNoteSlackMessage {resolution_note.id} updated with SlackChannel {slack_channel.id} " - f"(slack_id: {slack_id})." - ) - except SlackChannel.DoesNotExist: - missing_resolution_notes += 1 - logger.warning( - f"SlackChannel with slack_id {slack_id} and slack_team_identity {slack_team_identity} " - f"does not exist for ResolutionNoteSlackMessage {resolution_note.id}." - ) - - if resolution_notes_to_update: - ResolutionNoteSlackMessage.objects.bulk_update(resolution_notes_to_update, ["slack_channel"]) - logger.info( - f"Bulk updated {len(resolution_notes_to_update)} ResolutionNoteSlackMessage with their Slack channel." - ) - - logger.info( - f"Finished migration. Total resolution note slack messages processed: {total_resolution_notes}. " - f"Resolution note slack messages updated: {updated_resolution_notes}. " - f"Missing SlackChannels: {missing_resolution_notes}." - ) + with schema_editor.connection.cursor() as cursor: + cursor.execute(sql) + updated_rows = cursor.rowcount # Number of rows updated + logger.info(f"Bulk updated {updated_rows} ResolutionNoteSlackMessage records with their Slack channel.") + logger.info("Finished migration to populate slack_channel field.") class Migration(migrations.Migration): diff --git a/engine/apps/schedules/migrations/0019_auto_20241021_1735.py b/engine/apps/schedules/migrations/0019_auto_20241021_1735.py index 8fe349dd..edc89366 100644 --- a/engine/apps/schedules/migrations/0019_auto_20241021_1735.py +++ b/engine/apps/schedules/migrations/0019_auto_20241021_1735.py @@ -9,47 +9,26 @@ logger = logging.getLogger(__name__) def populate_slack_channel(apps, schema_editor): OnCallSchedule = apps.get_model("schedules", "OnCallSchedule") SlackChannel = apps.get_model("slack", "SlackChannel") + Organization = apps.get_model("user_management", "Organization") logger.info("Starting migration to populate slack_channel field.") - queryset = OnCallSchedule.objects.filter(channel__isnull=False, organization__slack_team_identity__isnull=False) - total_schedules = queryset.count() - updated_schedules = 0 - missing_channels = 0 - schedules_to_update = [] + sql = f""" + UPDATE {OnCallSchedule._meta.db_table} AS ocs + JOIN {Organization._meta.db_table} AS org ON org.id = ocs.organization_id + JOIN {SlackChannel._meta.db_table} AS sc ON sc.slack_id = ocs.channel + AND sc.slack_team_identity_id = org.slack_team_identity_id + SET ocs.slack_channel_id = sc.id + WHERE ocs.channel IS NOT NULL + AND org.slack_team_identity_id IS NOT NULL; + """ - logger.info(f"Total schedules to process: {total_schedules}") - - for schedule in queryset: - slack_id = schedule.channel - slack_team_identity = schedule.organization.slack_team_identity - - try: - slack_channel = SlackChannel.objects.get(slack_id=slack_id, slack_team_identity=slack_team_identity) - - schedule.slack_channel = slack_channel - schedules_to_update.append(schedule) - - updated_schedules += 1 - logger.info( - f"Schedule {schedule.id} updated with SlackChannel {slack_channel.id} (slack_id: {slack_id})." - ) - except SlackChannel.DoesNotExist: - missing_channels += 1 - logger.warning( - f"SlackChannel with slack_id {slack_id} and slack_team_identity {slack_team_identity} " - f"does not exist for Schedule {schedule.id}." - ) - - if schedules_to_update: - OnCallSchedule.objects.bulk_update(schedules_to_update, ["slack_channel"]) - logger.info(f"Bulk updated {len(schedules_to_update)} OnCallSchedules with their Slack channel.") - - logger.info( - f"Finished migration. Total schedules processed: {total_schedules}. " - f"Schedules updated: {updated_schedules}. Missing SlackChannels: {missing_channels}." - ) + with schema_editor.connection.cursor() as cursor: + cursor.execute(sql) + updated_rows = cursor.rowcount # Number of rows updated + logger.info(f"Bulk updated {updated_rows} OnCallSchedules with their Slack channel.") + logger.info("Finished migration to populate slack_channel field.") class Migration(migrations.Migration): diff --git a/engine/apps/user_management/migrations/0026_auto_20241017_1919.py b/engine/apps/user_management/migrations/0026_auto_20241017_1919.py index c0869217..df28b026 100644 --- a/engine/apps/user_management/migrations/0026_auto_20241017_1919.py +++ b/engine/apps/user_management/migrations/0026_auto_20241017_1919.py @@ -13,43 +13,21 @@ def populate_default_slack_channel(apps, schema_editor): logger.info("Starting migration to populate default_slack_channel field.") - queryset = Organization.objects.filter(general_log_channel_id__isnull=False, slack_team_identity__isnull=False) - total_orgs = queryset.count() - updated_orgs = 0 - missing_channels = 0 - organizations_to_update = [] + sql = f""" + UPDATE {Organization._meta.db_table} AS org + JOIN {SlackChannel._meta.db_table} AS sc ON sc.slack_id = org.general_log_channel_id + AND sc.slack_team_identity_id = org.slack_team_identity_id + SET org.default_slack_channel_id = sc.id + WHERE org.general_log_channel_id IS NOT NULL + AND org.slack_team_identity_id IS NOT NULL; + """ - logger.info(f"Total organizations to process: {total_orgs}") + with schema_editor.connection.cursor() as cursor: + cursor.execute(sql) + updated_rows = cursor.rowcount # Number of rows updated - for org in queryset: - slack_id = org.general_log_channel_id - slack_team_identity = org.slack_team_identity - - try: - slack_channel = SlackChannel.objects.get(slack_id=slack_id, slack_team_identity=slack_team_identity) - - org.default_slack_channel = slack_channel - organizations_to_update.append(org) - - updated_orgs += 1 - logger.info( - f"Organization {org.id} updated with SlackChannel {slack_channel.id} (slack_id: {slack_id})." - ) - except SlackChannel.DoesNotExist: - missing_channels += 1 - logger.warning( - f"SlackChannel with slack_id {slack_id} and slack_team_identity {slack_team_identity} " - f"does not exist for Organization {org.id}." - ) - - if organizations_to_update: - Organization.objects.bulk_update(organizations_to_update, ["default_slack_channel"]) - logger.info(f"Bulk updated {len(organizations_to_update)} organizations with their default Slack channel.") - - logger.info( - f"Finished migration. Total organizations processed: {total_orgs}. " - f"Organizations updated: {updated_orgs}. Missing SlackChannels: {missing_channels}." - ) + logger.info(f"Bulk updated {updated_rows} organizations with their default Slack channel.") + logger.info("Finished migration to populate default_slack_channel field.") class Migration(migrations.Migration):