Context Manager 를 활용 하여 SFTP 서버 커넥션 관리하기

안녕하세요. 데이터 관리 업무를 맡고 있는 주니어 개발자입니다. 최근 배치 프로세싱 속에서 SFTP 서버와 커넥션을 맺고 끊는 과정을 경험 하고, 컨텍스트 매니저를 통하여 커넥션을 언매니지드 한 상태로 관리하는 것이 아닌 매니지드 상태에서 관리하는 방법을 공유 해보려고 합니다.

기존 커넥션 연결 방식

class SftpConnection:
    def __init__(self):  
        self._ip = "127.0.0.1"  
        self._user = "test"  
        self._key_file = "test.pem"  
        self.client: SSHClient = SSHClient()  
  
        key = RSAKey.from_private_key_file(self._key_file)  
        self.client.set_missing_host_key_policy(AutoAddPolicy())  
        self.client.connect(self._ip, username=self._user, pkey=key)  
        self._sftp = self.client.open_sftp()  

		# SFTP 접속 후 처리 로직
 
 		# 해당 부분이 직접적으로 커넥션 종료 하는 부분
        self._sftp.close()  
        self.client.close()

위 소스 코드는 예시로 작성 된 코드이며 보다시피 커넥션을 맺은 후 비즈니스 로직이 끝남과 동시에 클라이언트 객체와 현재 접속중인 SFTP 커넥션이랑도 종료 해주는 모습을 볼 수 있습니다.

이렇게 SFTP 객체가 한 번만 사용 된다면 굳이 저 로직을 어떻게 처리하든 상관 없겠죠? 하지만, 이와 관련된 비슷한 방식으로 여러 군데에서 커넥션을 맺고 끊고 반복한다고 생각 했을 때 일일이 관리하는 것은 휴먼에러의 위험성이 증가할 것 같습니다.

그래서 컨텍스트 매니저를 통해 이 방법을 조금 더 편리하게 관리 해볼건데요. 대표적인 컨텍스트 매니저는 file I/O의 예제에서 많이 나타납니다.

with open("file.txt", "w") as f:
	f.write("context manager file")

이렇게 파일을 작성하고 닫지 않아도 with context 내부 로직이 끝나면 파일이 닫힌다고 모두가 알고 있습니다.

네, 이 방법으로 제가 만든 SFTP객체 또한 with context 내부 로직이 끝났을 때 커넥션이 종료 되게 유도 해보겠습니다.

컨텍스트 매니저를 활용한 SFTP 커넥션 객체 소스코드

방법은 두 가지가 존재 합니다. 1. 컨텍스트 매니저 데코레이터 사용 2. 클래스 매직 메서드 오버라이딩

이 중 2번째 방법인 클래스 매직 메서드를 오버라이딩 하여 사용해보겠습니다.

from argparse import Namespace  
from paramiko import SFTPClient  
from paramiko.client import AutoAddPolicy, SSHClient  
from paramiko.rsakey import RSAKey  
from paramiko.transport import Transport   
from pathlib import Path  
  
  
  
def sftp_download_error_handler(func):  
    """  
    sftp download exception handler    """
    def wrap(*args, **kwargs):  
        try:  
            return func(*args, **kwargs)  
        except OSError as e:  
            logger.error(e, exc_info=True)  
            logger.error("Remote file does not exist !!!")  
            return False  
        except Exception as e:  
            logger.error(e, exc_info=True)  
            logger.error("Download fail : ")  
            return False  
    return wrap  
  
  
class SftpConnection:
     """  
     sftp client context manager  
     """  
     def __init__(self, sftp_schema):  
  
         self._ip = sftp_schema["ip"]  
         self._port = sftp_schema["port"]  
         self._user = sftp_schema["user"]  
         self._password = sftp_schema["password"]  
  
     def __enter__(self):  
         self.transport: Transport = Transport((self._ip, nt(self._port)))  
         self.transport.connect(username=self._user, assword=self._password)  
         self._sftp: SFTPClient = FTPClient.from_transport(self.transport)  
         return self._sftp  
  
     def __exit__(self):  
         self._sftp.close()  
         self.transport.close()  

  
class SftpDownloader:  
    def __init__(self, sftp_config, sftp_args: Namespace):  
        self._args = sftp_args  
        self._schema = sftp_config
        self.sftp_connection = SftpConnection(self._schema)

	def _get(self, finfo):  
		remote = file_info["fname"]  
        with self.sftp_connection as sftp:  
			remote = file_info["remote"]  
            local = file_info["local"]
			sftp.get(remote_file_name, local_file_name)
            
            return str(remote)  # NOTE: PosixPath타입을 문자열로 반환  
  
    @sftp_download_error_handler  
    def download(self):
        success_download_files = list(map(self._get, download_file_info))  
        return success_download_files  
  
  
def file_download(schema, args):  
    sftp_downloader = SftpDownloader(schema, args)  
    print(f'sftp downloader init: {sftp_downloader}')  
    result = sftp_downloader.download()  
    print(f"file download result: {result}")  
    return result  
  
  
if __name__ == "__main__":  
	file_download(**parameters)

커넥션을 관리하기 전 코드와 별반 다를게 없어 보이지만 위에서 내렸던 가설대로 커넥션을 맺고 끊는 곳이 다른 모듈에서 많이 호출된다면 이렇게 컨텍스트 매니저 방식을 적용 했을 때 이점이 드러나는 것 같습니다.

우선 소스코드를 살펴 보면 큰 틀로 이런 모습을 나타냅니다.

  1. SFTP 커넥션 객체
  2. 다운로드 역할 객체
  3. 에러 핸들러
  4. 호출 메서드

위와 같은 구조를 갖고 있고 맨 윗 줄 부터 내용을 읽어보고 어떤 역할을 하는지 확인 해보겠습니다.

  1. sftp_download_error_handler
  1. SftpConnection
  1. SftpDownloader
def _get(self, finfo):  
        remote = file_info["fname"]  
        with self.sftp_connection as sftp:  
            remote = file_info["remote"]  
            local = file_info["local"]
            sftp.get(remote_file_name, local_file_name)

            return str(remote)  # NOTE: PosixPath타입을 문자열로 반환  

실제 사용은 위와 같이 사용 되고 있습니다.

  1. file_download

끝맺음

이렇게 사소하고 또 사소하지만 간단한 컨텍스트 매니저를 활용하여 언매니지드한 관리보다 매니지드한 관리 방식을 여러 방면으로 응용 할 수 있는 경험을 한 것 같아서 글을 남기게 되었습니다.

아쉬운 점이 있다면 직접 두 눈으로 커넥션을 무자비하게 맺고 끊고 반복 했을 때 관리 하는 것과 하지 않는 두 코드에서 어느 시점을 기준으로 성능 저하가 일어나는지에 대한 분석이나 이런 부분들이 있었다면 좋았을 것 같다고 생각이 드네요.

이상입니다 감사합니다.