How to trigger operator inside Python function using Airflow?

user_10163853 0
Viewed 3452

I have the following code:

def chunck_import(**kwargs):
    ...
    for i in range(1, num_pages + 1):
        start = lower +  chunks * i
        end = start + chunks
        if i>1:
             start = start + 1
        logging.info(start, end)
        if end > max_current:
            end = max_current
        where = 'where orders_id between {0} and {1}'.format(start,end)
        logging.info(where)
        import_orders_products_op = MySqlToGoogleCloudStorageOperator(
            task_id='import_orders_and_upload_to_storage_orders_products_{}'.format(i),
            mysql_conn_id='mysql_con',
            google_cloud_storage_conn_id='gcp_con',
            provide_context=True,
            approx_max_file_size_bytes = 100000000, #100MB per file
            sql = 'import_orders.sql',
            params={'WHERE': where},
            bucket=GCS_BUCKET_ID,
            filename=file_name_orders_products,
            dag=dag)

    start_task_op = DummyOperator(task_id='start_task', dag=dag)

    chunck_import_op = PythonOperator(
        task_id='chunck_import',
        provide_context=True,
        python_callable=chunck_import,
        dag=dag)

    start_task_op >>  chunck_import_op

This code uses PythonOperator to calculate how many runs I need from the MySqlToGoogleCloudStorageOperator and create the WHERE cluster of the SQL then it needs to execute it.

The problem is that the MySqlToGoogleCloudStorageOperator isn't being executed.

I can't actually do 

chunck_import_op >> import_orders_products_op

How can I make the MySqlToGoogleCloudStorageOperator be executed inside the PythonOperator?

airflow
1 Answers

I think at the end of your for loop, you'll want to call import_orders_products_op.execute(context=kwargs) possibly preceded by import_orders_products_op.pre_execute(context=kwargs). This is a bit complicated in that it skips the render_templates() call of the task_instance, and actually if you instead made a task_instance to put each of these tasks in, you could call run or _raw_run_task instead but these both require information from the dagrun (which you can get in the python callable's context like kwargs['dag_run'])

Looking at what you've passed to the operators it looks like as is you'll need the templating step to load the import_orders.sql file and fill in the WHERE parameter. Alternatively it's okay within the callable itself to load the file into a string, replace the {{ params.WHERE }} part (and any others) manually without Jinja2 (or you could spend time to figure out the right jinja2 calls), and then set the import_orders_products_op.sql=the_string_you_loaded before calling import_orders_products_op.pre_execute(context=kwargs) and import_orders_products_op.execute(context=kwargs).

dlamblin_4590
Related