def test_variable_setdefault_existing_json(self):
key = "tested_var_setdefault_2_id"
value = {"city": 'Paris', "Happiness": True}
Variable.set(key, value, serialize_json=True)
val = Variable.setdefault(key, value, deserialize_json=True)
# Check the returned value, and the stored value are handled correctly.
self.assertEqual(value, val)
self.assertEqual(value, Variable.get(key, deserialize_json=True))
def test_variable_setdefault_round_trip(self):
key = "tested_var_setdefault_1_id"
value = "Monday morning breakfast in Paris"
Variable.setdefault(key, value)
self.assertEqual(value, Variable.get(key))
def test_variable_setdefault_round_trip_json(self):
key = "tested_var_setdefault_2_id"
value = {"city": 'Paris', "Happiness": True}
Variable.setdefault(key, value, deserialize_json=True)
self.assertEqual(value, Variable.get(key, deserialize_json=True))
print("~" * 100)
# default_args = {
# "owner": "saumalya",
# "depends_on_past": False,
# "start_date": datetime(2020, 5, 8),
# "email": ["*****@*****.**"],
# "email_on_failure": False,
# "email_on_retry": False,
# "retries": 1,
# "retry_delay": timedelta(minutes=1)
# }
try:
config = Variable.setdefault(CONFIG_DB_KEY,
json.loads(json.dumps(dag_config,
default=str)),
deserialize_json=True)
except Exception as e:
print(
"Something went wrong while trying to set default variables during run time: "
+ str(e))
print("~" * 100)
traceback.print_exc(file=sys.stdout)
print("~" * 100)
with DAG(**dag_config) as dag:
# Declare pipeline start and end task
start_task = DummyOperator(task_id='pipeline_start')
end_task = DummyOperator(task_id='pipeline_end')
for account_details in pipeline_config['task_details']['accounts']:
'''
'''
---------------------------------------
READ THE CONFIG & LOAD THE VARIABLE
---------------------------------------
'''
CONFIG_FILE = '/usr/local/airflow/dags/config/config.yml'
CONFIG_DB_KEY = 'anuvaad_config'
with open(CONFIG_FILE) as file:
default_config = yaml.load(file, Loader=yaml.FullLoader)
default_config = json.dumps(default_config, default=str)
config = Variable.setdefault(CONFIG_DB_KEY,
default_config,
deserialize_json=True)
config = yaml.load(CONFIG_FILE, Loader=yaml.FullLoader)
'''
---------------------------------------
DEFAULT ARGS
---------------------------------------
'''
default_args = {
'owner': 'aravinth',
'depends_on_past': False,
'start_date': airflow.utils.dates.days_ago(0),
'email': ['*****@*****.**'],
'email_on_failure': True,
'email_on_retry': False,
'retries': 0,