<s> import logging <EOL> import re <EOL> import time <EOL> from django . utils . translation import ugettext as _ <EOL> from desktop . lib . exceptions_renderable import PopupException <EOL> from desktop . lib . i18n import force_unicode <EOL> from desktop . lib . rest . http_client import RestException <EOL> from desktop . models import DefaultConfiguration <EOL> from notebook . data_export import download as spark_download <EOL> from notebook . connectors . base import Api , QueryError , SessionExpired , _get_snippet_session <EOL> LOG = logging . getLogger ( __name__ ) <EOL> try : <EOL> from spark . conf import LIVY_SERVER_SESSION_KIND <EOL> from spark . job_server_api import get_api as get_spark_api <EOL> except ImportError , e : <EOL> LOG . exception ( '<STR_LIT>' ) <EOL> class SparkConfiguration ( object ) : <EOL> APP_NAME = '<STR_LIT>' <EOL> PROPERTIES = [ <EOL> { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT:value>" : [ ] , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT:value>" : [ ] , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT:value>" : [ ] , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:1>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:1>' , <EOL> } , <EOL> { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:1>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:2>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:1>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:1>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT:string>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT:value>" : '<STR_LIT:default>' , <EOL> } , { <EOL> "<STR_LIT:name>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" : _ ( "<STR_LIT>" ) , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT>" : True , <EOL> "<STR_LIT:value>" : [ ] , <EOL> } <EOL> ] <EOL> class SparkApi ( Api ) : <EOL> SPARK_UI_RE = re . compile ( "<STR_LIT>" ) <EOL> YARN_JOB_RE = re . compile ( "<STR_LIT>" ) <EOL> STANDALONE_JOB_RE = re . compile ( "<STR_LIT>" ) <EOL> @ staticmethod <EOL> def get_properties ( ) : <EOL> return SparkConfiguration . PROPERTIES <EOL> def create_session ( self , lang = '<STR_LIT>' , properties = None ) : <EOL> if not properties : <EOL> config = DefaultConfiguration . objects . get_configuration_for_user ( app = '<STR_LIT>' , user = self . user ) <EOL> if config is not None : <EOL> properties = config . properties_list <EOL> else : <EOL> properties = self . get_properties ( ) <EOL> props = dict ( [ ( p [ '<STR_LIT:name>' ] , p [ '<STR_LIT:value>' ] ) for p in properties ] ) if properties is not None else { } <EOL> props [ '<STR_LIT>' ] = lang <EOL> api = get_spark_api ( self . user ) <EOL> response = api . create_session ( ** props ) <EOL> status = api . get_session ( response [ '<STR_LIT:id>' ] ) <EOL> count = <NUM_LIT:0> <EOL> while status [ '<STR_LIT:state>' ] == '<STR_LIT>' and count < <NUM_LIT> : <EOL> status = api . get_session ( response [ '<STR_LIT:id>' ] ) <EOL> count += <NUM_LIT:1> <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> if status [ '<STR_LIT:state>' ] != '<STR_LIT>' : <EOL> info = '<STR_LIT:\n>' . join ( status [ '<STR_LIT>' ] ) if status [ '<STR_LIT>' ] else '<STR_LIT>' <EOL> raise QueryError ( _ ( '<STR_LIT>' ) % info ) <EOL> return { <EOL> '<STR_LIT:type>' : lang , <EOL> '<STR_LIT:id>' : response [ '<STR_LIT:id>' ] , <EOL> '<STR_LIT>' : properties <EOL> } <EOL> def execute ( self , notebook , snippet ) : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> try : <EOL> response = api . submit_statement ( session [ '<STR_LIT:id>' ] , snippet [ '<STR_LIT>' ] ) <EOL> return { <EOL> '<STR_LIT:id>' : response [ '<STR_LIT:id>' ] , <EOL> '<STR_LIT>' : True , <EOL> } <EOL> except Exception , e : <EOL> message = force_unicode ( str ( e ) ) . lower ( ) <EOL> if '<STR_LIT>' in message or '<STR_LIT>' in message or '<STR_LIT>' in message : <EOL> raise SessionExpired ( e ) <EOL> else : <EOL> raise e <EOL> def check_status ( self , notebook , snippet ) : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> cell = snippet [ '<STR_LIT:result>' ] [ '<STR_LIT>' ] [ '<STR_LIT:id>' ] <EOL> try : <EOL> response = api . fetch_data ( session [ '<STR_LIT:id>' ] , cell ) <EOL> return { <EOL> '<STR_LIT:status>' : response [ '<STR_LIT:state>' ] , <EOL> } <EOL> except Exception , e : <EOL> message = force_unicode ( str ( e ) ) . lower ( ) <EOL> if '<STR_LIT>' in message : <EOL> raise SessionExpired ( e ) <EOL> else : <EOL> raise e <EOL> def fetch_result ( self , notebook , snippet , rows , start_over ) : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> cell = snippet [ '<STR_LIT:result>' ] [ '<STR_LIT>' ] [ '<STR_LIT:id>' ] <EOL> try : <EOL> response = api . fetch_data ( session [ '<STR_LIT:id>' ] , cell ) <EOL> except Exception , e : <EOL> message = force_unicode ( str ( e ) ) . lower ( ) <EOL> if '<STR_LIT>' in message : <EOL> raise SessionExpired ( e ) <EOL> else : <EOL> raise e <EOL> content = response [ '<STR_LIT>' ] <EOL> if content [ '<STR_LIT:status>' ] == '<STR_LIT>' : <EOL> data = content [ '<STR_LIT:data>' ] <EOL> images = [ ] <EOL> try : <EOL> table = data [ '<STR_LIT>' ] <EOL> except KeyError : <EOL> try : <EOL> images = [ data [ '<STR_LIT>' ] ] <EOL> except KeyError : <EOL> images = [ ] <EOL> data = [ [ data [ '<STR_LIT>' ] ] ] <EOL> meta = [ { '<STR_LIT:name>' : '<STR_LIT>' , '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } ] <EOL> type = '<STR_LIT:text>' <EOL> else : <EOL> data = table [ '<STR_LIT:data>' ] <EOL> headers = table [ '<STR_LIT>' ] <EOL> meta = [ { '<STR_LIT:name>' : h [ '<STR_LIT:name>' ] , '<STR_LIT:type>' : h [ '<STR_LIT:type>' ] , '<STR_LIT>' : '<STR_LIT>' } for h in headers ] <EOL> type = '<STR_LIT>' <EOL> if not start_over : <EOL> data = [ ] <EOL> return { <EOL> '<STR_LIT:data>' : data , <EOL> '<STR_LIT>' : images , <EOL> '<STR_LIT>' : meta , <EOL> '<STR_LIT:type>' : type <EOL> } <EOL> elif content [ '<STR_LIT:status>' ] == '<STR_LIT:error>' : <EOL> tb = content . get ( '<STR_LIT>' , None ) <EOL> if tb is None : <EOL> msg = content . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> evalue = content . get ( '<STR_LIT>' ) <EOL> if evalue is not None : <EOL> msg = '<STR_LIT>' % ( msg , evalue ) <EOL> else : <EOL> msg = '<STR_LIT>' . join ( tb ) <EOL> raise QueryError ( msg ) <EOL> def download ( self , notebook , snippet , format ) : <EOL> try : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> cell = snippet [ '<STR_LIT:result>' ] [ '<STR_LIT>' ] [ '<STR_LIT:id>' ] <EOL> return spark_download ( api , session [ '<STR_LIT:id>' ] , cell , format ) <EOL> except Exception , e : <EOL> raise PopupException ( e ) <EOL> def cancel ( self , notebook , snippet ) : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> response = api . cancel ( session [ '<STR_LIT:id>' ] ) <EOL> return { '<STR_LIT:status>' : <NUM_LIT:0> } <EOL> def get_log ( self , notebook , snippet , startFrom = <NUM_LIT:0> , size = None ) : <EOL> api = get_spark_api ( self . user ) <EOL> session = _get_snippet_session ( notebook , snippet ) <EOL> return api . get_log ( session [ '<STR_LIT:id>' ] , startFrom = startFrom , size = size ) <EOL> def progress ( self , snippet , logs ) : <EOL> return <NUM_LIT:50> <EOL> def close_statement ( self , snippet ) : <EOL> pass <EOL> def close_session ( self , session ) : <EOL> api = get_spark_api ( self . user ) <EOL> if session [ '<STR_LIT:id>' ] is not None : <EOL> try : <EOL> api . close ( session [ '<STR_LIT:id>' ] ) <EOL> return { <EOL> '<STR_LIT>' : session [ '<STR_LIT:id>' ] , <EOL> '<STR_LIT:status>' : <NUM_LIT:0> <EOL> } <EOL> except RestException , e : <EOL> if e . code == <NUM_LIT> or e . code == <NUM_LIT> : <EOL> raise SessionExpired ( e ) <EOL> else : <EOL> return { '<STR_LIT:status>' : - <NUM_LIT:1> } <EOL> def get_jobs ( self , notebook , snippet , logs ) : <EOL> if self . _is_yarn_mode ( ) : <EOL> start_logs = self . get_log ( notebook , snippet , startFrom = <NUM_LIT:0> , size = <NUM_LIT:100> ) <EOL> return self . _get_yarn_jobs ( start_logs ) <EOL> else : <EOL> return self . _get_standalone_jobs ( logs ) <EOL> def _get_standalone_jobs ( self , logs ) : <EOL> job_ids = set ( [ ] ) <EOL> spark_ui_url = self . SPARK_UI_RE . search ( logs ) <EOL> if not spark_ui_url : <EOL> LOG . warn ( '<STR_LIT>' ) <EOL> return [ ] <EOL> else : <EOL> spark_ui_url = spark_ui_url . group ( <NUM_LIT:1> ) <EOL> for match in self . STANDALONE_JOB_RE . finditer ( logs ) : <EOL> job_id = match . group ( <NUM_LIT:1> ) <EOL> job_ids . add ( job_id ) <EOL> jobs = [ { <EOL> '<STR_LIT:name>' : job_id , <EOL> '<STR_LIT:url>' : '<STR_LIT>' % ( spark_ui_url , job_id ) <EOL> } for job_id in job_ids ] <EOL> return jobs <EOL> def _get_yarn_jobs ( self , logs ) : <EOL> tracking_urls = set ( [ ] ) <EOL> for match in self . YARN_JOB_RE . finditer ( logs ) : <EOL> url = match . group ( <NUM_LIT:1> ) <EOL> tracking_urls . add ( url ) <EOL> jobs = [ { <EOL> '<STR_LIT:name>' : url . strip ( '<STR_LIT:/>' ) . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] , <EOL> '<STR_LIT:url>' : url <EOL> } for url in tracking_urls ] <EOL> return jobs <EOL> def _is_yarn_mode ( self ) : <EOL> return LIVY_SERVER_SESSION_KIND . get ( ) == "<STR_LIT>" </s>
