add delay for a task until particular file is moved from bucket

user_13516482 0
Viewed 1044

I am new to Airflow. I have to check for a file which is generated from DAG (eg: sample.txt) is moved from bucket(in my case the file I have generated will be moved away from the bucket when picked up by other system, and then there won't be this output file in the bucket. It might take few minutes for the file to be removed from bucket)

How to add a task in the same DAG where it waits/retires till the file is moved away from the bucket and when the sample.txt file is moved away then proceed with the next task.

Is there any operator which satisfies the above criteria? please throw some light on how to proceed

airflow google-cloud-storage python
2 Answers

You can create a custom sensor based on the current GCSObjectExistenceSensor

The modification is simple:

from airflow.providers.google.cloud.sensors.gcs import GCSObjectExistenceSensor
class GCSObjectNotExistenceSensor(GCSObjectExistenceSensor):

    def poke(self, context: dict) -> bool:
        self.log.info('Sensor checks if : %s, %s does not exist', self.bucket, self.object)
        hook = GCSHook(
            gcp_conn_id=self.google_cloud_conn_id,
            delegate_to=self.delegate_to,
            impersonation_chain=self.impersonation_chain,
        )
        return not hook.exists(self.bucket, self.object)

Then use the sensor GCSObjectNotExistenceSensor in your code like:

gcs_object_does_not_exists = GCSObjectNotExistenceSensor(
    bucket=BUCKET_1,
    object=PATH_TO__FILE,
    mode='poke',
    task_id="gcs_object_does_not_exists_task",
)

The sensor will not let the pipeline to continue until the object PATH_TO__FILE is removed.

user_146244090

You can use airflow PythonOperator to achieve the task. Make the Python callable continuously poke GCS and check if the file is removed. Return from Python function when the file from GCS is removed.

from airflow.operators.python_operator import PythonOperator
from google.cloud import storage
import google.auth

def check_file_in_gcs():
    credentials, project = google.auth.default()
    storage_client = storage.Client('your_Project_id', credentials=credentials)
    name = 'sample.txt'   
    bucket_name = 'Your_Bucket_name'
    bucket = storage_client.bucket(bucket_name)
    while True:
        stats = storage.Blob(bucket=bucket, name=name).exists(storage_client)
        if not stats:
           print("Returning as file is removed!!!!")
           return

check_gcs_file_removal = PythonOperator(
            task_id='check_gcs_file_removal',
            python_callable= check_file_in_gcs,
            #op_kwargs={'params': xyz},
            #Pass bucket name and other details if needed by commentating above 
            dag=dag
        )

you might need to install Python packages for the google cloud libraries to work. Please install one from below. (Not sure which one to install exactly.Taken from my virtualenv)

google-api-core==1.16.0
google-api-python-client==1.8.0
google-auth==1.12.0
google-auth-httplib2==0.0.3
google-auth-oauthlib==0.4.1
google-cloud-core==1.3.0
google-cloud-storage==1.27.0
user_156314820
Related