From 150d09d3fe89bc85949bf22c30f4fffe672ee458 Mon Sep 17 00:00:00 2001
From: biellls <bielllobera@gmail.com>
Date: Wed, 17 Aug 2016 22:48:47 +0200
Subject: [PATCH 01/18] Added SFTPHook

---
 airflow/contrib/hooks/sftp_hook.py | 141 +++++++++++++++++++++++++++++
 1 file changed, 141 insertions(+)
 create mode 100644 airflow/contrib/hooks/sftp_hook.py

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
new file mode 100644
index 0000000000000..b2cb34156c9e3
--- /dev/null
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pysftp
+import logging
+import datetime
+from airflow.hooks.base_hook import BaseHook
+
+class SFTPHook(BaseHook):
+    """
+    Interact with SFTP. Aims to be interchangeable with FTPHook.
+
+    Pitfalls: - In contrast with FTPHook describe_directory only returns size and modify. It doesn't return unix.owner,
+                unix.mode, perm, unix.group, unique and type.
+              - retrieve_file and store_file only take a local full path and not a buffer.
+              - If no mode is passed to create_directory it will be created with 777 permissions.
+
+    Errors that may occur throughout but should be handled
+    downstream.
+    """
+
+    def __init__(self, ftp_conn_id):
+        self.ftp_conn_id = ftp_conn_id
+        self.conn = None
+
+    def get_conn(self):
+        """
+        Returns an SFTP connection object
+        """
+        if self.conn is None:
+            params = self.get_connection(self.ftp_conn_id)
+            self.conn = pysftp.Connection(host=params.host, port=params.port, username=params.login, password=params.password)
+
+        return self.conn
+
+    def close_conn(self):
+        """
+        Closes the connection. An error will occur if the
+        connection wasnt ever opened.
+        """
+        conn = self.conn
+        conn.close()
+
+    def describe_directory(self, path):
+        """
+        Returns a dictionary of {filename: {attributes}} for all files
+        on the remote system.
+        :param path: full path to the remote directory
+        :type path: str
+        """
+        conn = self.get_conn()
+        flist = conn.listdir_attr(path)
+        files = {}
+        for f in flist:
+            modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime('%Y%m%d%H%M%S')
+            files[f.filename] = {'size': f.st_size, 'modify': modify}
+        return files
+
+    def list_directory(self, path, nlst=False):
+        """
+        Returns a list of files on the remote system.
+        :param path: full path to the remote directory to list
+        :type path: str
+        """
+        conn = self.get_conn()
+        files = conn.listdir(path)
+        return files
+
+    def create_directory(self, path, mode=777):
+        """
+        Creates a directory on the remote system.
+        :param path: full path to the remote directory to create
+        :type path: str
+        :param mode: int representation of octal mode for directory
+        """
+        conn = self.get_conn()
+        conn.mkdir(path, mode)
+
+    def delete_directory(self, path):
+        """
+        Deletes a directory on the remote system.
+        :param path: full path to the remote directory to delete
+        :type path: str
+        """
+        conn = self.get_conn()
+        conn.rmdir(path)
+
+    def retrieve_file(self, remote_full_path, local_full_path):
+        """
+        Transfers the remote file to a local location.
+        If local_full_path is a string path, the file will be put
+        at that location
+        :param remote_full_path: full path to the remote file
+        :type remote_full_path: str
+        :param local_full_path: full path to the local file
+        :type local_full_path: str
+        """
+        conn = self.get_conn()
+        logging.info('Retrieving file from FTP: {}'.format(remote_full_path))
+        conn.get(remote_full_path, local_full_path)
+        logging.info('Finished retrieving file from FTP: {}'.format(remote_full_path))
+
+    def store_file(self, remote_full_path, local_full_path):
+        """
+        Transfers a local file to the remote location.
+        If local_full_path_or_buffer is a string path, the file will be read
+        from that location
+        :param remote_full_path: full path to the remote file
+        :type remote_full_path: str
+        :param local_full_path: full path to the local file
+        :type local_full_path: str
+        """
+        conn = self.get_conn()
+        conn.put(local_full_path, remote_full_path)
+
+    def delete_file(self, path):
+        """
+        Removes a file on the FTP Server
+        :param path: full path to the remote file
+        :type path: str
+        """
+        conn = self.get_conn()
+        conn.remove(path)
+
+    def get_mod_time(self, path):
+        conn = self.get_conn()
+        ftp_mdtm = conn.stat(path).st_mtime
+        return datetime.datetime.fromtimestamp(ftp_mdtm).strftime('%Y%m%d%H%M%S')
+

