Coverage for app / backend / src / couchers / servicers / gis.py: 81%

1import json 

2import logging 

3from typing import Any 

4 

5from google.protobuf import empty_pb2 

6from sqlalchemy import Function, select 

7from sqlalchemy.dialects.postgresql import JSON 

8from sqlalchemy.orm import Session 

9from sqlalchemy.sql import func 

10from sqlalchemy.sql.selectable import GenerativeSelect 

11 

12from couchers.context import CouchersContext 

13from couchers.materialized_views import ClusteredUser, LiteUser 

14from couchers.models import Node, Page, PageType, PageVersion 

15from couchers.proto import gis_pb2_grpc 

16from couchers.proto.google.api import httpbody_pb2 

17from couchers.sql import users_visible 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22def _build_geojson_select(statement: GenerativeSelect) -> Function[Any]: 

23 """ 

24 See usages below. 

25 """ 

26 # this is basically a translation of the postgis ST_AsGeoJSON example into sqlalchemy/geoalchemy2 

27 return func.json_build_object( 

28 "type", 

29 "FeatureCollection", 

30 "features", 

31 func.json_agg(func.ST_AsGeoJSON(statement.subquery(), maxdecimaldigits=5).cast(JSON)), 

32 ) 

33 

34 

35def _statement_to_geojson_response(session: Session, statement: GenerativeSelect) -> httpbody_pb2.HttpBody: 

36 json_dict = session.execute(select(_build_geojson_select(statement))).scalar_one_or_none() 

37 return httpbody_pb2.HttpBody( 

38 content_type="application/json", 

39 # json.dumps escapes non-ascii characters 

40 data=json.dumps(json_dict).encode("ascii"), 

41 ) 

42 

43 

44class GIS(gis_pb2_grpc.GISServicer): 

45 def GetUsers(self, request: empty_pb2.Empty, context: CouchersContext, session: Session) -> httpbody_pb2.HttpBody: 

46 # Build FeatureCollection from precomputed per-row GeoJSON in the materialized view, 

47 # assembling with string_agg in Postgres 

48 result = session.execute( 

49 select( 

50 func.concat( 

51 '{"type":"FeatureCollection","features":[', 

52 func.coalesce(func.string_agg(LiteUser.geojson, ","), ""), 

53 "]}", 

54 ) 

55 ).where(users_visible(context, table=LiteUser)) 

56 ).scalar_one() 

57 return httpbody_pb2.HttpBody( 

58 content_type="application/json", 

59 data=result.encode("ascii"), 

60 ) 

61 

62 def GetClusteredUsers( 

63 self, request: empty_pb2.Empty, context: CouchersContext, session: Session 

64 ) -> httpbody_pb2.HttpBody: 

65 return _statement_to_geojson_response(session, select(ClusteredUser.geom, ClusteredUser.count)) 

66 

67 def GetCommunities( 

68 self, request: empty_pb2.Empty, context: CouchersContext, session: Session 

69 ) -> httpbody_pb2.HttpBody: 

70 return _statement_to_geojson_response(session, select(Node).where(Node.geom != None)) 

71 

72 def GetPlaces(self, request: empty_pb2.Empty, context: CouchersContext, session: Session) -> httpbody_pb2.HttpBody: 

73 # need to do a subquery here so we get pages without a geom, not just versions without geom 

74 latest_pages = ( 

75 select(func.max(PageVersion.id).label("id")) 

76 .join(Page, Page.id == PageVersion.page_id) 

77 .where(Page.type == PageType.place) 

78 .group_by(PageVersion.page_id) 

79 .subquery() 

80 ) 

81 

82 statement = ( 

83 select(PageVersion.page_id.label("id"), PageVersion.slug.label("slug"), PageVersion.geom) 

84 .join(latest_pages, latest_pages.c.id == PageVersion.id) 

85 .where(PageVersion.geom != None) 

86 ) 

87 

88 return _statement_to_geojson_response(session, statement) 

89 

90 def GetGuides(self, request: empty_pb2.Empty, context: CouchersContext, session: Session) -> httpbody_pb2.HttpBody: 

91 latest_pages = ( 

92 select(func.max(PageVersion.id).label("id")) 

93 .join(Page, Page.id == PageVersion.page_id) 

94 .where(Page.type == PageType.guide) 

95 .group_by(PageVersion.page_id) 

96 .subquery() 

97 ) 

98 

99 statement = ( 

100 select(PageVersion.page_id.label("id"), PageVersion.slug.label("slug"), PageVersion.geom) 

101 .join(latest_pages, latest_pages.c.id == PageVersion.id) 

102 .where(PageVersion.geom != None) 

103 ) 

104 

105 return _statement_to_geojson_response(session, statement)