def test_sql_sensor_presto(self):
op = SqlSensor(
task_id='hdfs_sensor_check',
conn_id='presto_default',
sql="SELECT 'x' FROM airflow.static_babynames LIMIT 1;",
dag=self.dag,
)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_unsupported_conn_type(self):
op = SqlSensor(
task_id='sql_sensor_check',
conn_id='redis_default',
sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
dag=self.dag,
)
with pytest.raises(AirflowException):
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_sql_sensor_postgres_poke(self, mock_hook):
op = SqlSensor(
task_id='sql_sensor_check',
conn_id='postgres_default',
sql="SELECT 1",
)
mock_hook.get_connection('postgres_default').conn_type = "postgres"
mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records
mock_get_records.return_value = []
assert not op.poke(None)
mock_get_records.return_value = [[None]]
assert not op.poke(None)
mock_get_records.return_value = [['None']]
assert op.poke(None)
mock_get_records.return_value = [[0.0]]
assert not op.poke(None)
mock_get_records.return_value = [[0]]
assert not op.poke(None)
mock_get_records.return_value = [['0']]
assert op.poke(None)
mock_get_records.return_value = [['1']]
assert op.poke(None)
def test_sql_sensor_postgres_poke_fail_on_empty(self, mock_hook):
op = SqlSensor(
task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", fail_on_empty=True
)
mock_hook.get_connection('postgres_default').conn_type = "postgres"
mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records
mock_get_records.return_value = []
with pytest.raises(AirflowException):
op.poke(None)
def test_sql_sensor_postgres_poke_failure(self, mock_hook):
op = SqlSensor(
task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", failure=lambda x: x in [1]
)
mock_hook.get_connection('postgres_default').conn_type = "postgres"
mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records
mock_get_records.return_value = []
assert not op.poke(None)
mock_get_records.return_value = [[1]]
with pytest.raises(AirflowException):
op.poke(None)
def test_sql_sensor_postgres_poke_invalid_success(self, mock_hook):
op = SqlSensor(
task_id='sql_sensor_check',
conn_id='postgres_default',
sql="SELECT 1",
success=[1],
)
mock_hook.get_connection('postgres_default').conn_type = "postgres"
mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records
mock_get_records.return_value = [[1]]
with pytest.raises(AirflowException) as ctx:
op.poke(None)
assert "self.success is present, but not callable -> [1]" == str(ctx.value)
def test_sql_sensor_postgres_poke_success(self, mock_hook):
op = SqlSensor(
task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", success=lambda x: x in [1]
)
mock_hook.get_connection('postgres_default').conn_type = "postgres"
mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records
mock_get_records.return_value = []
assert not op.poke(None)
mock_get_records.return_value = [[1]]
assert op.poke(None)
mock_get_records.return_value = [['1']]
assert not op.poke(None)
def test_sql_sensor_postgres(self):
op1 = SqlSensor(
task_id='sql_sensor_check_1',
conn_id='postgres_default',
sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
dag=self.dag,
)
op1.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
op2 = SqlSensor(
task_id='sql_sensor_check_2',
conn_id='postgres_default',
sql="SELECT count(%s) FROM INFORMATION_SCHEMA.TABLES",
parameters=["table_name"],
dag=self.dag,
)
op2.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
dag = DAG('run_daily_census',
default_args=default_args,
catchup=False,
schedule_interval='35 5 * * *')
# deps = ExternalTaskSensor(
# external_dag_id='get_etl_deps',
# external_task_id='query_and_save_deps',
# task_id='wait_for_dependencies_file',
# dag=dag
# )
check_max_census_date = SqlSensor(
task_id='check_max_census_date',
conn_id='ebi_datamart',
sql=
'select CAST(MAX(EFFECTIVE_TIME) as DATE) from CLARITY_CLARITY_ADT where EVENT_TYPE_C = 6',
success=check_date,
dag=dag,
)
RDC = PythonOperator(
task_id='refresh_daily_census',
python_callable=refresh_tableau_extract,
op_kwargs={'datasource_id': '86a56fc4-d842-40ad-8898-2249e302c88d'},
dag=dag)
# deps >>
check_max_census_date >> RDC
call_date_sql = """
select MAX(CAST(rqi.[DateTime] as date)) as max_cisco_calldate
from ccehds_t_router_queue_interval rqi with (nolock)
where CAST(rqi.[DateTime] as date) < CAST(GETDATE() as date)
and rqi.precisionqueueid != 5082
having SUM(ISNULL(rqi.callsoffered, 0)) > 0
"""
check_max_call_date = SqlSensor(
task_id='check_max_call_date',
conn_id='ebi_datamart',
sql=call_date_sql,
success=check_date,
dag=dag,
)
ACC = PythonOperator(
task_id='refresh_avaya_call_center',
python_callable=refresh_tableau_extract,
op_kwargs={'datasource_id': '6EE81264-DFE1-4318-8397-9EC853DF7085'},
dag=dag)
ACCA = PythonOperator(
task_id='refresh_avaya_call_center_agents',
python_callable=refresh_tableau_extract,
op_kwargs={'datasource_id': '082E7473-5803-4681-A847-3ADAB4B1E8E7'},
dag=dag)