<s> import logging <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 libsolr . api import SolrApi as NativeSolrApi <EOL> from notebook . connectors . base import Api , QueryError <EOL> from notebook . models import escape_rows <EOL> LOG = logging . getLogger ( __name__ ) <EOL> def query_error_handler ( func ) : <EOL> def decorator ( * args , ** kwargs ) : <EOL> try : <EOL> return func ( * args , ** kwargs ) <EOL> except Exception , e : <EOL> message = force_unicode ( str ( e ) ) <EOL> raise QueryError ( message ) <EOL> return decorator <EOL> class SolrApi ( Api ) : <EOL> def __init__ ( self , user , interpreter = None ) : <EOL> Api . __init__ ( self , user , interpreter = interpreter ) <EOL> self . options = interpreter [ '<STR_LIT>' ] <EOL> @ query_error_handler <EOL> def execute ( self , notebook , snippet ) : <EOL> from search . conf import SOLR_URL <EOL> api = NativeSolrApi ( SOLR_URL . get ( ) , self . user . username ) <EOL> collection = self . options . get ( '<STR_LIT>' ) or snippet . get ( '<STR_LIT>' ) <EOL> if not collection or collection == '<STR_LIT:default>' : <EOL> collection = api . collections2 ( ) [ <NUM_LIT:0> ] <EOL> response = api . sql ( collection , snippet [ '<STR_LIT>' ] ) <EOL> info = response [ '<STR_LIT>' ] [ '<STR_LIT>' ] . pop ( - <NUM_LIT:1> ) <EOL> if info . get ( '<STR_LIT>' ) : <EOL> raise QueryError ( info [ '<STR_LIT>' ] ) <EOL> headers = [ ] <EOL> for row in response [ '<STR_LIT>' ] [ '<STR_LIT>' ] : <EOL> for col in row . keys ( ) : <EOL> if col not in headers : <EOL> headers . append ( col ) <EOL> data = [ [ doc . get ( col ) for col in headers ] for doc in response [ '<STR_LIT>' ] [ '<STR_LIT>' ] ] <EOL> has_result_set = bool ( data ) <EOL> return { <EOL> '<STR_LIT>' : True , <EOL> '<STR_LIT>' : has_result_set , <EOL> '<STR_LIT>' : <NUM_LIT:0> , <EOL> '<STR_LIT:result>' : { <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT:data>' : data if has_result_set else [ ] , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:name>' : col , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } for col in headers ] if has_result_set else [ ] , <EOL> '<STR_LIT:type>' : '<STR_LIT>' <EOL> } , <EOL> '<STR_LIT>' : <NUM_LIT:0> , <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT>' : <NUM_LIT:1> <EOL> } <EOL> @ query_error_handler <EOL> def check_status ( self , notebook , snippet ) : <EOL> return { '<STR_LIT:status>' : '<STR_LIT>' } <EOL> @ query_error_handler <EOL> def fetch_result ( self , notebook , snippet , rows , start_over ) : <EOL> return { <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT:data>' : [ ] , <EOL> '<STR_LIT>' : [ ] , <EOL> '<STR_LIT:type>' : '<STR_LIT>' <EOL> } <EOL> @ query_error_handler <EOL> def fetch_result_metadata ( self ) : <EOL> pass <EOL> @ query_error_handler <EOL> def cancel ( self , notebook , snippet ) : <EOL> return { '<STR_LIT:status>' : <NUM_LIT:0> } <EOL> @ query_error_handler <EOL> def get_log ( self , notebook , snippet , startFrom = None , size = None ) : <EOL> return '<STR_LIT>' <EOL> def download ( self , notebook , snippet , format ) : <EOL> raise PopupException ( '<STR_LIT>' ) <EOL> @ query_error_handler <EOL> def close_statement ( self , snippet ) : <EOL> return { '<STR_LIT:status>' : - <NUM_LIT:1> } <EOL> @ query_error_handler <EOL> def autocomplete ( self , snippet , database = None , table = None , column = None , nested = None ) : <EOL> from search . conf import SOLR_URL <EOL> api = NativeSolrApi ( SOLR_URL . get ( ) , self . user . username ) <EOL> assist = Assist ( self , self . user , api ) <EOL> response = { '<STR_LIT:status>' : - <NUM_LIT:1> } <EOL> if database is None : <EOL> response [ '<STR_LIT>' ] = [ self . options . get ( '<STR_LIT>' ) or snippet . get ( '<STR_LIT>' ) or '<STR_LIT:default>' ] <EOL> elif table is None : <EOL> tables_meta = [ ] <EOL> for t in assist . get_tables ( database ) : <EOL> tables_meta . append ( { '<STR_LIT:name>' : t , '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } ) <EOL> response [ '<STR_LIT>' ] = tables_meta <EOL> else : <EOL> columns = assist . get_columns ( database , table ) <EOL> response [ '<STR_LIT>' ] = [ col [ '<STR_LIT:name>' ] for col in columns ] <EOL> response [ '<STR_LIT>' ] = columns <EOL> response [ '<STR_LIT:status>' ] = <NUM_LIT:0> <EOL> return response <EOL> @ query_error_handler <EOL> def get_sample_data ( self , snippet , database = None , table = None , column = None ) : <EOL> from search . conf import SOLR_URL <EOL> db = NativeSolrApi ( SOLR_URL . get ( ) , self . user ) <EOL> assist = Assist ( self , self . user , db ) <EOL> response = { '<STR_LIT:status>' : - <NUM_LIT:1> } <EOL> sample_data = assist . get_sample_data ( database , table , column ) <EOL> if sample_data : <EOL> response [ '<STR_LIT:status>' ] = <NUM_LIT:0> <EOL> response [ '<STR_LIT>' ] = sample_data [ '<STR_LIT>' ] <EOL> response [ '<STR_LIT>' ] = sample_data [ '<STR_LIT>' ] <EOL> else : <EOL> response [ '<STR_LIT:message>' ] = _ ( '<STR_LIT>' ) <EOL> return response <EOL> class Assist ( ) : <EOL> def __init__ ( self , api , user , db ) : <EOL> self . api = api <EOL> self . user = user <EOL> self . db = db <EOL> def get_databases ( self ) : <EOL> return self . options [ '<STR_LIT>' ] . get ( '<STR_LIT>' ) or [ '<STR_LIT:default>' ] <EOL> def get_tables ( self , database , table_names = [ ] ) : <EOL> return self . db . collections2 ( ) <EOL> def get_columns ( self , database , table ) : <EOL> return [ { '<STR_LIT:name>' : field [ '<STR_LIT:name>' ] , '<STR_LIT:type>' : field [ '<STR_LIT:type>' ] , '<STR_LIT>' : '<STR_LIT>' } for field in self . db . schema_fields ( table ) [ '<STR_LIT>' ] ] <EOL> def get_sample_data ( self , database , table , column = None ) : <EOL> if column is None : <EOL> column = '<STR_LIT:U+002CU+0020>' . join ( [ col [ '<STR_LIT:name>' ] for col in self . get_columns ( database , table ) ] ) <EOL> snippet = { <EOL> '<STR_LIT>' : table , <EOL> '<STR_LIT>' : '<STR_LIT>' % ( column , table ) <EOL> } <EOL> res = self . api . execute ( None , snippet ) <EOL> response = { '<STR_LIT:status>' : - <NUM_LIT:1> } <EOL> if res : <EOL> response [ '<STR_LIT:status>' ] = <NUM_LIT:0> <EOL> response [ '<STR_LIT>' ] = [ col [ '<STR_LIT:name>' ] for col in res [ '<STR_LIT:result>' ] [ '<STR_LIT>' ] ] <EOL> response [ '<STR_LIT>' ] = escape_rows ( res [ '<STR_LIT:result>' ] [ '<STR_LIT:data>' ] , nulls_only = True ) <EOL> else : <EOL> response [ '<STR_LIT:message>' ] = _ ( '<STR_LIT>' ) <EOL> return response </s>
