class FastExportOperator(BaseOperator):
"""
Export a table from Teradata to csv file
:param sql_select_stmt: Select statament to export
:type sql_select_stmt: str
:param output_file: output file to export
:type output_file: str
"""
template_fields = ('sql','output_file',)
template_ext = ('.sql',)
ui_color = '#a8e4b1'
def __init__(
self,
*,
sql_select_stmt: str,
output_file: str,
delimiter: str = ';',
encoding: str = 'UTF8',
spool_mode: str = 'SPOOL',
xcom_push: bool = True,
ttu_conn_id: str = 'ttu_default',
**kwargs,
) -> None:
super().__init__(**kwargs)
self.sql = sql_select_stmt
self.output_file = output_file
self.delimiter = delimiter
self.encoding = encoding
self.spool_mode = spool_mode
self.xcom_push = xcom_push
self._hook = None
self.ttu_conn_id = ttu_conn_id
def execute(self, context):
"""
Call the function
"""
self._hook = TtuHook(ttu_conn_id=self.ttu_conn_id)
self._hook.execute_tptexport(sql=self.sql,
output_file=self.output_file,
delimiter=self.delimiter,
encoding=self.encoding,
spool_mode=self.spool_mode,
xcom_push_flag=self.xcom_push)
def on_kill(self):
self._hook.on_kill()
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()
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()