class BteqOperator(BaseOperator):
"""
Executes BTEQ code in a specific Teradata database
:param sql: the bteq code to be executed. (templated)
:type sql: Can receive a str representing a sql statement,
a list of str (sql statements), or reference to a template file.
Template reference are recognized by str ending in '.sql'
:param postgres_conn_id: reference to a specific postgres database
:type postgres_conn_id: str
:param autocommit: if True, each command is automatically committed.
(default value: False)
:type autocommit: bool
:param parameters: (optional) the parameters to render the SQL query with.
:type parameters: dict or iterable
:param database: name of database which overwrite defined one in connection
:type database: str
"""
template_fields = ('sql',)
template_ext = ('.sql', '.bteq',)
ui_color = '#ff976d'
@apply_defaults
def __init__(
self,
*,
bteq: str,
xcom_push: bool = True,
ttu_conn_id: str = 'ttu_default',
**kwargs
) -> None:
super().__init__(**kwargs)
self.sql = bteq
self._hook = None
self.xcom_push = xcom_push
self.ttu_conn_id = ttu_conn_id
def execute(self, context):
"""
Call execute_bteq method from TttuHook to run the provided BTEQ string
"""
self._hook = TtuHook(ttu_conn_id=self.ttu_conn_id)
self._hook.execute_bteq(self.sql, self.xcom_push)
def on_kill(self):
self._hook.on_kill()
def test_execute_bteq_runcmd(self, mock_tmpfile, mock_tmpdir, mock_popen):
# Given
mock_subprocess = MockSubProcess()
mock_subprocess.returncode = 0
mock_popen.return_value = mock_subprocess
mock_tmpdir.return_value.__enter__.return_value = '/tmp/airflowtmp_ttu_bteq'
mock_tmpfile.return_value.__enter__.return_value.name = 'test.bteq'
# When
hook = TtuHook(ttu_conn_id='ttu_default')
hook.execute_bteq(bteq="")
# Then
mock_popen.assert_called_with(['bteq'],
stdin=mock.ANY,
stdout=mock_subprocess.PIPE,
stderr=mock_subprocess.STDOUT,
cwd='/tmp/airflowtmp_ttu_bteq',
preexec_fn=mock.ANY)
def test_execute_bteq_runcmd_error_noraise(self, mock_tmpfile, mock_tmpdir,
mock_popen):
# Givens
mock_subprocess = MockSubProcess(
output=self._bteq_error_no_failure_subprocess_output)
mock_subprocess.returncode = 0
mock_popen.return_value = mock_subprocess
mock_tmpdir.return_value.__enter__.return_value = '/tmp/airflowtmp_ttu_bteq'
mock_tmpfile.return_value.__enter__.return_value.name = 'test.bteq'
# When
hook = TtuHook(ttu_conn_id='ttu_default')
# Then
with self.assertLogs(level="INFO") as cm:
hook.execute_bteq(bteq="")
self.assertEqual(
"INFO:airflow.providers.teradata.hooks.ttu.TtuHook:BTEQ command exited with return code 0",
cm.output[-1],
)
def test_execute_bteq_runcmd_error_raise(self, mock_tmpfile, mock_tmpdir,
mock_popen):
# Given
mock_subprocess = MockSubProcess(
output=self._bteq_failure_subprocess_output)
mock_subprocess.returncode = 311
mock_popen.return_value = mock_subprocess
mock_tmpdir.return_value.__enter__.return_value = '/tmp/airflowtmp_ttu_bteq'
mock_tmpfile.return_value.__enter__.return_value.name = 'test.bteq'
# When
hook = TtuHook(ttu_conn_id='ttu_default')
# Then
with self.assertRaises(AirflowException) as cm:
hook.execute_bteq(bteq="")
msg = (
"BTEQ command exited with return code 311 because of "
"*** Failure 3706 Syntax error: expected something between '(' and the string 'test'"
)
self.assertEqual(str(cm.exception), msg)
def test_execute_bteq_runcmd_return_last_line(self, mock_tmpfile,
mock_tmpdir, mock_popen):
# Givens
mock_subprocess = MockSubProcess(output=self._bteq_subprocess_output)
mock_subprocess.returncode = 0
mock_popen.return_value = mock_subprocess
mock_tmpdir.return_value.__enter__.return_value = '/tmp/airflowtmp_ttu_bteq'
mock_tmpfile.return_value.__enter__.return_value.name = 'test.bteq'
# When
hook = TtuHook(ttu_conn_id='ttu_default')
# Then
res = hook.execute_bteq(bteq="", xcom_push_flag=True)
self.assertEqual(
"*** RC (return code) = 0",
res,
)
class FastLoadOperator(BaseOperator):
"""
Load a CSV file (without header) to Teradata Database Table
:param input_file: output file to export
:type input_file: str
:param target_table: target table
:type target_table: str
:param delimiter: file delimiter
:type delimiter: str
:param working_database: database to create log tables
:type working_database: str
:param encoding: encoding of the file
:type encoding: str
:param xcom_push: True if return last log to xcom
:type xcom_push: Bool
"""
template_fields = (
'input_file',
'target_table',
'preoperator_bteq',
)
template_ext = (
'.sql',
'.bteq',
)
ui_color = '#4aa3ba'
@apply_defaults
def __init__(
self,
*,
input_file: str,
target_table: str,
working_database: str,
delimiter: str = ';',
encoding: str = 'UTF8',
preoperator_bteq: Optional[str],
raise_on_rows_error: bool = True,
raise_on_rows_duplicated: bool = True,
xcom_push: bool = True,
ttu_conn_id: str = 'ttu_default',
**kwargs,
) -> None:
super().__init__(**kwargs)
self.input_file = input_file
self.target_table = target_table
self.working_database = working_database
self.delimiter = delimiter
self.encoding = encoding
self.xcom_push = xcom_push
self._hook = None
self.ttu_conn_id = ttu_conn_id
self.preoperator_bteq = preoperator_bteq
self.raise_on_rows_error = raise_on_rows_error
self.raise_on_rows_duplicated = raise_on_rows_duplicated
def execute(self, context):
"""
Call the executable from teradata
"""
self._hook = TtuHook(ttu_conn_id=self.ttu_conn_id)
if self.preoperator_bteq:
logging.info('Executing preoperator BTEQ')
self._hook.execute_bteq(self.preoperator_bteq, self.xcom_push)
logging.info('Executing tdload command')
self._hook.execute_tdload(
input_file=self.input_file,
table=self.target_table,
working_database=self.working_database,
delimiter=self.delimiter,
encoding=self.encoding,
xcom_push_flag=self.xcom_push,
raise_on_rows_error=self.raise_on_rows_error,
raise_on_rows_duplicated=self.raise_on_rows_duplicated)
def on_kill(self):
self._hook.on_kill()