Repository: piskvorky/smart_open
Pull request #29 (2015-08-18 09:17:49)

Implemented possibility to read directly from WebHDFS using smart_open
This PR partially solves the issue https://github.com/piskvorky/smart_open/issues/23. It implements reading from WebHDFS based on API described at https://hadoop.apache.org/docs/r1.0.4/webhdfs.html.

For reading from WebHDFS you can use:
`webhdfs://host:port/path/file`

**NOTE:**
 - It does not support kerberos authentication yet (so it only works for WebHDFS that is not secured)
 - Please do not merge yet, I'll try to implement read and seek methods in this PR soon as well. 


Patch:
diff --git a/.travis.yml b/.travis.yml
index 6ce9767..a79f44e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,6 +8,6 @@ python:
 - "3.4"
 install:
 - if [[ $TRAVIS_PYTHON_VERSION == 2.6 ]]; then pip install ordereddict; fi
-- pip install mock moto
+- pip install mock moto responses
 - python setup.py install
 script: python -W ignore setup.py test
diff --git a/README.rst b/README.rst
index 9895cd2..8ec8b76 100644
--- a/README.rst
+++ b/README.rst
@@ -34,7 +34,11 @@ It is well tested (using `moto <https://github.com/spulec/moto>`_), well documen
 
   >>> # stream from HDFS
   >>> for line in smart_open.smart_open('hdfs://user/hadoop/my_file.txt'):
-  ...    print line
+  ...     print line
+
+  >>> # stream from WebHDFS
+  >>> for line in smart_open.smart_open('webhdfs://host:port/user/hadoop/my_file.txt'):
+  ...     print line
 
   >>> # stream content *into* S3 (write mode):
   >>> with smart_open.smart_open('s3://mybucket/mykey.txt', 'wb') as fout:
@@ -86,7 +90,7 @@ Or, if you prefer to install from the `source tar.gz <http://pypi.python.org/pyp
     python setup.py test  # run unit tests
     python setup.py install
 
-To run the unit tests (optional), you'll also need to install `mock <https://pypi.python.org/pypi/mock>`_ and `moto <https://github.com/spulec/moto>`_ (``pip install mock moto``). The tests are also run automatically with `Travis CI <https://travis-ci.org/piskvorky/smart_open>`_ on every commit push & pull request.
+To run the unit tests (optional), you'll also need to install `mock <https://pypi.python.org/pypi/mock>`_ , `moto <https://github.com/spulec/moto>`_ and `responses <https://github.com/getsentry/responses>` (``pip install mock moto responses``). The tests are also run automatically with `Travis CI <https://travis-ci.org/piskvorky/smart_open>`_ on every commit push & pull request.
 
 Todo
 ----
diff --git a/smart_open/smart_open_lib.py b/smart_open/smart_open_lib.py
index b21704d..abf1a71 100644
--- a/smart_open/smart_open_lib.py
+++ b/smart_open/smart_open_lib.py
@@ -26,6 +26,7 @@ import multiprocessing.pool
 import os
 import subprocess
 import sys
+import requests
 
 from boto.compat import BytesIO, urlsplit, six
 import boto.s3.key
@@ -93,6 +94,8 @@ def smart_open(uri, mode="rb"):
             return S3OpenRead(parsed_uri)
         elif parsed_uri.scheme in ("hdfs", ):
             return HdfsOpenRead(parsed_uri)
+        elif parsed_uri.scheme in ("webhdfs", ):
+            return WebHdfsOpenRead(parsed_uri)
         else:
             raise NotImplementedError("read mode not supported for %r scheme", parsed_uri.scheme)
     elif mode in ('w', 'wb'):
@@ -119,6 +122,8 @@ class ParseUri(object):
       * s3://my_bucket/my_key
       * s3://my_key:my_secret@my_bucket/my_key
       * hdfs:///path/file
+      * hdfs://path/file
+      * webhdfs://host:port/path/file
       * ./local/path/file
       * ./local/path/file.gz
       * file:///home/user/file
@@ -140,9 +145,15 @@ class ParseUri(object):
 
         if self.scheme == "hdfs":
             self.uri_path = parsed_uri.netloc + parsed_uri.path
+            self.uri_path = "/" + self.uri_path.lstrip("/")
 
             if not self.uri_path:
                 raise RuntimeError("invalid HDFS URI: %s" % uri)
+        elif self.scheme == "webhdfs":
+            self.uri_path = parsed_uri.netloc + "/webhdfs/v1" + parsed_uri.path
+
+            if not self.uri_path:
+                raise RuntimeError("invalid WebHDFS URI: %s" % uri)
         elif self.scheme in ("s3", "s3n"):
             self.bucket_id = (parsed_uri.netloc + parsed_uri.path).split('@')
             self.key_id = None
@@ -191,6 +202,7 @@ class S3OpenRead(object):
         self.read_key = s3_connection.get_bucket(parsed_uri.bucket_id).lookup(parsed_uri.key_id)
         if self.read_key is None:
             raise KeyError(parsed_uri.key_id)
