diff --git a/cartoframes/data/services/__init__.py b/cartoframes/data/services/__init__.py index c50a0871b..5dc850752 100644 --- a/cartoframes/data/services/__init__.py +++ b/cartoframes/data/services/__init__.py @@ -1,7 +1,9 @@ from .geocoding import Geocoding from .isolines import Isolines +from .bq_datasets import BQUserDataset __all__ = [ 'Geocoding', - 'Isolines' + 'Isolines', + 'BQUserDataset', ] diff --git a/cartoframes/data/services/bq_datasets.py b/cartoframes/data/services/bq_datasets.py new file mode 100644 index 000000000..04342cb87 --- /dev/null +++ b/cartoframes/data/services/bq_datasets.py @@ -0,0 +1,50 @@ +import requests +from carto.utils import ResponseStream +# from carto.auth import APIKeyAuthClient + +from carto.exceptions import CartoException + +# TODO: this shouldn't be hardcoded +DO_ENRICHMENT_API_URL = 'http://localhost:7070/bq' + + +class BQDataset: + + def __init__(self, name_id): + self.name = name_id + # TODO fix this crap + self.session = requests.Session() + self.api_key = 'my_valid_api_key' + + def download(self): + url = DO_ENRICHMENT_API_URL + '/datasets/' + self.name + params = {'api_key': self.api_key} + + try: + response = self.session.get(url, + params=params, + stream=True) + response.raise_for_status() + except requests.HTTPError as e: + if 400 <= response.status_code < 500: + # Client error, provide better reason + reason = response.json()['error'][0] + error_msg = u'%s Client Error: %s' % (response.status_code, + reason) + raise CartoException(error_msg) + else: + raise CartoException(e) + except Exception as e: + raise CartoException(e) + + return response + + def download_stream(self): + return ResponseStream(self.download()) + + +class BQUserDataset: + + @staticmethod + def name(name_id): + return BQDataset(name_id) diff --git a/tests/e2e/data/services/test_bq_datasets.py b/tests/e2e/data/services/test_bq_datasets.py new file mode 100644 index 000000000..4e8566d53 --- /dev/null +++ b/tests/e2e/data/services/test_bq_datasets.py @@ -0,0 +1,40 @@ +import unittest +import pandas +import geopandas +from shapely import wkt + +from cartoframes.data.services import BQUserDataset + + +EXPECTED_CSV_SAMPLE = """state_fips_code,county_fips_code,geo_id,tract_name,internal_point_geo +60,10,60010950100,9501.0,POINT (-170.5618796 -14.2587411) +60,10,60010950200,9502.0,POINT (-170.5589852 -14.2859572) +60,10,60010950300,9503.0,POINT (-170.6310985 -14.2760947) +60,10,60010950500,9505.0,POINT (-170.6651925 -14.2713653) +60,10,60010950600,9506.0,POINT (-170.701028 -14.252446) +""" + + +class TestBQDataset(unittest.TestCase): + + def test_can_download_to_dataframe(self): + result = BQUserDataset.name('census_tracts_american_samoa').download_stream() + df = pandas.read_csv(result) + + self.assertEqual(df.shape, (18, 13)) + + # do some checks on the contents + sample = pandas.DataFrame( + df.head(), + columns=( + 'state_fips_code', + 'county_fips_code', + 'geo_id', + 'tract_name', + 'internal_point_geo' + ) + ) + sample['internal_point_geo'] = df['internal_point_geo'].apply(wkt.loads) + geosample = geopandas.GeoDataFrame(sample, geometry='internal_point_geo') + + self.assertEqual(geosample.to_csv(index=False), EXPECTED_CSV_SAMPLE)