Merge pull request #1744 from CartoDB/feature/ch175987/economicalinsu…

…rance-cartoframes-datasets

Add cartodbfy step for create_table_from_query and copy_table functions

cartoframes/io/carto.py

@@ -275,7 +275,7 @@ def rename_table(table_name, new_table_name, credentials=None, if_exists='fail',
         log.info('Success! Table "{0}" renamed to table "{1}" correctly'.format(table_name, new_table_name))
 
 
-def copy_table(table_name, new_table_name, credentials=None, if_exists='fail', log_enabled=True):
+def copy_table(table_name, new_table_name, credentials=None, if_exists='fail', log_enabled=True, cartodbfy=True):
     """Copy a table into a new table in the CARTO account.
 
     Args:
@@ -285,7 +285,8 @@ def copy_table(table_name, new_table_name, credentials=None, if_exists='fail', l
             instance of Credentials (username, api_key, etc).
         if_exists (str, optional): 'fail', 'replace', 'append'. Default is 'fail'.
         log_enabled (bool, optional): enable the logging mechanism. Default is True.
-
+        cartodbfy (bool, optional): convert the table to CARTO format. Default True. More info
+            `here <https://carto.com/developers/sql-api/guides/creating-tables/#create-tables>`.
     Raises:
         ValueError: if the table names provided are wrong or the if_exists param is not valid.
 
@@ -303,13 +304,19 @@ def copy_table(table_name, new_table_name, credentials=None, if_exists='fail', l
     query = 'SELECT * FROM {}'.format(table_name)
 
     context_manager = ContextManager(credentials)
-    new_table_name = context_manager.create_table_from_query(query, new_table_name, if_exists)
+    new_table_name = context_manager.create_table_from_query(query, new_table_name, if_exists, cartodbfy)
 
     if log_enabled:
         log.info('Success! Table "{0}" copied to table "{1}" correctly'.format(table_name, new_table_name))
 
 
-def create_table_from_query(query, new_table_name, credentials=None, if_exists='fail', log_enabled=True):
+def create_table_from_query(
+        query,
+        new_table_name,
+        credentials=None,
+        if_exists='fail',
+        log_enabled=True,
+        cartodbfy=True):
     """Create a new table from an SQL query in the CARTO account.
 
     Args:
@@ -319,7 +326,8 @@ def create_table_from_query(query, new_table_name, credentials=None, if_exists='
             instance of Credentials (username, api_key, etc).
         if_exists (str, optional): 'fail', 'replace', 'append'. Default is 'fail'.
         log_enabled (bool, optional): enable the logging mechanism. Default is True.
-
+        cartodbfy (bool, optional): convert the table to CARTO format. Default True. More info
+            `here <https://carto.com/developers/sql-api/guides/creating-tables/#create-tables>`.
     Raises:
         ValueError: if the query or table name provided is wrong or the if_exists param is not valid.
 
@@ -335,7 +343,7 @@ def create_table_from_query(query, new_table_name, credentials=None, if_exists='
             ', '.join(IF_EXISTS_OPTIONS)))
 
     context_manager = ContextManager(credentials)
-    new_table_name = context_manager.create_table_from_query(query, new_table_name, if_exists)
+    new_table_name = context_manager.create_table_from_query(query, new_table_name, if_exists, cartodbfy)
 
     if log_enabled:
         log.info('Success! Table "{0}" created correctly'.format(new_table_name))

cartoframes/io/managers/context_manager.py

@@ -122,7 +122,7 @@ def copy_from(self, gdf, table_name, if_exists='fail', cartodbfy=True,
 
         return table_name
 
-    def create_table_from_query(self, query, table_name, if_exists):
+    def create_table_from_query(self, query, table_name, if_exists, cartodbfy=True):
         schema = self.get_schema()
         table_name = self.normalize_table_name(table_name)
 
@@ -140,6 +140,10 @@ def create_table_from_query(self, query, table_name, if_exists):
         else:
             self._drop_create_table_from_query(table_name, schema, query)
 
+        if cartodbfy is True:
+            cartodbfy_query = _cartodbfy_query(table_name, schema)
+            self.execute_long_running_query(cartodbfy_query)
+
         return table_name
 
     def list_tables(self, schema=None):

tests/unit/io/managers/test_context_manager.py

@@ -278,3 +278,29 @@ def __init__(self):
 
         with pytest.raises(CartoRateLimitException):
             test_function(retry_times=0)
+
+    def test_create_table_from_query_cartodbfy(self, mocker):
+        # Given
+        mocker.patch.object(ContextManager, 'has_table', return_value=False)
+        mocker.patch.object(ContextManager, 'get_schema', return_value='schema')
+        mock = mocker.patch.object(ContextManager, 'execute_long_running_query')
+
+        # When
+        cm = ContextManager(self.credentials)
+        cm.create_table_from_query('SELECT * FROM table_name', '__new_table_name__', if_exists='fail', cartodbfy=True)
+
+        # Then
+        mock.assert_called_with("SELECT CDB_CartodbfyTable('schema', '__new_table_name__')")
+
+    def test_create_table_from_query_cartodbfy_default(self, mocker):
+        # Given
+        mocker.patch.object(ContextManager, 'has_table', return_value=False)
+        mocker.patch.object(ContextManager, 'get_schema', return_value='schema')
+        mock = mocker.patch.object(ContextManager, 'execute_long_running_query')
+
+        # When
+        cm = ContextManager(self.credentials)
+        cm.create_table_from_query('SELECT * FROM table_name', '__new_table_name__', if_exists='fail')
+
+        # Then
+        mock.assert_called_with("SELECT CDB_CartodbfyTable('schema', '__new_table_name__')")

tests/unit/io/test_carto.py

@@ -500,6 +500,39 @@ def test_copy_table_wrong_if_exists(mocker):
     assert str(e.value) == 'Wrong option for the `if_exists` param. You should provide: fail, replace, append.'
 
 
+def test_copy_table_no_cartodbfy(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    copy_table('__table_name__', '__new_table_name__', CREDENTIALS, cartodbfy=False)
+
+    # Then
+    assert cm_mock.call_args[0][3] is False
+
+
+def test_copy_table_cartodbfy(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    copy_table('__table_name__', '__new_table_name__', CREDENTIALS, cartodbfy=True)
+
+    # Then
+    assert cm_mock.call_args[0][3] is True
+
+
+def test_copy_table_cartodbfy_default(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    copy_table('__table_name__', '__new_table_name__', CREDENTIALS)
+
+    # Then
+    assert cm_mock.call_args[0][3] is True
+
+
 def test_create_table_from_query_wrong_query(mocker):
     # When
     with pytest.raises(ValueError) as e:
@@ -535,3 +568,36 @@ def test_create_table_from_query_wrong_if_exists(mocker):
 
     # Then
     assert str(e.value) == 'Wrong option for the `if_exists` param. You should provide: fail, replace, append.'
+
+
+def test_create_table_from_query_no_cartodbfy(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    create_table_from_query('SELECT * FROM table', '__new_table_name__', CREDENTIALS, cartodbfy=False)
+
+    # Then
+    assert cm_mock.call_args[0][3] is False
+
+
+def test_create_table_from_query_cartodbfy(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    create_table_from_query('SELECT * FROM table', '__new_table_name__', CREDENTIALS, cartodbfy=True)
+
+    # Then
+    assert cm_mock.call_args[0][3] is True
+
+
+def test_create_table_from_query_cartodbfy_default(mocker):
+    # Given
+    cm_mock = mocker.patch.object(ContextManager, 'create_table_from_query')
+
+    # When
+    create_table_from_query('SELECT * FROM table', '__new_table_name__', CREDENTIALS)
+
+    # Then
+    assert cm_mock.call_args[0][3] is True