def test_execute_mysql_preoperator(self, mock_hive_hook, mock_mysql_hook):
self.kwargs.update(dict(mysql_preoperator='preoperator'))
HiveToMySqlTransfer(**self.kwargs).execute(context={})
mock_mysql_hook.return_value.run.assert_called_once_with(
self.kwargs['mysql_preoperator'])
def test_execute_with_hive_conf(self, mock_hive_hook, mock_mysql_hook):
context = {}
self.kwargs.update(
dict(hive_conf={'mapreduce.job.queuename': 'fake_queue'}))
HiveToMySqlTransfer(**self.kwargs).execute(context=context)
hive_conf = context_to_airflow_vars(context)
hive_conf.update(self.kwargs['hive_conf'])
mock_hive_hook.return_value.get_records.assert_called_once_with(
self.kwargs['sql'], hive_conf=hive_conf)
def test_execute(self, mock_hive_hook, mock_mysql_hook):
HiveToMySqlTransfer(**self.kwargs).execute(context={})
mock_hive_hook.assert_called_once_with(
hiveserver2_conn_id=self.kwargs['hiveserver2_conn_id'])
mock_hive_hook.return_value.get_records.assert_called_once_with(
'sql', hive_conf={})
mock_mysql_hook.assert_called_once_with(
mysql_conn_id=self.kwargs['mysql_conn_id'])
mock_mysql_hook.return_value.insert_rows.assert_called_once_with(
table=self.kwargs['mysql_table'],
rows=mock_hive_hook.return_value.get_records.return_value)
def test_execute_bulk_load(self, mock_hive_hook, mock_tmp_file,
mock_mysql_hook):
type(mock_tmp_file).name = PropertyMock(return_value='tmp_file')
context = {}
self.kwargs.update(dict(bulk_load=True))
HiveToMySqlTransfer(**self.kwargs).execute(context=context)
mock_tmp_file.assert_called_once_with()
mock_hive_hook.return_value.to_csv.assert_called_once_with(
self.kwargs['sql'],
mock_tmp_file.return_value.name,
delimiter='\t',
lineterminator='\n',
output_header=False,
hive_conf=context_to_airflow_vars(context))
mock_mysql_hook.return_value.bulk_load.assert_called_once_with(
table=self.kwargs['mysql_table'],
tmp_file=mock_tmp_file.return_value.name)
mock_tmp_file.return_value.close.assert_called_once_with()
def test_hive_to_mysql(self):
op = HiveToMySqlTransfer(
mysql_conn_id='airflow_db',
task_id='hive_to_mysql_check',
create=True,
sql="""
SELECT name
FROM airflow.static_babynames
LIMIT 100
""",
mysql_table='test_static_babynames',
mysql_preoperator=[
'DROP TABLE IF EXISTS test_static_babynames;',
'CREATE TABLE test_static_babynames (name VARCHAR(500))',
],
dag=self.dag)
op.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
op.run(start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE,
ignore_ti_state=True)