Coverage for src/couchers/servicers/gis.py: 77%

1import json 

2import logging 

3 

4from sqlalchemy.dialects.postgresql import JSON 

5from sqlalchemy.sql import func 

6 

7from couchers.materialized_views import clustered_users, lite_users 

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

9from couchers.sql import couchers_select as select 

10from proto import gis_pb2_grpc 

11from proto.google.api import httpbody_pb2 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16def _build_geojson_select(statement): 

17 """ 

18 See usages below. 

19 """ 

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

21 return func.json_build_object( 

22 "type", 

23 "FeatureCollection", 

24 "features", 

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

26 ) 

27 

28 

29def _statement_to_geojson_response(session, statement): 

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

31 return httpbody_pb2.HttpBody( 

32 content_type="application/json", 

33 # json.dumps escapes non-ascii characters 

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

35 ) 

36 

37 

38class GIS(gis_pb2_grpc.GISServicer): 

39 def GetUsers(self, request, context, session): 

40 statement = select(lite_users.c.id, lite_users.c.geom, lite_users.c.has_completed_profile).where_users_visible( 

41 context, table=lite_users.c 

42 ) 

43 return _statement_to_geojson_response(session, statement) 

44 

45 def GetClusteredUsers(self, request, context, session): 

46 return _statement_to_geojson_response(session, select(clustered_users.c.geom, clustered_users.c.count)) 

47 

48 def GetCommunities(self, request, context, session): 

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

50 

51 def GetPlaces(self, request, context, session): 

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

53 latest_pages = ( 

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

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

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

57 .group_by(PageVersion.page_id) 

58 .subquery() 

59 ) 

60 

61 statement = ( 

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

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

64 .where(PageVersion.geom != None) 

65 ) 

66 

67 return _statement_to_geojson_response(session, statement) 

68 

69 def GetGuides(self, request, context, session): 

70 latest_pages = ( 

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

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

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

74 .group_by(PageVersion.page_id) 

75 .subquery() 

76 ) 

77 

78 statement = ( 

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

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

81 .where(PageVersion.geom != None) 

82 ) 

83 

84 return _statement_to_geojson_response(session, statement)