Skip to content

fix(hc): Introduce Organization.get_bulk_owner_profiles #68756

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Apr 22, 2024
57 changes: 56 additions & 1 deletion src/sentry/models/organization.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
from sentry.roles.manager import Role
from sentry.services.hybrid_cloud.notifications import notifications_service
from sentry.services.hybrid_cloud.organization_mapping import organization_mapping_service
from sentry.services.hybrid_cloud.user import RpcUser
from sentry.services.hybrid_cloud.user import RpcUser, RpcUserProfile
from sentry.services.hybrid_cloud.user.service import user_service
from sentry.types.organization import OrganizationAbsoluteUrlMixin
from sentry.utils.http import is_using_customer_domain
Expand Down Expand Up @@ -329,6 +329,61 @@ def default_owner_id(self):
self._default_owner_id = owners[0].id
return self._default_owner_id

@classmethod
def _get_bulk_owner_ids(cls, organizations: Collection[Organization]) -> dict[int, int]:
"""Find user IDs of the default owners of multiple organization.

The returned table maps organization ID to user ID.
"""
from sentry.models.organizationmember import OrganizationMember

owner_id_table: dict[int, int] = {}
org_ids_to_query: list[int] = []
for org in organizations:
default_owner = getattr(org, "_default_owner", None)
if default_owner:
owner_id_table[org.id] = default_owner.id
else:
org_ids_to_query.append(org.id)

if org_ids_to_query:
queried_owner_ids = OrganizationMember.objects.filter(
organization_id__in=org_ids_to_query, role=roles.get_top_dog().id
).values_list("organization_id", "user_id")

for (org_id, user_id) in queried_owner_ids:
# An org may have multiple owners. Here we mimic the behavior of
# `get_default_owner`, which is to use the first one in the query
# result's iteration order.
if org_id not in owner_id_table:
owner_id_table[org_id] = user_id

return owner_id_table

@classmethod
def get_bulk_owner_profiles(
cls, organizations: Collection[Organization]
) -> dict[int, RpcUserProfile]:
"""Query for profile data of owners of multiple organizations.

The returned table is keyed by organization ID and shows the default owner.
An organization may have multiple owners, in which case only the default
owner is shown. Organization IDs may be absent from the returned table if no
owner was found.
"""

owner_id_table = cls._get_bulk_owner_ids(organizations)
owner_ids = list(owner_id_table.values())

profiles = user_service.get_many_profiles(filter=dict(user_ids=owner_ids))
profile_table = {c.id: c for c in profiles}

return {
org_id: profile_table[user_id]
for (org_id, user_id) in owner_id_table.items()
if user_id in profile_table
}

def has_single_owner(self):
owners = list(
self.get_members_with_org_roles([roles.get_top_dog().id])[:2].values_list(
Expand Down
16 changes: 15 additions & 1 deletion tests/sentry/models/test_organization.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
from sentry.testutils.helpers.features import with_feature
from sentry.testutils.hybrid_cloud import HybridCloudTestMixin
from sentry.testutils.outbox import outbox_runner
from sentry.testutils.silo import assume_test_silo_mode
from sentry.testutils.silo import assume_test_silo_mode, assume_test_silo_mode_of


class OrganizationTest(TestCase, HybridCloudTestMixin):
Expand Down Expand Up @@ -418,6 +418,20 @@ def test_absolute_url_with_customer_domain(self):
url = org.absolute_url("/organizations/acme/issues/", query="?project=123", fragment="#ref")
assert url == "http://acme.testserver/issues/?project=123#ref"

def test_get_bulk_owner_profiles(self):
u1, u2, u3 = (self.create_user() for _ in range(3))
o1, o2, o3 = (self.create_organization(owner=u) for u in (u1, u2, u3))
o2.get_default_owner() # populate _default_owner
with assume_test_silo_mode_of(User):
u3.delete()

bulk_owner_profiles = Organization.get_bulk_owner_profiles([o1, o2, o3])
assert set(bulk_owner_profiles.keys()) == {o1.id, o2.id}
assert bulk_owner_profiles[o1.id].id == u1.id
assert bulk_owner_profiles[o2.id].id == u2.id
assert bulk_owner_profiles[o2.id].name == u2.name
assert bulk_owner_profiles[o2.id].email == u2.email


class OrganizationDeletionTest(TestCase):
def add_org_notification_settings(self, org: Organization, user: User):
Expand Down
Loading