基于URI是否含key_file字段修改retrieve_ssh_key_var函数

需求说明 #

当SSH/SFTP连接URI包含key_file参数（示例：ssh://user@host.com?key_file=%2Ftmp%2-user-ssh-key）时，提取并返回key文件的basename；当URI不含key_file参数（示例：sftp://user:password@host.com:22）时，函数不执行任何操作。

原函数代码 #

def retrieve_ssh_key_var(self):
    key_file = SSHHook(ssh_conn_id=self.sftp_conn_id).key_file
    return basename(key_file)

用户尝试的无效代码 #

def retrieve_ssh_key_var(self):
    if 'key_file' in self.sftp_conn_id
        try:
            key_file = SSHHook(ssh_conn_id=self.sftp_conn_id).key_file
            return basename(key_file)
        except Exception:
            pass

无效代码问题分析 #

1. 语法错误：if语句末尾缺少冒号:
2. 判断逻辑错误：self.sftp_conn_id是连接ID而非原始URI，直接用字符串包含判断无法正确识别是否存在key_file参数
3. 未处理key_file为空的情况，且无默认返回值

正确实现代码 #

from os.path import basename
from airflow.providers.ssh.hooks.ssh import SSHHook

def retrieve_ssh_key_var(self):
    try:
        ssh_hook = SSHHook(ssh_conn_id=self.sftp_conn_id)
        # 检查key_file是否存在且非空
        if ssh_hook.key_file:
            return basename(ssh_hook.key_file)
    except Exception:
        # 捕获连接初始化或key文件处理异常，不执行任何操作
        pass
    # 无key_file参数或异常时返回None（或根据需求调整）
    return None

代码说明 #

• 先初始化SSHHook实例，直接通过实例的key_file属性判断是否存在密钥文件参数
• 仅当key_file非空时，返回其basename
• 用try-except捕获可能的异常（如连接ID无效、key文件路径错误等），异常时直接跳过操作
• 无有效key_file时返回None，符合“不执行任何操作”的需求

内容的提问来源于stack exchange，提问作者NavySeal2026