def test_query_function_called_only_if_needed(self) -> None:
hamlet = self.example_user("hamlet")
# Get the user cached:
get_user_profile_by_id(hamlet.id)
class CustomException(Exception):
pass
def query_function(ids: List[int]) -> List[UserProfile]:
raise CustomException("The query function was called")
# query_function shouldn't be called, because the only requested object
# is already cached:
result: Dict[int, UserProfile] = bulk_cached_fetch(
cache_key_function=user_profile_by_id_cache_key,
query_function=query_function,
object_ids=[hamlet.id],
id_fetcher=get_user_id,
)
self.assertEqual(result, {hamlet.id: hamlet})
with self.assertLogs(level="INFO") as info_log:
flush_cache(Mock())
self.assertEqual(info_log.output, ["INFO:root:Clearing memcached cache after migrations"])
# With the cache flushed, the query_function should get called:
with self.assertRaises(CustomException):
result = bulk_cached_fetch(
cache_key_function=user_profile_by_id_cache_key,
query_function=query_function,
object_ids=[hamlet.id],
id_fetcher=get_user_id,
)
def test_query_function_called_only_if_needed(self) -> None:
# Get the user cached:
hamlet = get_user_profile_by_email(self.example_email("hamlet"))
class CustomException(Exception):
pass
def query_function(emails: List[str]) -> List[UserProfile]:
raise CustomException("The query function was called")
# query_function shouldn't be called, because the only requested object
# is already cached:
result: Dict[str, UserProfile] = bulk_cached_fetch(
cache_key_function=user_profile_by_email_cache_key,
query_function=query_function,
object_ids=[self.example_email("hamlet")],
id_fetcher=get_user_email,
)
self.assertEqual(result, {hamlet.delivery_email: hamlet})
flush_cache(Mock())
# With the cache flushed, the query_function should get called:
with self.assertRaises(CustomException):
result = bulk_cached_fetch(
cache_key_function=user_profile_by_email_cache_key,
query_function=query_function,
object_ids=[self.example_email("hamlet")],
id_fetcher=get_user_email,
)
def user_ids_to_users(user_ids: Sequence[int],
realm: Realm) -> List[UserProfile]:
# TODO: Consider adding a flag to control whether deactivated
# users should be included.
def fetch_users_by_id(user_ids: List[int]) -> List[UserProfile]:
return list(
UserProfile.objects.filter(id__in=user_ids).select_related())
user_profiles_by_id: Dict[int, UserProfile] = bulk_cached_fetch(
cache_key_function=user_profile_by_id_cache_key,
query_function=fetch_users_by_id,
object_ids=user_ids,
id_fetcher=get_user_id,
)
found_user_ids = user_profiles_by_id.keys()
missed_user_ids = [
user_id for user_id in user_ids if user_id not in found_user_ids
]
if missed_user_ids:
raise JsonableError(
_("Invalid user ID: {}").format(missed_user_ids[0]))
user_profiles = list(user_profiles_by_id.values())
for user_profile in user_profiles:
if user_profile.realm != realm:
raise JsonableError(
_("Invalid user ID: {}").format(user_profile.id))
return user_profiles
def bulk_get_user_profile_by_id(
uids: List[int]) -> Dict[int, UserDisplayRecipient]:
return bulk_cached_fetch(
# Use a separate cache key to protect us from conflicts with
# the get_user_profile_by_id cache.
# (Since we fetch only several fields here)
cache_key_function=display_recipient_bulk_get_users_by_id_cache_key,
query_function=lambda ids: list(
UserProfile.objects.filter(id__in=ids).values(
*display_recipient_fields)),
object_ids=uids,
id_fetcher=user_dict_id_fetcher,
)
def test_empty_object_ids_list(self) -> None:
class CustomException(Exception):
pass
def cache_key_function(email: str) -> str: # nocoverage -- this is just here to make sure it's not called
raise CustomException("The cache key function was called")
def query_function(emails: List[str]) -> List[UserProfile]: # nocoverage -- this is just here to make sure it's not called
raise CustomException("The query function was called")
# query_function and cache_key_function shouldn't be called, because
# objects_ids is empty, so there's nothing to do.
result: Dict[str, UserProfile] = bulk_cached_fetch(
cache_key_function=cache_key_function,
query_function=query_function,
object_ids=[],
id_fetcher=get_user_email,
)
self.assertEqual(result, {})
def bulk_get_users(
emails: List[str],
realm: Optional[Realm],
base_query: Optional[QuerySet[UserProfile]] = None
) -> Dict[str, UserProfile]:
if base_query is None:
assert realm is not None
query = UserProfile.objects.filter(realm=realm, is_active=True)
realm_id = realm.id
else:
# WARNING: Currently, this code path only really supports one
# version of `base_query` being used (because otherwise,
# they'll share the cache, which can screw up the filtering).
# If you're using this flow, you'll need to re-do any filters
# in base_query in the code itself; base_query is just a perf
# optimization.
query = base_query
realm_id = 0
def fetch_users_by_email(emails: List[str]) -> QuerySet[UserProfile]:
# This should be just
#
# UserProfile.objects.select_related("realm").filter(email__iexact__in=emails,
# realm=realm)
#
# But chaining __in and __iexact doesn't work with Django's
# ORM, so we have the following hack to construct the relevant where clause
where_clause = "upper(zerver_userprofile.email::text) IN (SELECT upper(email) FROM unnest(%s) AS email)"
return query.select_related("realm").extra(where=[where_clause],
params=(emails, ))
def user_to_email(user_profile: UserProfile) -> str:
return user_profile.email.lower()
return bulk_cached_fetch(
# Use a separate cache key to protect us from conflicts with
# the get_user cache.
lambda email: "bulk_get_users:" + user_profile_cache_key_id(
email, realm_id),
fetch_users_by_email,
[email.lower() for email in emails],
id_fetcher=user_to_email,
)
