Coverage for api/utils/tallies.py: 94%

1from collections import defaultdict 

2from datetime import datetime, timedelta 

3 

4import django_redis 

5import structlog 

6from django_redis.client.default import Redis 

7from redis.exceptions import ConnectionError 

8 

9 

10logger = structlog.get_logger(__name__) 

11 

12 

13def get_weekly_timestamp() -> str: 

14 """Get a timestamp for the Monday of any given week.""" 

15 now = datetime.now() 

16 monday = now - timedelta(days=now.weekday()) 

17 return monday.strftime("%Y-%m-%d") 

18 

19 

20def get_monthly_timestamp() -> str: 

21 """Get a timestamp for the month.""" 

22 now = datetime.now() 

23 return now.strftime("%Y-%m") 

24 

25 

26def count_provider_occurrences(results: list[dict], index: str) -> None: 

27 # Use ``get_redis_connection`` rather than Django's caches 

28 # so that we can open a pipeline rather than sending off ``n`` 

29 # writes and because the RedisPy client's ``incr`` method 

30 # is safe by default rather than Django's handspun method which: 

31 # 1. Takes two requests to execute; and 

32 # 2. Raises a ``ValueError`` if the key doesn't exist rather than 

33 # just initialising the key to the value like Redis's behaviour. 

34 tallies: Redis = django_redis.get_redis_connection("tallies") 

35 

36 provider_occurrences = defaultdict(int) 

37 for result in results: 

38 provider_occurrences[result["provider"]] += 1 

39 

40 week = get_weekly_timestamp() 

41 with tallies.pipeline() as pipe: 

42 for provider, occurrences in provider_occurrences.items(): 

43 pipe.incr(f"provider_occurrences:{index}:{week}:{provider}", occurrences) 

44 pipe.incr(f"provider_appeared_in_searches:{index}:{week}:{provider}", 1) 

45 try: 

46 pipe.execute() 

47 except ConnectionError: 

48 logger.warning("Redis connect failed, cannot increment provider tallies.")