def test_parquet_with_delimited_hint_hint_first(self, mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
@contextmanager
def mysource(records_format=None):
self.assertIsInstance(records_format, ParquetRecordsFormat)
yield mock_source
@contextmanager
def mytarget(db_engine):
yield mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {
'compression': 'GZIP',
'format': 'parquet',
},
'target': {
'db_name': 'foo'
},
}
with self.assertRaises(NotImplementedError):
run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
def test_variant_with_no_hints(self, mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
@contextmanager
def mysource(initial_hints=None, records_format=None):
self.assertIsNone(initial_hints)
self.assertEqual(records_format.variant, 'vertica')
yield mock_source
@contextmanager
def mytarget(db_engine):
yield mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {
'variant': 'vertica',
},
'target': {
'db_name': 'foo'
},
}
out = run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
mock_records.move.assert_called()
self.assertEqual(out, mock_records.move.return_value)
def test_parquet(self, mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
@contextmanager
def mysource(records_format=None):
self.assertIsInstance(records_format, ParquetRecordsFormat)
yield mock_source
@contextmanager
def mytarget(db_engine):
yield mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {
'format': 'parquet'
},
'target': {
'db_name': 'foo'
},
}
out = run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
mock_records.move.assert_called()
self.assertEqual(out, mock_records.move.return_value)
def test_run_records_mover_job(self,
mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
def mysource(a: int,
schema_name: str,
google_cloud_creds,
existing_table_handling=None,
spectrum_base_url=None):
self.assertEqual(a, 1)
self.assertEqual(google_cloud_creds, mock_session.creds.google_sheets.return_value)
self.assertEqual(schema_name, 'myschema')
self.assertEqual(spectrum_base_url, 'spectrumschema')
return mock_source
def mytarget(b: int, db_engine):
assert b == 2
self.assertEqual(db_engine, mock_session.get_db_engine.return_value)
return mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_db_facts = {
'redshift_spectrum_base_url_myschema': 'spectrumschema'
}
mock_session.creds.db_facts.return_value = mock_db_facts
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {
'a': 1,
'db_name': 'mydbname',
'gcp_creds_name': 'mygcpcreds',
'existing_table': 'drop_and_recreate',
'schema_name': 'myschema'
},
'target': {
'b': 2,
'db_name': 'foo',
}
}
out = run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
mock_records.move.assert_called()
self.assertEqual(out, mock_records.move.return_value)
def test_hints_with_no_variant(self, mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
@contextmanager
def mysource(initial_hints=None, records_format=None):
self.assertTrue(initial_hints is not None)
self.assertEqual(',', initial_hints['field-delimiter'])
self.assertEqual('GZIP', initial_hints['compression'])
self.assertEqual('all', initial_hints['quoting'])
self.assertIsNone(records_format)
yield mock_source
@contextmanager
def mytarget(db_engine):
yield mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {
'field-delimiter': ',',
'compression': 'GZIP',
'quoting': 'all',
},
'target': {
'db_name': 'foo'
}
}
out = run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
mock_records.move.assert_called()
self.assertEqual(out, mock_records.move.return_value)
def test_run_records_mover_hint_and_then_variant(self,
mock_Session):
mock_source = Mock(name='source')
mock_target = Mock(name='target')
def mysource():
return mock_source
def mytarget(records_format: Optional[BaseRecordsFormat] = None):
assert isinstance(records_format, DelimitedRecordsFormat)
self.assertEqual(records_format.variant, 'csv')
self.assertIsNone(records_format.hints['escape'])
return mock_target
mock_job_name = Mock(name='job_name')
mock_session = mock_Session.return_value
mock_db_facts = {
'redshift_spectrum_base_url_myschema': 'spectrumschema'
}
mock_session.creds.db_facts.return_value = mock_db_facts
mock_records = mock_session.records
mock_records.sources.mysource = mysource
mock_records.targets.mytarget = mytarget
config = {
'fail_if_dont_understand': False,
'source': {},
'target': {
'compression': None,
'variant': 'csv'
}
}
out = run_records_mover_job(source_method_name='mysource',
target_method_name='mytarget',
job_name=mock_job_name,
config=config)
mock_records.move.assert_called()
self.assertEqual(out, mock_records.move.return_value)