[Airflow] Nifi Trigger Operator 만들기

airflow에서 nifi processor로 trigger를 보내야 했다.

사실 nifi에서 restAPI를 지원해주기 때문에 NIFI ETL 자체를 시작하는 것은 큰 문제가 안됐지만 nifi 작업 이후에 airflow에서 후속 task가 있다면 ETL이 어느 시점에 끝나는 지를 확인하는 게 어려웠다.

NIFI는 processor 단위로 데이터를 처리하며, flowfile이 processor에 인풋되면 바로 처리 후 output stream으로 내보내기 때문에 종료 시점을 명확히 알 수가 없었다.

 

아래 블로그에서 airflow에서 nifi 를 트리거할 수 있는 방법을 소개했지만 시작 노드를 `GenerateFlowFile`로 설정해야 하는 등 제약조건이 있었다.

https://towardsdatascience.com/interconnecting-airflow-with-a-nifi-etl-pipeline-8abea0667b8a

 

예를 들면 ListSFTP 로 시작하는 작업이 있다면, ListSFTP는 incoming relationship을 허용하지 않으므로 블로그에서 소개하는 방법을 사용할 수 없었다.

Input requirement:
This component does not allow an incoming relationship.

출처: 도큐먼트

 

그리고 예시에서는 Nifi cluster가 single node를 가정한 것이라 여러개의 노드일 경우 failover를 추가로 고려해야 했고, flowfile이 한 개인 경우에만 사용할 수 있었다 ㅠ

그래도 방법은 이해가 가서 조금 더 범용적으로 쓸 수 있도록 블로그 코드를 참고해서 살짝 바꿨다.

 

1. 컨셉

 

2. Airflow Connection 설정

HttpHook을 상속해서 사용할 것이기 때문에 Connection Type은 HTTP로 설정해주었다.

 

`Connection Id`: nifi1_api

`Connection Type`: HTTP

`Host`: http://x.x.x.x:8080

 

이런식으로 만들어준다.

노드가 여러개라면 여러 개의 커넥션을 만든다.

 

3. Nifi ETL 설정

필요조건

마지막 End Node로 `UpdateAttribute`를 사용해야 한다.

clusterNode의 상태변화를 확인하기 위함이다.

 

예시로 아래와 같이 간단한 작업을 만들어 보았다.

원격지에서 List -> Fetch 후 HDFS에 넣고 End Node로 `UpdateAttribute`를 사용하는 모습이다.

Start Node만 켰다 껐다 하면서 사용할거라 Start Node는 중지상태, 다른 나머지는 Running 상태로 두면 된다.

Start Node 의 스케쥴링은 timer driven으로 두되, penalty ducation은 충분히 큰 값으로 두어 running -> stopped 상태가 될 때 무수히 많은 flowfile을 생성하지 못하도록 (한 번만 실행되도록) 해주어야 한다.

 

End Node 의 properties는 아래와 같이 설정해준다.

Store State는 "Store state locally"로 설정하고,

initial value는 아무거나 두고 `last_tms`라는 properties를 추가해서 값은 현재 시간 `${now()}`으로 작성해준다.

이렇게 되면 Flowfile이 processor를 지날 때마다 ClusterNode 내에 있는 키 값인 last_tms의 value는 flowfile이 나간 시점으로 업데이트 된다.

 

2. 구현: nifiutil

nifi restAPI는 공식문서에 정리되어 있다.

프로세서 정보 가져오기

def get_processor(url_nifi_api, processor_id):
	header = {"Content-Type": "application/json"}
    
    response = requests.get(url_nifi_api + f"processors/{processor_id}", headers=header)
    return json.loads(response.content)


def get_processor_state(url_nifi_api, processor_id, token):
    header = {"Content-Type": "application/json"}
    
    response = requests.get(url_nifi_api + f"processors/{processor_id}/state", headers=header)
    return json.loads(response.content)

 

`get_processor`: 프로세서를 가져온다.

`get_processor_state`: 프로세서의 현재 상태를 가져온다.

 

프로세서 상태 정보 파싱

노드가 두 개인 nifi 의 프로세서에 대해 `get_processor_state`로 상태를 가져오면 아래와 같은 json을 반환한다. 

