diff --git a/airflow/contrib/sensors/hdfs_sensor.py b/airflow/contrib/sensors/hdfs_sensor.py deleted file mode 100644 index 832b81b8e5f25..0000000000000 --- a/airflow/contrib/sensors/hdfs_sensor.py +++ /dev/null @@ -1,78 +0,0 @@ -# -*- coding: utf-8 -*- -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you 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 airflow.sensors.hdfs_sensor import HdfsSensor - - -class HdfsSensorRegex(HdfsSensor): - def __init__(self, - regex, - *args, - **kwargs): - super(HdfsSensorRegex, self).__init__(*args, **kwargs) - self.regex = regex - - def poke(self, context): - """ - poke matching files in a directory with self.regex - - :return: Bool depending on the search criteria - """ - sb = self.hook(self.hdfs_conn_id).get_conn() - self.log.info( - 'Poking for {self.filepath} to be a directory ' - 'with files matching {self.regex.pattern}'. - format(**locals()) - ) - result = [f for f in sb.ls([self.filepath], include_toplevel=False) if - f['file_type'] == 'f' and - self.regex.match(f['path'].replace('%s/' % self.filepath, ''))] - result = self.filter_for_ignored_ext(result, self.ignored_ext, - self.ignore_copying) - result = self.filter_for_filesize(result, self.file_size) - return bool(result) - - -class HdfsSensorFolder(HdfsSensor): - def __init__(self, - be_empty=False, - *args, - **kwargs): - super(HdfsSensorFolder, self).__init__(*args, **kwargs) - self.be_empty = be_empty - - def poke(self, context): - """ - poke for a non empty directory - - :return: Bool depending on the search criteria - """ - sb = self.hook(self.hdfs_conn_id).get_conn() - result = [f for f in sb.ls([self.filepath], include_toplevel=True)] - result = self.filter_for_ignored_ext(result, self.ignored_ext, - self.ignore_copying) - result = self.filter_for_filesize(result, self.file_size) - if self.be_empty: - self.log.info('Poking for filepath {self.filepath} to a empty directory' - .format(**locals())) - return len(result) == 1 and result[0]['path'] == self.filepath - else: - self.log.info('Poking for filepath {self.filepath} to a non empty directory' - .format(**locals())) - result.pop(0) - return bool(result) and result[0]['file_type'] == 'f' diff --git a/airflow/hooks/hdfs_hook.py b/airflow/hooks/hdfs_hook.py index 597b7c4f7ec83..2ef51d873282a 100644 --- a/airflow/hooks/hdfs_hook.py +++ b/airflow/hooks/hdfs_hook.py @@ -17,85 +17,117 @@ # specific language governing permissions and limitations # under the License. -from six import PY2 +import hdfs3 +from hdfs3.utils import MyNone from airflow import configuration -from airflow.exceptions import AirflowException from airflow.hooks.base_hook import BaseHook -snakebite_imported = False -if PY2: - from snakebite.client import Client, HAClient, Namenode, AutoConfigClient - snakebite_imported = True +class HdfsHook(BaseHook): + """Hook for interacting with HDFS using the hdfs3 library. + By default hdfs3 loads its configuration from `core-site.xml` and + `hdfs-site.xml` if these files can be found in any of the typical + locations. The hook loads `host` and `port` parameters from the + hdfs connection (if given) and extra configuration parameters can be + supplied in the connections extra JSON as follows: -class HDFSHookException(AirflowException): - pass + { + "pars": { + "dfs.domain.socket.path": "/var/lib/hadoop-hdfs/dn_socket" + }, + "ha": { + "host": "nameservice1", + "conf": { + "dfs.nameservices": "ns1", + "dfs.ha.namenodes.ns1": "nn1,nn2", + "dfs.namenode.rpc-address.ns1.nn1": "host1:8020", + "dfs.namenode.rpc-address.ns1.nn2": "host2:8020", + "dfs.namenode.http-address.ns1.nn1": "host1:50070", + "dfs.namenode.http-address.ns1.nn2": "host2:50070" + } + }, + "autoconf": True, + "token": "...", + "ticket_cache": "..." + } + Here `pars` can be used to supply configuration options with the same key + names as typically contained in the XML config files, which will take + precedence over any parameters loaded from files. The `ha` configuration + section can be used to supply options for using hdfs3 in high-availability + mode (see the hdfs3 documentation for more details). The `autoconf` option + controls whether hdfs3 uses the available hadoop config files for + its configuration, whilst the `token` and `ticket_cache` options are + used for configuring kerberos. -class HDFSHook(BaseHook): - """ - Interact with HDFS. This class is a wrapper around the snakebite library. - - :param hdfs_conn_id: Connection id to fetch connection info - :type hdfs_conn_id: str - :param proxy_user: effective user for HDFS operations - :type proxy_user: str - :param autoconfig: use snakebite's automatically configured client - :type autoconfig: bool + Security modes can also be configured by defining appropriate value for + the `hadoop.security.authentication` key in `pars`. Kerberos is used + automatically if Airflow has been configured to use kerberos. + + :param str hdfs_conn_id: Connection ID to fetch parameters from. + :param bool autoconf: Whether to use autoconfig to discover + configuration options from the hdfs XML configuration files. """ - def __init__(self, hdfs_conn_id='hdfs_default', proxy_user=None, - autoconfig=False): - if not snakebite_imported: - raise ImportError( - 'This HDFSHook implementation requires snakebite, but ' - 'snakebite is not compatible with Python 3 ' - '(as of August 2015). Please use Python 2 if you require ' - 'this hook -- or help by submitting a PR!') + + def __init__(self, hdfs_conn_id=None): + super(HdfsHook, self).__init__(None) + self.hdfs_conn_id = hdfs_conn_id - self.proxy_user = proxy_user - self.autoconfig = autoconfig + self._conn = None def get_conn(self): - """ - Returns a snakebite HDFSClient object. - """ - # When using HAClient, proxy_user must be the same, so is ok to always - # take the first. - effective_user = self.proxy_user - autoconfig = self.autoconfig - use_sasl = configuration.conf.get('core', 'security') == 'kerberos' - - try: - connections = self.get_connections(self.hdfs_conn_id) - - if not effective_user: - effective_user = connections[0].login - if not autoconfig: - autoconfig = connections[0].extra_dejson.get('autoconfig', - False) - hdfs_namenode_principal = connections[0].extra_dejson.get( - 'hdfs_namenode_principal') - except AirflowException: - if not autoconfig: - raise - - if autoconfig: - # will read config info from $HADOOP_HOME conf files - client = AutoConfigClient(effective_user=effective_user, - use_sasl=use_sasl) - elif len(connections) == 1: - client = Client(connections[0].host, connections[0].port, - effective_user=effective_user, use_sasl=use_sasl, - hdfs_namenode_principal=hdfs_namenode_principal) - elif len(connections) > 1: - nn = [Namenode(conn.host, conn.port) for conn in connections] - client = HAClient(nn, effective_user=effective_user, - use_sasl=use_sasl, - hdfs_namenode_principal=hdfs_namenode_principal) - else: - raise HDFSHookException("conn_id doesn't exist in the repository " - "and autoconfig is not specified") - - return client + if self._conn is None: + hdfs_params = {} + + # Configure kerberos security if used by Airflow. + if configuration.conf.get("core", "security") == "kerberos": + hdfs_params["hadoop.security.authentication"] = "kerberos" + + if self.hdfs_conn_id is None: + self._conn = hdfs3.HDFileSystem(pars=hdfs_params, autoconf=True) + else: + conn_params = self.get_connection(self.hdfs_conn_id) + conn_extra_params = conn_params.extra_dejson + + # Extract hadoop parameters from extra. + hdfs_params.update(conn_extra_params.get("pars", {})) + + # Extract high-availability config if given. + ha_params = conn_extra_params.get("ha", {}) + hdfs_params.update(ha_params.get("conf", {})) + + # Collect extra parameters to pass to kwargs. Note that we + # avoid passing empty parameters, as hdfs3 seems to do + # funky things with defining its own None variable. + extra_kws = { + "user": conn_params.login or None, + "ticket_cache": conn_extra_params.get("ticket_cache"), + "token": conn_extra_params.get("token"), + } + extra_kws = {k: v for k, v in extra_kws.items() if v} + + # Build connection. + self._conn = hdfs3.HDFileSystem( + host=ha_params.get("host") or conn_params.host or MyNone, + port=conn_params.port or MyNone, + pars=hdfs_params, + autoconf=conn_extra_params.get("autoconf", True), + **extra_kws + ) + + return self._conn + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + def close(self): + """Closes the HDFSHook and any underlying connections.""" + + if self._conn is not None: + self._conn.disconnect() + self._conn = None diff --git a/airflow/sensors/hdfs_sensor.py b/airflow/sensors/hdfs_sensor.py index 3eb5145ec9668..155903ed3607a 100644 --- a/airflow/sensors/hdfs_sensor.py +++ b/airflow/sensors/hdfs_sensor.py @@ -17,107 +17,231 @@ # specific language governing permissions and limitations # under the License. -import re -import sys -from builtins import str +import posixpath from airflow import settings -from airflow.hooks.hdfs_hook import HDFSHook +from airflow.hooks.hdfs_hook import HdfsHook from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.utils.decorators import apply_defaults -from airflow.utils.log.logging_mixin import LoggingMixin -class HdfsSensor(BaseSensorOperator): - """ - Waits for a file or folder to land in HDFS +class HdfsFileSensor(BaseSensorOperator): + """Sensor that waits for files matching a specific (glob) pattern to land in HDFS. + + :param str file_pattern: Glob pattern to match. + :param str conn_id: Connection to use. + :param Iterable[FilePathFilter] filters: Optional list of filters that can be + used to apply further filtering to any file paths matching the glob pattern. + Any files that fail a filter are dropped from consideration. + :param int min_size: Minimum size (in MB) for files to be considered. Can be used + to filter any intermediate files that are below the expected file size. + :param Set[str] ignore_exts: File extensions to ignore. By default, files with + a '_COPYING_' extension are ignored, as these represent temporary files. """ - template_fields = ('filepath',) - ui_color = settings.WEB_COLORS['LIGHTBLUE'] + + template_fields = ("_pattern",) + ui_color = settings.WEB_COLORS["LIGHTBLUE"] @apply_defaults - def __init__(self, - filepath, - hdfs_conn_id='hdfs_default', - ignored_ext=None, - ignore_copying=True, - file_size=None, - hook=HDFSHook, - *args, - **kwargs): - super(HdfsSensor, self).__init__(*args, **kwargs) - if ignored_ext is None: - ignored_ext = ['_COPYING_'] - self.filepath = filepath - self.hdfs_conn_id = hdfs_conn_id - self.file_size = file_size - self.ignored_ext = ignored_ext - self.ignore_copying = ignore_copying - self.hook = hook - - @staticmethod - def filter_for_filesize(result, size=None): - """ - Will test the filepath result and test if its size is at least self.filesize - - :param result: a list of dicts returned by Snakebite ls - :param size: the file size in MB a file should be at least to trigger True - :return: (bool) depending on the matching criteria - """ - if size: - log = LoggingMixin().log - log.debug( - 'Filtering for file size >= %s in files: %s', - size, map(lambda x: x['path'], result) - ) - size *= settings.MEGABYTE - result = [x for x in result if x['length'] >= size] - log.debug('HdfsSensor.poke: after size filter result is %s', result) - return result - - @staticmethod - def filter_for_ignored_ext(result, ignored_ext, ignore_copying): - """ - Will filter if instructed to do so the result to remove matching criteria - - :param result: list of dicts returned by Snakebite ls - :type result: list[dict] - :param ignored_ext: list of ignored extensions - :type ignored_ext: list - :param ignore_copying: shall we ignore ? - :type ignore_copying: bool - :return: list of dicts which were not removed - :rtype: list[dict] - """ - if ignore_copying: - log = LoggingMixin().log - regex_builder = r"^.*\.(%s$)$" % '$|'.join(ignored_ext) - ignored_extensions_regex = re.compile(regex_builder) - log.debug( - 'Filtering result for ignored extensions: %s in files %s', - ignored_extensions_regex.pattern, map(lambda x: x['path'], result) - ) - result = [x for x in result if not ignored_extensions_regex.match(x['path'])] - log.debug('HdfsSensor.poke: after ext filter result is %s', result) - return result + def __init__( + self, + pattern, + conn_id="hdfs_default", + filters=None, + min_size=None, + ignore_exts=("_COPYING_",), + **kwargs + ): + super(HdfsFileSensor, self).__init__(**kwargs) + + # Min-size and ignore-ext filters are added via + # arguments for backwards compatibility. + filters = list(filters or []) + + if min_size: + filters.append(SizeFilter(min_size=min_size)) + + if ignore_exts: + filters.append(ExtFilter(exts=ignore_exts)) + + self._pattern = pattern + self._conn_id = conn_id + self._filters = filters def poke(self, context): - sb = self.hook(self.hdfs_conn_id).get_conn() - self.log.info('Poking for file {self.filepath}'.format(**locals())) - try: - # IMOO it's not right here, as there no raise of any kind. - # if the filepath is let's say '/data/mydirectory', - # it's correct but if it is '/data/mydirectory/*', - # it's not correct as the directory exists and sb does not raise any error - # here is a quick fix - result = [f for f in sb.ls([self.filepath], include_toplevel=False)] - self.log.debug('HdfsSensor.poke: result is %s', result) - result = self.filter_for_ignored_ext( - result, self.ignored_ext, self.ignore_copying + with HdfsHook(self._conn_id) as hook: + conn = hook.get_conn() + + # Fetch files matching glob pattern. + self.log.info("Poking for file pattern %s", self._pattern) + + try: + file_paths = [ + fp for fp in conn.glob(self._pattern) if not conn.isdir(fp) + ] + except IOError: + # File path doesn't exist yet. + file_paths = [] + + self.log.info("Files matching pattern: %s", file_paths) + + # Filter using any provided filters. + for filter_ in self._filters: + file_paths = filter_(file_paths, hook) + file_paths = list(file_paths) + + self.log.info("Filters after filtering: %s", file_paths) + + return bool(file_paths) + + +class HdfsFolderSensor(BaseSensorOperator): + """Waits for folder(s) matching the given pattern to be created in HDFS. + + By default, the sensor does not check whether any matched folders contain + files. If folders are required to (not) be empty, this behaviour can be modified + using the `require_empty` and `require_not_empty` parameters. If `require_empty` + is True, the sensor fails if any matched folders contain files. Similarly, if + `require_not_empty` is True, the sensor fails if any matched folders do not contain + files. The list of file paths considered in these checks can be modified using the + `sub_pattern` and `sub_filters` parameters. + + :param str pattern: Glob pattern to match. + :param str conn_id: Connection to use. + :param bool require_empty: Whether folders are required to be empty. If true, + the sensor fails if any of the matched directories is not empty. + :param bool require_not_empty: Whether folders are required to be NOT empty. + If true, the sensor fails if any of the matched directories is empty. + :param str sub_pattern: Glob pattern to filter nested file paths on when + checking whether directories are (not) empty. + :param Iterable[FilePathFilter] sub_filters: File path filters that should be used + to filter nested file paths on when checking whether directories are (not) + empty. See `HdfsFileSensor` for more details. + """ + + template_fields = ("_pattern",) + ui_color = settings.WEB_COLORS["LIGHTBLUE"] + + def __init__( + self, + pattern, + conn_id="hdfs_default", + require_empty=False, + require_not_empty=False, + sub_pattern=None, + sub_filters=None, + **kwargs + ): + super(HdfsFolderSensor, self).__init__(**kwargs) + + if require_empty and require_not_empty: + raise ValueError( + "Either require_empty or require_not_empty must be false, " + "as the two conditions are mutually exclusive." ) - result = self.filter_for_filesize(result, self.file_size) - return bool(result) - except Exception: - e = sys.exc_info() - self.log.debug("Caught an exception !: %s", str(e)) - return False + + self._pattern = pattern + self._conn_id = conn_id + + self._require_empty = require_empty + self._require_not_empty = require_not_empty + + self._sub_pattern = sub_pattern or "*" + self._sub_filters = list(sub_filters or []) + + def poke(self, context): + with HdfsHook(self._conn_id) as hook: + conn = hook.get_conn() + + # Try to expand glob pattern, returns single dir if not glob. + self.log.info("Poking for directories matching %s", self._pattern) + + try: + dir_paths = [ + path_ for path_ in conn.glob(self._pattern) if conn.isdir(path_) + ] + except IOError: + # File path doesn't exist yet. + dir_paths = [] + + self.log.info("Directories matching pattern: %s", dir_paths) + + if self._require_empty or self._require_not_empty: + self.log.info( + "Checking for files or subdirectories matching pattern: %s", + self._sub_pattern, + ) + + # Check if directories do/don't contain files. Returns False + # if any of the directories fail the required condition. + for dir_path in dir_paths: + sub_pattern = posixpath.join(dir_path, self._sub_pattern) + sub_paths = conn.glob(sub_pattern) + + for filter_ in self._sub_filters: + sub_paths = filter_(sub_paths, hook) + + sub_paths = list(sub_paths) + + self.log.info( + "Sub-directories/files matching pattern: %s", sub_paths + ) + + is_empty = not sub_paths + if (self._require_empty and not is_empty) or ( + self._require_not_empty and is_empty + ): + return False + + return bool(dir_paths) + + +class FilePathFilter: + """Base file path filter class. + + Allows a given list of file paths to be filtered on a given criteria. + Examples include a size filter (which filters files for a given minimum file + size) and a extension filter (which filters files with a given extension). + + Filters are required to implement a single __call__ method, which takes the + file system hook and a set of file paths to filter. As a result, this method + should return the set of filtered file paths. + """ + + def __call__(self, file_paths, hook): + raise NotImplementedError() + + +class SizeFilter(FilePathFilter): + """Filter that drops any file paths below the given file size. + + :param int min_size: Minimum file size in megabytes. + """ + + def __init__(self, min_size): + self._min_size = min_size + + def __call__(self, file_paths, hook): + min_size_mb = self._min_size * settings.MEGABYTE + + conn = hook.get_conn() + for file_path in file_paths: + info = conn.info(file_path) + + if info["kind"] == "file" and info["size"] > min_size_mb: + yield file_path + + +class ExtFilter(FilePathFilter): + """Filter that drops any file paths with given extensions. + + :param Set[str] exts: Set of extensions to filter for. + """ + + def __init__(self, exts): + self._exts = set(exts) + + def __call__(self, file_paths, hook): + for file_path in file_paths: + if posixpath.splitext(file_path)[1][1:] not in self._exts: + yield file_path diff --git a/setup.py b/setup.py index 1cbe29b64ad05..23278036419f2 100644 --- a/setup.py +++ b/setup.py @@ -183,7 +183,7 @@ def write_version(filename=os.path.join(*['airflow', ] github_enterprise = ['Flask-OAuthlib>=0.9.1'] google_auth = ['Flask-OAuthlib>=0.9.1'] -hdfs = ['snakebite>=2.7.8'] +hdfs = ['hdfs3>=0.3.0'] hive = [ 'hmsclient>=0.1.0', 'pyhive>=0.6.0', diff --git a/tests/contrib/sensors/test_hdfs_sensor.py b/tests/contrib/sensors/test_hdfs_sensor.py index b03b738686ed8..f9afaa16ce003 100644 --- a/tests/contrib/sensors/test_hdfs_sensor.py +++ b/tests/contrib/sensors/test_hdfs_sensor.py @@ -16,238 +16,190 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import logging -import unittest +import datetime as dt import re -from datetime import timedelta +import unittest +import warnings -from airflow.contrib.sensors.hdfs_sensor import HdfsSensorFolder, HdfsSensorRegex -from airflow.exceptions import AirflowSensorTimeout +from airflow import models +from airflow.contrib.sensors.hdfs_sensor import (HdfsSensorFolder, + HdfsSensorRegex, + HdfsRegexFileSensor) +from tests.sensors.test_hdfs_sensor import MockHdfs3Client + + +class HdfsRegexFileSensorTests(unittest.TestCase): + """Tests for the HdfsRegexFileSensor class.""" -class HdfsSensorFolderTests(unittest.TestCase): - def setUp(self): - from tests.core import FakeHDFSHook - self.hook = FakeHDFSHook - self.log = logging.getLogger() - self.log.setLevel(logging.DEBUG) - - def test_should_be_empty_directory(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - task = HdfsSensorFolder(task_id='Should_be_empty_directory', - filepath='/datadirectory/empty_directory', - be_empty=True, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - task.execute(None) - - # Then - # Nothing happens, nothing is raised exec is ok - - def test_should_be_empty_directory_fail(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - task = HdfsSensorFolder(task_id='Should_be_empty_directory_fail', - filepath='/datadirectory/not_empty_directory', - be_empty=True, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) - - def test_should_be_a_non_empty_directory(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - task = HdfsSensorFolder(task_id='Should_be_non_empty_directory', - filepath='/datadirectory/not_empty_directory', - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - task.execute(None) - - # Then - # Nothing happens, nothing is raised exec is ok - - def test_should_be_non_empty_directory_fail(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - task = HdfsSensorFolder(task_id='Should_be_empty_directory_fail', - filepath='/datadirectory/empty_directory', - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) - - -class HdfsSensorRegexTests(unittest.TestCase): def setUp(self): - from tests.core import FakeHDFSHook - self.hook = FakeHDFSHook - self.log = logging.getLogger() - self.log.setLevel(logging.DEBUG) + file_details = [ + { + 'kind': 'directory', + 'name': '/data', + 'size': 0 + }, + { + 'kind': 'file', + 'name': '/data/test1file', + 'size': 2000000 + }, + { + 'kind': 'file', + 'name': '/data/copying._COPYING_', + 'size': 2000000 + } + ] + + self._mock_client, self._mock_params = \ + MockHdfs3Client.from_file_details(file_details, test_instance=self) + + self._default_task_kws = { + 'timeout': 1, + 'retry_delay': dt.timedelta(seconds=1), + 'poke_interval': 1 + } def test_should_match_regex(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - compiled_regex = re.compile("test[1-2]file") - task = HdfsSensorRegex(task_id='Should_match_the_regex', - filepath='/datadirectory/regex_dir', - regex=compiled_regex, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - task.execute(None) - - # Then - # Nothing happens, nothing is raised exec is ok + """Tests example where files should match regex.""" + + regex = re.compile("test[1-2]file") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex', + pattern='/data/*', + regex=regex, + **self._default_task_kws) + self.assertTrue(task.poke(context={})) + + def test_should_match_regex_dir(self): + """Tests example where files should match regex with dir path.""" + + regex = re.compile("test[1-2]file") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex', + pattern='/data', + regex=regex, + **self._default_task_kws) + self.assertTrue(task.poke(context={})) def test_should_not_match_regex(self): - """ - test the empty directory behaviour - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - compiled_regex = re.compile("^IDoNotExist") - task = HdfsSensorRegex(task_id='Should_not_match_the_regex', - filepath='/datadirectory/regex_dir', - regex=compiled_regex, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) - - def test_should_match_regex_and_filesize(self): - """ - test the file size behaviour with regex - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - compiled_regex = re.compile("test[1-2]file") - task = HdfsSensorRegex(task_id='Should_match_the_regex_and_filesize', - filepath='/datadirectory/regex_dir', - regex=compiled_regex, - ignore_copying=True, - ignored_ext=['_COPYING_', 'sftp'], - file_size=10, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - task.execute(None) - - # Then - # Nothing happens, nothing is raised exec is ok - - def test_should_match_regex_but_filesize(self): - """ - test the file size behaviour with regex - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - compiled_regex = re.compile("test[1-2]file") - task = HdfsSensorRegex(task_id='Should_match_the_regex_but_filesize', - filepath='/datadirectory/regex_dir', - regex=compiled_regex, - file_size=20, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) - - def test_should_match_regex_but_copyingext(self): - """ - test the file size behaviour with regex - :return: - """ - # Given - self.log.debug('#' * 10) - self.log.debug('Running %s', self._testMethodName) - self.log.debug('#' * 10) - compiled_regex = re.compile(r"copying_file_\d+.txt") - task = HdfsSensorRegex(task_id='Should_match_the_regex_but_filesize', - filepath='/datadirectory/regex_dir', - regex=compiled_regex, - ignored_ext=['_COPYING_', 'sftp'], - file_size=20, - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) + """Tests example where files should match regex.""" + + regex = re.compile("^IDoNotExist") + task = HdfsRegexFileSensor( + task_id='should_not_match_the_regex', + pattern='/data/*', + regex=regex, + **self._default_task_kws) + self.assertFalse(task.poke(context={})) + + def test_should_match_regex_and_size(self): + """Tests example with matching regex and sufficient file size.""" + + regex = re.compile("test[1-2]file") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex_and_size', + pattern='/data/*', + regex=regex, + min_size=1, + **self._default_task_kws) + self.assertTrue(task.poke(context={})) + + def test_should_match_regex_not_size(self): + """Tests example with matching regex but too small file size.""" + + regex = re.compile("test[1-2]file") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex_but_not_size', + pattern='/data/*', + regex=regex, + min_size=10, + **self._default_task_kws) + self.assertFalse(task.poke(context={})) + + def test_should_match_regex_not_ext(self): + """Tests example with matching regex but wrong ext.""" + + regex = re.compile("test[1-2]file") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex_but_not_size', + pattern='/data/*', + regex=regex, + min_size=10, + **self._default_task_kws) + self.assertFalse(task.poke(context={})) + + compiled_regex = re.compile("copying.*") + task = HdfsRegexFileSensor( + task_id='should_match_the_regex_but_not_ext', + pattern='/data/*', + regex=compiled_regex, + ignore_exts=['_COPYING_'], + **self._default_task_kws) + self.assertFalse(task.poke(context={})) + + def test_calling_old_class(self): + """Tests call to old class.""" + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + regex = re.compile("test[1-2]file") + task = HdfsSensorRegex( + task_id='should_match_the_regex_old', + pattern='/data', + regex=regex, + **self._default_task_kws) + self.assertTrue(task.poke(context={})) + + +class HdfsSensorFolderTests(unittest.TestCase): + """Tests for the (deprecated) HdfsSensorFolder class.""" + + def setUp(self): + file_details = [ + {"kind": "directory", "name": "/empty", "size": 0}, + {"kind": "directory", "name": "/not_empty", "size": 0}, + {"kind": "file", "name": "/not_empty/a.tmp", "size": 2000000}, + {"kind": "directory", "name": "/nested", "size": 0}, + {"kind": "directory", "name": "/nested/a", "size": 0} + ] + + self._mock_client, self._mock_params = \ + MockHdfs3Client.from_file_details(file_details, test_instance=self) + + self._default_task_kws = { + "timeout": 1, + "retry_delay": dt.timedelta(seconds=1), + "poke_interval": 1, + } + + self._mock_client = MockHdfs3Client(file_details) + self._mock_params = models.Connection(conn_id="hdfs_default") + + def test_empty_directory(self): + """Tests example with an empty directory.""" + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + task = HdfsSensorFolder( + task_id="test_empty_directory", + pattern="/empty") + self.assertTrue(task.poke(context={})) + + def test_non_empty_directory(self): + """Tests example with a non-empty directory.""" + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + task = HdfsSensorFolder( + task_id="test_non_empty_directory", + pattern="/empty") + self.assertTrue(task.poke(context={})) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/core.py b/tests/core.py index 0fbe29952e66d..62a5817001670 100644 --- a/tests/core.py +++ b/tests/core.py @@ -2647,158 +2647,6 @@ def check_for_path(self, hdfs_path): return hdfs_path -class FakeSnakeBiteClientException(Exception): - pass - - -class FakeSnakeBiteClient(object): - - def __init__(self): - self.started = True - - def ls(self, path, include_toplevel=False): - """ - the fake snakebite client - :param path: the array of path to test - :param include_toplevel: to return the toplevel directory info - :return: a list for path for the matching queries - """ - if path[0] == '/datadirectory/empty_directory' and not include_toplevel: - return [] - elif path[0] == '/datadirectory/datafile': - return [{ - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 0, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/datafile' - }] - elif path[0] == '/datadirectory/empty_directory' and include_toplevel: - return [{ - 'group': u'supergroup', - 'permission': 493, - 'file_type': 'd', - 'access_time': 0, - 'block_replication': 0, - 'modification_time': 1481132141540, - 'length': 0, - 'blocksize': 0, - 'owner': u'hdfs', - 'path': '/datadirectory/empty_directory' - }] - elif path[0] == '/datadirectory/not_empty_directory' and include_toplevel: - return [{ - 'group': u'supergroup', - 'permission': 493, - 'file_type': 'd', - 'access_time': 0, - 'block_replication': 0, - 'modification_time': 1481132141540, - 'length': 0, - 'blocksize': 0, - 'owner': u'hdfs', - 'path': '/datadirectory/empty_directory' - }, { - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 0, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/not_empty_directory/test_file' - }] - elif path[0] == '/datadirectory/not_empty_directory': - return [{ - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 0, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/not_empty_directory/test_file' - }] - elif path[0] == '/datadirectory/not_existing_file_or_directory': - raise FakeSnakeBiteClientException - elif path[0] == '/datadirectory/regex_dir': - return [{ - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, 'length': 12582912, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/regex_dir/test1file' - }, { - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 12582912, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/regex_dir/test2file' - }, { - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 12582912, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/regex_dir/test3file' - }, { - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 12582912, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/regex_dir/copying_file_1.txt._COPYING_' - }, { - 'group': u'supergroup', - 'permission': 420, - 'file_type': 'f', - 'access_time': 1481122343796, - 'block_replication': 3, - 'modification_time': 1481122343862, - 'length': 12582912, - 'blocksize': 134217728, - 'owner': u'hdfs', - 'path': '/datadirectory/regex_dir/copying_file_3.txt.sftp' - }] - else: - raise FakeSnakeBiteClientException - - -class FakeHDFSHook(object): - def __init__(self, conn_id=None): - self.conn_id = conn_id - - def get_conn(self): - client = FakeSnakeBiteClient() - return client - - class ConnectionTest(unittest.TestCase): def setUp(self): configuration.load_test_config() @@ -2901,6 +2749,7 @@ def test_init_proxy_user(self): self.assertEqual('someone', c.proxy_user) +<<<<<<< HEAD HDFSHook = None if six.PY2: from airflow.hooks.hdfs_hook import HDFSHook @@ -2950,6 +2799,8 @@ def test_get_ha_client(self, mock_get_connections): self.assertIsInstance(client, snakebite.client.HAClient) +======= +>>>>>>> Apply kerberos even without connection. send_email_test = mock.Mock() diff --git a/tests/hooks/test_hdfs_hook.py b/tests/hooks/test_hdfs_hook.py new file mode 100644 index 0000000000000..b7e9a22a7c025 --- /dev/null +++ b/tests/hooks/test_hdfs_hook.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 + +import mock +from hdfs3.utils import MyNone + +from airflow.hooks.hdfs_hook import HdfsHook, hdfs3, configuration + + +class TestHDFSHook(unittest.TestCase): + """ + Tests for the Hdfs3Hook class. + + Note that the HDFileSystem class is mocked in most of these tests + to avoid the requirement of having a local HDFS instance for testing. + """ + + def setUp(self): + self._mock_fs = mock.Mock() + + self._mocked_hook = HdfsHook() + self._mocked_hook._conn = self._mock_fs + + @mock.patch.object(hdfs3, 'HDFileSystem') + @mock.patch.object(HdfsHook, 'get_connection') + def test_get_conn(self, conn_mock, hdfs3_mock): + """Tests get_conn call without ID.""" + + with HdfsHook() as hook: + hook.get_conn() + + conn_mock.assert_not_called() + hdfs3_mock.assert_called_once_with(autoconf=True, pars={}) + + @mock.patch.object(hdfs3, 'HDFileSystem') + @mock.patch.object(HdfsHook, 'get_connection') + def test_get_conn_with_conn(self, conn_mock, hdfs3_mock): + """Tests get_conn call with specified connection.""" + + conn_mock.return_value = mock.Mock( + host='namenode', + login='hdfs_user', + port=8020, + extra_dejson={ + 'pars': { + 'dfs.namenode.logging.level': 'info' + }, + 'autoconf': False, + 'ticket_cache': '/path/to/cache' + }) + + with HdfsHook(hdfs_conn_id='hdfs_default') as hook: + hook.get_conn() + + conn_mock.assert_called_once_with('hdfs_default') + + hdfs3_mock.assert_called_once_with( + host='namenode', + port=8020, + pars={'dfs.namenode.logging.level': 'info'}, + user='hdfs_user', + autoconf=False, + ticket_cache='/path/to/cache') + + @mock.patch.object(hdfs3, 'HDFileSystem') + @mock.patch.object(HdfsHook, 'get_connection') + def test_get_conn_with_conn_ha(self, conn_mock, hdfs3_mock): + """Tests get_conn call with connection containing HA config.""" + + conn_mock.return_value = mock.Mock( + host='namenode', + login='hdfs_user', + port=8020, + extra_dejson={ + 'pars': { + 'dfs.namenode.logging.level': 'info' + }, + "ha": { + "host": "ns1", + "conf": { + "dfs.nameservices": "ns1", + "dfs.ha.namenodes.ns1": "nn1,nn2", + "dfs.namenode.rpc-address.ns1.nn1": "host1:8020", + "dfs.namenode.rpc-address.ns1.nn2": "host2:8020", + "dfs.namenode.http-address.ns1.nn1": "host1:50070", + "dfs.namenode.http-address.ns1.nn2": "host2:50070" + } + } + }) + + with HdfsHook(hdfs_conn_id='hdfs_default') as hook: + hook.get_conn() + + conn_mock.assert_called_once_with('hdfs_default') + + hdfs3_mock.assert_called_once_with( + host='ns1', + port=8020, + pars={ + 'dfs.namenode.logging.level': 'info', + 'dfs.nameservices': 'ns1', + 'dfs.ha.namenodes.ns1': 'nn1,nn2', + 'dfs.namenode.rpc-address.ns1.nn1': 'host1:8020', + 'dfs.namenode.rpc-address.ns1.nn2': 'host2:8020', + 'dfs.namenode.http-address.ns1.nn1': 'host1:50070', + 'dfs.namenode.http-address.ns1.nn2': 'host2:50070' + }, + user='hdfs_user', + autoconf=True) + + @mock.patch.object(configuration.conf, 'get') + @mock.patch.object(hdfs3, 'HDFileSystem') + @mock.patch.object(HdfsHook, 'get_connection') + def test_kerberos(self, conn_mock, hdfs3_mock, conf_mock): + """Tests setting kerberos auth from Airflow config.""" + + conn_mock.return_value = mock.Mock( + host='namenode', + login='hdfs_user', + port=8020, + extra_dejson={}) + + conf_mock.return_value = 'kerberos' + + with HdfsHook(hdfs_conn_id='hdfs_default') as hook: + hook.get_conn() + + hdfs3_mock.assert_called_once_with( + host='namenode', + port=8020, + user='hdfs_user', + pars={'hadoop.security.authentication': 'kerberos'}, + autoconf=True) + + @mock.patch.object(configuration.conf, 'get') + @mock.patch.object(hdfs3, 'HDFileSystem') + def test_kerberos_wo_conn(self, hdfs3_mock, conf_mock): + """Tests setting kerberos auth from Airflow config without conn.""" + + conf_mock.return_value = 'kerberos' + + with HdfsHook() as hook: + hook.get_conn() + + hdfs3_mock.assert_called_once_with( + pars={'hadoop.security.authentication': 'kerberos'}, + autoconf=True + ) + + @mock.patch.object(hdfs3, 'HDFileSystem') + @mock.patch.object(HdfsHook, 'get_connection') + def test_get_conn_with_empty_conn(self, conn_mock, hdfs3_mock): + """Tests get_conn call with empty connection.""" + + conn_mock.return_value = mock.Mock( + host='', + login='', + port='', + extra_dejson={}) + + with HdfsHook(hdfs_conn_id='hdfs_default') as hook: + hook.get_conn() + + conn_mock.assert_called_once_with('hdfs_default') + + hdfs3_mock.assert_called_once_with( + host=MyNone, + port=MyNone, + pars={}, + autoconf=True) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/sensors/test_hdfs_sensor.py b/tests/sensors/test_hdfs_sensor.py index 26adeaa3ab60b..db2ac5092e571 100644 --- a/tests/sensors/test_hdfs_sensor.py +++ b/tests/sensors/test_hdfs_sensor.py @@ -16,76 +16,359 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import unittest from datetime import timedelta +import fnmatch +import posixpath +import unittest + +import mock + +from airflow import models +from airflow.sensors.hdfs_sensor import HdfsFileSensor, HdfsFolderSensor, HdfsHook + + +class MockHdfs3Client(object): + """Mock hdfs3 client for testing purposes.""" + + def __init__(self, file_details): + self._file_details = {entry["name"]: entry for entry in file_details} + + @classmethod + def from_file_details(cls, file_details, test_instance, conn_id="hdfs_default"): + """Builds mock hdsf3 client using file_details.""" + + mock_client = cls(file_details) + mock_params = models.Connection(conn_id=conn_id) + + # Setup mock for get_connection. + patcher = mock.patch.object( + HdfsHook, "get_connection", return_value=mock_params + ) + test_instance.addCleanup(patcher.stop) + patcher.start() -from airflow import configuration -from airflow.exceptions import AirflowSensorTimeout -from airflow.sensors.hdfs_sensor import HdfsSensor -from airflow.utils.timezone import datetime -from tests.core import FakeHDFSHook + # Setup mock for get_conn. + patcher = mock.patch.object(HdfsHook, "get_conn", return_value=mock_client) + test_instance.addCleanup(patcher.stop) + patcher.start() -configuration.load_test_config() + return mock_client, mock_params -DEFAULT_DATE = datetime(2015, 1, 1) -TEST_DAG_ID = 'unit_test_dag' + def glob(self, pattern): + """Returns glob of files matching pattern.""" + # Implements a non-recursive glob on file names. + pattern_dir = posixpath.dirname(pattern) + pattern_base = posixpath.basename(pattern) -class HdfsSensorTests(unittest.TestCase): + # Ensure pattern_dir ends with '/'. + pattern_dir = posixpath.join(pattern_dir, "") + + for file_path in self._file_details.keys(): + file_basename = posixpath.basename(file_path) + if file_path.startswith(pattern_dir) and \ + fnmatch.fnmatch(file_basename, pattern_base): + yield file_path + + def isdir(self, path): + """Returns true if path is a directory.""" + + try: + details = self._file_details[path] + except KeyError: + raise IOError() + + return details["kind"] == "directory" + + def info(self, path): + """Returns info for given file path.""" + + try: + return self._file_details[path] + except KeyError: + raise IOError() + + +class HdfsFileSensorTests(unittest.TestCase): + """Tests for the HdfsFileSensor class.""" def setUp(self): - self.hook = FakeHDFSHook + file_details = [ + {"kind": "directory", "name": "/data/empty", "size": 0}, + {"kind": "directory", "name": "/data/not_empty", "size": 0}, + {"kind": "file", "name": "/data/not_empty/small.txt", "size": 10}, + {"kind": "file", "name": "/data/not_empty/large.txt", "size": 10000000}, + {"kind": "file", "name": "/data/not_empty/file.txt._COPYING_"}, + ] - def test_legacy_file_exist(self): - """ - Test the legacy behaviour - :return: - """ - # When - task = HdfsSensor(task_id='Should_be_file_legacy', - filepath='/datadirectory/datafile', - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - task.execute(None) - - # Then - # Nothing happens, nothing is raised exec is ok - - def test_legacy_file_exist_but_filesize(self): + self._mock_client, self._mock_params = MockHdfs3Client.from_file_details( + file_details, test_instance=self + ) + + self._default_task_kws = { + "timeout": 1, + "retry_delay": timedelta(seconds=1), + "poke_interval": 1, + } + + def test_existing_file(self): + """Tests poking for existing file.""" + + task = HdfsFileSensor( + task_id="existing_file", + pattern="/data/not_empty/small.txt", + **self._default_task_kws + ) + self.assertTrue(task.poke(context={})) + + def test_existing_file_glob(self): + """Tests poking for existing file with glob.""" + + task = HdfsFileSensor( + task_id="existing_file", + pattern="/data/not_empty/*.txt", + **self._default_task_kws + ) + self.assertTrue(task.poke(context={})) + + def test_nonexisting_file(self): + """Tests poking for non-existing file.""" + + task = HdfsFileSensor( + task_id="nonexisting_file", + pattern="/data/not_empty/random.txt", + **self._default_task_kws + ) + self.assertFalse(task.poke(context={})) + + def test_nonexisting_file_glob(self): + """Tests poking for non-existing file with glob.""" + + task = HdfsFileSensor( + task_id="existing_file", + pattern="/data/not_empty/*.xml", + **self._default_task_kws + ) + self.assertFalse(task.poke(context={})) + + def test_nonexisting_path(self): + """Tests poking for non-existing path.""" + + task = HdfsFileSensor( + task_id="nonexisting_path", + pattern="/data/not_empty/small.txt", + **self._default_task_kws + ) + + with mock.patch.object(self._mock_client, "glob", side_effect=IOError): + self.assertFalse(task.poke(context={})) + + def test_file_filter_size_small(self): + """Tests poking for file while filtering for file size (too small).""" + + task = HdfsFileSensor( + task_id="existing_file_too_small", + pattern="/data/not_empty/small.txt", + min_size=1, + **self._default_task_kws + ) + self.assertFalse(task.poke(context={})) + + def test_file_filter_size_large(self): + """Tests poking for file while filtering for file size (large).""" + + task = HdfsFileSensor( + task_id="existing_file_large", + pattern="/data/not_empty/large.txt", + min_size=1, + **self._default_task_kws + ) + self.assertTrue(task.poke(context={})) + + def test_file_filter_ext(self): + """Tests poking for file while filtering for extension.""" + + task = HdfsFileSensor( + task_id="existing_file_large", + pattern="/data/not_empty/f*", + ignore_exts=("_COPYING_",), + **self._default_task_kws + ) + self.assertFalse(task.poke(context={})) + + +class HdfsFolderSensorTests(unittest.TestCase): + def setUp(self): + file_details = [ + {"kind": "directory", "name": "/empty", "size": 0}, + {"kind": "directory", "name": "/not_empty", "size": 0}, + {"kind": "file", "name": "/not_empty/a.tmp", "size": 2000000}, + {"kind": "directory", "name": "/nested", "size": 0}, + {"kind": "directory", "name": "/nested/a", "size": 0} + ] + + self._mock_client, self._mock_params = MockHdfs3Client.from_file_details( + file_details, test_instance=self + ) + + self._default_task_kws = { + "timeout": 1, + "retry_delay": timedelta(seconds=1), + "poke_interval": 1, + } + + self._mock_client = MockHdfs3Client(file_details) + self._mock_params = models.Connection(conn_id="hdfs_default") + + def test_empty_directory(self): + """Tests example with an empty directory.""" + + task = HdfsFolderSensor(task_id="test_empty_directory", pattern="/empty") + self.assertTrue(task.poke(context={})) + + def test_non_empty_directory(self): + """Tests example with a non-empty directory.""" + + task = HdfsFolderSensor(task_id="test_non_empty_directory", pattern="/empty") + self.assertTrue(task.poke(context={})) + + def test_required_empty_directory(self): + """Tests example with an empty directory and require_empty = True.""" + + task = HdfsFolderSensor( + task_id="test_empty_directory", pattern="/empty", require_empty=True + ) + self.assertTrue(task.poke(context={})) + + def test_required_empty_directory_fail(self): + """Tests example with a non-empty directory and require_empty = True.""" + + task = HdfsFolderSensor( + task_id="test_non_empty_directory", pattern="/not_empty", require_empty=True + ) + self.assertFalse(task.poke(context={})) + + def test_required_not_empty_directory(self): + """Tests example with a non-empty directory and + require_not_empty = True. """ - Test the legacy behaviour with the filesize - :return: + + task = HdfsFolderSensor( + task_id="test_required_not_empty_directory", + pattern="/not_empty", + require_not_empty=True, + ) + self.assertTrue(task.poke(context={})) + + def test_required_not_empty_directory_fail(self): + """Tests example with an empty directory and + require_not_empty = True. """ - # When - task = HdfsSensor(task_id='Should_be_file_legacy', - filepath='/datadirectory/datafile', - timeout=1, - file_size=20, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) - - def test_legacy_file_does_not_exists(self): + + task = HdfsFolderSensor( + task_id="test_required_not_empty_directory", + pattern="/empty", + require_not_empty=True, + ) + self.assertFalse(task.poke(context={})) + + def test_glob(self): + """Tests globbing for directories.""" + + task = HdfsFolderSensor(task_id="test_glob", pattern="/*") + self.assertTrue(task.poke(context={})) + + def test_glob_non_existing(self): + """Tests globbing for non-existing directories.""" + + task = HdfsFolderSensor( + task_id="test_glob_non_existing", pattern="/non-existing/*" + ) + self.assertFalse(task.poke(context={})) + + def test_glob_require_empty(self): + """Tests globbing with only empty dir and require_empty = True.""" + + task = HdfsFolderSensor( + task_id="test_glob_require_empty", pattern="/e*", require_empty=True + ) + self.assertTrue(task.poke(context={})) + + def test_glob_require_empty_fail(self): + """Tests globbing with non-empty dir and require_empty = True.""" + + task = HdfsFolderSensor( + task_id="test_glob_require_empty_fail", pattern="/*", require_empty=True + ) + self.assertFalse(task.poke(context={})) + + def test_glob_require_not_empty(self): + """Tests globbing with only non-empty dir and + require_not_empty = True. """ - Test the legacy behaviour - :return: + + task = HdfsFolderSensor( + task_id="test_glob_require_not_empty", pattern="/n*", require_not_empty=True + ) + self.assertTrue(task.poke(context={})) + + def test_glob_require_not_empty_fail(self): + """Tests globbing with empty dir and require_not_empty = True.""" + + task = HdfsFolderSensor( + task_id="test_glob_require_not_empty_fail", + pattern="/*", + require_not_empty=True, + ) + self.assertFalse(task.poke(context={})) + + def test_sub_pattern(self): + """Tests filtering directory with matching sub_pattern.""" + + task = HdfsFolderSensor( + task_id="test_sub_pattern", + pattern="/not_empty", + require_not_empty=True, + sub_pattern="*.tmp" + ) + self.assertTrue(task.poke(context={})) + + def test_sub_pattern_no_match(self): + """Tests filtering directory with sub_pattern that doesn't match.""" + + task = HdfsFolderSensor( + task_id="test_sub_pattern_no_match", + pattern="/not_empty", + require_not_empty=True, + sub_pattern="*.txt" + ) + self.assertFalse(task.poke(context={})) + + def test_sub_pattern_subdir(self): + """Tests filtering directory with sub_pattern that matches subdir.""" + + task = HdfsFolderSensor( + task_id="test_sub_pattern_subdir", + pattern="/nested", + require_not_empty=True, + sub_pattern="a" + ) + self.assertTrue(task.poke(context={})) + + def test_sub_pattern_subdir_no_match(self): + """Tests filtering directory with sub_pattern that doesn't + match subdir. """ - task = HdfsSensor(task_id='Should_not_be_file_legacy', - filepath='/datadirectory/not_existing_file_or_directory', - timeout=1, - retry_delay=timedelta(seconds=1), - poke_interval=1, - hook=self.hook) - - # When - # Then - with self.assertRaises(AirflowSensorTimeout): - task.execute(None) + + task = HdfsFolderSensor( + task_id="test_sub_pattern_subdir_no_match", + pattern="/nested", + require_not_empty=True, + sub_pattern="b" + ) + self.assertFalse(task.poke(context={})) + + +if __name__ == "__main__": + unittest.main()