From 16196822dee94f2b0fb8139eb375dd3a7bb64248 Mon Sep 17 00:00:00 2001 From: Joey Orlando Date: Wed, 1 Feb 2023 12:07:32 +0100 Subject: [PATCH] Add utility function to get readonly db key if defined (#1264) # What this PR does This is a minor refactor before implementing https://github.com/grafana/oncall-private/issues/1558. Additionally, it cleans up a few spots where we do this: ``` # Re-take in case we are in the readonly db context. ``` We currently don't read anything from a read-only database, so this should be not necessary. ## Checklist - [x] Tests updated - [ ] Documentation added (N/A) - [ ] `CHANGELOG.md` updated (N/A) --- engine/apps/slack/models/slack_message.py | 6 +----- .../apps/slack/models/slack_user_identity.py | 18 +++++----------- engine/common/database.py | 17 +++++++++++++++ engine/common/tests/test_database.py | 21 +++++++++++++++++++ 4 files changed, 44 insertions(+), 18 deletions(-) create mode 100644 engine/common/database.py create mode 100644 engine/common/tests/test_database.py diff --git a/engine/apps/slack/models/slack_message.py b/engine/apps/slack/models/slack_message.py index 614ef5d2..f373611c 100644 --- a/engine/apps/slack/models/slack_message.py +++ b/engine/apps/slack/models/slack_message.py @@ -68,12 +68,8 @@ class SlackMessage(models.Model): f"It is strange!" ) return None - # Re-take object to switch connection from readonly db to master. self._slack_team_identity = self.organization.slack_team_identity - - _self = SlackMessage.objects.get(pk=self.pk) - _self._slack_team_identity = _self.organization.slack_team_identity - _self.save() + self.save() return self._slack_team_identity def get_alert_group(self): diff --git a/engine/apps/slack/models/slack_user_identity.py b/engine/apps/slack/models/slack_user_identity.py index c9288baf..e988d9f7 100644 --- a/engine/apps/slack/models/slack_user_identity.py +++ b/engine/apps/slack/models/slack_user_identity.py @@ -149,7 +149,6 @@ class SlackUserIdentity(models.Model): @property def slack_login(self): if self.cached_slack_login is None or self.cached_slack_login == "slack_token_revoked_unable_to_cache_login": - _self = SlackUserIdentity.objects.get(pk=self.pk) # Re-take in case we are in the readonly db context. sc = SlackClientWithErrorHandling(self.slack_team_identity.bot_access_token) try: result = sc.api_call( @@ -158,20 +157,17 @@ class SlackUserIdentity(models.Model): team=self.slack_team_identity, ) self.cached_slack_login = result["user"]["name"] - _self.cached_slack_login = result["user"]["name"] - _self.save() + self.save() except SlackAPITokenException as e: logger.warning("Unable to get slack login: token revoked\n" + str(e)) self.cached_slack_login = "slack_token_revoked_unable_to_cache_login" - _self.cached_slack_login = "slack_token_revoked_unable_to_cache_login" - _self.save() + self.save() return "slack_token_revoked_unable_to_cache_login" except SlackAPIException as e: if e.response["error"] == "user_not_found": logger.warning("user_not_found " + str(e)) self.cached_slack_login = "user_not_found" - _self.cached_slack_login = "user_not_found" - _self.save() + self.save() elif e.response["error"] == "invalid_auth": return "no_enough_permissions_to_retrieve" else: @@ -182,7 +178,6 @@ class SlackUserIdentity(models.Model): @property def timezone(self): if self.cached_timezone is None or self.cached_timezone == "None": - _self = SlackUserIdentity.objects.get(pk=self.pk) # Re-take in case we are in the readonly db context. sc = SlackClientWithErrorHandling(self.slack_team_identity.bot_access_token) try: result = sc.api_call( @@ -194,8 +189,7 @@ class SlackUserIdentity(models.Model): if tz_from_slack == "None" or tz_from_slack is None: tz_from_slack = "UTC" self.cached_timezone = tz_from_slack - _self.cached_timezone = tz_from_slack - _self.save(update_fields=["cached_timezone"]) + self.save(update_fields=["cached_timezone"]) except SlackAPITokenException as e: print("Token revoked: " + str(e)) except requests.exceptions.Timeout: @@ -207,13 +201,11 @@ class SlackUserIdentity(models.Model): @property def im_channel_id(self): if self.cached_im_channel_id is None: - _self = SlackUserIdentity.objects.get(pk=self.pk) # Re-take in case we are in the readonly db context. sc = SlackClientWithErrorHandling(self.slack_team_identity.bot_access_token) try: result = sc.api_call("conversations.open", users=self.slack_id, return_im=True) self.cached_im_channel_id = result["channel"]["id"] - _self.cached_im_channel_id = result["channel"]["id"] - _self.save() + self.save() except SlackAPIException as e: if e.response["error"] == "cannot_dm_bot": logger.warning("Trying to DM bot " + str(e)) diff --git a/engine/common/database.py b/engine/common/database.py new file mode 100644 index 00000000..dd40a39e --- /dev/null +++ b/engine/common/database.py @@ -0,0 +1,17 @@ +import random + +from django.conf import settings + + +def get_random_readonly_database_key_if_present_otherwise_default() -> str: + """ + This function returns a string, representing a key in the DATABASES django settings. + If settings.READONLY_DATABASES is set, and non-empty, it randomly chooses one of the read-only databases, + otherwise it falls back to "default". + + This is primarily intended to be used for django's QuerySet.using() function + """ + using_db = "default" + if hasattr(settings, "READONLY_DATABASES") and len(settings.READONLY_DATABASES) > 0: + using_db = random.choice(list(settings.READONLY_DATABASES.keys())) + return using_db diff --git a/engine/common/tests/test_database.py b/engine/common/tests/test_database.py new file mode 100644 index 00000000..d426eae4 --- /dev/null +++ b/engine/common/tests/test_database.py @@ -0,0 +1,21 @@ +from django.test import override_settings + +from common.database import get_random_readonly_database_key_if_present_otherwise_default + +MOCK_READ_ONLY_DATABASES = { + "foo": "asdfkdjkdfjkdf", + "bar": "nmcvnmcvmnvc", +} + + +class TestGetRandomReadOnlyDatabaseKeyIfPresentOtherwiseDefault: + @override_settings(READONLY_DATABASES=MOCK_READ_ONLY_DATABASES) + def test_it_randomly_chooses_a_readonly_database(self) -> None: + assert get_random_readonly_database_key_if_present_otherwise_default() in MOCK_READ_ONLY_DATABASES + + @override_settings(READONLY_DATABASES={}) + def test_it_falls_back_to_default_if_readonly_databases_is_set_but_empty(self) -> None: + assert get_random_readonly_database_key_if_present_otherwise_default() == "default" + + def test_it_falls_back_to_default_if_readonly_databases_is_not_set(self) -> None: + assert get_random_readonly_database_key_if_present_otherwise_default() == "default"