get table metadata and manage dataset privacy

cartoframes/dataset.py

@@ -8,6 +8,7 @@
 
 from carto.exceptions import CartoException, CartoRateLimitException
 from .geojson import load_geojson
+from .dataset_info import DatasetInfo
 
 # avoid _lock issue: https://github.com/tqdm/tqdm/issues/457
 tqdm(disable=True, total=0)  # initialise internal lock
@@ -26,9 +27,9 @@ class Dataset(object):
     REPLACE = 'replace'
     APPEND = 'append'
 
-    PRIVATE = 'private'
-    PUBLIC = 'public'
-    LINK = 'link'
+    PRIVATE = DatasetInfo.PRIVATE
+    PUBLIC = DatasetInfo.PUBLIC
+    LINK = DatasetInfo.LINK
 
     STATE_LOCAL = 'local'
     STATE_REMOTE = 'remote'
@@ -55,6 +56,7 @@ def __init__(self, table_name=None, schema='public',
         self._cc = context or default_context
         self._state = state
         self._is_saved_in_carto = is_saved_in_carto
+        self._dataset_info = None
 
         self._normalized_column_names = None
 
@@ -105,6 +107,29 @@ def set_geodataframe(self, gdf):
     def get_table_name(self):
         return self._table_name
 
+    def get_dataset_info(self):
+        if self._is_saved_in_carto:
+            if self._dataset_info is None:
+                self._dataset_info = self._get_dataset_info()
+
+            return self._dataset_info
+        else:
+            raise CartoException('Your data is not synchronized with CARTO.'
+                                 'First of all, you should call upload method to save your data in CARTO.')
+
+    def set_dataset_info(self, dataset_info=None, privacy=None, name=None):
+        if self._is_saved_in_carto:
+            if self._dataset_info is None:
+                self._dataset_info = self._get_dataset_info()
+
+            if dataset_info:
+                self._dataset_info.update(privacy=dataset_info.privacy)
+            else:
+                self._dataset_info.update(privacy=privacy, name=name)
+        else:
+            raise CartoException('Your data is not synchronized with CARTO.'
+                                 'First of all, you should call upload method to save your data in CARTO.')
+
     def upload(self, with_lonlat=None, if_exists=FAIL, table_name=None, schema=None, context=None):
         if table_name:
             self._table_name = normalize_name(table_name)
@@ -384,6 +409,8 @@ def _map_geom_type(self, geom_type):
             'MultiPolygon': Dataset.GEOM_TYPE_POLYGON
         }[geom_type]
 
+    def _get_dataset_info(self):
+        return DatasetInfo(self._cc, self._table_name)
 
 def recursive_read(context, query, retry_times=Dataset.DEFAULT_RETRY_TIMES):
     try:

cartoframes/dataset_info.py

@@ -0,0 +1,64 @@
+from carto.datasets import DatasetManager
+from carto.exceptions import CartoException
+
+
+class DatasetInfo():
+    PRIVATE = 'PRIVATE'
+    PUBLIC = 'PUBLIC'
+    LINK = 'LINK'
+
+    def __init__(self, carto_context, table_name):
+        self._get_metadata(carto_context, table_name)
+
+        if self._metadata is not None:
+            self.privacy = self._metadata.privacy
+            self.name = self._metadata.name
+        else:
+            raise CartoException('Something goes wrong accessing the table metadata.')
+
+    def update(self, privacy=None, name=None):
+        modified = False
+        if privacy and self._validate_privacy(privacy):
+            self.privacy = privacy.upper()
+            modified = True
+        if name and self._validate_name(name):
+            self.name = name
+            modified = True
+
+        if modified:
+            self._save_metadata()
+
+    def _get_metadata(self, carto_context, table_name, retries=6, retry_wait_time=1):
+        ds_manager = DatasetManager(carto_context.auth_client)
+        try:
+            self._metadata = ds_manager.get(table_name)
+        except Exception as e:
+            if type(e).__name__ == 'NotFoundException' and retries > 0:
+                # if retry_wait_time > 7: # it should be after more than 15 seconds
+                    # warn('We are still procesing the CARTO table. Sorry for the delay.')
+                time.sleep(retry_wait_time)
+                self._get_metadata(carto_context= carto_context, table_name=table_name,
+                                   retries=retries-1, retry_wait_time=retry_wait_time*2)
+            else:
+                return None
+
+    def _save_metadata(self):
+        self._metadata.privacy = self.privacy
+        self._metadata.name = self.name
+        self._metadata.save()
+
+    def _validate_privacy(self, privacy):
+        if privacy.upper() not in [self.PRIVATE, self.PUBLIC, self.LINK]:
+            raise ValueError('Wrong privacy. The privacy: {p} is not valid. You can use: {o1}, {o2}, {o3}'.format(
+                p=privacy, o1=self.PRIVATE, o2=self.PUBLIC, o3=self.LINK))
+
+        if privacy != self.privacy:
+            return True
+
+        return False
+
+    def _validate_name(self, name):
+        if name != self.name:
+            return True
+
+        return False

