def test_temporary_loadable_directory_load_with_no_s3_temp_bucket_configured(self):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = None
vertica_unloader = VerticaUnloader(db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
with self.assertRaises(NoTemporaryBucketConfiguration):
with vertica_unloader.temporary_loadable_directory_loc():
pass
def test_unload_with_no_aws_creds(self, mock_open):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = MagicMock(name='s3_temp_base_loc')
mock_target_records_format = Mock(name='target_records_format',
spec=DelimitedRecordsFormat)
mock_target_records_format.hints = {}
mock_out = mock_db.execute.return_value
mock_out.fetchall.return_value = ['awslib']
unloader = VerticaUnloader(db=mock_db,
s3_temp_base_loc=mock_s3_temp_base_loc)
mock_table = Mock(name='table')
mock_unload_plan = Mock(name='unload_plan')
mock_unload_plan.records_format = Mock(spec=DelimitedRecordsFormat)
mock_unload_plan.records_format.hints = {}
mock_schema = Mock(name='schema')
mock_directory = Mock(name='directory')
mock_unload_plan.records_format = Mock(spec=DelimitedRecordsFormat)
mock_unload_plan.records_format.hints = {}
mock_unload_plan.records_format.format_type = 'delimited'
mock_column = Mock(name='column')
mock_db.dialect.get_columns.return_value = [mock_column]
with self.assertRaises(NotImplementedError):
unloader.unload(schema=mock_schema,
table=mock_table,
unload_plan=mock_unload_plan,
directory=mock_directory)
def test_no_temp_bucket_known_supported_records_formats_for_unload(self):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = None
unloader = VerticaUnloader(db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
out = unloader.known_supported_records_formats_for_unload()
self.assertEqual([], out)
def test_known_supported_records_formats_for_unload(self):
mock_db = Mock(name='db')
mock_source_records_format = Mock(name='source_records_format',
spec=DelimitedRecordsFormat)
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
vertica_unloader = VerticaUnloader(
db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
mock_source_records_format.hints = {}
out = vertica_unloader.known_supported_records_formats_for_unload()
self.assertEqual(out, [DelimitedRecordsFormat(variant='vertica')])
def test_no_temp_bucket_can_unload_this_format_true(self,
mock_RecordsUnloadPlan):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = None
mock_target_records_format = Mock(name='target_records_format', spec=DelimitedRecordsFormat)
mock_target_records_format.hints = {}
unloader = VerticaUnloader(db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
out = unloader.can_unload_this_format(mock_target_records_format)
self.assertFalse(out)
def test_can_unload_format_true(self, mock_vertica_export_options):
mock_db = Mock(name='db')
mock_source_records_format = Mock(name='source_records_format',
spec=DelimitedRecordsFormat)
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
vertica_unloader = VerticaUnloader(
db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
mock_source_records_format.hints = {}
out = vertica_unloader.can_unload_format(mock_source_records_format)
mock_vertica_export_options.assert_called_with(set(), ANY)
self.assertEqual(True, out)
def test_s3_export_available_false_no_awslib(self):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
mock_target_records_format = Mock(name='target_records_format',
spec=DelimitedRecordsFormat)
mock_target_records_format.hints = {}
mock_out = mock_db.execute.return_value
mock_out.fetchall.return_value = []
unloader = VerticaUnloader(db=mock_db,
s3_temp_base_loc=mock_s3_temp_base_loc)
self.assertEqual(False, unloader.s3_export_available())
def test_can_unload_to_scheme_s3_but_with_s3_export_true(self):
mock_db = Mock(name='db')
mock_resultset = Mock(name='resultset')
mock_db.execute.return_value = mock_resultset
mock_resultset.fetchall.return_value = ['awslib']
mock_s3_temp_base_loc = None
vertica_unloader = VerticaUnloader(
db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
self.assertTrue(vertica_unloader.can_unload_to_scheme('s3'))
mock_db.execute.\
assert_called_with("SELECT lib_name from user_libraries where lib_name = 'awslib'")
def test_known_supported_records_formats_for_unload(self):
mock_db = Mock(name='db')
mock_source_records_format = Mock(name='source_records_format',
spec=DelimitedRecordsFormat)
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
vertica_unloader = VerticaUnloader(
db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
mock_resultset = Mock(name='resultset')
mock_db.execute.return_value = mock_resultset
mock_resultset.fetchall.return_value = ['awslib']
mock_source_records_format.hints = {}
out = vertica_unloader.known_supported_records_formats_for_unload()
mock_db.execute.\
assert_called_with("SELECT lib_name from user_libraries where lib_name = 'awslib'")
self.assertEqual(out, [DelimitedRecordsFormat(variant='vertica')])
def test_can_unload_this_format_with_s3_true(self,
mock_vertica_export_options):
mock_db = Mock(name='db')
mock_source_records_format = Mock(name='source_records_format',
spec=DelimitedRecordsFormat)
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
vertica_unloader = VerticaUnloader(
db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
mock_resultset = Mock(name='resultset')
mock_db.execute.return_value = mock_resultset
mock_resultset.fetchall.return_value = ['awslib']
mock_source_records_format.hints = {}
out = vertica_unloader.can_unload_this_format(
mock_source_records_format)
mock_db.execute.\
assert_called_with("SELECT lib_name from user_libraries where lib_name = 'awslib'")
mock_vertica_export_options.assert_called_with(set(), ANY)
self.assertEqual(True, out)
def test_no_temp_bucket_unload(self):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = None
unloader = VerticaUnloader(db=mock_db, s3_temp_base_loc=mock_s3_temp_base_loc)
mock_schema = Mock(name='schema')
mock_table = Mock(name='table')
mock_unload_plan = Mock(name='unload_plan')
mock_unload_plan.records_format = Mock(spec=DelimitedRecordsFormat)
mock_unload_plan.records_format.hints = {
'compression': 'GZIP'
}
mock_directory = Mock(name='directory')
with self.assertRaises(NotImplementedError):
unloader.unload(schema=mock_schema,
table=mock_table,
unload_plan=mock_unload_plan,
directory=mock_directory)
def test_unload_to_s3_directory_with_token(self):
mock_db = Mock(name='db')
mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
mock_out = mock_db.execute.return_value
mock_out.fetchall.return_value = []
unloader = VerticaUnloader(db=mock_db,
s3_temp_base_loc=mock_s3_temp_base_loc)
mock_table = Mock(name='table')
mock_unload_plan = Mock(name='unload_plan')
mock_schema = Mock(name='schema')
mock_directory = Mock(name='directory')
mock_unload_plan.records_format = Mock(spec=DelimitedRecordsFormat)
mock_unload_plan.records_format.hints = {}
mock_unload_plan.records_format.format_type = 'delimited'
with self.assertRaises(CredsDoNotSupportS3Export):
unloader.unload_to_s3_directory(schema=mock_schema,
table=mock_table,
unload_plan=mock_unload_plan,
directory=mock_directory)
def test_s3_temp_bucket_available_false(self):
mock_db = Mock(name='db')
vertica_unloader = VerticaUnloader(db=mock_db, s3_temp_base_loc=None)
self.assertFalse(vertica_unloader.s3_temp_bucket_available())