27

I need to update a variable I have made in Airflow programmatically but I can not find the answer on how to do that with code.

I have retrieved my variable with this code:

column_number = Variable.get('column_number')

At the end of the function, I would like to increment the column_number by one

I have tried this: Variable.set_val("column_number", int(column_number) + 1)

And it does not work.

Here is the full code for reference:

import airflow
from datetime import datetime, timedelta
from random import randint
from airflow import DAG
from airflow.hooks.postgres_hook import PostgresHook
from airflow.models import Variable
from airflow.operators.python_operator import PythonOperator

args = {
    'owner': 'besteman',
    'start_date': datetime.utcnow(),
    'retries': 1,
    'retry_delay': timedelta(minutes=30)
}

dag = DAG(dag_id='test-postgres', default_args=args, schedule_interval='@hourly')

def add_columns_and_values():

    column_number = Variable.get('column_number')

    pg_hook = PostgresHook(postgres_conn_id='airflow-test')

    add_columns = f'ALTER TABLE students ADD COLUMN test{column_number} smallint;'

    pg_hook.run(add_columns) 

    for i in range(8):
        add_values = f"UPDATE students SET test{column_number} = '{randint(50, 100)}' WHERE id = {i+1};"
        pg_hook.run(add_values)

    Variable.set_val("column_number", int(column_number) + 1)


t1 = PythonOperator(task_id='add_columns_values',
    python_callable=add_columns_and_values,
    dag=dag)
Justin Besteman
  • 378
  • 1
  • 6
  • 17

2 Answers2

40

Use Variable.set instead of Variable.set_val. set_val() is a setter for the val attribute and not intended for outside use. This should do what you want:

Variable.set("column_number", int(column_number) + 1)

It will make the actual update to the database, along with handling session and serialization for you if needed.

Reference: https://github.com/apache/incubator-airflow/blob/1.10.1/airflow/models.py#L4558-L4569

Daniel Huang
  • 6,238
  • 34
  • 33
6

In Airflow 2.0 you should use Variable.update:

Variable.update(key=your_key, value=new_value)

Otherwise you will get a Unique Constraint violation exception stating that the key already exists.

Joaquín L. Robles
  • 6,261
  • 10
  • 66
  • 96
  • 1
    This maybe is out of date. Check here on the source code on 2.5.0: https://airflow.apache.org/docs/apache-airflow/stable/_modules/airflow/models/variable.html#Variable.set. The new implementation overwrites the variable. In fact, `Variable.update` will raise KeyError if the variable does not exist yet. – bvan Jan 10 '23 at 08:53