21
21
import pysftp
22
22
import logging
23
23
import datetime
24
- from airflow .hooks .base_hook import BaseHook
24
+ from airflow .contrib . hooks .ssh_hook import SSHHook
25
25
26
26
27
- class SFTPHook (BaseHook ):
27
+ class SFTPHook (SSHHook ):
28
28
"""
29
+ This hook is inherited from SSH hook. Please refer to SSH hook for the input
30
+ arguments.
31
+
29
32
Interact with SFTP. Aims to be interchangeable with FTPHook.
30
33
31
34
Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and
@@ -39,32 +42,66 @@ class SFTPHook(BaseHook):
39
42
Errors that may occur throughout but should be handled downstream.
40
43
"""
41
44
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
+
44
49
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' )
45
82
46
83
def get_conn (self ):
47
84
"""
48
85
Returns an SFTP connection object
49
86
"""
50
87
if self .conn is None :
51
- params = self .get_connection (self .ftp_conn_id )
52
88
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 :
55
90
cnopts .hostkeys = None
91
+ cnopts .compression = self .compress
56
92
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 ,
60
96
'cnopts' : cnopts
61
97
}
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
+
68
105
self .conn = pysftp .Connection (** conn_params )
69
106
return self .conn
70
107
0 commit comments