base.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738
  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. from builtins import object
  18. import json
  19. import logging
  20. import re
  21. import time
  22. import uuid
  23. from django.utils.translation import ugettext as _
  24. from django.utils.encoding import smart_str
  25. from desktop.auth.backend import is_admin
  26. from desktop.conf import TASK_SERVER, has_connectors
  27. from desktop.lib import export_csvxls
  28. from desktop.lib.exceptions_renderable import PopupException
  29. from desktop.lib.i18n import smart_unicode
  30. from desktop.models import get_cluster_config
  31. from metadata.optimizer.base import get_api as get_optimizer_api
  32. from notebook.conf import get_ordered_interpreters
  33. from notebook.sql_utils import get_current_statement
  34. LOG = logging.getLogger(__name__)
  35. class SessionExpired(Exception):
  36. pass
  37. class QueryExpired(Exception):
  38. def __init__(self, message=None):
  39. super(QueryExpired, self).__init__()
  40. self.message = message
  41. class AuthenticationRequired(Exception):
  42. def __init__(self, message=None):
  43. super(AuthenticationRequired, self).__init__()
  44. self.message = message
  45. class OperationTimeout(Exception):
  46. def __str__(self):
  47. return 'OperationTimeout'
  48. class OperationNotSupported(Exception):
  49. pass
  50. class QueryError(Exception):
  51. def __init__(self, message, handle=None):
  52. super(QueryError, self).__init__(message)
  53. self.message = message or _('No error message, please check the logs.')
  54. self.handle = handle
  55. self.extra = {}
  56. def __unicode__(self):
  57. return smart_unicode(self.message)
  58. class Notebook(object):
  59. def __init__(self, document=None, **options):
  60. self.document = None
  61. if document is not None:
  62. self.data = document.data
  63. self.document = document
  64. else:
  65. _data = {
  66. 'name': 'My Notebook',
  67. 'uuid': str(uuid.uuid4()),
  68. 'description': '',
  69. 'type': 'notebook',
  70. 'isSaved': False,
  71. 'isManaged': False, # Aka isTask
  72. 'skipHistorify': False,
  73. 'sessions': [],
  74. 'snippets': [],
  75. }
  76. _data.update(options)
  77. self.data = json.dumps(_data)
  78. def get_json(self):
  79. _data = self.get_data()
  80. return json.dumps(_data)
  81. def get_data(self):
  82. _data = json.loads(self.data)
  83. if self.document is not None:
  84. _data['id'] = self.document.id
  85. _data['is_history'] = self.document.is_history
  86. return _data
  87. def get_str(self, from_oozie_action=False):
  88. return '\n\n\n'.join([
  89. 'USE %s;\n\n%s' % (
  90. snippet['database'],
  91. snippet['statement_raw'] if from_oozie_action else Notebook.statement_with_variables(snippet)
  92. )
  93. for snippet in self.get_data()['snippets']
  94. ]
  95. )
  96. @staticmethod
  97. def statement_with_variables(snippet):
  98. statement_raw = snippet['statement_raw']
  99. hasCurlyBracketParameters = snippet['type'] != 'pig'
  100. variables = {}
  101. for variable in snippet['variables']:
  102. variables[variable['name']] = variable
  103. if variables:
  104. variables_names = []
  105. for variable in snippet['variables']:
  106. variables_names.append(variable['name'])
  107. variablesString = '|'.join(variables_names)
  108. def replace(match):
  109. p1 = match.group(1)
  110. p2 = match.group(2)
  111. variable = variables[p2]
  112. value = smart_str(variable['value'])
  113. return smart_str(p1) + smart_str(value if value is not None else variable['meta'].get('placeholder',''))
  114. return re.sub(
  115. "([^\\\\])\\$" + (
  116. "{(" if hasCurlyBracketParameters else "(") + variablesString + ")(=[^}]*)?" + ("}"
  117. if hasCurlyBracketParameters else ""
  118. ),
  119. replace,
  120. smart_str(statement_raw)
  121. )
  122. return statement_raw
  123. def add_hive_snippet(self, database, sql):
  124. _data = json.loads(self.data)
  125. _data['snippets'].append(self._make_snippet({
  126. 'status': 'running',
  127. 'statement_raw': sql,
  128. 'statement': sql,
  129. 'type': 'hive',
  130. 'properties': {
  131. 'files': [],
  132. 'functions': [],
  133. 'settings': [],
  134. },
  135. 'database': database,
  136. }))
  137. self._add_session(_data, 'hive')
  138. self.data = json.dumps(_data)
  139. def add_java_snippet(self, clazz, app_jar, arguments, files):
  140. _data = json.loads(self.data)
  141. _data['snippets'].append(self._make_snippet({
  142. u'type': u'java',
  143. u'status': u'running',
  144. u'properties': {
  145. u'files': files,
  146. u'class': clazz,
  147. u'app_jar': app_jar,
  148. u'arguments': arguments,
  149. u'archives': [],
  150. }
  151. }))
  152. self._add_session(_data, 'java')
  153. self.data = json.dumps(_data)
  154. def add_sqoop_snippet(self, statement, arguments, files):
  155. _data = json.loads(self.data)
  156. _data['snippets'].append(self._make_snippet({
  157. u'type': u'sqoop1',
  158. u'status': u'running',
  159. u'properties': {
  160. u'files': files,
  161. u'arguments': arguments,
  162. u'archives': [],
  163. u'statement': statement
  164. }
  165. }))
  166. self._add_session(_data, 'java')
  167. self.data = json.dumps(_data)
  168. def add_spark_snippet(self, clazz, jars, arguments, files):
  169. _data = json.loads(self.data)
  170. _data['snippets'].append(self._make_snippet({
  171. u'type': u'spark',
  172. u'status': u'running',
  173. u'properties': {
  174. u'files': files,
  175. u'class': clazz,
  176. u'app_jar': jars,
  177. u'arguments': arguments,
  178. u'archives': [],
  179. u'spark_opts': ''
  180. }
  181. }))
  182. self._add_session(_data, 'spark')
  183. self.data = json.dumps(_data)
  184. def add_shell_snippet(self, shell_command, arguments=None, archives=None, files=None, env_var=None, last_executed=None, capture_output=True):
  185. _data = json.loads(self.data)
  186. if arguments is None:
  187. arguments = []
  188. if archives is None:
  189. archives = []
  190. if files is None:
  191. files = []
  192. if env_var is None:
  193. env_var = []
  194. _data['snippets'].append(self._make_snippet({
  195. u'type': u'shell',
  196. u'status': u'running',
  197. u'properties': {
  198. u'files': files,
  199. u'shell_command': shell_command,
  200. u'arguments': arguments,
  201. u'archives': archives,
  202. u'env_var': env_var,
  203. u'command_path': shell_command,
  204. u'capture_output': capture_output
  205. },
  206. u'lastExecuted': last_executed
  207. }))
  208. self._add_session(_data, 'shell')
  209. self.data = json.dumps(_data)
  210. def _make_snippet(self, _snippet):
  211. return {
  212. 'status': _snippet.get('status', 'ready'),
  213. 'id': str(uuid.uuid4()),
  214. 'statement_raw': _snippet.get('statement', ''),
  215. 'statement': _snippet.get('statement', ''),
  216. 'type': _snippet.get('type'),
  217. 'properties': _snippet['properties'],
  218. 'name': _snippet.get('name', '%(type)s snippet' % _snippet),
  219. 'database': _snippet.get('database'),
  220. 'result': {},
  221. 'variables': [],
  222. 'lastExecuted': _snippet.get('lastExecuted'),
  223. 'capture_output': _snippet.get('capture_output', True)
  224. }
  225. def _add_session(self, data, snippet_type):
  226. from notebook.connectors.hiveserver2 import HS2Api # Cyclic dependency
  227. if snippet_type not in [_s['type'] for _s in data['sessions']]:
  228. data['sessions'].append({
  229. 'type': snippet_type,
  230. 'properties': HS2Api.get_properties(snippet_type),
  231. 'id': None
  232. }
  233. )
  234. def execute(self, request, batch=False):
  235. from notebook.api import _execute_notebook # Cyclic dependency
  236. notebook_data = self.get_data()
  237. snippet = notebook_data['snippets'][0]
  238. snippet['wasBatchExecuted'] = batch
  239. return _execute_notebook(request, notebook_data, snippet)
  240. def execute_and_wait(self, request, timeout_sec=30.0, sleep_interval=1, include_results=False):
  241. """
  242. Run query and check status until it finishes or timeouts.
  243. Check status until it finishes or timeouts.
  244. """
  245. handle = self.execute(request, batch=False)
  246. if handle['status'] != 0:
  247. raise QueryError(e, message='SQL statement failed.', handle=handle)
  248. operation_id = handle['history_uuid']
  249. curr = time.time()
  250. end = curr + timeout_sec
  251. handle = self.check_status(request, operation_id=operation_id)
  252. while curr <= end:
  253. if handle['status'] == 0 and handle['query_status']['status'] not in ('waiting', 'running'):
  254. if include_results and handle['query_status']['status'] == 'available':
  255. handle.update(
  256. self.fetch_result_data(request.user, operation_id=operation_id)
  257. )
  258. # TODO: close
  259. return handle
  260. handle = self.check_status(request, operation_id=operation_id)
  261. time.sleep(sleep_interval)
  262. curr = time.time()
  263. # TODO
  264. # msg = "The query timed out after %(timeout)d seconds, canceled query." % {'timeout': timeout_sec}
  265. # LOG.warning(msg)
  266. # try:
  267. # self.cancel_operation(handle)
  268. # # get_api(request, snippet).cancel(notebook, snippet)
  269. # except Exception as e:
  270. # msg = "Failed to cancel query."
  271. # LOG.warning(msg)
  272. # self.close_operation(handle)
  273. # raise QueryServerException(e, message=msg)
  274. raise OperationTimeout()
  275. def check_status(self, request, operation_id):
  276. from notebook.api import _check_status
  277. return _check_status(request, operation_id=operation_id)
  278. def fetch_result_data(self, user, operation_id):
  279. from notebook.api import _fetch_result_data
  280. return _fetch_result_data(user, operation_id=operation_id, rows=100, start_over=False, nulls_only=True)
  281. def get_interpreter(connector_type, user=None):
  282. interpreter = [
  283. interpreter for interpreter in get_ordered_interpreters(user) if connector_type == interpreter['type']
  284. ]
  285. if not interpreter:
  286. if connector_type == 'hbase': # TODO move to connectors
  287. interpreter = [{
  288. 'name': 'hbase',
  289. 'type': 'hbase',
  290. 'interface': 'hbase',
  291. 'options': {},
  292. 'is_sql': False
  293. }]
  294. elif connector_type == 'kafka':
  295. interpreter = [{
  296. 'name': 'kafka',
  297. 'type': 'kafka',
  298. 'interface': 'kafka',
  299. 'options': {},
  300. 'is_sql': False
  301. }]
  302. elif connector_type == 'solr':
  303. interpreter = [{
  304. 'name': 'solr',
  305. 'type': 'solr',
  306. 'interface': 'solr',
  307. 'options': {},
  308. 'is_sql': False
  309. }]
  310. else:
  311. raise PopupException(_('Snippet type %s is not configured.') % connector_type)
  312. elif len(interpreter) > 1:
  313. raise PopupException(_('Snippet type %s matching more than one interpreter: %s') % (connector_type, len(interpreter)))
  314. return interpreter[0]
  315. def patch_snippet_for_connector(snippet):
  316. """
  317. Connector backward compatibility switcher.
  318. # TODO Connector unification
  319. """
  320. if snippet.get('connector') and snippet['connector'].get('type'):
  321. snippet['type'] = snippet['connector']['type'] # To rename to 'id'
  322. snippet['dialect'] = snippet['connector']['dialect']
  323. else:
  324. snippet['dialect'] = snippet['type']
  325. def get_api(request, snippet):
  326. from notebook.connectors.oozie_batch import OozieApi
  327. if snippet.get('wasBatchExecuted') and not TASK_SERVER.ENABLED.get():
  328. return OozieApi(user=request.user, request=request)
  329. if snippet.get('type') == 'report':
  330. snippet['type'] = 'impala'
  331. patch_snippet_for_connector(snippet)
  332. connector_name = snippet['type']
  333. if has_connectors() and snippet.get('type') == 'hello' and is_admin(request.user):
  334. interpreter = snippet.get('interpreter')
  335. else:
  336. interpreter = get_interpreter(connector_type=connector_name, user=request.user)
  337. interface = interpreter['interface']
  338. if get_cluster_config(request.user)['has_computes']:
  339. compute = json.loads(request.POST.get('cluster', '""')) # Via Catalog autocomplete API or Notebook create sessions.
  340. if compute == '""' or compute == 'undefined':
  341. compute = None
  342. if not compute and snippet.get('compute'): # Via notebook.ko.js
  343. interpreter['compute'] = snippet['compute']
  344. LOG.debug('Selected interpreter %s interface=%s compute=%s' % (
  345. interpreter['type'],
  346. interface,
  347. interpreter.get('compute') and interpreter['compute']['name'])
  348. )
  349. if interface == 'hiveserver2' or interface == 'hms':
  350. from notebook.connectors.hiveserver2 import HS2Api
  351. return HS2Api(user=request.user, request=request, interpreter=interpreter)
  352. elif interface == 'oozie':
  353. return OozieApi(user=request.user, request=request)
  354. elif interface == 'livy':
  355. from notebook.connectors.spark_shell import SparkApi
  356. return SparkApi(request.user)
  357. elif interface == 'livy-batch':
  358. from notebook.connectors.spark_batch import SparkBatchApi
  359. return SparkBatchApi(request.user)
  360. elif interface == 'text' or interface == 'markdown':
  361. from notebook.connectors.text import TextApi
  362. return TextApi(request.user)
  363. elif interface == 'rdbms':
  364. from notebook.connectors.rdbms import RdbmsApi
  365. return RdbmsApi(request.user, interpreter=snippet['type'], query_server=snippet.get('query_server'))
  366. elif interface == 'jdbc':
  367. if interpreter['options'] and interpreter['options'].get('url', '').find('teradata') >= 0:
  368. from notebook.connectors.jdbc_teradata import JdbcApiTeradata
  369. return JdbcApiTeradata(request.user, interpreter=interpreter)
  370. if interpreter['options'] and interpreter['options'].get('url', '').find('awsathena') >= 0:
  371. from notebook.connectors.jdbc_athena import JdbcApiAthena
  372. return JdbcApiAthena(request.user, interpreter=interpreter)
  373. elif interpreter['options'] and interpreter['options'].get('url', '').find('presto') >= 0:
  374. from notebook.connectors.jdbc_presto import JdbcApiPresto
  375. return JdbcApiPresto(request.user, interpreter=interpreter)
  376. elif interpreter['options'] and interpreter['options'].get('url', '').find('clickhouse') >= 0:
  377. from notebook.connectors.jdbc_clickhouse import JdbcApiClickhouse
  378. return JdbcApiClickhouse(request.user, interpreter=interpreter)
  379. elif interpreter['options'] and interpreter['options'].get('url', '').find('vertica') >= 0:
  380. from notebook.connectors.jdbc_vertica import JdbcApiVertica
  381. return JdbcApiVertica(request.user, interpreter=interpreter)
  382. else:
  383. from notebook.connectors.jdbc import JdbcApi
  384. return JdbcApi(request.user, interpreter=interpreter)
  385. elif interface == 'teradata':
  386. from notebook.connectors.jdbc import JdbcApiTeradata
  387. return JdbcApiTeradata(request.user, interpreter=interpreter)
  388. elif interface == 'athena':
  389. from notebook.connectors.jdbc import JdbcApiAthena
  390. return JdbcApiAthena(request.user, interpreter=interpreter)
  391. elif interface == 'presto':
  392. from notebook.connectors.jdbc_presto import JdbcApiPresto
  393. return JdbcApiPresto(request.user, interpreter=interpreter)
  394. elif interface == 'sqlalchemy':
  395. from notebook.connectors.sql_alchemy import SqlAlchemyApi
  396. return SqlAlchemyApi(request.user, interpreter=interpreter)
  397. elif interface == 'solr':
  398. from notebook.connectors.solr import SolrApi
  399. return SolrApi(request.user, interpreter=interpreter)
  400. elif interface == 'hbase':
  401. from notebook.connectors.hbase import HBaseApi
  402. return HBaseApi(request.user)
  403. elif interface == 'ksql':
  404. from notebook.connectors.ksql import KSqlApi
  405. return KSqlApi(request.user, interpreter=interpreter)
  406. elif interface == 'flink':
  407. from notebook.connectors.flink_sql import FlinkSqlApi
  408. return FlinkSqlApi(request.user, interpreter=interpreter)
  409. elif interface == 'kafka':
  410. from notebook.connectors.kafka import KafkaApi
  411. return KafkaApi(request.user)
  412. elif interface == 'pig':
  413. return OozieApi(user=request.user, request=request) # Backward compatibility until Hue 4
  414. else:
  415. raise PopupException(_('Notebook connector interface not recognized: %s') % interface)
  416. def _get_snippet_session(notebook, snippet):
  417. session = [session for session in notebook['sessions'] if session['type'] == snippet['type']]
  418. if not session:
  419. raise SessionExpired()
  420. else:
  421. return session[0]
  422. # Base API
  423. class Api(object):
  424. def __init__(self, user, interpreter=None, request=None, query_server=None):
  425. self.user = user
  426. self.interpreter = interpreter
  427. self.request = request
  428. self.query_server = query_server
  429. def create_session(self, lang, properties=None):
  430. return {
  431. 'type': lang,
  432. 'id': None,
  433. 'properties': properties if not None else []
  434. }
  435. def close_session(self, session):
  436. pass
  437. def fetch_result(self, notebook, snippet, rows, start_over):
  438. pass
  439. def can_start_over(self, notebook, snippet):
  440. return False
  441. def fetch_result_size(self, notebook, snippet):
  442. raise OperationNotSupported()
  443. def download(self, notebook, snippet, file_format='csv'):
  444. from beeswax import data_export #TODO: Move to notebook?
  445. from beeswax import conf
  446. result_wrapper = ExecutionWrapper(self, notebook, snippet)
  447. max_rows = conf.DOWNLOAD_ROW_LIMIT.get()
  448. max_bytes = conf.DOWNLOAD_BYTES_LIMIT.get()
  449. content_generator = data_export.DataAdapter(result_wrapper, max_rows=max_rows, max_bytes=max_bytes)
  450. return export_csvxls.create_generator(content_generator, file_format)
  451. def get_log(self, notebook, snippet, startFrom=None, size=None):
  452. return 'No logs'
  453. def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
  454. return {}
  455. def progress(self, notebook, snippet, logs=None):
  456. return 50
  457. def get_jobs(self, notebook, snippet, logs):
  458. return []
  459. def get_sample_data(self, snippet, database=None, table=None, column=None, is_async=False, operation=None): raise NotImplementedError()
  460. def export_data_as_hdfs_file(self, snippet, target_file, overwrite): raise NotImplementedError()
  461. def export_data_as_table(self, notebook, snippet, destination, is_temporary=False, location=None): raise NotImplementedError()
  462. def export_large_data_to_hdfs(self, notebook, snippet, destination): raise NotImplementedError()
  463. def statement_risk(self, interface, notebook, snippet):
  464. response = self._get_current_statement(notebook, snippet)
  465. query = response['statement']
  466. client = get_optimizer_api(self.user, interface)
  467. patch_snippet_for_connector(snippet)
  468. return client.query_risk(query=query, source_platform=snippet['dialect'], db_name=snippet.get('database') or 'default')
  469. def statement_compatibility(self, interface, notebook, snippet, source_platform, target_platform):
  470. response = self._get_current_statement(notebook, snippet)
  471. query = response['statement']
  472. client = get_optimizer_api(self.user, interface)
  473. return client.query_compatibility(source_platform, target_platform, query)
  474. def statement_similarity(self, interface, notebook, snippet, source_platform):
  475. response = self._get_current_statement(notebook, snippet)
  476. query = response['statement']
  477. client = get_optimizer_api(self.user, interface)
  478. return client.similar_queries(source_platform, query)
  479. def describe(self, notebook, snippet, database=None, table=None, column=None):
  480. if column:
  481. response = self.describe_column(notebook, snippet, database=database, table=table, column=column)
  482. elif table:
  483. response = {
  484. 'status': 0,
  485. 'name': table or '',
  486. 'partition_keys': [],
  487. 'cols': [],
  488. 'path_location': '',
  489. 'hdfs_link': '',
  490. 'comment': '',
  491. 'is_view': False,
  492. 'properties': [],
  493. 'details': {'properties': {'table_type': ''}, 'stats': {}},
  494. 'stats': []
  495. }
  496. describe_table = self.describe_table(notebook, snippet, database, table)
  497. response.update(describe_table)
  498. else:
  499. response = {
  500. 'status': 0,
  501. 'owner_name': '',
  502. 'owner_type': '',
  503. 'parameters': '',
  504. 'hdfs_link': '',
  505. 'message': ''
  506. }
  507. describe_database = self.describe_database(notebook, snippet, database)
  508. response.update(describe_database)
  509. return response
  510. def describe_column(self, notebook, snippet, database=None, table=None, column=None):
  511. return []
  512. def describe_table(self, notebook, snippet, database=None, table=None):
  513. response = {}
  514. autocomplete = self.autocomplete(snippet, database=database, table=table)
  515. response['cols'] = autocomplete['extended_columns'] if autocomplete and autocomplete.get('extended_columns') else [],
  516. return response
  517. def describe_database(self, notebook, snippet, database=None):
  518. return {}
  519. def close_statement(self, notebook, snippet): pass
  520. def _get_current_statement(self, notebook, snippet):
  521. should_close, resp = get_current_statement(snippet)
  522. if should_close:
  523. try:
  524. self.close_statement(notebook, snippet) # Close all the time past multi queries
  525. except:
  526. LOG.warn('Could not close previous multiquery query')
  527. return resp
  528. def get_log_is_full_log(self, notebook, snippet):
  529. return True
  530. def _get_snippet_name(notebook, unique=False, table_format=False):
  531. name = (('%(name)s' + ('-%(id)s' if unique else '') if notebook.get('name') else '%(type)s-%(id)s') % notebook)
  532. if table_format:
  533. name = re.sub('[-|\s:]', '_', name)
  534. return name
  535. class ExecutionWrapper(object):
  536. def __init__(self, api, notebook, snippet, callback=None):
  537. self.api = api
  538. self.notebook = notebook
  539. self.snippet = snippet
  540. self.callback = callback
  541. self.should_close = False
  542. def fetch(self, handle, start_over=None, rows=None):
  543. if start_over:
  544. if not self.snippet['result'].get('handle') \
  545. or not self.snippet['result']['handle'].get('guid') \
  546. or not self.api.can_start_over(self.notebook, self.snippet):
  547. start_over = False
  548. handle = self.api.execute(self.notebook, self.snippet)
  549. self.snippet['result']['handle'] = handle
  550. if self.callback and hasattr(self.callback, 'on_execute'):
  551. self.callback.on_execute(handle)
  552. self.should_close = True
  553. self._until_available()
  554. if self.snippet['result']['handle'].get('sync', False):
  555. result = self.snippet['result']['handle']['result']
  556. else:
  557. result = self.api.fetch_result(self.notebook, self.snippet, rows, start_over)
  558. return ResultWrapper(result.get('meta'), result.get('data'), result.get('has_more'))
  559. def _until_available(self):
  560. if self.snippet['result']['handle'].get('sync', False):
  561. return # Request is already completed
  562. count = 0
  563. sleep_seconds = 1
  564. check_status_count = 0
  565. get_log_is_full_log = self.api.get_log_is_full_log(self.notebook, self.snippet)
  566. while True:
  567. response = self.api.check_status(self.notebook, self.snippet)
  568. if self.callback and hasattr(self.callback, 'on_status'):
  569. self.callback.on_status(response['status'])
  570. if self.callback and hasattr(self.callback, 'on_log'):
  571. log = self.api.get_log(self.notebook, self.snippet, startFrom=count)
  572. if get_log_is_full_log:
  573. log = log[count:]
  574. self.callback.on_log(log)
  575. count += len(log)
  576. if response['status'] not in ['waiting', 'running', 'submitted']:
  577. break
  578. check_status_count += 1
  579. if check_status_count > 5:
  580. sleep_seconds = 5
  581. elif check_status_count > 10:
  582. sleep_seconds = 10
  583. time.sleep(sleep_seconds)
  584. def close(self, handle):
  585. if self.should_close:
  586. self.should_close = False
  587. self.api.close_statement(self.notebook, self.snippet)
  588. class ResultWrapper(object):
  589. def __init__(self, cols, rows, has_more):
  590. self._cols = cols
  591. self._rows = rows
  592. self.has_more = has_more
  593. def full_cols(self):
  594. return self._cols
  595. def rows(self):
  596. return self._rows