Repository: piskvorky/smart_open
Pull request #31 (2015-08-26 09:39:25)

Added writing to webhdfs support
This PR adds ability to write directly to webhdfs.

**NOTE:**
Please do not merge yet

**TODO:**
 - tests
 - ~~investigation of `close`,` __enter__` and `__exit__` methods~~


Patch:
diff --git a/README.rst b/README.rst
index 8ec8b76..33b4576 100644
--- a/README.rst
+++ b/README.rst
@@ -45,6 +45,11 @@ It is well tested (using `moto <https://github.com/spulec/moto>`_), well documen
   ...     for line in ['first line', 'second line', 'third line']:
   ...          fout.write(line + '\n')
 
+  >>> # stream content *into* WebHDFS (write mode):
+  >>> with smart_open.smart_open('webhdfs://host:port/user/hadoop/my_file.txt', 'wb') as fout:
+  ...     for line in ['first line', 'second line', 'third line']:
+  ...          fout.write(line + '\n')
+
   >>> # stream from/to local compressed files:
   >>> for line in smart_open.smart_open('./foo.txt.gz'):
   ...    print line
diff --git a/smart_open/smart_open_lib.py b/smart_open/smart_open_lib.py
index f458321..bbe11e1 100644
--- a/smart_open/smart_open_lib.py
+++ b/smart_open/smart_open_lib.py
@@ -27,6 +27,10 @@ import os
 import subprocess
 import sys
 import requests
+if sys.version_info[0] == 2:
+    import httplib
+elif sys.version_info[0] == 3:
+    import http.client as httplib
 
 from boto.compat import BytesIO, urlsplit, six
 import boto.s3.key
@@ -34,6 +38,7 @@ import boto.s3.key
 logger = logging.getLogger(__name__)
 
 S3_MIN_PART_SIZE = 50 * 1024**2  # minimum part size for S3 multipart uploads
+WEBHDFS_MIN_PART_SIZE = 50 * 1024**2  # minimum part size for HDFS multipart uploads
 
 
 def smart_open(uri, mode="rb", **kw):
@@ -105,6 +110,8 @@ def smart_open(uri, mode="rb", **kw):
             outkey = boto.s3.key.Key(outbucket)
             outkey.key = parsed_uri.key_id
             return S3OpenWrite(outbucket, outkey, **kw)
+        elif parsed_uri.scheme in ("webhdfs", ):
+            return WebHdfsOpenWrite(parsed_uri, **kw)
         else:
             raise NotImplementedError("write mode not supported for %r scheme", parsed_uri.scheme)
     else:
@@ -462,6 +469,85 @@ class S3OpenWrite(object):
             raise
 
 
