Reading sql file with jinja templates in custom operator in Airflow

Question

I am trying to read sql file that contains query with jinja templates in the custom operator in Airflow. I have already achieved it using PythonOperator that calls function where I used

def execute_query(**kwargs)
    sql_query = open('my_sql_query.sql').read()   #(SELECT * FROM my_table WHERE date > {})
    sql_query.format(kwargs['ds'])

but I would prefer use this syntax {{ ds }} directly in the query like SELECT * FROM my_table WHERE date > {{ ds }}

What I have done:

I created CustomOperator with template_fields and template_ext

class SQLOperator(BaseOperator):

   template_fields = ['sql']
   template_ext = ('.sql',)

   @apply_defaults
   def __init__(
       self,
       name = None,
       sql = None,
       *args, **kwargs) -> None:
       super().__init__(*args, **kwargs)
       self.name = name
       self.sql = sql

    def execute(self, context):
       print("Name", name) # <- works
       print("Query", sql) # <- doesn't work and I don't know how to get the sql file content

Dag

default_args = {...}

dag = DAG(
    'sql_operator_test',
     schedule_interval='0 0 * * *',
     template_searchpath=['/Users/username/airflow/dags/sql/test/'],
     default_args=default_args)

sql_task = SQLOperator(
    task_id='sql_process',
    name="Aaa",
    sql="/Users/username/airflow/dags/sql/test.sql",
    dag=dag)

SQL query

SELECT * FROM my_table WHERE date > {{ ds }}

I am running out of ideas. Is there any option to pass the file to the operator or get its rendered content?

user_14624409 · Answer

Your approach is OK. I took your code and created a working example showing {{ ds }} is being templated as required:

Create a .py file as:

from datetime import datetime, timedelta
from airflow import DAG
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults


class SQLOperator(BaseOperator):

   template_fields = ['sql']
   template_ext = ('.sql',)

   @apply_defaults
   def __init__(
           self,
           name = None,
           sql = None,
           *args,
           **kwargs
   ) -> None:

       super().__init__(**kwargs)
       self.name = name
       self.sql = sql

   def execute(self, context):
       print("Name", self.name)  # <- works
       print("Query", self.sql)  # <- Also works :)

default_args = {
    'owner': 'a',
    'start_date': datetime(2020, 3, 24, 2, 0, 0),
}


dag = DAG(
    'sql_operator_test',
     schedule_interval=None,
     default_args=default_args)

sql_task = SQLOperator(
    task_id='sql_process',
    name="Aaa",
    sql="test.sql",
    dag=dag)

a .sql file as:

SELECT * FROM my_table WHERE date > {{ ds }}

Running it gives:

And from the task log:

Reading sql file with jinja templates in custom operator in Airflow

1 Answers