Coverage for src/couchers/materialized_views.py: 100%

1import logging 

2 

3from google.protobuf import empty_pb2 

4from sqlalchemy import Index 

5from sqlalchemy.sql import func 

6from sqlalchemy.sql import select as sa_select 

7from sqlalchemy_utils import create_materialized_view, refresh_materialized_view 

8 

9from couchers.db import session_scope 

10from couchers.models import Base, ClusterRole, ClusterSubscription, Upload, User 

11 

12logger = logging.getLogger(__name__) 

13 

14cluster_subscription_counts_selectable = ( 

15 sa_select( 

16 ClusterSubscription.cluster_id.label("cluster_id"), 

17 func.count().label("count"), 

18 ) 

19 .select_from(ClusterSubscription) 

20 .outerjoin(User, User.id == ClusterSubscription.user_id) 

21 .where(User.is_visible) 

22 .group_by(ClusterSubscription.cluster_id) 

23) 

24 

25cluster_subscription_counts = create_materialized_view( 

26 "cluster_subscription_counts", 

27 cluster_subscription_counts_selectable, 

28 Base.metadata, 

29 [ 

30 Index( 

31 "uq_cluster_subscription_counts_cluster_id", 

32 cluster_subscription_counts_selectable.c.cluster_id, 

33 unique=True, 

34 ) 

35 ], 

36) 

37 

38cluster_admin_counts_selectable = ( 

39 sa_select( 

40 ClusterSubscription.cluster_id.label("cluster_id"), 

41 func.count().label("count"), 

42 ) 

43 .select_from(ClusterSubscription) 

44 .outerjoin(User, User.id == ClusterSubscription.user_id) 

45 .where(ClusterSubscription.role == ClusterRole.admin) 

46 .where(User.is_visible) 

47 .group_by(ClusterSubscription.cluster_id) 

48) 

49 

50cluster_admin_counts = create_materialized_view( 

51 "cluster_admin_counts", 

52 cluster_admin_counts_selectable, 

53 Base.metadata, 

54 [Index("uq_cluster_admin_counts_cluster_id", cluster_admin_counts_selectable.c.cluster_id, unique=True)], 

55) 

56 

57lite_users_selectable = ( 

58 sa_select( 

59 User.id.label("id"), 

60 User.username.label("username"), 

61 User.name.label("name"), 

62 User.city.label("city"), 

63 User.age.label("age"), 

64 func.ST_Y(User.geom).label("lat"), 

65 func.ST_X(User.geom).label("lng"), 

66 User.geom_radius.label("radius"), 

67 User.is_visible.label("is_visible"), 

68 Upload.filename.label("avatar_filename"), 

69 ) 

70 .select_from(User) 

71 .outerjoin(Upload, Upload.key == User.avatar_key) 

72) 

73 

74lite_users = create_materialized_view( 

75 "lite_users", 

76 lite_users_selectable, 

77 Base.metadata, 

78 [Index("uq_lite_users_id", lite_users_selectable.c.id, unique=True)], 

79) 

80 

81 

82def refresh_materialized_views(payload: empty_pb2.Empty): 

83 logger.info("Refreshing materialized views") 

84 with session_scope() as session: 

85 refresh_materialized_view(session, "cluster_subscription_counts", concurrently=True) 

86 refresh_materialized_view(session, "cluster_admin_counts", concurrently=True) 

87 

88 

89def refresh_materialized_views_rapid(payload: empty_pb2.Empty): 

90 logger.info("Refreshing materialized views (rapid)") 

91 with session_scope() as session: 

92 refresh_materialized_view(session, "lite_users", concurrently=True)