From 04f9446eac1cb5d35131a1dc8a05b50c45a1c255 Mon Sep 17 00:00:00 2001
From: Biel Llobera <biellls@Biels-MacBook-Pro.local>
Date: Sun, 13 Nov 2016 20:18:49 +0100
Subject: [PATCH 02/18] Yapf refactor

---
 airflow/contrib/hooks/sftp_hook.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index b2cb34156c9e3..0c273cd168bbd 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -18,6 +18,7 @@
 import datetime
 from airflow.hooks.base_hook import BaseHook
 
+
 class SFTPHook(BaseHook):
     """
     Interact with SFTP. Aims to be interchangeable with FTPHook.
@@ -41,7 +42,11 @@ def get_conn(self):
         """
         if self.conn is None:
             params = self.get_connection(self.ftp_conn_id)
-            self.conn = pysftp.Connection(host=params.host, port=params.port, username=params.login, password=params.password)
+            self.conn = pysftp.Connection(
+                host=params.host,
+                port=params.port,
+                username=params.login,
+                password=params.password)
 
         return self.conn
 
@@ -64,7 +69,8 @@ def describe_directory(self, path):
         flist = conn.listdir_attr(path)
         files = {}
         for f in flist:
-            modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime('%Y%m%d%H%M%S')
+            modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime(
+                '%Y%m%d%H%M%S')
             files[f.filename] = {'size': f.st_size, 'modify': modify}
         return files
 
@@ -110,7 +116,8 @@ def retrieve_file(self, remote_full_path, local_full_path):
         conn = self.get_conn()
         logging.info('Retrieving file from FTP: {}'.format(remote_full_path))
         conn.get(remote_full_path, local_full_path)
-        logging.info('Finished retrieving file from FTP: {}'.format(remote_full_path))
+        logging.info('Finished retrieving file from FTP: {}'.format(
+            remote_full_path))
 
     def store_file(self, remote_full_path, local_full_path):
         """
@@ -137,5 +144,6 @@ def delete_file(self, path):
     def get_mod_time(self, path):
         conn = self.get_conn()
         ftp_mdtm = conn.stat(path).st_mtime
-        return datetime.datetime.fromtimestamp(ftp_mdtm).strftime('%Y%m%d%H%M%S')
+        return datetime.datetime.fromtimestamp(ftp_mdtm).strftime(
+            '%Y%m%d%H%M%S')
 

From 6aba1d39038d5a005ef6363f6651b0203b3c8ebf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Tue, 18 Jul 2017 18:56:28 +0100
Subject: [PATCH 03/18] Cleaning SFTP conn attribute after closing conn.

---
 airflow/contrib/hooks/sftp_hook.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index b2cb34156c9e3..8c5cbfe6fe7cb 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -18,6 +18,7 @@
 import datetime
 from airflow.hooks.base_hook import BaseHook
 
+
 class SFTPHook(BaseHook):
     """
     Interact with SFTP. Aims to be interchangeable with FTPHook.
@@ -52,6 +53,7 @@ def close_conn(self):
         """
         conn = self.conn
         conn.close()
+        self.conn = None
 
     def describe_directory(self, path):
         """

From b3d7c64c5224a4d9d8ce870f34522605a043a0f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Tue, 18 Jul 2017 18:58:05 +0100
Subject: [PATCH 04/18] Cleaning FTP conn attribute after closing conn.

