I realized that SSHOperator class calls a function from SSHHook class which actually makes the ssh connection to given details.

Problem

There are lots of Dags (classes) which refers to function. How to add retry mechanism without modifying each and every class

Solution

Monkey patching is first thing which came to my mind. However, when I applied the monkey patching it created a recursion, as I had to call function again for retry. It took a while to fix this issue. However, finally I was able to fix it by creating another reference of original function is same class I am patching.

plugins/hooks/ssh_hook_with_retry.py

# -*- coding: utf-8 -*-

# Author: Gaurang Shah
#
from airflow.contrib.hooks.ssh_hook import SSHHook
import time 

max_retry = 3
retry = 0


def get_conn_with_retry(self):
    try:
		# call the original function 
        return SSHHook.original_func(self)
    except Exception as e:
        global retry
        if retry < max_retry:
            retry += 1
            print("tried %s times, failed to connect retrying again" %retry)
			time.sleep(60)
			# call the patched function which inturn calls get_conn_with_retry function 
            self.get_conn()
        else:
            raise e

plugins/hooks/__init__.py

from airflow.contrib.hooks import ssh_hook

# add new reference to original function to avoid recursion 
ssh_hook.SSHHook.original_func = ssh_hook.SSHHook.get_conn
from hooks.ssh_hook_with_retry import get_conn_with_retry

# path get_conn function with our function 
ssh_hook.SSHHook.get_conn = get_conn_with_retry

# when someone calls SSHHook.get_conn then ssh_hook_with_retry.get_conn_with_retry will be called
# and to call orignal get_conn we will have to call  SSHHook.original_func 

The post Monkey patching to decorate your library function appeared first on Big Data Processing .