Support YARN endpoints protected by Kerberos/SPNEGO

Author: lresende

.travis.yml

@@ -13,7 +13,11 @@ matrix:
 
 
 install:
-  - pip install tox coveralls
+  - pip install --upgrade setuptools pip
+  - pip install --upgrade requests requests_mock requests_kerberos tox coveralls
+  - python setup.py bdist_wheel
+  - pip install dist/*.whl
+  - pip freeze
 
 script:
   tox

setup.py

@@ -19,17 +19,17 @@ def find_version(*file_paths):
         return version_match.group(1)
     raise RuntimeError("Unable to find version string.")
 
-
-install_requires = []
-
 setup(
     name = 'yarn-api-client',
     version = find_version('yarn_api_client', '__init__.py'),
     description='Python client for Hadoop® YARN API',
     long_description=read('README.rst'),
-    packages = find_packages(exclude=['tests']),
+    packages = find_packages(exclude=['tests','itests']),
 
-    install_requires = install_requires,
+    install_requires = [
+        'requests>=2.7,<3.0',
+        'requests-kerberos==0.12.0',
+    ],
     entry_points = {
         'console_scripts': [
             'yarn_client = yarn_api_client.main:main',

tests/test_base.py

@@ -4,56 +4,66 @@
 except ImportError:
     from http.client import OK
 
-from mock import patch
-from tests import TestCase
+import json
+import requests
+import requests_mock
 
+from tests import TestCase
 from yarn_api_client import base
 from yarn_api_client.errors import APIError, ConfigurationError
 
 
 class BaseYarnAPITestCase(TestCase):
-    def test_request(self):
-        client = self.get_client()
-        with patch('yarn_api_client.base.HTTPConnection') as http_conn_mock:
-            with patch('yarn_api_client.base.json'):
-                http_conn_mock().getresponse().status = OK
+    @staticmethod
+    def success_response():
+        return {
+            'status':'success'
+        }
 
-                client.request('/ololo', foo='bar')
+    def test_valid_request(self):
+        with requests_mock.mock() as requests_get_mock:
+            requests_get_mock.get('/ololo', text=json.dumps(BaseYarnAPITestCase.success_response()))
 
-                http_conn_mock().request.assert_called_with('GET', '/ololo?foo=bar')
+            client = self.get_client()
+            response = client.request('/ololo', foo='bar')
 
-                http_conn_mock.reset_mock()
-                client.request('/ololo')
+            assert requests_get_mock.called
+            self.assertIn(response.data['status'], 'success')
+
+
+    def test_valid_request_with_parameters(self):
+        with requests_mock.mock() as requests_get_mock:
+            requests_get_mock.get('/ololo?foo=bar', text=json.dumps(BaseYarnAPITestCase.success_response()))
 
-                http_conn_mock()
+            client = self.get_client()
+            response = client.request('/ololo', foo='bar')
 
-                http_conn_mock().request.assert_called_with('GET', '/ololo')
+            assert requests_get_mock.called
+            self.assertIn(response.data['status'], 'success')
 
     def test_bad_request(self):
-        client = self.get_client()
-        with patch('yarn_api_client.base.HTTPConnection') as http_conn_mock:
-            http_conn_mock().getresponse().status = 404
+        with requests_mock.mock() as requests_get_mock:
+            requests_get_mock.get('/ololo', status_code=404)
 
+            client = self.get_client()
             with self.assertRaises(APIError):
                 client.request('/ololo')
-            
-    def test_http_configuration(self):
-        client = self.get_client()
-        client.address = None
-        client.port = 80
 
-        with self.assertRaises(ConfigurationError):
-            conn = client.http_conn
+    def test_http_configuration(self):
+        with requests_mock.mock() as requests_get_mock:
+            requests_get_mock.get('/ololo', text=json.dumps(BaseYarnAPITestCase.success_response()))
 
-        client.address = 'localhost'
-        client.port = None
+            client = self.get_client()
+            client.address = None
+            client.port = 80
 
-        with self.assertRaises(ConfigurationError):
-            conn = client.http_conn
+            with self.assertRaises(ConfigurationError):
+                client.request('/ololo')
 
     def get_client(self):
         client = base.BaseYarnAPI()
         client.address = 'example.com'
         client.port = 80
         client.timeout = 0
+        client.kerberos_enabled = False
         return client

tox.ini

@@ -1,15 +1,11 @@
 [tox]
-envlist = py26,py27,py34,py35
+envlist = py27,py35,py36
 
 [testenv]
 deps =
+     requests
+     requests-kerberos
+     requests_mock
      coverage
      mock
 commands = coverage run --source=yarn_api_client setup.py test
-
-[testenv:py26]
-deps =
-     argparse
-     coverage
-     mock
-     unittest2

yarn_api_client/application_master.py

@@ -18,9 +18,10 @@ class ApplicationMaster(BaseYarnAPI):
     :param str address: Proxy HTTP address
     :param int port: Proxy HTTP port
     :param int timeout: API connection timeout in seconds
+    :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
-    def __init__(self, address=None, port=8088, timeout=30):
-        self.address, self.port, self.timeout = address, port, timeout
+    def __init__(self, address=None, port=8088, timeout=30, kerberos_enabled=False):
+        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
         if address is None:
             self.logger.debug('Get configuration from hadoop conf dir')
             address, port = get_webproxy_host_port()

yarn_api_client/base.py

@@ -1,60 +1,72 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
-try:
-    from httplib import HTTPConnection, OK
-except ImportError:
-    from http.client import HTTPConnection, OK
+
 import json
 import logging
-try:
-    from urllib import urlencode
-except ImportError:
-    from urllib.parse import urlencode
+import requests
 
+from requests_kerberos import HTTPKerberosAuth
 from .errors import APIError, ConfigurationError
 
 
 class Response(object):
-    def __init__(self, http_response):
-        self.data = json.load(http_response)
+    def __init__(self, response):
+        self.data = response.json()
 
 
 class BaseYarnAPI(object):
+    __logger = None
     response_class = Response
 
+    def _validate_configuration(self):
+        if self.address is None:
+            raise ConfigurationError('API address is not set')
+        elif self.port is None:
+            raise ConfigurationError('API port is not set')
+
     def request(self, api_path, **query_args):
-        params = urlencode(query_args)
-        if params:
-            path = api_path + '?' + params
+        params = query_args
+        api_endpoint = 'http://{}:{}{}'.format(self.address, self.port, api_path)
+
+        self.logger.info('API Endpoint {}'.format(api_endpoint))
+
+        self._validate_configuration()
+
+        response = None
+        if self.kerberos_enabled:
+            response = requests.get(api_endpoint, params, auth=HTTPKerberosAuth())
         else:
-            path = api_path
+            response = requests.get(api_endpoint, params)
 
-        self.logger.info('Request http://%s:%s%s', self.address, self.port, path)
-        
-        http_conn = self.http_conn
-        http_conn.request('GET', path)
-        response = http_conn.getresponse()
+        if response.status_code == requests.codes.ok:
+            return self.response_class(response)
+        else:
+            msg = 'Response finished with status: %s. Details: %s' % (response.status_code, response.text)
+            raise APIError(msg)
 
-        if response.status == OK:
+    def update(self, api_path, data):
+        api_endpoint = 'http://{}:{}{}'.format(self.address, self.port, api_path)
+
+        self.logger.info('API Endpoint {}'.format(api_endpoint))
+
+        self._validate_configuration()
+
+        response = None
+        if self.kerberos_enabled:
+            response = requests.put(api_endpoint, data=data, auth=HTTPKerberosAuth())
+        else:
+            response = requests.put(api_endpoint, data=data)
+
+        if response.status_code == requests.codes.ok:
             return self.response_class(response)
         else:
-            explanation = response.read()
-            msg = 'Response finished with status: %s. Details: %s' % (response.status, explanation)
+            msg = 'Response finished with status: %s. Details: %s' % (response.status_code, response.text)
             raise APIError(msg)
 
     def construct_parameters(self, arguments):
         params = dict((key, value) for key, value in arguments if value is not None)
         return params
 
-    @property
-    def http_conn(self):
-        if self.address is None:
-            raise ConfigurationError('API address is not set')
-        elif self.port is None:
-            raise ConfigurationError('API port is not set')
-        return HTTPConnection(self.address, self.port, timeout=self.timeout)
-
-    __logger = None
     @property
     def logger(self):
         if self.__logger is None:

yarn_api_client/history_server.py

@@ -18,9 +18,10 @@ class HistoryServer(BaseYarnAPI):
     :param str address: HistoryServer HTTP address
     :param int port: HistoryServer HTTP port
     :param int timeout: API connection timeout in seconds
+    :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
-    def __init__(self, address=None, port=19888, timeout=30):
-        self.address, self.port, self.timeout = address, port, timeout
+    def __init__(self, address=None, port=19888, timeout=30, kerberos_enabled=False):
+        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
         if address is None:
             self.logger.debug('Get information from hadoop conf dir')
             address, port = get_jobhistory_host_port()

yarn_api_client/main.py

@@ -256,7 +256,7 @@ def main():
         method_args = [getattr(opts, arg) for arg in opts.method_args]
     else:
         method_args = []
-    # Construce key arguments for method
+    # Construct key arguments for method
     if 'method_kwargs' in opts:
         method_kwargs = dict((key, getattr(opts, key)) for key in opts.method_kwargs)
     else:

yarn_api_client/node_manager.py

@@ -12,9 +12,10 @@ class NodeManager(BaseYarnAPI):
     :param str address: NodeManager HTTP address
     :param int port: NodeManager HTTP port
     :param int timeout: API connection timeout in seconds
+    :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
-    def __init__(self, address=None, port=8042, timeout=30):
-        self.address, self.port, self.timeout = address, port, timeout
+    def __init__(self, address=None, port=8042, timeout=30, kerberos_enabled=False):
+        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
 
     def node_information(self):
         """

yarn_api_client/resource_manager.py

@@ -19,9 +19,10 @@ class ResourceManager(BaseYarnAPI):
     :param str address: ResourceManager HTTP address
     :param int port: ResourceManager HTTP port
     :param int timeout: API connection timeout in seconds
+    :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
-    def __init__(self, address=None, port=8088, timeout=30):
-        self.address, self.port, self.timeout = address, port, timeout
+    def __init__(self, address=None, port=8088, timeout=30, kerberos_enabled=False):
+        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
         if address is None:
             self.logger.debug('Get configuration from hadoop conf dir')
             address, port = get_resource_manager_host_port()
@@ -132,7 +133,7 @@ def cluster_application_statistics(self, state_list=None,
             comma-separated list. If states is not provided, the API will
             enumerate all application states and return the counts of them.
         :param list application_type_list: types of the applications,
-            specified as a comma-separated list. If applicationTypes is not
+            specified as a comma-separated list. If application_types is not
             provided, the API will count the applications of any application
             type. In this case, the response shows * to indicate any
             application type. Note that we only support at most one
@@ -146,13 +147,13 @@ def cluster_application_statistics(self, state_list=None,
         # TODO: validate state argument
         states = ','.join(state_list) if state_list is not None else None
         if application_type_list is not None:
-            applicationTypes = ','.join(application_type_list)
+            application_types = ','.join(application_type_list)
         else:
-            applicationTypes = None
+            application_types = None
 
         loc_args = (
             ('states', states),
-            ('applicationTypes', applicationTypes))
+            ('applicationTypes', application_types))
         params = self.construct_parameters(loc_args)
 
         return self.request(path, **params)
@@ -184,6 +185,36 @@ def cluster_application_attempts(self, application_id):
 
         return self.request(path)
 
+    def cluster_application_state(self, application_id):
+        """
+        With the application state API, you can obtain the current
+        state of an application.
+
+        :param str application_id: The application id
+        :returns: API response object with JSON data
+        :rtype: :py:class:`yarn_api_client.base.Response`
+        """
+        path = '/ws/v1/cluster/apps/{appid}/state'.format(
+            appid=application_id)
+
+        return self.request(path)
+
+    def cluster_application_kill(self, application_id):
+        """
+        With the application kill API, you can kill an application
+        that is not in FINISHED or FAILED state.
+
+        :param str application_id: The application id
+        :returns: API response object with JSON data
+        :rtype: :py:class:`yarn_api_client.base.Response`
+        """
+
+        data = '{"state": "KILLED"}'
+        path = '/ws/v1/cluster/apps/{appid}/state'.format(
+            appid=application_id)
+
+        return self.put(path, data)
+
     def cluster_nodes(self, state=None, healthy=None):
         """
         With the Nodes API, you can obtain a collection of resources, each of