---
 airflow/contrib/hooks/ftp_hook.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/airflow/contrib/hooks/ftp_hook.py b/airflow/contrib/hooks/ftp_hook.py
index 148811ff1e4f3..540c07021f005 100644
--- a/airflow/contrib/hooks/ftp_hook.py
+++ b/airflow/contrib/hooks/ftp_hook.py
@@ -90,6 +90,7 @@ def close_conn(self):
         """
         conn = self.conn
         conn.quit()
+        self.conn = None
 
     def describe_directory(self, path):
         """

From 97af689d818420d7cb73c2f0b0048b6c5a59ed4f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Wed, 19 Jul 2017 23:46:21 +0100
Subject: [PATCH 05/18] Adding legacy imports.

---
 airflow/contrib/hooks/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/airflow/contrib/hooks/__init__.py b/airflow/contrib/hooks/__init__.py
index 977c2cee2eef4..0de409e8b1aa8 100644
--- a/airflow/contrib/hooks/__init__.py
+++ b/airflow/contrib/hooks/__init__.py
@@ -36,6 +36,7 @@
     'ftps_hook': ['FTPSHook'],
     'vertica_hook': ['VerticaHook'],
     'ssh_hook': ['SSHHook'],
+    'sftp_hook': ['SFTPHook'],
     'bigquery_hook': ['BigQueryHook'],
     'qubole_hook': ['QuboleHook'],
     'gcs_hook': ['GoogleCloudStorageHook'],

From 960a5ca05f67a2f7eb14851d8da175936deacf2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Wed, 19 Jul 2017 23:46:48 +0100
Subject: [PATCH 06/18] Adding requirements.

---
 scripts/ci/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/ci/requirements.txt b/scripts/ci/requirements.txt
index 0e5dbaf20bfaa..0804523781371 100644
--- a/scripts/ci/requirements.txt
+++ b/scripts/ci/requirements.txt
@@ -70,6 +70,7 @@ pygments
 pyhive
 pykerberos
 PyOpenSSL
+pysftp
 PySmbClient
 python-daemon
 python-dateutil

From a6b16532625b0d1bf2c2ea2a0630426558258b75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Wed, 19 Jul 2017 23:47:28 +0100
Subject: [PATCH 07/18] Adding default SFTP connection.

---
 airflow/contrib/hooks/sftp_hook.py | 3 +--
 airflow/utils/db.py                | 4 ++++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index 8c5cbfe6fe7cb..1d44e26723b22 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
 
 import pysftp
 import logging
