Skip to content

Commit 4bf9478

Browse files
ckljohnashb
authored andcommitted
[AIRFLOW-3112] Make SFTP hook to inherit SSH hook (#3945)
This is to aline the arguments of SFTP hook with SSH hook
1 parent dbb6659 commit 4bf9478

File tree

1 file changed

+53
-16
lines changed

1 file changed

+53
-16
lines changed

airflow/contrib/hooks/sftp_hook.py

+53-16
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,14 @@
2121
import pysftp
2222
import logging
2323
import datetime
24-
from airflow.hooks.base_hook import BaseHook
24+
from airflow.contrib.hooks.ssh_hook import SSHHook
2525

2626

27-
class SFTPHook(BaseHook):
27+
class SFTPHook(SSHHook):
2828
"""
29+
This hook is inherited from SSH hook. Please refer to SSH hook for the input
30+
arguments.
31+
2932
Interact with SFTP. Aims to be interchangeable with FTPHook.
3033
3134
Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and
@@ -39,32 +42,66 @@ class SFTPHook(BaseHook):
3942
Errors that may occur throughout but should be handled downstream.
4043
"""
4144

42-
def __init__(self, ftp_conn_id='sftp_default'):
43-
self.ftp_conn_id = ftp_conn_id
45+
def __init__(self, ftp_conn_id='sftp_default', *args, **kwargs):
46+
kwargs['ssh_conn_id'] = ftp_conn_id
47+
super(SFTPHook, self).__init__(*args, **kwargs)
48+
4449
self.conn = None
50+
self.private_key_pass = None
51+
52+
if self.ssh_conn_id is not None:
53+
conn = self.get_connection(self.ssh_conn_id)
54+
if conn.extra is not None:
55+
extra_options = conn.extra_dejson
56+
if 'private_key_pass' in extra_options:
57+
self.private_key_pass = extra_options.get('private_key_pass', None)
58+
59+
# For backward compatibility
60+
# TODO: remove in Airflow 2.1
61+
import warnings
62+
if 'ignore_hostkey_verification' in extra_options \
63+
and str(extra_options["ignore_hostkey_verification"])\
64+
.lower() == 'false':
65+
warnings.warn(
66+
'Extra option `ignore_hostkey_verification` is deprecated.'
67+
'Please use `no_host_key_check` instead.'
68+
'This option will be removed in Airflow 2.1',
69+
DeprecationWarning,
70+
stacklevel=2,
71+
)
72+
self.no_host_key_check = False
73+
if 'private_key' in extra_options:
74+
warnings.warn(
75+
'Extra option `private_key` is deprecated.'
76+
'Please use `key_file` instead.'
77+
'This option will be removed in Airflow 2.1',
78+
DeprecationWarning,
79+
stacklevel=2,
80+
)
81+
self.key_file = extra_options.get('private_key')
4582

4683
def get_conn(self):
4784
"""
4885
Returns an SFTP connection object
4986
"""
5087
if self.conn is None:
51-
params = self.get_connection(self.ftp_conn_id)
5288
cnopts = pysftp.CnOpts()
53-
if ('ignore_hostkey_verification' in params.extra_dejson and
54-
params.extra_dejson['ignore_hostkey_verification']):
89+
if self.no_host_key_check:
5590
cnopts.hostkeys = None
91+
cnopts.compression = self.compress
5692
conn_params = {
57-
'host': params.host,
58-
'port': params.port,
59-
'username': params.login,
93+
'host': self.remote_host,
94+
'port': self.port,
95+
'username': self.username,
6096
'cnopts': cnopts
6197
}
62-
if params.password is not None:
63-
conn_params['password'] = params.password
64-
if 'private_key' in params.extra_dejson:
65-
conn_params['private_key'] = params.extra_dejson['private_key']
66-
if 'private_key_pass' in params.extra_dejson:
67-
conn_params['private_key_pass'] = params.extra_dejson['private_key_pass']
98+
if self.password and self.password.strip():
99+
conn_params['password'] = self.password
100+
if self.key_file:
101+
conn_params['private_key'] = self.key_file
102+
if self.private_key_pass:
103+
conn_params['private_key_pass'] = self.private_key_pass
104+
68105
self.conn = pysftp.Connection(**conn_params)
69106
return self.conn
70107

0 commit comments

Comments
 (0)