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

31 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-11-21 04:21 +0000

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 = ( 

41 select(lite_users.c.id, lite_users.c.geom, lite_users.c.has_completed_profile) 

42 .where_users_visible(context, table=lite_users.c) 

43 .where(lite_users.c.geom != None) 

44 ) 

45 return _statement_to_geojson_response(session, statement) 

46 

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

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

49 

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

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

52 

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

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

55 latest_pages = ( 

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

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

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

59 .group_by(PageVersion.page_id) 

60 .subquery() 

61 ) 

62 

63 statement = ( 

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

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

66 .where(PageVersion.geom != None) 

67 ) 

68 

69 return _statement_to_geojson_response(session, statement) 

70 

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

72 latest_pages = ( 

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

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

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

76 .group_by(PageVersion.page_id) 

77 .subquery() 

78 ) 

79 

80 statement = ( 

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

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

83 .where(PageVersion.geom != None) 

84 ) 

85 

86 return _statement_to_geojson_response(session, statement)