def test_column_info_without_geom(self):
df = DataFrame([['Gran Vía 46', 'Madrid'], ['Ebro 1', 'Sevilla']],
columns=['Address', 'City'])
dataframe_columns_info = get_dataframe_columns_info(df)
assert dataframe_columns_info == [
ColumnInfo('Address', 'address', 'text', False),
ColumnInfo('City', 'city', 'text', False)
]
def test_column_info_basic_troubled_names(self):
gdf = GeoDataFrame(
[[1, 'POINT (1 1)', 'fake_geom']],
columns=['cartodb_id', 'the_geom', 'the_geom_webmercator'])
set_geometry(gdf, 'the_geom', inplace=True)
dataframe_columns_info = get_dataframe_columns_info(gdf)
assert dataframe_columns_info == [
ColumnInfo('cartodb_id', 'cartodb_id', 'bigint', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Point, 4326)', True)
]
def test_column_info_with_geom(self):
gdf = GeoDataFrame(
[['Gran Vía 46', 'Madrid', 'POINT (0 0)'], ['Ebro 1', 'Sevilla', 'POINT (1 1)']],
columns=['Address', 'City', 'the_geom'])
set_geometry(gdf, 'the_geom', inplace=True)
dataframe_columns_info = get_dataframe_columns_info(gdf)
assert dataframe_columns_info == [
ColumnInfo('Address', 'address', 'text', False),
ColumnInfo('City', 'city', 'text', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Point, 4326)', True)
]
def test_column_info_geometry_troubled_names(self):
gdf = GeoDataFrame(
[['POINT (0 0)', 'POINT (1 1)', 'POINT (2 2)']],
columns=['Geom', 'the_geom', 'g-e-o-m-e-t-r-y'])
set_geometry(gdf, 'the_geom', inplace=True)
dataframe_columns_info = get_dataframe_columns_info(gdf)
assert dataframe_columns_info == [
ColumnInfo('Geom', 'geom', 'text', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Point, 4326)', True),
ColumnInfo('g-e-o-m-e-t-r-y', 'g_e_o_m_e_t_r_y', 'text', False)
]
def test_converters(self):
columns = [
ColumnInfo('cartodb_id', 'cartodb_id', 'integer', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Geometry, 4326)', True),
ColumnInfo('name', 'name', 'text', False),
ColumnInfo('flag', 'flag', 'boolean', False),
ColumnInfo('number', 'number', 'double precision', False)
]
converters = obtain_converters(columns)
assert type(converters) == dict
assert converters['cartodb_id'] == _convert_int
assert converters['the_geom'] == _convert_generic
assert converters['name'] == _convert_generic
assert converters['flag'] == _convert_bool
assert converters['number'] == _convert_float
def test_column_info_sort(self):
columns = [
ColumnInfo('cartodb_id', 'cartodb_id', 'integer', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Geometry, 4326)',
True),
ColumnInfo('Name', 'name', 'text', False),
ColumnInfo('flag', 'flag', 'boolean', False),
ColumnInfo('NUMBER', 'number', 'double precision', False)
]
columns.sort()
assert columns == [
ColumnInfo('cartodb_id', 'cartodb_id', 'integer', False),
ColumnInfo('flag', 'flag', 'boolean', False),
ColumnInfo('Name', 'name', 'text', False),
ColumnInfo('NUMBER', 'number', 'double precision', False),
ColumnInfo('the_geom', 'the_geom', 'geometry(Geometry, 4326)',
True)
]
def test_internal_copy_from(self, mocker):
# Given
from shapely.geometry import Point
mocker.patch(
'cartoframes.io.managers.context_manager._create_auth_client')
mock = mocker.patch.object(CopySQLClient, 'copyfrom')
gdf = GeoDataFrame({'A': [1, 2], 'B': [Point(0, 0), Point(1, 1)]})
columns = [
ColumnInfo('A', 'a', 'bigint', False),
ColumnInfo('B', 'b', 'geometry', True)
]
# When
cm = ContextManager(self.credentials)
cm._copy_from(gdf, 'table_name', columns)
# Then
assert mock.call_args[0][0] == '''
COPY table_name("a","b") FROM stdin WITH (FORMAT csv, DELIMITER '|', NULL '__null');
'''.strip()
assert list(mock.call_args[0][1]) == [
b'1|0101000020E610000000000000000000000000000000000000\n',
b'2|0101000020E6100000000000000000F03F000000000000F03F\n'
]
def test_copy_from(self, mocker):
# Given
mocker.patch(
'cartoframes.io.managers.context_manager._create_auth_client')
mocker.patch.object(ContextManager, 'has_table', return_value=False)
mock = mocker.patch.object(ContextManager, '_copy_from')
df = DataFrame({'A': [1]})
columns = [ColumnInfo('A', 'a', 'bigint', False)]
# When
cm = ContextManager(self.credentials)
cm.copy_from(df, 'TABLE NAME')
# Then
mock.assert_called_once_with(df, 'table_name', columns)
def test_copy_to(self, mocker):
# Given
query = '__query__'
columns = [ColumnInfo('A', 'a', 'bigint', False)]
mocker.patch.object(ContextManager,
'compute_query',
return_value=query)
mocker.patch.object(ContextManager,
'_get_query_columns_info',
return_value=columns)
mock = mocker.patch.object(ContextManager, '_copy_to')
# When
cm = ContextManager(self.credentials)
cm.copy_to(query)
# Then
mock.assert_called_once_with('SELECT "A" FROM (__query__) _q', columns,
3)
def test_copy_from_exists_replace_truncate_and_drop_add_columns(
self, mocker):
# Given
mocker.patch(
'cartoframes.io.managers.context_manager._create_auth_client')
mocker.patch.object(ContextManager, 'has_table', return_value=True)
mocker.patch.object(ContextManager,
'get_schema',
return_value='schema')
mock = mocker.patch.object(ContextManager,
'_truncate_and_drop_add_columns')
df = DataFrame({'A': [1]})
columns = [ColumnInfo('A', 'a', 'bigint', False)]
# When
cm = ContextManager(self.credentials)
cm.copy_from(df, 'TABLE NAME', 'replace')
# Then
mock.assert_called_once_with('table_name', 'schema', columns, [], True)
def test_copy_from(self, mocker):
# Given
mocker.patch(
'cartoframes.io.managers.context_manager._create_auth_client')
mocker.patch.object(ContextManager, 'has_table', return_value=False)
mocker.patch.object(ContextManager,
'get_schema',
return_value='schema')
mock_create_table = mocker.patch.object(ContextManager,
'execute_long_running_query')
mock = mocker.patch.object(ContextManager, '_copy_from')
df = DataFrame({'A': [1]})
columns = [ColumnInfo('A', 'a', 'bigint', False)]
# When
cm = ContextManager(self.credentials)
cm.copy_from(df, 'TABLE NAME')
# Then
mock_create_table.assert_called_once_with('''
BEGIN; CREATE TABLE table_name ("a" bigint); SELECT CDB_CartodbfyTable(\'schema\', \'table_name\'); COMMIT;
'''.strip())
mock.assert_called_once_with(df, 'table_name', columns,
DEFAULT_RETRY_TIMES)
def test_column_info_compare(self):
column = ColumnInfo('cartodb_id', 'cartodb_id', 'integer', False)
assert column == ColumnInfo('CARTODB_ID', 'cartodb_id', 'integer',
False)