<s> import os <EOL> import re <EOL> import locale <EOL> import logging <EOL> import subprocess <EOL> from io import open <EOL> import click <EOL> import sqlparse <EOL> from . import export <EOL> from . main import special_command , NO_QUERY , PARSED_QUERY <EOL> from . favoritequeries import favoritequeries <EOL> from . utils import handle_cd_command <EOL> TIMING_ENABLED = False <EOL> use_expanded_output = False <EOL> PAGER_ENABLED = True <EOL> @ export <EOL> def set_timing_enabled ( val ) : <EOL> global TIMING_ENABLED <EOL> TIMING_ENABLED = val <EOL> @ export <EOL> def set_pager_enabled ( val ) : <EOL> global PAGER_ENABLED <EOL> PAGER_ENABLED = val <EOL> @ export <EOL> def is_pager_enabled ( ) : <EOL> return PAGER_ENABLED <EOL> @ export <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , arg_type = PARSED_QUERY , aliases = ( '<STR_LIT>' , ) , case_sensitive = True ) <EOL> def set_pager ( arg , ** _ ) : <EOL> if arg : <EOL> os . environ [ '<STR_LIT>' ] = arg <EOL> msg = '<STR_LIT>' % arg <EOL> set_pager_enabled ( True ) <EOL> else : <EOL> if '<STR_LIT>' in os . environ : <EOL> msg = '<STR_LIT>' % os . environ [ '<STR_LIT>' ] <EOL> else : <EOL> msg = '<STR_LIT>' <EOL> set_pager_enabled ( True ) <EOL> return [ ( None , None , None , msg ) ] <EOL> @ export <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> arg_type = NO_QUERY , aliases = ( '<STR_LIT>' , ) , case_sensitive = True ) <EOL> def disable_pager ( ) : <EOL> set_pager_enabled ( False ) <EOL> return [ ( None , None , None , '<STR_LIT>' ) ] <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , arg_type = NO_QUERY , aliases = ( '<STR_LIT>' , ) , case_sensitive = True ) <EOL> def toggle_timing ( ) : <EOL> global TIMING_ENABLED <EOL> TIMING_ENABLED = not TIMING_ENABLED <EOL> message = "<STR_LIT>" <EOL> message += "<STR_LIT>" if TIMING_ENABLED else "<STR_LIT>" <EOL> return [ ( None , None , None , message ) ] <EOL> @ export <EOL> def is_timing_enabled ( ) : <EOL> return TIMING_ENABLED <EOL> @ export <EOL> def set_expanded_output ( val ) : <EOL> global use_expanded_output <EOL> use_expanded_output = val <EOL> @ export <EOL> def is_expanded_output ( ) : <EOL> return use_expanded_output <EOL> def quit ( * args ) : <EOL> raise NotImplementedError <EOL> def stub ( * args ) : <EOL> raise NotImplementedError <EOL> _logger = logging . getLogger ( __name__ ) <EOL> @ export <EOL> def editor_command ( command ) : <EOL> """<STR_LIT>""" <EOL> return command . strip ( ) . endswith ( '<STR_LIT>' ) or command . strip ( ) . startswith ( '<STR_LIT>' ) <EOL> @ export <EOL> def get_filename ( sql ) : <EOL> if sql . strip ( ) . startswith ( '<STR_LIT>' ) : <EOL> command , _ , filename = sql . partition ( '<STR_LIT:U+0020>' ) <EOL> return filename . strip ( ) or None <EOL> @ export <EOL> def open_external_editor ( filename = None , sql = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> sql = sql . strip ( ) <EOL> pattern = re . compile ( '<STR_LIT>' ) <EOL> while pattern . search ( sql ) : <EOL> sql = pattern . sub ( '<STR_LIT>' , sql ) <EOL> message = None <EOL> filename = filename . strip ( ) . split ( '<STR_LIT:U+0020>' , <NUM_LIT:1> ) [ <NUM_LIT:0> ] if filename else None <EOL> MARKER = '<STR_LIT>' <EOL> query = click . edit ( sql + '<STR_LIT>' + MARKER , filename = filename , <EOL> extension = '<STR_LIT>' ) <EOL> if filename : <EOL> try : <EOL> with open ( filename , encoding = '<STR_LIT:utf-8>' ) as f : <EOL> query = f . read ( ) <EOL> except IOError : <EOL> message = '<STR_LIT>' % filename <EOL> if query is not None : <EOL> query = query . split ( MARKER , <NUM_LIT:1> ) [ <NUM_LIT:0> ] . rstrip ( '<STR_LIT:\n>' ) <EOL> else : <EOL> query = sql <EOL> return ( query , message ) <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , arg_type = PARSED_QUERY , case_sensitive = True ) <EOL> def execute_favorite_query ( cur , arg ) : <EOL> """<STR_LIT>""" <EOL> if arg == '<STR_LIT>' : <EOL> for result in list_favorite_queries ( ) : <EOL> yield result <EOL> query = favoritequeries . get ( arg ) <EOL> if query is None : <EOL> message = "<STR_LIT>" % ( arg ) <EOL> yield ( None , None , None , message ) <EOL> else : <EOL> for sql in sqlparse . split ( query ) : <EOL> sql = sql . rstrip ( '<STR_LIT:;>' ) <EOL> title = '<STR_LIT>' % ( sql ) <EOL> cur . execute ( sql ) <EOL> if cur . description : <EOL> headers = [ x [ <NUM_LIT:0> ] for x in cur . description ] <EOL> yield ( title , cur , headers , None ) <EOL> else : <EOL> yield ( title , None , None , None ) <EOL> def list_favorite_queries ( ) : <EOL> """<STR_LIT>""" <EOL> headers = [ "<STR_LIT:Name>" , "<STR_LIT>" ] <EOL> rows = [ ( r , favoritequeries . get ( r ) ) for r in favoritequeries . list ( ) ] <EOL> if not rows : <EOL> status = '<STR_LIT>' + favoritequeries . usage <EOL> else : <EOL> status = '<STR_LIT>' <EOL> return [ ( '<STR_LIT>' , rows , headers , status ) ] <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def save_favorite_query ( arg , ** _ ) : <EOL> """<STR_LIT>""" <EOL> usage = '<STR_LIT>' + favoritequeries . usage <EOL> if not arg : <EOL> return [ ( None , None , None , usage ) ] <EOL> name , _ , query = arg . partition ( '<STR_LIT:U+0020>' ) <EOL> if ( not name ) or ( not query ) : <EOL> return [ ( None , None , None , <EOL> usage + '<STR_LIT>' ) ] <EOL> favoritequeries . save ( name , query ) <EOL> return [ ( None , None , None , "<STR_LIT>" ) ] <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def delete_favorite_query ( arg , ** _ ) : <EOL> """<STR_LIT>""" <EOL> usage = '<STR_LIT>' + favoritequeries . usage <EOL> if not arg : <EOL> return [ ( None , None , None , usage ) ] <EOL> status = favoritequeries . delete ( arg ) <EOL> return [ ( None , None , None , status ) ] <EOL> @ special_command ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def execute_system_command ( arg , ** _ ) : <EOL> """<STR_LIT>""" <EOL> usage = "<STR_LIT>" <EOL> if not arg : <EOL> return [ ( None , None , None , usage ) ] <EOL> try : <EOL> command = arg . strip ( ) <EOL> if command . startswith ( '<STR_LIT>' ) : <EOL> ok , error_message = handle_cd_command ( arg ) <EOL> if not ok : <EOL> return [ ( None , None , None , error_message ) ] <EOL> return [ ( None , None , None , '<STR_LIT>' ) ] <EOL> args = arg . split ( '<STR_LIT:U+0020>' ) <EOL> process = subprocess . Popen ( args , stdout = subprocess . PIPE , stderr = subprocess . PIPE ) <EOL> output , error = process . communicate ( ) <EOL> response = output if not error else error <EOL> if isinstance ( response , bytes ) : <EOL> encoding = locale . getpreferredencoding ( False ) <EOL> response = response . decode ( encoding ) <EOL> return [ ( None , None , None , response ) ] <EOL> except OSError as e : <EOL> return [ ( None , None , None , '<STR_LIT>' % e . strerror ) ] </s>