{
  "componentState": {
    "componentId": "...",
    "stateDescription": "...",
    "clusterState": {},
    "localState": {
      "scope": "LOCAL",
      "state": [
        {
          "key": "last_tms",
          "value": "0",
          "clusterNodeId": "...",
          "clusterNodeAddress": "nifinode1:8080"
        },
        {
          "key": "last_tms",
          "value": "0",
          "clusterNodeId": "...",
          "clusterNodeAddress": "nifinode2:8080"
        }
      ],
      "totalEntryCount": 2
    }
  }
}

1에서  stateful variable inital value를 0으로 설정했다면 last_tms 의 value를 0으로 줄 것이다.

def parse_state(json_obj, state_key):
	result = {}
    states = json_obj["componentState"]["localState"]["state"]
    for state in states:
    	if state["key"] == state_key:
        	result[state["clusterNodeId"]] = state["value"]
    return result

last_tms 와 clusterNodeId만 필요하기 때문에 위 함수와 같이 파싱하여 {clusterNodeId:value} 형태로 맵을 만든다.

결과:

{"node1의 nodeid":"0","node2의 nodeid":"0"}

 

프로세서 상태 변경

def update_processor_status(processor_id, new_state, url_nifi_api):
	processor = get_processor(url_nifi_api, processor_id)
    
    put_dict = {
    	"revision": processor["revision"],
        "state": new_state,
        "disconnectedNodeAcknowledged": True,
    }
    
    payload = json.dumps(put_dict).encode("utf8")
    header = {
    	"Content-Type": "application/json"
    }
    response = requests.put(
    	url_nifi_api + f"processors/{processor_id}/run-status",
    	headers=header,
        data=payload,
    )
    return response

 

프로세서 상태 변경 검사

def wait_for_update(url_nifi_api, processor_id, state_key, initial_state):
	processor_state = get_processor_state(url_nifi_api, processor_id, None)
    while True:
    	processor_state = get_processor_state(url_nifi_api, processor_id, None)
        value_current = parse_state(processor_state, state_key)
        
        if initial_state == value_current:
        	pause(1)
        else:
        	pause(10)
            processor_state = get_processor_state(url_nifi_api, processor_id, None)
            second_current = parse_state(processor_state, state_key)
            if second_current == value_current:
            	break
            else:
            	continue

End Node를 통해 ETL이 종료되었는지 종료되지 않았는지 판단할 때 사용할 함수이다.

1초에 한 번씩 프로세서의 state를 확인해 초기 값의 last_tms와 현재의 last_tms값에 변경이 있는지 검사한다.

만약 변경이 있다면 바로 종료하지 않고 10초간 더 기다린 후에도 상태가 동일한지 검사한다 (두번째 if 부분) .

두 번 검사하는 이유는:

예를들어 100개의 flowfile이 round robin으로 두 개의 노드에서 처리된다고 가정하면, 두 개의 노드 중 가장 먼저 첫 번째 flowfile을 처리한 노드는 local state를 업데이트 하게된다.

last_tms 변경이 감지되었지만 사실은 ETL이 끝난 것이 아니다 (98개의 flowfile이 남아있음).

이러한 오작동을 피하기 위해 일정시간 이후에도 last_tms 값의 변경이 있었는지 확인하고, 없었을 때 ETL이 종료되는 것으로 작성했다.

 

3. 구현: nifihookHook

2에서 만든 nifiutl을 import해서 사용한다.

import requests
from airflow.exceptions import AirflowException
from airflow.providers.http.hooks.http import HttpHook
from some.nifi.path.nifiutil import update_processor_status, wait_for_update, get_processor_state, parse_state