@@ -32,7 +31,7 @@ class SFTPHook(BaseHook):
     downstream.
     """
 
-    def __init__(self, ftp_conn_id):
+    def __init__(self, ftp_conn_id='sftp_default'):
         self.ftp_conn_id = ftp_conn_id
         self.conn = None
 
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 04b151232f9de..ac32367d3fc70 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -185,6 +185,10 @@ def initdb():
         models.Connection(
             conn_id='ssh_default', conn_type='ssh',
             host='localhost'))
+    merge_conn(
+        models.Connection(
+            conn_id='sftp_default', conn_type='sftp',
+            host='localhost', port=22, login='root'))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',

From 00874cbbb9dce530ec772f398002ca76045e09ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Wed, 19 Jul 2017 23:48:01 +0100
Subject: [PATCH 08/18] Adding SFTP tests.

---
 tests/contrib/hooks/test_sftp_hook.py | 97 +++++++++++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 tests/contrib/hooks/test_sftp_hook.py

diff --git a/tests/contrib/hooks/test_sftp_hook.py b/tests/contrib/hooks/test_sftp_hook.py
new file mode 100644
index 0000000000000..bced0ce2556d8
--- /dev/null
+++ b/tests/contrib/hooks/test_sftp_hook.py
@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+import shutil
+import os
+import pysftp
+
+from airflow import configuration
+from airflow.contrib.hooks.sftp_hook import SFTPHook
+
+TMP_PATH = '/tmp'
+TMP_DIR_FOR_TESTS = 'tests_sftp_hook_dir'
+TMP_FILE_FOR_TESTS = 'test_file.txt'
+
+
+class SFTPHookTest(unittest.TestCase):
+    def setUp(self):
+        configuration.load_test_config()
+        self.hook = SFTPHook()
+        os.makedirs(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        with open(os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS), 'a') as f:
+                f.write('Test file')
+
+    def test_get_conn(self):
+        output = self.hook.get_conn()
+        self.assertEqual(type(output), pysftp.Connection)
+
+    def test_close_conn(self):
+        self.hook.conn = self.hook.get_conn()
+        self.assertTrue(self.hook.conn is not None)
+        self.hook.close_conn()
+        self.assertTrue(self.hook.conn is None)
+
+    def test_describe_directory(self):
+        output = self.hook.describe_directory(TMP_PATH)
+        self.assertTrue(TMP_DIR_FOR_TESTS in output)
+
+    def test_list_directory(self):
+        output = self.hook.list_directory(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        self.assertEqual(output, [])
+
+    def test_create_and_delete_directory(self):
+        new_dir_name = 'new_dir'
+        self.hook.create_directory(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, new_dir_name))
+        output = self.hook.describe_directory(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        self.assertTrue(new_dir_name in output)
+        self.hook.delete_directory(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, new_dir_name))
+        output = self.hook.describe_directory(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        self.assertTrue(new_dir_name not in output)
+
+    def test_store_retrieve_and_delete_file(self):
+        self.hook.store_file(
+                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+                local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
+        )
+        output = self.hook.list_directory(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        self.assertEqual(output, [TMP_FILE_FOR_TESTS])
+        retrieved_file_name = 'retrieved.txt'
+        self.hook.retrieve_file(
+                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+                local_full_path=os.path.join(TMP_PATH, retrieved_file_name)
+        )
+        self.assertTrue(retrieved_file_name in os.listdir(TMP_PATH))
+        os.remove(os.path.join(TMP_PATH, retrieved_file_name))
+        self.hook.delete_file(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS))
+        output = self.hook.list_directory(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        self.assertEqual(output, [])
+
+    def test_get_mod_time(self):
+        self.hook.store_file(
+                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+                local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
+        )
+        output = self.hook.get_mod_time(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS))
+        self.assertEqual(len(output), 14)
+
+    def tearDown(self):
+        shutil.rmtree(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
+        os.remove(os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS))
+
+
+if __name__ == '__main__':
+    unittest.main()

From 1139ad04eb48f37a5c8c3955b924a1499af1d661 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Wed, 19 Jul 2017 23:48:14 +0100
Subject: [PATCH 09/18] Adding SFTPHook to docs.

---
 docs/code.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/code.rst b/docs/code.rst
index b17c3fef80a89..cf8399e3318ac 100644
--- a/docs/code.rst
+++ b/docs/code.rst
@@ -239,6 +239,7 @@ Community contributed hooks
         VerticaHook,
         FTPHook,
         SSHHook,
+        SFTPHook,
         CloudantHook
 
 .. autoclass:: airflow.contrib.hooks.gcs_hook.GoogleCloudStorageHook

From abde25805cc9c9d47dc4eeeea2cd61261895bde5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Thu, 20 Jul 2017 00:13:51 +0100
Subject: [PATCH 10/18] Fixing indentation

---
 tests/contrib/hooks/test_sftp_hook.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/contrib/hooks/test_sftp_hook.py b/tests/contrib/hooks/test_sftp_hook.py
index bced0ce2556d8..058ca07581b1b 100644
--- a/tests/contrib/hooks/test_sftp_hook.py
+++ b/tests/contrib/hooks/test_sftp_hook.py
@@ -33,7 +33,7 @@ def setUp(self):
         self.hook = SFTPHook()
         os.makedirs(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
         with open(os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS), 'a') as f:
-                f.write('Test file')
+            f.write('Test file')
 
     def test_get_conn(self):
         output = self.hook.get_conn()
@@ -64,15 +64,15 @@ def test_create_and_delete_directory(self):
 
     def test_store_retrieve_and_delete_file(self):
         self.hook.store_file(
-                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
-                local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
+            remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+            local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
         )
         output = self.hook.list_directory(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
         self.assertEqual(output, [TMP_FILE_FOR_TESTS])
         retrieved_file_name = 'retrieved.txt'
         self.hook.retrieve_file(
-                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
-                local_full_path=os.path.join(TMP_PATH, retrieved_file_name)
+            remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+            local_full_path=os.path.join(TMP_PATH, retrieved_file_name)
         )
         self.assertTrue(retrieved_file_name in os.listdir(TMP_PATH))
         os.remove(os.path.join(TMP_PATH, retrieved_file_name))
@@ -82,8 +82,8 @@ def test_store_retrieve_and_delete_file(self):
 
     def test_get_mod_time(self):
         self.hook.store_file(
-                remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
-                local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
+            remote_full_path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS),
+            local_full_path=os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS)
         )
         output = self.hook.get_mod_time(path=os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, TMP_FILE_FOR_TESTS))
         self.assertEqual(len(output), 14)

From 8f4fa1ba409a3eea6939df7abb2c3f3706088cb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Thu, 20 Jul 2017 22:50:06 +0100
Subject: [PATCH 11/18] [AIRFLOW-442] SFTP connection improvements

---
 airflow/contrib/hooks/sftp_hook.py | 19 ++++++++++++++++---
 airflow/utils/db.py                |  3 ++-
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index 1d44e26723b22..d85017d59132e 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -41,8 +41,22 @@ def get_conn(self):
         """
         if self.conn is None:
             params = self.get_connection(self.ftp_conn_id)
