airflow-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From bo...@apache.org
Subject incubator-airflow git commit: [AIRFLOW-866] Add FTPSensor
Date Sun, 19 Feb 2017 08:52:31 GMT
Repository: incubator-airflow
Updated Branches:
  refs/heads/master 0ed608dca -> 5f87f8a68


[AIRFLOW-866] Add FTPSensor

Closes #2070 from s7anley/ftp-sensor


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/5f87f8a6
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/5f87f8a6
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/5f87f8a6

Branch: refs/heads/master
Commit: 5f87f8a686550ade1c13f133a97523e8ebc3619f
Parents: 0ed608d
Author: Ján Koščo <3k.stanley@gmail.com>
Authored: Sun Feb 19 09:52:16 2017 +0100
Committer: Bolke de Bruin <bolke@xs4all.nl>
Committed: Sun Feb 19 09:52:19 2017 +0100

----------------------------------------------------------------------
 airflow/contrib/sensors/ftp_sensor.py | 64 +++++++++++++++++++++++++++++
 tests/contrib/sensors/ftp_sensor.py   | 66 ++++++++++++++++++++++++++++++
 2 files changed, 130 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5f87f8a6/airflow/contrib/sensors/ftp_sensor.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/sensors/ftp_sensor.py b/airflow/contrib/sensors/ftp_sensor.py
new file mode 100644
index 0000000..4a9428b
--- /dev/null
+++ b/airflow/contrib/sensors/ftp_sensor.py
@@ -0,0 +1,64 @@
+# -*- 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 ftplib
+import logging
+
+from airflow.contrib.hooks.ftp_hook import FTPHook, FTPSHook
+from airflow.operators.sensors import BaseSensorOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class FTPSensor(BaseSensorOperator):
+    """
+    Waits for a file or directory to be present on FTP.
+
+    :param path: Remote file or directory path
+    :type path: str
+    :param ftp_conn_id: The connection to run the sensor against
+    :type ftp_conn_id: str
+    """
+    template_fields = ('path',)
+
+    @apply_defaults
+    def __init__(self, path, ftp_conn_id='ftp_default', *args, **kwargs):
+        super(FTPSensor, self).__init__(*args, **kwargs)
+
+        self.path = path
+        self.ftp_conn_id = ftp_conn_id
+
+    def _create_hook(self):
+        """Return connection hook."""
+        return FTPHook(ftp_conn_id=self.ftp_conn_id)
+
+    def poke(self, context):
+        with self._create_hook() as hook:
+            logging.info('Poking for %s', self.path)
+            try:
+                hook.get_mod_time(self.path)
+            except ftplib.error_perm as e:
+                error = str(e).split(None, 1)
+                if error[1] != "Can't check for file existence":
+                    raise e
+
+                return False
+
+            return True
+
+
+class FTPSSensor(FTPSensor):
+    """Waits for a file or directory to be present on FTP over SSL."""
+    def _create_hook(self):
+        """Return connection hook."""
+        return FTPSHook(ftp_conn_id=self.ftp_conn_id)

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5f87f8a6/tests/contrib/sensors/ftp_sensor.py
----------------------------------------------------------------------
diff --git a/tests/contrib/sensors/ftp_sensor.py b/tests/contrib/sensors/ftp_sensor.py
new file mode 100644
index 0000000..50f8b8b
--- /dev/null
+++ b/tests/contrib/sensors/ftp_sensor.py
@@ -0,0 +1,66 @@
+# -*- 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 unittest
+from ftplib import error_perm
+
+from mock import MagicMock
+
+from airflow.contrib.hooks.ftp_hook import FTPHook
+from airflow.contrib.sensors.ftp_sensor import FTPSensor
+
+
+class TestFTPSensor(unittest.TestCase):
+    def setUp(self):
+        super(TestFTPSensor, self).setUp()
+        self._create_hook_orig = FTPSensor._create_hook
+        self.hook_mock = MagicMock(spec=FTPHook)
+
+        def _create_hook_mock(sensor):
+            mock = MagicMock()
+            mock.__enter__ = lambda x: self.hook_mock
+
+            return mock
+
+        FTPSensor._create_hook = _create_hook_mock
+
+    def tearDown(self):
+        FTPSensor._create_hook = self._create_hook_orig
+        super(TestFTPSensor, self).tearDown()
+
+    def test_poke(self):
+        op = FTPSensor(path="foobar.json", ftp_conn_id="bob_ftp",
+                       task_id="test_task")
+
+        self.hook_mock.get_mod_time.side_effect = \
+            [error_perm("550: Can't check for file existence"), None]
+
+        self.assertFalse(op.poke(None))
+        self.assertTrue(op.poke(None))
+
+    def test_poke_fails_due_error(self):
+        op = FTPSensor(path="foobar.json", ftp_conn_id="bob_ftp",
+                       task_id="test_task")
+
+        self.hook_mock.get_mod_time.side_effect = \
+            error_perm("530: Login authentication failed")
+
+        with self.assertRaises(error_perm) as context:
+            op.execute(None)
+
+        self.assertTrue("530" in str(context.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()


Mime
View raw message