ashb commented on a change in pull request #6090: [AIRFLOW-5470] Add Apache
Livy REST operator
URL: https://github.com/apache/airflow/pull/6090#discussion_r333927564
##########
File path: airflow/contrib/hooks/livy_hook.py
##########
@@ -47,66 +46,98 @@ class BatchState(Enum):
SUCCESS = 'success'
-TERMINAL_STATES = {
- BatchState.SUCCESS,
- BatchState.DEAD,
- BatchState.KILLED,
- BatchState.ERROR,
-}
-
-
-class LivyHook(BaseHook, LoggingMixin):
+class LivyHook(HttpHook, LoggingMixin):
"""
Hook for Apache Livy through the REST API.
- For more information about the API refer to
- https://livy.apache.org/docs/latest/rest-api.html
-
:param livy_conn_id: reference to a pre-defined Livy Connection.
:type livy_conn_id: str
+
+ .. seealso::
+ For more details refer to the Apache Livy API reference:
+ https://livy.apache.org/docs/latest/rest-api.html
"""
- def __init__(self, livy_conn_id='livy_default'):
- super(LivyHook, self).__init__(livy_conn_id)
- self._livy_conn_id = livy_conn_id
- self._build_base_url()
- def _build_base_url(self):
- """
- Build connection URL
- """
- params = self.get_connection(self._livy_conn_id)
+ TERMINAL_STATES = {
+ BatchState.SUCCESS,
+ BatchState.DEAD,
+ BatchState.KILLED,
+ BatchState.ERROR,
+ }
+
+ _def_headers = {
+ 'Content-Type': 'application/json',
+ 'Accept': 'application/json'
+ }
- base_url = params.host
+ def __init__(self, livy_conn_id='livy_default'):
+ super(LivyHook, self).__init__(http_conn_id=livy_conn_id)
- if not base_url:
- raise AirflowException("Missing Livy endpoint hostname")
+ def get_conn(self, headers=None):
+ """
+ Returns http session for use with requests
- if '://' not in base_url:
- base_url = '{}://{}'.format('http', base_url)
- if not re.search(r':\d+$', base_url):
- base_url = '{}:{}'.format(base_url, str(params.port or 8998))
+ :param headers: additional headers to be passed through as a dictionary
+ :type headers: dict
+ :return: requests session
+ :rtype: requests.Session
+ """
+ tmp_headers = self._def_headers.copy() # setting default headers
+ if headers:
+ tmp_headers.update(headers)
+ return super(LivyHook, self).get_conn(tmp_headers)
- self._base_url = base_url
+ def run_method(self, method='GET', endpoint=None, data=None, headers=None,
extra_options=None):
+ """
+ Wrapper for HttpHook, allows to change method on the same HttpHook
+
+ :param method: http method
+ :type method: str
+ :param endpoint: endpoint
+ :type endpoint: str
+ :param data: request payload
+ :type data: dict
+ :param headers: headers
+ :type headers: dict
+ :param extra_options: extra options
+ :type extra_options: dict
+ :return: http response
+ :rtype: requests.Response
+ """
+ if method not in ('GET', 'POST', 'PUT', 'DELETE', 'HEAD'):
+ raise AirflowException("Invalid http method '{}'".format(method))
- def get_conn(self):
- pass
+ back_method = self.method
+ self.method = method
+ result = self.run(endpoint, data, headers, extra_options)
+ self.method = back_method
Review comment:
You should wrap this in a `try`/`finally`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
