flink_sql.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539
  1. #!/usr/bin/env python
  2. # Licensed to Cloudera, Inc. under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. Cloudera, Inc. licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing, software
  13. # distributed under the License is distributed on an "AS IS" BASIS,
  14. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15. # See the License for the specific language governing permissions and
  16. # limitations under the License.
  17. import re
  18. import json
  19. import time
  20. import logging
  21. import posixpath
  22. from desktop.auth.backend import rewrite_user
  23. from desktop.lib.i18n import force_unicode
  24. from desktop.lib.rest.http_client import HttpClient, RestException
  25. from desktop.lib.rest.resource import Resource
  26. from notebook.connectors.base import Api, QueryError
  27. LOG = logging.getLogger()
  28. _JSON_CONTENT_TYPE = 'application/json'
  29. _API_VERSION = 'v3'
  30. SESSION_KEY = '%(username)s-%(connector_name)s'
  31. OPERATION_TOKEN = '%(username)s-%(connector_name)s' + '-operation-token'
  32. def query_error_handler(func):
  33. def decorator(*args, **kwargs):
  34. try:
  35. return func(*args, **kwargs)
  36. except RestException as e:
  37. try:
  38. message = force_unicode(json.loads(e.message)['errors'])
  39. except Exception:
  40. message = e.message
  41. message = force_unicode(message)
  42. raise QueryError(parse_error(message))
  43. except Exception as e:
  44. message = force_unicode(str(e))
  45. raise QueryError(message)
  46. return decorator
  47. def parse_error(error):
  48. lines = re.split(r'\\n', error)
  49. caused_by = [line for line in lines if 'Caused by:' in line]
  50. if len(caused_by) == 0:
  51. return error
  52. elif len(caused_by) >= 1:
  53. return caused_by[-1]
  54. class FlinkSqlApi(Api):
  55. def __init__(self, user, interpreter=None):
  56. Api.__init__(self, user, interpreter=interpreter)
  57. self.options = interpreter['options']
  58. api_url = self.options['url']
  59. self.db = FlinkSqlClient(user=user, api_url=api_url)
  60. @query_error_handler
  61. def create_session(self, lang=None, properties=None):
  62. session = self._get_session()
  63. response = {
  64. 'type': lang,
  65. 'id': session['id']
  66. }
  67. return response
  68. def _get_session_key(self):
  69. return SESSION_KEY % {
  70. 'username': self.user.username if hasattr(self.user, 'username') else self.user,
  71. 'connector_name': self.interpreter['name']
  72. }
  73. def _get_session_info_from_user(self):
  74. self.user = rewrite_user(self.user)
  75. session_key = self._get_session_key()
  76. if self.user.profile.data.get(session_key):
  77. return self.user.profile.data[session_key]
  78. def _set_session_info_to_user(self, session_info):
  79. self.user = rewrite_user(self.user)
  80. session_key = self._get_session_key()
  81. self.user.profile.update_data({session_key: session_info})
  82. self.user.profile.save()
  83. def _remove_session_info_from_user(self):
  84. self.user = rewrite_user(self.user)
  85. session_key = self._get_session_key()
  86. operation_token_key = self._get_operation_token_key()
  87. if self.user.profile.data.get(session_key):
  88. json_data = self.user.profile.data
  89. json_data.pop(session_key)
  90. json_data.pop(operation_token_key)
  91. self.user.profile.json_data = json.dumps(json_data)
  92. self.user.profile.save()
  93. def _get_operation_token_key(self):
  94. return OPERATION_TOKEN % {
  95. 'username': self.user.username if hasattr(self.user, 'username') else self.user,
  96. 'connector_name': self.interpreter['name']
  97. }
  98. def _get_operation_token_info_from_user(self, operation_handle):
  99. self.user = rewrite_user(self.user)
  100. operation_token_key = self._get_operation_token_key()
  101. if self.user.profile.data.get(operation_token_key):
  102. return self.user.profile.data[operation_token_key][operation_handle]
  103. def _set_operation_token_info_to_user(self, operation_handle, token):
  104. self.user = rewrite_user(self.user)
  105. operation_token_key = self._get_operation_token_key()
  106. json_data = self.user.profile.data
  107. if self.user.profile.data.get(operation_token_key) is None:
  108. json_data[operation_token_key] = {}
  109. json_data[operation_token_key][operation_handle] = token
  110. self.user.profile.update_data(json_data)
  111. self.user.profile.save()
  112. def _remove_operation_token_info_from_user(self, operation_handle):
  113. self.user = rewrite_user(self.user)
  114. operation_token_key = self._get_operation_token_key()
  115. if self.user.profile.data.get(operation_token_key):
  116. json_data = self.user.profile.data
  117. json_data[operation_token_key].pop(operation_handle)
  118. self.user.profile.json_data = json.dumps(json_data)
  119. self.user.profile.save()
  120. def _get_session(self):
  121. session = self._get_session_info_from_user()
  122. if not session:
  123. session = self.db.create_session()
  124. try:
  125. self.db.session_heartbeat(session_handle=session['sessionHandle'])
  126. except Exception as e:
  127. if "Session '%(sessionHandle)s' does not exist" % session in str(e):
  128. LOG.warning('Session %(sessionHandle)s does not exist, opening a new one' % session)
  129. session = self.db.create_session()
  130. else:
  131. raise e
  132. session['id'] = session['sessionHandle']
  133. self._set_session_info_to_user(session)
  134. return session
  135. @query_error_handler
  136. def execute(self, notebook, snippet):
  137. session = self._get_session()
  138. session_handle = session['id']
  139. statement = snippet['statement'].strip().rstrip(';')
  140. # TODO: Operations such as add, alter, create, drop, use, load, unload can be executed using simple path via
  141. # /sessions/:session_handle/configure-session
  142. operation_handle = self.db.execute_statement(session_handle=session_handle, statement=statement)
  143. self._set_operation_token_info_to_user(operation_handle['operationHandle'], 0)
  144. return {
  145. 'has_result_set': True,
  146. 'guid': operation_handle['operationHandle'],
  147. }
  148. def _is_sync_statement(self, statement):
  149. return bool(re.match(r'^(add|alter|create|drop|load|unload|use)\b', statement, re.IGNORECASE))
  150. @query_error_handler
  151. def check_status(self, notebook, snippet):
  152. response = {}
  153. session = self._get_session()
  154. status = 'expired'
  155. if snippet.get('result'):
  156. statement_id = snippet['result']['handle']['guid']
  157. if session:
  158. if not statement_id: # Sync result
  159. status = 'available'
  160. else:
  161. try:
  162. resp = self.db.fetch_status(session['id'], statement_id)
  163. if resp.get('status') == 'RUNNING':
  164. status = 'running'
  165. elif resp.get('status') == 'FINISHED':
  166. status = 'available'
  167. elif resp.get('status') == 'CANCELED':
  168. status = 'expired'
  169. elif resp.get('status') == 'CLOSED':
  170. status = 'closed'
  171. elif resp.get('status') == 'ERROR':
  172. status = 'error'
  173. self._remove_operation_token_info_from_user(statement_id)
  174. result_resp = self.db.fetch_results(session['id'], statement_id, 0)
  175. raise QueryError(parse_error(result_resp['errors'][-1]))
  176. except Exception as e:
  177. if 'Can not find the submitted operation in the OperationManager with the %s' % statement_id in str(e):
  178. LOG.warning('Operation Handle: %s does not exist' % statement_id)
  179. else:
  180. raise e
  181. response['status'] = status
  182. return response
  183. @query_error_handler
  184. def fetch_result(self, notebook, snippet, rows, start_over):
  185. session = self._get_session()
  186. statement_id = snippet['result']['handle']['guid']
  187. token = self._get_operation_token_info_from_user(statement_id)
  188. # Is race condition between cancel and fetch possible?
  189. resp = self.db.fetch_results(session['id'], operation_handle=statement_id, token=token)
  190. if resp['resultType'] == 'EOS':
  191. next_result = None
  192. else:
  193. next_result = resp.get('nextResultUri') if resp else None
  194. if next_result:
  195. # nextResultUri format:
  196. # /sessions/:session_handle/operations/:operation_handle/result/:token?rowFormat=JSON
  197. # Step 1: Drop URL query part ("?rowFormat=JSON")
  198. url_path = next_result.rsplit('?', 1)[0]
  199. # Step 2: Extract "token" from URL path
  200. n = int(url_path.rsplit('/', 1)[-1])
  201. self._set_operation_token_info_to_user(statement_id, n)
  202. data = [db['fields'] for db in resp['results']['data'] if resp and resp['results'] and resp['results']['data']]
  203. if not bool(next_result):
  204. # This will not be required if close_statement one will start working
  205. self._remove_operation_token_info_from_user(statement_id)
  206. return {
  207. 'has_more': bool(next_result),
  208. 'data': data, # No escaping...
  209. 'meta': [{
  210. 'name': column['name'],
  211. 'type': column['logicalType']['type'],
  212. 'comment': column['comment']
  213. }
  214. for column in resp['results']['columns'] if resp
  215. ],
  216. 'type': 'table'
  217. }
  218. @query_error_handler
  219. def autocomplete(self, snippet, database=None, table=None, column=None, nested=None, operation=None):
  220. response = {}
  221. if operation == 'functions':
  222. response['functions'] = self._show_functions(database)
  223. elif database is None:
  224. response['databases'] = self._show_databases()
  225. elif table is None:
  226. response['tables_meta'] = self._show_tables(database)
  227. elif column is None:
  228. columns = self._get_columns(database, table)
  229. response['columns'] = [col['name'] for col in columns]
  230. response['extended_columns'] = [{
  231. 'comment': col.get('comment'),
  232. 'name': col.get('name'),
  233. 'type': col['type']
  234. }
  235. for col in columns
  236. ]
  237. return response
  238. @query_error_handler
  239. def get_sample_data(self, snippet, database=None, table=None, column=None, nested=False, is_async=False,
  240. operation=None):
  241. if operation == 'hello':
  242. snippet['statement'] = "SELECT 'Hello World!'"
  243. else:
  244. snippet['statement'] = "SELECT * FROM `%(database)s`.`%(table)s` LIMIT 25;" % {
  245. 'database': database,
  246. 'table': table
  247. }
  248. session = self._get_session()
  249. session_id = session['id']
  250. operation_handle = self.db.execute_statement(session_handle=session_id, statement=snippet['statement'])
  251. statement_id = operation_handle['operationHandle']
  252. resp = self.db.fetch_results(session_id, statement_id, 0)
  253. while resp['resultType'] == 'NOT_READY':
  254. time.sleep(0.1)
  255. resp = self.db.fetch_results(session_id, statement_id, 0)
  256. sample = [db['fields'] for db in resp['results']['data'] if resp and resp['results'] and resp['results']['data']]
  257. n = 0
  258. while resp['resultType'] != 'EOS':
  259. resp = self.db.fetch_results(session_id, statement_id, n)
  260. if resp['resultType'] == 'PAYLOAD':
  261. n = n + 1
  262. sample += [db['fields'] for db in resp['results']['data'] if resp and resp['results'] and resp['results']['data']]
  263. time.sleep(1)
  264. if len(sample) > 0:
  265. break
  266. return {
  267. 'status': 0,
  268. 'result': {
  269. 'handle': {
  270. 'guid': statement_id
  271. }
  272. },
  273. 'rows': sample,
  274. 'full_headers': [
  275. {
  276. 'name': column['name'],
  277. 'type': column['logicalType']['type'],
  278. 'comment': column['comment']
  279. }
  280. for column in resp['results']['columns'] if resp
  281. ]
  282. }
  283. @query_error_handler
  284. def cancel(self, notebook, snippet):
  285. session = self._get_session()
  286. operation_handle = snippet['result']['handle']['guid']
  287. try:
  288. self.db.close_statement(session['id'], operation_handle)
  289. except Exception as e:
  290. message = force_unicode(str(e)).lower()
  291. LOG.debug(message)
  292. return {'status': 0}
  293. @query_error_handler
  294. def close_statement(self, notebook, snippet):
  295. session = self._get_session()
  296. statement_id = snippet['result']['handle']['guid']
  297. try:
  298. if session and statement_id:
  299. self.db.close_statement(session_handle=session['id'], operation_handle=statement_id)
  300. # self._remove_operation_token_info_from_user(statement_id) ## Needs to check why Hue db not getting updated
  301. else:
  302. return {'status': -1} # missing operation ids
  303. except Exception as e:
  304. if 'does not exist in current session:' in str(e):
  305. return {'status': -1} # skipped
  306. else:
  307. raise e
  308. return {'status': 0}
  309. def close_session(self, session):
  310. if self._get_session_info_from_user():
  311. self._remove_session_info_from_user()
  312. self.db.close_session(session['id'])
  313. return {
  314. 'status': 0,
  315. 'session': session['id']
  316. }
  317. def _check_status_and_fetch_result(self, session_handle, operation_handle):
  318. resp = self.db.fetch_results(session_handle, operation_handle, 0)
  319. while resp['resultType'] == 'NOT_READY':
  320. resp = self.db.fetch_results(session_handle, operation_handle, 0)
  321. data = [i['fields'] for i in resp['results']['data'] if resp and resp['results'] and resp['results']['data']]
  322. return data
  323. def _show_databases(self):
  324. session = self._get_session()
  325. session_handle = session['id']
  326. operation_handle = self.db.execute_statement(session_handle=session_handle, statement='SHOW DATABASES')
  327. db_list = self._check_status_and_fetch_result(session_handle, operation_handle['operationHandle'])
  328. return [db[0] for db in db_list]
  329. def _show_tables(self, database):
  330. session = self._get_session()
  331. session_handle = session['id']
  332. operation_handle = self.db.execute_statement(session_handle=session_handle,
  333. statement='SHOW TABLES IN `%(database)s`' % {'database': database})
  334. table_list = self._check_status_and_fetch_result(session_handle, operation_handle['operationHandle'])
  335. return [{
  336. 'name': table[0],
  337. 'type': 'Table',
  338. 'comment': '',
  339. }
  340. for table in table_list
  341. ]
  342. def _get_columns(self, database, table):
  343. session = self._get_session()
  344. session_handle = session['id']
  345. operation_handle = self.db.execute_statement(
  346. session_handle=session_handle,
  347. statement='DESCRIBE `%(database)s`.`%(table)s`' % {'database': database, 'table': table})
  348. column_list = self._check_status_and_fetch_result(session_handle, operation_handle['operationHandle'])
  349. return [{
  350. 'name': col[0],
  351. 'type': col[1], # Types to unify
  352. 'comment': '',
  353. }
  354. for col in column_list
  355. ]
  356. def _show_functions(self, database):
  357. session = self._get_session()
  358. session_handle = session['id']
  359. operation_handle = self.db.execute_statement(
  360. session_handle=session_handle,
  361. statement='SHOW FUNCTIONS IN `%(database)s`' % {'database': database})
  362. function_list = self._check_status_and_fetch_result(session_handle, operation_handle['operationHandle'])
  363. return [{'name': function[0]} for function in function_list]
  364. class FlinkSqlClient:
  365. """
  366. Implements https://nightlies.apache.org/flink/flink-docs-master/docs/dev/table/sql-gateway/rest/.
  367. Could be a pip module or sqlalchemy dialect in the future.
  368. """
  369. def __init__(self, user, api_url):
  370. self.user = user
  371. self._url = posixpath.join(api_url.rstrip('/') + '/' + _API_VERSION + '/')
  372. self._client = HttpClient(self._url, logger=LOG)
  373. self._root = Resource(self._client)
  374. def __str__(self):
  375. return "FlinkClient at %s" % (self._url,)
  376. def info(self):
  377. return self._root.get('info')
  378. def create_session(self, **properties):
  379. data = {
  380. "sessionName": self.user.username + "-flink-sql",
  381. }
  382. data.update(properties)
  383. return self._root.post('sessions', data=json.dumps(data), contenttype=_JSON_CONTENT_TYPE)
  384. def close_session(self, session_handle):
  385. return self._root.delete('sessions/%(session_handle)s' % {'session_handle': session_handle})
  386. def get_session_conf(self, session_handle):
  387. return self._root.get('sessions/%(session_handle)s' % {'session_handle': session_handle})
  388. def session_heartbeat(self, session_handle):
  389. return self._root.post('sessions/%(session_handle)s/heartbeat' % {'session_handle': session_handle})
  390. def configure_session(self, session_handle, statement):
  391. data = {
  392. "statement": statement,
  393. }
  394. json_data = json.dumps(data)
  395. path = 'sessions/%(session_handle)s/configure-session' % {'session_handle': session_handle}
  396. self._root.post(path, data=json_data, contenttype=_JSON_CONTENT_TYPE)
  397. def execute_statement(self, session_handle, statement):
  398. data = {
  399. "statement": statement,
  400. }
  401. json_data = json.dumps(data)
  402. path = 'sessions/%(session_handle)s/statements' % {'session_handle': session_handle}
  403. return self._root.post(path, data=json_data, contenttype=_JSON_CONTENT_TYPE)
  404. def fetch_status(self, session_handle, operation_handle):
  405. return self._root.get(
  406. 'sessions/%(session_handle)s/operations/%(operation_handle)s/status' % {
  407. 'session_handle': session_handle,
  408. 'operation_handle': operation_handle,
  409. }
  410. )
  411. def fetch_results(self, session_handle, operation_handle, token=0):
  412. return self._root.get(
  413. 'sessions/%(session_handle)s/operations/%(operation_handle)s/result/%(token)s' % {
  414. 'session_handle': session_handle,
  415. 'operation_handle': operation_handle,
  416. 'token': token
  417. })
  418. def close_statement(self, session_handle, operation_handle):
  419. return self._root.delete(
  420. 'sessions/%(session_handle)s/operations/%(operation_handle)s/close' % {
  421. 'session_handle': session_handle,
  422. 'operation_handle': operation_handle,
  423. }
  424. )
  425. def cancel(self, session_handle, operation_handle):
  426. return self._root.post(
  427. 'sessions/%(session_handle)s/operations/%(operation_handle)s/cancel' % {
  428. 'session_handle': session_handle,
  429. 'operation_handle': operation_handle
  430. }
  431. )