Use Datacommons V2 client for geojson fetch

cloudbuild.py.yaml

@@ -15,6 +15,8 @@ steps:
   entrypoint: /bin/sh
   waitFor:
   - python_install
+  secretEnv:
+    - DC_API_KEY
   args:
     - -c
     - |
@@ -41,3 +43,8 @@ steps:
 #  - python_install
 
 # TODO(rsned): Add any other useful code health tools.
+
+availableSecrets:
+  secretManager:
+    - versionName: projects/datcom-ci/secrets/dc-api-key/versions/latest
+      env: DC_API_KEY

util/latlng_recon_geojson.py

@@ -16,17 +16,27 @@
 See latlng_recon_geojson_test.py for usage example.
 """
 
-from shapely import geometry
-import datacommons as dc
 import json
 import logging
-import time
-import urllib
+from pathlib import Path
+from shapely import geometry
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+
+from util.dc_api_wrapper import dc_api_batched_wrapper
+from util.dc_api_wrapper import dc_api_wrapper
+from util.dc_api_wrapper import get_datacommons_client
 
 _WORLD = 'Earth'
 _USA = 'country/USA'
 _MAX_RETRIES = 3
 _RETRY_DELAY = 15
+_DC_API_CONFIG = {
+    'dc_api_retries': _MAX_RETRIES,
+    'dc_api_retry_secs': _RETRY_DELAY,
+}
 
 _GJ_PROP = {
     'Country': 'geoJsonCoordinatesDP2',
@@ -36,28 +46,75 @@
 }
 
 
-def _get_geojsons(place_type, parent_place, retry=0):
-    try:
-        places = dc.get_places_in([parent_place], place_type)[parent_place]
-        resp = dc.get_property_values(places, _GJ_PROP[place_type])
-        geojsons = {}
-        for p, gj in resp.items():
-            if not gj:
-                continue
-            geojsons[p] = geometry.shape(json.loads(gj[0]))
-        return geojsons
-    except urllib.error.URLError:
-        if retry > _MAX_RETRIES:
-            logging.error("Exceeded max retries(%s)" % str(_MAX_RETRIES))
-            raise RuntimeError
-        else:
-            # retry after a small delay
-            time.sleep(_RETRY_DELAY)
-            return _get_geojsons(place_type, parent_place, retry + 1)
+def _get_dc_client():
+    return get_datacommons_client(_DC_API_CONFIG)
+
+
+def _get_geojsons(place_type, parent_place):
+    client = _get_dc_client()
+    places_response = dc_api_wrapper(
+        function=client.node.fetch_place_children,
+        args={
+            'place_dcids': [parent_place],
+            'children_type': place_type,
+            'as_dict': True,
+        },
+        retries=_MAX_RETRIES,
+        retry_secs=_RETRY_DELAY,
+    )
+    if not places_response or parent_place not in places_response:
+        response_keys = None
+        if isinstance(places_response, dict):
+            response_keys = sorted(places_response.keys())
+        logging.error(
+            'Failed to fetch place children. place_type=%s parent_place=%s '
+            'response_type=%s response_keys=%s',
+            place_type,
+            parent_place,
+            type(places_response).__name__,
+            response_keys,
+        )
+        raise RuntimeError
+    places = [
+        node.get('dcid')
+        for node in places_response.get(parent_place, [])
+        if node.get('dcid')
+    ]
+    resp = dc_api_batched_wrapper(function=client.node.fetch_property_values,
+                                  dcids=places,
+                                  args={'properties': _GJ_PROP[place_type]},
+                                  dcid_arg_kw='node_dcids',
+                                  config=_DC_API_CONFIG)
+    geojsons = {}
+    for place in places:
+        nodes = (resp.get(place, {}).get('arcs',
+                                         {}).get(_GJ_PROP[place_type],
+                                                 {}).get('nodes', []))
+        if not nodes:
+            continue
+        geojson = nodes[0].get('value')
+        if not geojson:
+            continue
+        geojsons[place] = geometry.shape(json.loads(geojson))
+    return geojsons
 
 
 def _get_continent_map(countries):
-    return dc.get_property_values(countries, 'containedInPlace')
+    client = _get_dc_client()
+    resp = dc_api_batched_wrapper(function=client.node.fetch_property_values,
+                                  dcids=countries,
+                                  args={'properties': 'containedInPlace'},
+                                  dcid_arg_kw='node_dcids',
+                                  config=_DC_API_CONFIG)
+    continent_map = {}
+    for country in countries:
+        nodes = (resp.get(country, {}).get('arcs',
+                                           {}).get('containedInPlace',
+                                                   {}).get('nodes', []))
+        continent_map[country] = [
+            node.get('dcid') for node in nodes if node.get('dcid')
+        ]
+    return continent_map
 
 
 class LatLng2Places:

util/latlng_recon_geojson_test.py

@@ -14,16 +14,15 @@
 """Tests for util.latlng_recon_geojson"""
 
 import json
-import os
+from pathlib import Path
 import sys
 import unittest
 from shapely import geometry
 from unittest import mock
 
-sys.path.append(
-    os.path.dirname(os.path.dirname(os.path.dirname(
-        os.path.abspath(__file__)))))
-import latlng_recon_geojson
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+from util import latlng_recon_geojson
 
 _SC_COUNTY_GJ_STR = """
 {"type": "Polygon", "coordinates": [[[-122.202653, 37.363046], [-122.026107, 37.16681], [-121.575402, 36.893033], [-121.488949, 36.983148], [-121.215406, 36.961248], [-121.23711, 37.157204], [-121.399019, 37.150135], [-121.45575, 37.24944], [-121.409075, 37.380672], [-121.472952, 37.482333], [-122.115161, 37.46628], [-122.202653, 37.363046]]]}
@@ -46,8 +45,8 @@ def _mock_get_gj(place_type, parent_place):
 
 class LatlngReconGeojsonTest(unittest.TestCase):
 
-    @mock.patch('latlng_recon_geojson._get_geojsons')
-    @mock.patch('latlng_recon_geojson._get_continent_map')
+    @mock.patch('util.latlng_recon_geojson._get_geojsons')
+    @mock.patch('util.latlng_recon_geojson._get_continent_map')
     def test_main(self, mock_cmap, mock_gj):
         mock_cmap.return_value = {'country/USA': ['northamerica']}
         mock_gj.side_effect = _mock_get_gj