contextmanagerを使ったSFTPのconnection管理 (Paramiko)

はじめに

Zennとダブルポストです。

contextmanagerの実践例 の続きです。今回はSFTPが題材です。SSH/SFTPのクライアントライブラリとしてよく使われるParamikoを使います。

Paramikoでの実践例

パスフレーズ付きの秘密鍵を使った、公開鍵認証での例です。

from paramiko import Transport
from paramiko import RSAKey
from paramiko import SFTPClient

class SFTPService:

   def __init__(self, host, port, user, private_key, pass_phrase):
       self.host = host
       self.port = port
       self.private_key = private_key
       self.pass_phrase = pass_phrase       

   @contextlib.contextmanager
   def connect(self):
       try:
         transport = Transport((self.host, self.port))
         pkey = RSAKey.from_private_key_file(self.private_key, password=self.pass_phrase)
         transport.connect(self.user, pkey=pkey)
         self.client = SFTPClient.from_transport(transport)  
         yield self
       finally:
         self.client.close()
         
         

   def put(local_file: Path, remote_file: Path):
       self.client.put(local_file.as_posix(), remote_file.as_posix()) 
         

複数回のsftpコマンドを実行している間は、ひとつのコネクションでやることを想定して、この形に落ち着きました。
このクラスのインスタンス生成を、with構文を使わないとまともにできないような形も検討したのですが、テストコードが書きづらくなるので、コネクションをつくるメソッド + contextmanager という形にしました。

SFTPClient#closeは、インスタンス生成時に渡されたTransportのcloseを呼び出しています。
なので、clientのcloseだけ呼び出せば十分です。

    def __init__(self, sock):
        """
        Create an SFTP client from an existing `.Channel`.  The channel
        should already have requested the ``"sftp"`` subsystem.
        An alternate way to create an SFTP client context is by using
        `from_transport`.
        :param .Channel sock: an open `.Channel` using the ``"sftp"`` subsystem
        :raises:
            `.SSHException` -- if there's an exception while negotiating sftp
        """
        BaseSFTP.__init__(self)
        self.sock = sock
        ...


    def close(self):
        """
        Close the SFTP session and its underlying channel.
        .. versionadded:: 1.4
        """
        self._log(INFO, "sftp session closed.")
        self.sock.close()

使い方としては、このようになります。　

sftp_service = SFTPService(...)
src_dir = Path('tmp', 'src')
with sftp_service.connect() as sftp:
     for file in src_dir.iterdir():
         remote_file = Path('dst').joinpath(file.name)
         sftp.put(file, remote_file) 

これで、コネクション周りのコードを、ビジネスロジックを書くところから隔離しやすくなります。

参考