+class WebHdfsOpenWrite(object):
+    """
+    Context manager for writing into webhdfs files
+
+    """
+    def __init__(self, parsed_uri, min_part_size=WEBHDFS_MIN_PART_SIZE):
+        if parsed_uri.scheme not in ("webhdfs"):
+            raise TypeError("can only process WebHDFS files")
+        self.parsed_uri = parsed_uri
+        self.closed = False
+        self.min_part_size = min_part_size
+        # creating empty file first
+        payload = {"op": "CREATE", "overwrite": True}
+        init_response = requests.put("http://" + self.parsed_uri.uri_path, params=payload, allow_redirects=False)
+        if not init_response.status_code == httplib.TEMPORARY_REDIRECT:
+            raise WebHdfsException(str(init_response.status_code) + "\n" + init_response.content)
+        uri = init_response.headers['location']
+        response = requests.put(uri, data="", headers={'content-type': 'application/octet-stream'})
+        if not response.status_code == httplib.CREATED:
+            raise WebHdfsException(str(response.status_code) + "\n" + response.content)
+        self.lines = []
+        self.parts = 0
+        self.chunk_bytes = 0
+        self.total_size = 0
+
+    def upload(self, data):
+        payload = {"op": "APPEND"}
+        init_response = requests.post("http://" + self.parsed_uri.uri_path, params=payload, allow_redirects=False)
+        if not init_response.status_code == httplib.TEMPORARY_REDIRECT:
+            raise WebHdfsException(str(init_response.status_code) + "\n" + init_response.content)
+        uri = init_response.headers['location']
+        response = requests.post(uri, data=data, headers={'content-type': 'application/octet-stream'})
+        if not response.status_code == httplib.OK:
+            raise WebHdfsException(str(response.status_code) + "\n" + response.content)
+
+    def write(self, b):
+        """
+        Write the given bytes (binary string) into the WebHDFS file from constructor.
+
+        """
+        if self.closed:
+            raise ValueError("I/O operation on closed file")
+        if isinstance(b, six.text_type):
+            # not part of API: also accept unicode => encode it as utf8
+            b = b.encode('utf8')
+
+        if not isinstance(b, six.binary_type):
+            raise TypeError("input must be a binary string")
+
+        self.lines.append(b)
+        self.chunk_bytes += len(b)
+        self.total_size += len(b)
+
+        if self.chunk_bytes >= self.min_part_size:
+            buff = b"".join(self.lines)
+            logger.info("uploading part #%i, %i bytes (total %.3fGB)" % (self.parts, len(buff), self.total_size / 1024.0 ** 3))
+            self.upload(buff)
+            logger.debug("upload of part #%i finished" % self.parts)
+            self.parts += 1
+            self.lines, self.chunk_bytes = [], 0
+
+    def seek(self, offset, whence=None):
+        raise NotImplementedError("seek() not implemented yet")
+
+    def close(self):
+        buff = b"".join(self.lines)
+        if buff:
+            logger.info("uploading last part #%i, %i bytes (total %.3fGB)" % (self.parts, len(buff), self.total_size / 1024.0 ** 3))
+            self.upload(buff)
+            logger.debug("upload of last part #%i finished" % self.parts)
+        self.closed = True
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+
 def s3_iter_bucket_process_key(key):
     """
     Conceptually part of `s3_iter_bucket`, but must remain top-level method because
@@ -552,3 +638,9 @@ def s3_iter_lines(key):
     # process the last line, too
     if buf:
         yield buf
+
+
+class WebHdfsException(Exception):
+    def __init__(self, msg=str()):
+        self.msg = msg
+        super(WebHdfsException, self).__init__(self.msg)


Test patch:
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index 3b09d09..1cdc143 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -420,6 +420,52 @@ class S3OpenWriteTest(unittest.TestCase):
         self.assertEqual(output, [b"testtest\n", b"test"])
 
 
+class WebHdfsWriteTest(unittest.TestCase):
+    """
+    Test writing into webhdfs files.
+
+    """
+
+    @responses.activate
+    def test_initialize_write(self):
+        def request_callback(request):
+            resp_body = ""
+            headers = {'location': 'http://127.0.0.1:8440/file'}
+            return (307, headers, resp_body)
+        responses.add_callback(responses.PUT, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
+        responses.add(responses.PUT, "http://127.0.0.1:8440/file", status=201)
+        smart_open_object = smart_open.WebHdfsOpenWrite(smart_open.ParseUri("webhdfs://127.0.0.1:8440/path/file"))
+        assert len(responses.calls) == 2
+        path, params = responses.calls[0].request.url.split("?")
+        assert path == "http://127.0.0.1:8440/webhdfs/v1/path/file"
+        assert params == "overwrite=True&op=CREATE" or params == "op=CREATE&overwrite=True"
+        assert responses.calls[1].request.url == "http://127.0.0.1:8440/file"
+
+    @responses.activate
+    def test_write(self):
+        def request_callback(request):
+            resp_body = ""
+            headers = {'location': 'http://127.0.0.1:8440/file'}
+            return (307, headers, resp_body)
+
+        responses.add_callback(responses.PUT, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
+        responses.add(responses.PUT, "http://127.0.0.1:8440/file", status=201)
+        smart_open_object = smart_open.WebHdfsOpenWrite(smart_open.ParseUri("webhdfs://127.0.0.1:8440/path/file"))
+
+        def write_callback(request):
+            assert request.body == u"žluťoučký koníček".encode('utf8')
+            headers = {}
+            return (200, headers, "")
+        test_string = u"žluťoučký koníček".encode('utf8')
+        responses.add_callback(responses.POST, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
+        responses.add_callback(responses.POST, "http://127.0.0.1:8440/file", callback=write_callback)
+        smart_open_object.write(test_string)
+        smart_open_object.close()
+        assert len(responses.calls) == 4
+        assert responses.calls[2].request.url == "http://127.0.0.1:8440/webhdfs/v1/path/file?op=APPEND"
+        assert responses.calls[3].request.url == "http://127.0.0.1:8440/file"
+
+
 class S3IterBucketTest(unittest.TestCase):
     """
     Test parallel iteration of given bucket.