-            self.conn = pysftp.Connection(host=params.host, port=params.port, username=params.login, password=params.password)
-
+            cnopts = pysftp.CnOpts()
+            if 'ignore_hostkey_verification' in params.extra_dejson and params.extra_dejson['ignore_hostkey_verification']:
+                cnopts.hostkeys = None
+            conn_params = {
+                'host': params.host,
+                'port': params.port,
+                'username': params.login,
+                'cnopts': cnopts
+            }
+            if params.password is not None:
+                conn_params['password'] = params.password
+            if 'private_key' in params.extra_dejson:
+                conn_params['private_key'] = params.extra_dejson['private_key']
+            if 'private_key_pass' in params.extra_dejson:
+                conn_params['private_key_pass'] = params.extra_dejson['private_key_pass']
+            self.conn = pysftp.Connection(**conn_params)
         return self.conn
 
     def close_conn(self):
@@ -139,4 +153,3 @@ def get_mod_time(self, path):
         conn = self.get_conn()
         ftp_mdtm = conn.stat(path).st_mtime
         return datetime.datetime.fromtimestamp(ftp_mdtm).strftime('%Y%m%d%H%M%S')
-
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index ac32367d3fc70..639314f337234 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -188,7 +188,8 @@ def initdb():
     merge_conn(
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
-            host='localhost', port=22, login='root'))
+            host='localhost', port=22, login='root',
+            extra='{"ignore_hostkey_verification": true}'))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',

From 6f0d6670d290c25cbeae99cf39fc33de89810791 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Thu, 20 Jul 2017 23:33:12 +0100
Subject: [PATCH 12/18] [AIRFLOW-442] Add private key to SFTP default conn

---
 airflow/utils/db.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 639314f337234..4f2dadf301007 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -189,7 +189,7 @@ def initdb():
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
             host='localhost', port=22, login='root',
-            extra='{"ignore_hostkey_verification": true}'))
+            extra='{"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}'))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',

From b4cc98c17d39e31b09868943deaa37a72888b767 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Fri, 21 Jul 2017 18:17:07 +0100
Subject: [PATCH 13/18] [AIRFLOW-442] Changing SFTP default conn to use it
 within unit tests

