Airflow: How to SSH and run BashOperator from a di

Question

Is there a way to ssh to different server and run BashOperator using Airbnb's Airflow? I am trying to run a hive sql command with Airflow but I need to SSH to a different box in order to run the hive shell. My tasks should look like this:

1. SSH to server1
2. start Hive shell
3. run Hive command

Thanks!

Answer 1

I think that I just figured it out:

1. Create a SSH connection in UI under Admin > Connection. Note: the connection will be deleted if you reset the database

2. In the Python file add the following

   from airflow.contrib.hooks import SSHHook
   sshHook = SSHHook(conn_id=<YOUR CONNECTION ID FROM THE UI>)
   

3. Add the SSH operator task

   t1 = SSHExecuteOperator(
       task_id="task1",
       bash_command=<YOUR COMMAND>,
       ssh_hook=sshHook,
       dag=dag)
   

Thanks!

Answer 2

One thing to note with Anton's answer is that the argument is actually ssh_conn_id, not conn_id for the SSHOperator object. At least in version 1.10.

A quick example would look like

from datetime import timedelta, datetime
import airflow
from airflow import DAG
from airflow.contrib.operators.ssh_operator import SSHOperator
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email': ['[e-mail]'],
    'email_on_failure': False,
    'email_on_retry': False,
    'start_date': datetime.now() - timedelta(minutes=20),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}
dag = DAG(dag_id='testing_stuff',
          default_args=default_args,
          schedule_interval='0,10,20,30,40,50 * * * *',
          dagrun_timeout=timedelta(seconds=120))
# Step 1 - Dump data from postgres databases
t1_bash = """
echo 'Hello World'
"""
t1 = SSHOperator(
    ssh_conn_id='ssh_default',
    task_id='test_ssh_operator',
    command=t1_bash,
    dag=dag)