def test_parse_bucket_key(self, key, bucket, parsed_key, parsed_bucket): op = S3KeySensor( task_id='s3_key_sensor', bucket_key=key, bucket_name=bucket, ) self.assertEqual(op.bucket_key, parsed_key) self.assertEqual(op.bucket_name, parsed_bucket)
def test_bucket_name_none_and_bucket_key_as_relative_path(self): """ Test if exception is raised when bucket_name is None and bucket_key is provided as relative path rather than s3:// url. :return: """ with self.assertRaises(AirflowException): S3KeySensor(task_id='s3_key_sensor', bucket_key="file_in_bucket")
def test_bucket_name_provided_and_bucket_key_is_s3_url(self): """ Test if exception is raised when bucket_name is provided while bucket_key is provided as a full s3:// url. :return: """ with self.assertRaises(AirflowException): S3KeySensor(task_id='s3_key_sensor', bucket_key="s3://test_bucket/file", bucket_name='test_bucket')
def test_poke(self, mock_hook): op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file') mock_check_for_key = mock_hook.return_value.check_for_key mock_check_for_key.return_value = False self.assertFalse(op.poke(None)) mock_check_for_key.assert_called_once_with(op.bucket_key, op.bucket_name) mock_hook.return_value.check_for_key.return_value = True self.assertTrue(op.poke(None))
def test_poke_wildcard(self, mock_hook): op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file', wildcard_match=True) mock_check_for_wildcard_key = mock_hook.return_value.check_for_wildcard_key mock_check_for_wildcard_key.return_value = False assert not op.poke(None) mock_check_for_wildcard_key.assert_called_once_with( op.bucket_key, op.bucket_name) mock_check_for_wildcard_key.return_value = True assert op.poke(None)
def test_parse_bucket_key(self, key, bucket, parsed_key, parsed_bucket, mock_hook): mock_hook.return_value.check_for_key.return_value = False op = S3KeySensor( task_id='s3_key_sensor', bucket_key=key, bucket_name=bucket, ) op.poke(None) self.assertEqual(op.bucket_key, parsed_key) self.assertEqual(op.bucket_name, parsed_bucket)
def Sensor_0(config): sensor = S3KeySensor( bucket_key= "s3://abinitio-spark-redshift-testing/prophecy-libs-assembly-2.jar", wildcard_match=False, aws_conn_id="aws_default_pankaj", verify=False, soft_fail=False, poke_interval=60, timeout=604800, mode="poke", task_id="Sensor_0", exponential_backoff=False) return sensor
def check_minio_2(): @task() def log(): print("Something happened in bucket") sensor = S3KeySensor( task_id="check_s3_for_file_in_s3", bucket_key="s3://my-first-bucket/my-test-file-2.txt-*", wildcard_match=True, aws_conn_id="local_minio", ) aaa = log() sensor >> aaa
def Wait_for_Data(config): sensor = S3KeySensor( bucket_key = "s3://abinitio-spark-redshift-testing/gates/2021-05-03", wildcard_match = False, aws_conn_id = "aws_default", verify = False, soft_fail = False, poke_interval = 60, timeout = 604800, mode = "poke", task_id = "Wait_for_Data", exponential_backoff = False ) return sensor
def test_parse_bucket_key_from_jinja(self, mock_hook): mock_hook.return_value.check_for_key.return_value = False Variable.set("test_bucket_key", "s3://bucket/key") execution_date = datetime(2020, 1, 1) dag = DAG("test_s3_key", start_date=execution_date) op = S3KeySensor( task_id='s3_key_sensor', bucket_key='{{ var.value.test_bucket_key }}', bucket_name=None, dag=dag, ) ti = TaskInstance(task=op, execution_date=execution_date) context = ti.get_template_context() ti.render_templates(context) op.poke(None) self.assertEqual(op.bucket_key, "key") self.assertEqual(op.bucket_name, "bucket")
"Value": "Advertiser MR PID Airflow Project" }, ], } cluster_creator = EmrCreateJobFlowOperator( task_id="create_emr_cluster", aws_conn_id="aws_default", emr_conn_id="emr_default", job_flow_overrides=JOB_FLOW_OVERRIDES, dag=dag, ) sensor_stage1_key = S3KeySensor( task_id="s3_sensor_stage1_key", bucket_name="{{ dag_run.conf['metaBucketName'] }}", bucket_key="{{ dag_run.conf['instanceId'] }}/step_1_meta_enc_kc/_SUCCESS", ) SPARK_STEP_1 = [{ "Name": "adv-mr-pid-stage1", "ActionOnFailure": "TERMINATE_JOB_FLOW", "HadoopJarStep": { "Jar": "command-runner.jar", "Args": [ "spark-submit", "--deploy-mode", "cluster", "--master", "yarn",
from airflow.providers.amazon.aws.sensors.s3_key import S3KeySensor default_args = { "owner": "iotoi", "start_date": days_ago(1), } dag = DAG( "check_minio_1", default_args=default_args, schedule_interval="@once", tags=["iotoi-samples"], ) t1 = BashOperator( task_id="bash_test", bash_command='echo "hello, it should work" > s3_conn_test.txt', dag=dag, ) sensor = S3KeySensor( task_id="check_s3_for_file_in_s3", bucket_key="*", bucket_name="my-first-bucket", wildcard_match=True, aws_conn_id="local_minio", poke_interval=10, dag=dag, ) t1.set_upstream(sensor)
dag=dag, ) stage1_checker = EmrStepSensor( task_id="watch_stage1", job_flow_id= "{{ task_instance.xcom_pull(task_ids='create_emr_cluster', key='return_value') }}", step_id= "{{ task_instance.xcom_pull(task_ids='add_stage_1', key='return_value')[0] }}", aws_conn_id="aws_default", dag=dag, ) sensor_stage2_key = S3KeySensor( task_id="s3_sensor_stage2_key", bucket_name="{{ dag_run.conf['advBucketName'] }}", bucket_key= "{{ dag_run.conf['instanceId'] }}/step_1_meta_enc_kc_kp/_SUCCESS", ) SPARK_STEP_2 = [{ "Name": "meta-mr-pid-stage2", "ActionOnFailure": "TERMINATE_JOB_FLOW", "HadoopJarStep": { "Jar": "command-runner.jar", "Args": [ "spark-submit", "--deploy-mode", "cluster", "--master", "yarn",