Coverage for api/views/health_views.py: 66%

1from django.conf import settings 

2from django.db import connection 

3from django.db.utils import OperationalError 

4from rest_framework import status 

5from rest_framework.exceptions import APIException 

6from rest_framework.request import Request 

7from rest_framework.response import Response 

8from rest_framework.views import APIView 

9 

10from api.utils.throttle import ExemptOAuth2IdRateThrottle, HealthcheckAnonRateThrottle 

11 

12 

13class HealthCheckException(APIException): 

14 status_code = status.HTTP_503_SERVICE_UNAVAILABLE 

15 

16 

17class HealthCheck(APIView): 

18 """ 

19 Return a "200 OK" response if the server is running normally, 503 otherwise. 

20 

21 This endpoint is used in production to ensure that the server should receive 

22 traffic. If no response is provided, the server is deregistered from the 

23 load balancer and destroyed. 

24 """ 

25 

26 throttle_classes = [HealthcheckAnonRateThrottle, ExemptOAuth2IdRateThrottle] 

27 schema = None # Hide this view from the OpenAPI schema. 

28 

29 @staticmethod 

30 def _check_db() -> None: 

31 """ 

32 Check that the database is available. 

33 

34 Returns nothing if everything is OK, throws error otherwise. 

35 """ 

36 try: 

37 connection.ensure_connection() 

38 except OperationalError as err: 

39 raise HealthCheckException(f"postgres: {err}") 

40 

41 @staticmethod 

42 def _check_es() -> None: 

43 """ 

44 Check Elasticsearch cluster health. 

45 

46 Raises an exception if ES is not healthy. 

47 """ 

48 es_health = settings.ES.cluster.health(timeout="5s") 

49 

50 if es_health["timed_out"]: 

51 raise HealthCheckException("elasticsearch: es_timed_out") 

52 

53 if (es_status := es_health["status"]) != "green": 

54 raise HealthCheckException(f"elasticsearch: es_status_{es_status}") 

55 

56 def get(self, request: Request): 

57 if "check_es" in request.query_params: 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 self._check_es() 

59 self._check_db() 

60 

61 return Response({"status": "200 OK"}, status=200)