Schedule the shell script in different servers using same DAG file

Viewed 312

I am completely new to Apache Airflow. I have a situation. The code I have used is

import airflow
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta
from airflow.contrib.operators.ssh_operator import SSHOperator
from airflow.contrib.hooks.ssh_hook import SSHHook

default_args = {
                 'owner': 'john',
                 'depends_on_past': False,
                 'email': [''],
                 'email_on_failure': False,
                 'email_on_retry': False,
                 'retries': 0,
                 'retry_delay': timedelta(minutes=5)
               }
  
 dag = DAG(
            'tutorial',
            default_args = default_args,
            description='A simple tutorial DAG',
            schedule_interval=None)

bash_tutorial = """
  echo "Execute shell file: /A/B/server/tutorial.ksh"
  echo "{{macros.ds_format(ds, "%Y-%m-%d", "%m-%d-%Y"}}"
  source /home/johnbs/.profile
  /A/B/server/tutorial.ksh
 """

t1 = SSHOperator(
                ssh_conn_id='dev'
                task_id='tutorial.ksh'
                command=bash_tutorial,
                dag=dag
        )

Using airflow, I want to trigger a ksh script in different servers like dev and test servers i.e.

tutorial.ksh is present in dev server(conn_id is 'dev') of the path (/A/B/C/tutorial.ksh) and in test server(conn_id is 'test') of the path (/A/B/D/tutorial.ksh)...Here you can see C folder from dev and D folder from test...Which area should I update the code?

1 Answers

Each instance of SSHOperator execute a command on a single server. You will need to define each connection separately as explained in the docs and then you can do:

server_connection_ids = ['dev', 'test']
start_op = DummyOperator(task_id="start_task", dag=dag)
for conn in server_connection_ids:
    bash_tutorial = f"""
      echo "Execute shell file: /A/B/server/tutorial.ksh"
      echo "{{macros.ds_format(ds, "%Y-%m-%d", "%m-%d-%Y"}}"
      source /home/johnbs/.profile
      /A/B/{conn}/tutorial.ksh
     """
    ssh_op = SSHOperator(
        ssh_conn_id=f'{conn}',
        task_id=f'ssh_{conn}_task',
        command=bash_tutorial,
        dag=dag
    )

This will create a task per server.

Related