---
 airflow/utils/db.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 4f2dadf301007..c676ade536842 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -188,7 +188,7 @@ def initdb():
     merge_conn(
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
-            host='localhost', port=22, login='root',
+            host='localhost', port=22, login='travis',
             extra='{"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}'))
     merge_conn(
         models.Connection(

From 8b0db3edded7fd8c889de6f13d079af69fefdaf5 Mon Sep 17 00:00:00 2001
From: Biel Llobera <biellls@Biels-MacBook-Pro.local>
Date: Sat, 22 Jul 2017 11:11:26 +0200
Subject: [PATCH 14/18] Improvements to sftp hook

---
 airflow/contrib/hooks/sftp_hook.py | 34 ++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index 0c273cd168bbd..a12f228b90dbd 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+import stat
 import pysftp
 import logging
 import datetime
@@ -42,11 +43,23 @@ def get_conn(self):
         """
         if self.conn is None:
             params = self.get_connection(self.ftp_conn_id)
-            self.conn = pysftp.Connection(
-                host=params.host,
-                port=params.port,
-                username=params.login,
-                password=params.password)
+
+            if params.password is not None:
+                self.conn = pysftp.Connection(
+                    host=params.host,
+                    port=params.port,
+                    username=params.login,
+                    password=params.password)
+            else:
+                private_key = db_conn.extra_dejson['private_key'] #Todo test not specified
+                private_key_pass = db_conn.extra_dejson['private_key_pass'] if 'private_key_pass' in db_conn.extra_dejson else None #Todo test not specified
+                self.conn = pysftp.Connection(
+                    host=params.host,
+                    port=params.port,
+                    username=params.login,
+                    private_key=private_key,
+                    private_key_pass=private_key_pass)
+
 
         return self.conn
 
@@ -57,11 +70,12 @@ def close_conn(self):
         """
         conn = self.conn
         conn.close()
+        self.conn = None
 
     def describe_directory(self, path):
         """
         Returns a dictionary of {filename: {attributes}} for all files
-        on the remote system.
+        on the remote system (where the MLSD command is supported).
         :param path: full path to the remote directory
         :type path: str
         """
@@ -69,9 +83,11 @@ def describe_directory(self, path):
         flist = conn.listdir_attr(path)
         files = {}
         for f in flist:
-            modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime(
-                '%Y%m%d%H%M%S')
-            files[f.filename] = {'size': f.st_size, 'modify': modify}
+            modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime('%Y%m%d%H%M%S')
+            files[f.filename] = {
+                'size': f.st_size,
+                'type': 'dir' if stat.S_ISDIR(f.st_mode) else 'file',
+                'modify': modify}
         return files
 
     def list_directory(self, path, nlst=False):

From 817eb6ba411b63cc69ce278b1d9d3283c5134bf5 Mon Sep 17 00:00:00 2001
From: Biel Llobera <biellls@Biels-MacBook-Pro.local>
Date: Sat, 22 Jul 2017 11:21:57 +0200
Subject: [PATCH 15/18] Updated docstring

---
 airflow/contrib/hooks/sftp_hook.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index a12f228b90dbd..1644b052642a8 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -24,8 +24,8 @@ class SFTPHook(BaseHook):
     """
     Interact with SFTP. Aims to be interchangeable with FTPHook.
 
-    Pitfalls: - In contrast with FTPHook describe_directory only returns size and modify. It doesn't return unix.owner,
-                unix.mode, perm, unix.group, unique and type.
+    Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and modify. It doesn't return unix.owner,
+                unix.mode, perm, unix.group and unique.
               - retrieve_file and store_file only take a local full path and not a buffer.
               - If no mode is passed to create_directory it will be created with 777 permissions.
 

From 72c7cc8f3de7811470b64ca460630aee6b85caee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Sat, 22 Jul 2017 15:48:09 +0100
Subject: [PATCH 16/18] [AIRFLOW-442] Fix code quality checks

---
 airflow/contrib/hooks/sftp_hook.py | 2 +-
 airflow/utils/db.py                | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index 45d45ec9126ee..61af7655355ec 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -87,7 +87,7 @@ def describe_directory(self, path):
                 'modify': modify}
         return files
 
-    def list_directory(self, path, nlst=False):
+    def list_directory(self, path):
         """
         Returns a list of files on the remote system.
         :param path: full path to the remote directory to list
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index c676ade536842..7865399479ca6 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -189,7 +189,8 @@ def initdb():
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
             host='localhost', port=22, login='travis',
-            extra='{"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}'))
+            extra='{"private_key": "~/.ssh/id_rsa", 
+            "ignore_hostkey_verification": true}'))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',

