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_r333927159
##########
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)
Review comment:
Master is python3, so for new code please make this
```suggestion
return super().get_conn(tmp_headers)
```
(and we'll fix this up when we back-port it.)
----------------------------------------------------------------
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
