def execute(self, context=None):
self.log.info('Executing SQL: {}'.format(self.sql))
hook = VerticaHook(vertica_conn_id=self.vertica_conn_id)
records = len(hook.get_records(self.sql))
self.log.info('Returned rows: {}'.format(records))
if records == 2 and not self.reverse:
raise AirflowException('Row counts are not equal')
elif records == 1 and self.reverse:
raise AirflowException('Row counts are equal')
elif records not in (1, 2):
raise AirflowException(
'Expected 1 or 2 rows, got {}'.format(records))
self.log.info('Success!')
class VerticaSensor(BaseSensorOperator):
def __init__(self, vertica_conn_id='vertica_default',
checker=lambda x: True,
sql='select 1',
*args, **kwargs):
super().__init__(*args, **kwargs)
self.hook = VerticaHook(vertica_conn_id=vertica_conn_id)
self.checker = checker
self.sql = sql
def poke(self, context):
self.log.info('Poking: {}'.format(self.sql))
try:
response = self.hook.get_records(self.sql)
return self.checker(response)
except Exception as e:
return False
raise e
return True
def execute(self, context):
self.log.info('Running SQL query {}'.format(self.sql))
hook = VerticaHook(vertica_conn_id=self.vertica_conn_id)
result_set = hook.get_records(self.sql)
condition = self.test(result_set, context=context)
if condition:
self.log.info('Proceeding with downstream tasks...')
return
self.log.info('Skipping downstream tasks...')
downstream_tasks = context['task'].get_flat_relatives(upstream=False)
self.log.debug("Downstream task_ids %s", downstream_tasks)
if downstream_tasks:
self.skip(context['dag_run'], context['ti'].execution_date,
downstream_tasks)
self.log.info("Done.")
def __init__(self,
source,
target,
date_column=None,
truncate_date=False,
direct=False,
column_mapping=None,
force_introspection=False,
exclude=None,
sql=None,
*args,
**kwargs):
if column_mapping is None or force_introspection:
exclude = exclude or []
vertica_conn_id = kwargs.get('vertica_conn_id', 'vertica_default')
hook = VerticaHook(vertica_conn_id=vertica_conn_id)
columns = {
c: c
for (c,
) in hook.get_records(GET_TABLE_COLUMNS, dict(
table=target)) if c not in exclude
}
self.log.debug('Target columns discovered: {}'.format(
list(columns)))
column_mapping = column_mapping or {}
columns.update(column_mapping)
source_columns, target_columns = list(zip(*columns.items()))
_sql = sql or INSERT
params = dict(source=source,
target=target,
date_column=date_column,
truncate_date=truncate_date,
direct=direct,
source_columns=', '.join(source_columns),
target_columns=', '.join(target_columns))
super().__init__(sql=_sql, params=params, *args, **kwargs)
