diff --git a/cloudbuild.py.yaml b/cloudbuild.py.yaml index 8832658acc..a030466a7a 100644 --- a/cloudbuild.py.yaml +++ b/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 diff --git a/util/latlng_recon_geojson.py b/util/latlng_recon_geojson.py index 9a301e543e..7217e577d5 100644 --- a/util/latlng_recon_geojson.py +++ b/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: diff --git a/util/latlng_recon_geojson_test.py b/util/latlng_recon_geojson_test.py index 8f810ce6b6..54c6fa3c02 100644 --- a/util/latlng_recon_geojson_test.py +++ b/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