From ba610a8c70e5ee59c42b9cd2a7fcd7d9f7d09418 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Sat, 22 Jul 2017 16:00:24 +0100
Subject: [PATCH 17/18] [AIRFLOW-442] Fix SFTP default conn

---
 airflow/utils/db.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 7865399479ca6..c676ade536842 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -189,8 +189,7 @@ def initdb():
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
             host='localhost', port=22, login='travis',
-            extra='{"private_key": "~/.ssh/id_rsa", 
-            "ignore_hostkey_verification": true}'))
+            extra='{"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}'))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',

From 2fbe059e7bb336ec0d5160887ee3d2453d6c1fd6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Di=CC=81az=20Bautista?= <s.diazbautista@gmail.com>
Date: Sat, 22 Jul 2017 17:13:33 +0100
Subject: [PATCH 18/18] [AIRFLOW-442] Fix code quality checks

---
 airflow/contrib/hooks/sftp_hook.py | 19 +++++++++++--------
 airflow/utils/db.py                |  4 +++-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/airflow/contrib/hooks/sftp_hook.py b/airflow/contrib/hooks/sftp_hook.py
index 61af7655355ec..7340146625aa9 100644
--- a/airflow/contrib/hooks/sftp_hook.py
+++ b/airflow/contrib/hooks/sftp_hook.py
@@ -23,13 +23,15 @@ class SFTPHook(BaseHook):
     """
     Interact with SFTP. Aims to be interchangeable with FTPHook.
 
-    Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and modify. It doesn't return unix.owner,
-                unix.mode, perm, unix.group and unique.
-              - retrieve_file and store_file only take a local full path and not a buffer.
-              - If no mode is passed to create_directory it will be created with 777 permissions.
-
-    Errors that may occur throughout but should be handled
-    downstream.
+    Pitfalls: - In contrast with FTPHook describe_directory only returns size, type and
+                modify. It doesn't return unix.owner, unix.mode, perm, unix.group and
+                unique.
+              - retrieve_file and store_file only take a local full path and not a
+                buffer.
+              - If no mode is passed to create_directory it will be created with 777
+                permissions.
+
+    Errors that may occur throughout but should be handled downstream.
     """
 
     def __init__(self, ftp_conn_id='sftp_default'):
@@ -43,7 +45,8 @@ def get_conn(self):
         if self.conn is None:
             params = self.get_connection(self.ftp_conn_id)
             cnopts = pysftp.CnOpts()
-            if 'ignore_hostkey_verification' in params.extra_dejson and params.extra_dejson['ignore_hostkey_verification']:
+            if ('ignore_hostkey_verification' in params.extra_dejson
+                    and params.extra_dejson['ignore_hostkey_verification']):
                 cnopts.hostkeys = None
             conn_params = {
                 'host': params.host,
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index c676ade536842..d4a50de436361 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -189,7 +189,9 @@ def initdb():
         models.Connection(
             conn_id='sftp_default', conn_type='sftp',
             host='localhost', port=22, login='travis',
-            extra='{"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}'))
+            extra='''
+                {"private_key": "~/.ssh/id_rsa", "ignore_hostkey_verification": true}
+            '''))
     merge_conn(
         models.Connection(
             conn_id='fs_default', conn_type='fs',