test/mocks/context_mock.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+
+
+class CredsMock():
+    def __init__(self, key=None, username=None):
+        self._key = key
+        self._username = username
+
+    def username(self):
+        return self._username
+
+
+class ContextMock():
+    def __init__(self, username, api_key):
+        self.is_org = True
+        self.creds = CredsMock(key=api_key, username=username)

test/mocks/dataset_mock.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+
+import pandas as pd
+
+from cartoframes.dataset import Dataset, DatasetInfo
+
+
+class MetadataMock():
+    def __init__(self):
+        self.privacy = DatasetInfo.PRIVATE
+        self.name = None
+
+
+class DatasetInfoMock(DatasetInfo):
+    def _get_metadata(self, _1, _2):
+        self._metadata = MetadataMock()
+
+    def _save_metadata(self):
+        return True
+
+
+class DatasetMock(Dataset):
+    def download(self):
+        self._df = pd.DataFrame({'column_name': [1]})
+        return self._df
+
+    def _copyfrom(self, _):
+        return True
+
+    def _create_table(self, _):
+        return True
+
+    def _create_table_from_query(self):
+        return True
+
+    def exists(self):
+        return False
+
+    def _get_dataset_info(self):
+        return DatasetInfoMock(self._cc, self._table_name)

test/test_dataset.py

@@ -13,6 +13,8 @@
 from cartoframes.dataset import Dataset, _decode_geom
 from cartoframes.columns import normalize_name
 from cartoframes.geojson import load_geojson
+from mocks.dataset_mock import DatasetMock
+from mocks.context_mock import ContextMock
 
 from utils import _UserUrlLoader
 
@@ -483,3 +485,41 @@ def assertNotExistsTable(self, table_name):
                 '''.format(table=table_name))
         except CartoException as e:
             self.assertTrue('relation "{}" does not exist'.format(table_name) in str(e))
+
+
+class TestDatasetInfo(unittest.TestCase):
+    def setUp(self):
+        self.username = 'fake_username'
+        self.api_key = 'fake_api_key'
+        self.context = ContextMock(username=self.username, api_key=self.api_key)
+
+    def test_dataset_get_privacy_from_new_table(self):
+        query = 'SELECT 1'
+        dataset = DatasetMock.from_query(query=query, context=self.context)
+        dataset.upload(table_name='fake_table')
+        self.assertEqual(dataset.get_dataset_info().privacy, Dataset.PRIVATE)
+
+    def test_dataset_get_privacy_from_not_sync(self):
+        query = 'SELECT 1'
+        dataset = DatasetMock.from_query(query=query, context=self.context)
+        error_msg = ('Your data is not synchronized with CARTO.'
+                     'First of all, you should call upload method to save your data in CARTO.')
+        with self.assertRaises(CartoException, msg=error_msg):
+            dataset.get_dataset_info()
+
+    def test_dataset_set_privacy_to_new_table(self):
+        query = 'SELECT 1'
+        dataset = DatasetMock.from_query(query=query, context=self.context)
+        dataset.upload(table_name='fake_table')
+        dataset.set_dataset_info(privacy=Dataset.PUBLIC)
+        self.assertEqual(dataset.get_dataset_info().privacy, Dataset.PUBLIC)
+
+    def test_dataset_set_privacy_with_wrong_parameter(self):
+        query = 'SELECT 1'
+        dataset = DatasetMock.from_query(query=query, context=self.context)
+        dataset.upload(table_name='fake_table')
+        wrong_privacy = 'wrong_privacy'
+        error_msg = 'Wrong privacy. The privacy: {p} is not valid. You can use: {o1}, {o2}, {o3}'.format(
+                        p=wrong_privacy, o1=Dataset.PRIVATE, o2=Dataset.PUBLIC, o3=Dataset.LINK)
+        with self.assertRaises(ValueError, msg=error_msg):
+            dataset.set_dataset_info(privacy=wrong_privacy)