+        self.line_generator = s3_iter_lines(self.read_key)
 
     def __iter__(self):
         s3_connection = boto.connect_s3(
@@ -246,7 +258,7 @@ class HdfsOpenRead(object):
         self.parsed_uri = parsed_uri
 
     def __iter__(self):
-        hdfs = subprocess.Popen(["hdfs", "dfs", "-cat", os.path.join("/", self.parsed_uri.uri_path)], stdout=subprocess.PIPE)
+        hdfs = subprocess.Popen(["hdfs", "dfs", "-cat", self.parsed_uri.uri_path], stdout=subprocess.PIPE)
         return hdfs.stdout
 
     def read(self, size=None):
@@ -262,6 +274,61 @@ class HdfsOpenRead(object):
         pass
 
 
+class WebHdfsOpenRead(object):
+    """
+    Implement streamed reader from WebHDFS, as an iterable & context manager.
+    NOTE: it does not support kerberos authentication yet
+
+    """
+    def __init__(self, parsed_uri):
+        if parsed_uri.scheme not in ("webhdfs"):
+            raise TypeError("can only process WebHDFS files")
+        self.parsed_uri = parsed_uri
+        self.offset = 0
+
+    def __iter__(self):
+        payload = {"op": "OPEN"}
+        response = requests.get("http://" + self.parsed_uri.uri_path, params=payload, stream=True)
+        return response.iter_lines()
+
+    def read(self, size=None):
+        """
+        Read the specific number of bytes from the file
+
+        Note read() and line iteration (`for line in self: ...`) each have their
+        own file position, so they are independent. Doing a `read` will not affect
+        the line iteration, and vice versa.
+        """
+        if not size or size < 0:
+            payload = {"op": "OPEN", "offset": self.offset}
+            self.offset = 0
+        else:
+            payload = {"op": "OPEN", "offset": self.offset, "length": size}
+            self.offset = self.offset + size
+        response = requests.get("http://" + self.parsed_uri.uri_path, params=payload, stream=True)
+        return response.content
+
+    def seek(self, offset, whence=0):
+        """
+        Seek to the specified position.
+
+        Only seeking to the beginning (offset=0) supported for now.
+
+        """
+        if whence == 0 and offset == 0:
+            self.offset = 0
+        elif whence == 0:
+            self.offset = offset
+        else:
+            raise NotImplementedError("operations with whence not implemented yet")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        pass
+
+
 def make_closing(base, **attrs):
     """
     Add support for `with Base(attrs) as fout:` to the base class if it's missing.


Test patch:
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index 58eb2c8..3b09d09 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -16,6 +16,7 @@ import os
 import boto
 import mock
 from moto import mock_s3
+import responses
 
 import smart_open
 from smart_open import smart_open_lib
@@ -77,6 +78,12 @@ class ParseUriTest(unittest.TestCase):
         # incorrect uri - only one '@' in uri is allowed
         self.assertRaises(RuntimeError, smart_open.ParseUri, "s3://access_id@access_secret@mybucket/mykey")
 
+    def test_webhdfs_uri(self):
+        """Do webhdfs URIs parse correctly"""
+        parsed_uri = smart_open.ParseUri("webhdfs://host:port/path/file")
+        self.assertEqual(parsed_uri.scheme, "webhdfs")
+        self.assertEqual(parsed_uri.uri_path, "host:port/webhdfs/v1/path/file")
+
 
 class SmartOpenReadTest(unittest.TestCase):
     """
@@ -92,7 +99,6 @@ class SmartOpenReadTest(unittest.TestCase):
         # called with the correct path?
         mock_smart_open.assert_called_with("/tmp/test.txt", "rb")
 
-
     # couldn't find any project for mocking up HDFS data
     # TODO: we want to test also a content of the files, not just fnc call params
     @mock.patch('smart_open.smart_open_lib.subprocess')
@@ -109,6 +115,22 @@ class SmartOpenReadTest(unittest.TestCase):
         smart_open_object.__iter__()
         mock_subprocess.Popen.assert_called_with(["hdfs", "dfs", "-cat", "/tmp/test.txt"], stdout=mock_subprocess.PIPE)
 
+    @responses.activate
+    def test_webhdfs(self):
+        """Is webhdfs line iterator called correctly"""
+        responses.add(responses.GET, "http://127.0.0.1:8440/webhdfs/v1/path/file", body='line1\nline2')
+        smart_open_object = smart_open.WebHdfsOpenRead(smart_open.ParseUri("webhdfs://127.0.0.1:8440/path/file"))
+        iterator = iter(smart_open_object)
+        self.assertEqual(next(iterator).decode("utf-8"), "line1")
+        self.assertEqual(next(iterator).decode("utf-8"), "line2")
+
+    @responses.activate
+    def test_webhdfs_read(self):
+        """Does webhdfs read method work correctly"""
+        responses.add(responses.GET, "http://127.0.0.1:8440/webhdfs/v1/path/file", body='line1\nline2')
+        smart_open_object = smart_open.WebHdfsOpenRead(smart_open.ParseUri("webhdfs://127.0.0.1:8440/path/file"))
+        self.assertEqual(smart_open_object.read().decode("utf-8"), "line1\nline2")
+
     @mock.patch('smart_open.smart_open_lib.boto')
     @mock.patch('smart_open.smart_open_lib.s3_iter_lines')
     def test_s3_boto(self, mock_s3_iter_lines, mock_boto):