diff --git a/cartoframes/auth/credentials.py b/cartoframes/auth/credentials.py index 9cd61ba68..c02cc76d0 100644 --- a/cartoframes/auth/credentials.py +++ b/cartoframes/auth/credentials.py @@ -109,15 +109,25 @@ def session(self, session): """Set session""" self._session = session + @property + def me_data(self): + me_data = {} + + try: + me_data = self.get_api_key_auth_client().send(ME_SERVICE, 'get').json() + except Exception: + pass + + return me_data + @property def user_id(self): """Credentials user ID""" if not self._user_id: log.debug('Getting `user_id` for {}'.format(self._username)) - api_key_auth_client = self.get_api_key_auth_client() try: - user_me = api_key_auth_client.send(ME_SERVICE, 'get').json() + user_me = self.me_data() user_data = user_me.get('user_data') if user_data: self._user_id = user_data.get('id') diff --git a/cartoframes/io/carto.py b/cartoframes/io/carto.py index 8abf343b4..9575eb2a9 100644 --- a/cartoframes/io/carto.py +++ b/cartoframes/io/carto.py @@ -18,6 +18,7 @@ MAX_UPLOAD_SIZE_BYTES = 2000000000 # 2GB SAMPLE_ROWS_NUMBER = 100 +CSV_TO_CARTO_RATIO = 1.4 @send_metrics('data_downloaded') @@ -74,7 +75,8 @@ def read_carto(source, credentials=None, limit=None, retry_times=3, schema=None, @send_metrics('data_uploaded') def to_carto(dataframe, table_name, credentials=None, if_exists='fail', geom_col=None, index=False, index_label=None, - cartodbfy=True, log_enabled=True, retry_times=3, max_upload_size=MAX_UPLOAD_SIZE_BYTES): + cartodbfy=True, log_enabled=True, retry_times=3, max_upload_size=MAX_UPLOAD_SIZE_BYTES, + skip_quota_warning=False): """Upload a DataFrame to CARTO. The geometry's CRS must be WGS 84 (EPSG:4326) so you can use it on CARTO. Args: @@ -89,6 +91,9 @@ def to_carto(dataframe, table_name, credentials=None, if_exists='fail', geom_col uses the name of the index from the dataframe. cartodbfy (bool, optional): convert the table to CARTO format. Default True. More info `here `. + skip_quota_warning (bool, optional): skip the quota exceeded check and force the upload. + (The upload will still fail if the size of the dataset exceeds the remaining DB quota). + Default is False. retry_times (int, optional): Number of time to retry the upload in case it fails. Default is 3. @@ -114,6 +119,19 @@ def to_carto(dataframe, table_name, credentials=None, if_exists='fail', geom_col context_manager = ContextManager(credentials) + if not skip_quota_warning: + me_data = context_manager.credentials.me_data + if me_data is not None and me_data.get('user_data'): + n = min(SAMPLE_ROWS_NUMBER, len(dataframe)) + estimated_byte_size = len(dataframe.sample(n=n).to_csv(header=False)) * len(dataframe) \ + / n / CSV_TO_CARTO_RATIO + remaining_byte_quota = me_data.get('user_data').get('remaining_byte_quota') + + if remaining_byte_quota is not None and estimated_byte_size > remaining_byte_quota: + raise CartoException('DB Quota will be exceeded. ' + 'The remaining quota is {} bytes and the dataset size is {} bytes.'.format( + remaining_byte_quota, estimated_byte_size)) + gdf = GeoDataFrame(dataframe, copy=True) if index: diff --git a/tests/unit/io/test_carto.py b/tests/unit/io/test_carto.py index 2c9cd83fa..ae642f494 100644 --- a/tests/unit/io/test_carto.py +++ b/tests/unit/io/test_carto.py @@ -8,6 +8,7 @@ from shapely.geometry.base import BaseGeometry from shapely import wkt +from carto.exceptions import CartoException from cartoframes.auth import Credentials from cartoframes.io.managers.context_manager import ContextManager from cartoframes.io.carto import read_carto, to_carto, copy_table, create_table_from_query @@ -249,6 +250,54 @@ def test_to_carto(mocker): assert norm_table_name == table_name +def test_to_carto_quota_warning(mocker): + class NoQuotaCredentials(Credentials): + @property + def me_data(self): + return { + 'user_data': { + 'remaining_byte_quota': 0 + } + } + + # Given + table_name = '__table_name__' + cm_mock = mocker.patch.object(ContextManager, 'copy_from') + cm_mock.return_value = table_name + df = GeoDataFrame({'geometry': [Point([0, 0])]}) + + # Then + with pytest.raises(CartoException): + to_carto(df, table_name, NoQuotaCredentials('fake_user', 'fake_api_key')) + + +def test_to_carto_quota_warning_skip(mocker): + class NoQuotaCredentials(Credentials): + @property + def me_data(self): + return { + 'user_data': { + 'remaining_byte_quota': 0 + } + } + + # Given + table_name = '__table_name__' + cm_mock = mocker.patch.object(ContextManager, 'copy_from') + cm_mock.return_value = table_name + df = GeoDataFrame({'geometry': [Point([0, 0])]}) + + # When + norm_table_name = to_carto(df, table_name, NoQuotaCredentials('fake_user', 'fake_api_key'), + skip_quota_warning=True) + + # Then + assert cm_mock.call_args[0][1] == table_name + assert cm_mock.call_args[0][2] == 'fail' + assert cm_mock.call_args[0][3] is True + assert norm_table_name == table_name + + def test_to_carto_chunks(mocker): # Given table_name = '__table_name__'