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)