class NifihookHook(HttpHook):
	def __init__(
    	self,
        http_conn_ids: list,
        start_processor_id = '',
        end_processor_id = '',
        state_key = '',
        *args,
        **kwargs,
    ):
    	super().__init__(*args, **kwargs)
        self.http_conn_ids = http_conn_ids,
        self.conn = None,
        self.start_processor_id = start_processor_id
        self.end_processor_id = end_processor_id
        self.state_key = state_key

	def _find_stable_nodes(self) -> None:
    	for conn in self.http_conn_ids:
        	conn_url = self.get_connection(conn).host
            r = requests.post(conn_url)
            if r.status_code == 200:
            	if conn_url[-1] == '/':
                	self.nifi_api = conn_url + "nifi-api/"
                else:
                	self.nifi_api = conn_url + "/nifi-api/"
                break

	def execute(self) -> None:
    	self._find_stable_nodes()
        
        initial_state = parse_state(get_processor_state(self.nifi_api, self.end_processor_id, None), self.state_key)
        r = update_processor_status(self.start_processor_id, "RUNNING", self.nifi_api)
        pause(15)
        r = update_processor_status(self.start_processor_id, "STOPPED", self.nifi_api)
        if length(self.end_processor_id)>0:
        	wait_for_update(self.nifi_api, self.end_processor_id, self.state_key, initial_state)

 

 

`http_conn_ids` 는 list[string] 형태로 받아서 nifi 여러 노드의 정보를 받는다.

`_find_stable_nodes()`에서는 list로 받아온 nifi 노드에 요청을 보내면서 응답코드 200을 리턴해주는 노드를 찾는다.

만약 정상 응답을 받는 conn_id가 없다면 `self.nifi_api` 는 정의되지 않고 NifihookHook은 실패된다.

`execute()`에서는 아래의 순서로 실행한다.

1. End node의 현재 상태를 가져온다. 
2. Start node를 시작한다.
3. 일정 시간 이후 Start node를 중지한다.
4. 만약 end_processor_id 값이 있다면 End Node를 검사하며 끝났다고 판단될때까지 확인한다. 없다면 이부분은 스킵한다.

 

4. 구현: NifiTriggerOperator

from typing import Any, Dict, Optional
from airflow.providers.http.operators.http import SimpleHttpOperator
from airflow.utils.decorators import apply_defaults
from some.nifi.path.nifihookHook import NifihookHook

class NifiTriggerOperator(SimpleHttpOperator):
	@apply_defaults
    def __init__(
    	self,
        *,
        http_conn_ids: list,
        start_processor_id: str='',
        end_processor_id: str='',
        state_key: str='',
        **kwargs,
    ) -> None:
    	super().__init__(**kwargs)
        self.http_conn_ids = http_conn_ids
        self.start_processor_id = start_processor_id
        self.end_processor_id = end_processor_id
        self.state_key = state_key
        self.hook: Optional[NifihookHook] = None

	def execute(self, context: Dict[str, Any]) -> None:
    	self.hook = NifihookHook(
        	self.http_conn_ids,
            self.start_processor_id,
            self.end_processor_id,
            self.state_key
        )
        self.hook.execute()

그냥 인자값 받아서 NifihookHook init하고 execute 함수 실행하는 부분이다.

 

5. 사용

test = NifiTriggerOperator(
	task_id = "test",
    http_conn_ids = ["nifi1_api", "nifi2_api"],
    start_processor_id = f"{start_node_id}",
    end_processor_id = f"{end_node_id}",
    state_key = f"{my_state_key}"
)

위에서 정의해준 변수값들을 인자로 넣어주면 된다.

state_key 같은 경우는 3에서처럼 property 이름을 last_tms로 설정했다면 `state_key = "last_tms"`로 작성하면 된다.

 

 

추가 고려 필요 사항

1, ETL이 너무 오래 걸리면 `wait_for_update`에서 계속 while loop를 돌 수 있다.

납득가능한 시간을 두고 연결을 끊는 부분을 추가해야할 것 같다.

2. start node running -> stopped 상태변화의 텀을 좀 고민해봐야할 것 같다.

일례로 예전에 listsftp를 쓰는데 파일이 너무 많아서 리스팅 하는데만 1분이 걸린 적이 있었다..;

operator 사용 시에 텀을 몇초로 둘 건지 값을 받는것도 생각해봤는데 사용자가 이 오퍼레이터를 사용할 때 그런것까지 고려해야 하나.. 싶기도 하다.

3. 실패에 대한 케이스를 만들어서 task가 실패했을 때 로그를 보고 직관적으로 알 수 있게 추가되면 좋을 것 같다.