Clean up code

getsentry · Jun 19, 2023 · 375a8d1 · 375a8d1
1 parent d0c1aab
commit 375a8d1
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 18 deletions.
diff --git a/snuba/admin/auth.py b/snuba/admin/auth.py
@@ -45,7 +45,7 @@ def _is_member_of_group(user: AdminUser, group: str) -> bool:
     return google_api.check_group_membership(group_email=group, member=user.email)
 
 
-def get_iam_roles_from_file(user: AdminUser) -> Sequence[str]:
+def get_iam_roles_from_user(user: AdminUser) -> Sequence[str]:
     iam_roles = []
     try:
         with open(settings.ADMIN_IAM_POLICY_FILE, "r") as policy_file:
@@ -69,16 +69,37 @@ def get_iam_roles_from_file(user: AdminUser) -> Sequence[str]:
     return iam_roles
 
 
+def get_cached_iam_roles(user: AdminUser) -> Sequence[str]:
+    iam_roles_str = redis_client.get(f"roles-{user.email}")
+    if not iam_roles_str:
+        return []
+
+    iam_roles = rapidjson.loads(iam_roles_str)
+    if isinstance(iam_roles, list):
+        return iam_roles
+
+    return []
+
+
 def _set_roles(user: AdminUser) -> AdminUser:
     # todo: depending on provider convert user email
     # to subset of DEFAULT_ROLES based on IAM roles
-    key = f"roles-{user.email}"
-    iam_roles_str = redis_client.get(key)
-    iam_roles = rapidjson.loads(iam_roles_str) if iam_roles_str else None
+    iam_roles: Sequence[str] = []
+    try:
+        iam_roles = get_cached_iam_roles(user)
+    except Exception as e:
+        logger.exception("Failed to load roles from cache", exception=e)
+
     if not iam_roles:
-        iam_roles = get_iam_roles_from_file(user)
-        redis_client.set(key, rapidjson.dumps(iam_roles))
-        redis_client.expire(key, settings.ADMIN_ROLES_REDIS_TTL)
+        iam_roles = get_iam_roles_from_user(user)
+        try:
+            redis_client.set(
+                f"roles-{user.email}",
+                rapidjson.dumps(iam_roles),
+                ex=settings.ADMIN_ROLES_REDIS_TTL,
+            )
+        except Exception as e:
+            logger.exception(e)
 
     user.roles = [*[ROLES[role] for role in iam_roles if role in ROLES], *DEFAULT_ROLES]
     return user

diff --git a/tests/admin/clickhouse_migrations/test_api.py b/tests/admin/clickhouse_migrations/test_api.py
@@ -365,6 +365,81 @@ def test_get_iam_roles(caplog: Any) -> None:
             ]
 
         iam_file.close()
+
+        with patch(
+            "snuba.admin.auth.settings.ADMIN_IAM_POLICY_FILE", "file_not_exists.json"
+        ):
+            log = CapturingLogger()
+            with patch("snuba.admin.auth.logger", log):
+                user3 = AdminUser(email="test_user3@sentry.io", id="unknown")
+                _set_roles(user3)
+                assert "IAM policy file not found file_not_exists.json" in str(
+                    log.calls
+                )
+
+
+@pytest.mark.redis_db
+def test_get_iam_roles_cache() -> None:
+    system_role = generate_migration_test_role("system", "all")
+    tool_role = generate_tool_test_role("snql-to-sql")
+    with patch(
+        "snuba.admin.auth.DEFAULT_ROLES",
+        [system_role, tool_role],
+    ):
+        iam_file = tempfile.NamedTemporaryFile()
+        iam_file.write(
+            json.dumps(
+                {
+                    "bindings": [
+                        {
+                            "members": [
+                                "group:team-sns@sentry.io",
+                                "user:test_user1@sentry.io",
+                            ],
+                            "role": "roles/NonBlockingMigrationsExecutor",
+                        },
+                        {
+                            "members": [
+                                "group:team-sns@sentry.io",
+                                "user:test_user1@sentry.io",
+                                "user:test_user2@sentry.io",
+                            ],
+                            "role": "roles/TestMigrationsExecutor",
+                        },
+                        {
+                            "members": [
+                                "group:team-sns@sentry.io",
+                                "user:test_user1@sentry.io",
+                                "user:test_user2@sentry.io",
+                            ],
+                            "role": "roles/owner",
+                        },
+                        {
+                            "members": [
+                                "group:team-sns@sentry.io",
+                                "user:test_user1@sentry.io",
+                            ],
+                            "role": "roles/AllTools",
+                        },
+                    ]
+                }
+            ).encode("utf-8")
+        )
+
+        iam_file.flush()
+        with patch("snuba.admin.auth.settings.ADMIN_IAM_POLICY_FILE", iam_file.name):
+
+            user1 = AdminUser(email="test_user1@sentry.io", id="unknown")
+            _set_roles(user1)
+
+            assert user1.roles == [
+                ROLES["NonBlockingMigrationsExecutor"],
+                ROLES["TestMigrationsExecutor"],
+                ROLES["AllTools"],
+                system_role,
+                tool_role,
+            ]
+
         iam_file = tempfile.NamedTemporaryFile()
         iam_file.write(json.dumps({"bindings": []}).encode("utf-8"))
         iam_file.flush()
@@ -388,14 +463,3 @@ def test_get_iam_roles(caplog: Any) -> None:
                 system_role,
                 tool_role,
             ]
-
-        with patch(
-            "snuba.admin.auth.settings.ADMIN_IAM_POLICY_FILE", "file_not_exists.json"
-        ):
-            log = CapturingLogger()
-            with patch("snuba.admin.auth.logger", log):
-                user3 = AdminUser(email="test_user3@sentry.io", id="unknown")
-                _set_roles(user3)
-                assert "IAM policy file not found file_not_exists.json" in str(
-                    log.calls
-                )