yql

#!/usr/bin/env python3

try:
    import json
except ImportError:
    import simplejson as json

import readline, os, sys, getopt

from urllib.request import urlopen
from urllib.parse import quote, urlencode
from urllib.error import URLError
from pprint import pprint
from io import StringIO
import traceback

from help import *
from issue import *
from pager import *

setup_issue('dryfish', 'python-yql')

InteractiveGreeting = """\
YQL endpoint {url}

YQL statements may be on multiple lines.
End your YQL statements with `;' to have the executed.
Send EOF (Unix: Control-D DOS: Control-Z) to quit.
"""

ClientCommandHelp = """\
Client commands:

\\use   lists the external tables being used.
\\env   lists the external environments being used.
\\help  tries to give usage help for YQL keywords.
"""

ResetUses = False

class YQL(object):
    PublicURL = 'http://query.yahooapis.com/v1/public/yql'
    Timeout = 5

    def __init__(_, public=True, diagnostics=True, env=[]):
        assert(public == True) # Private authentication not yet supported.
        _.qs = {
            'format': 'json',
            'callback': '',
            'diagnostics': {True: 'true', False:'False'}[diagnostics],
        }
        _.env = env
        _.url = _.PublicURL

    def __call__(_, query):
        try:
            _.qs['q'] = query
            if _.env:
                _.qs['env'] = _.env
            q = urlencode(_.qs, True)
            r = urlopen(_.url + '?' + q, None, _.Timeout)
        except URLError as e:
            print("Error:", e)
            r = None
        if r is None:
            return None
        return json.loads(str(r.read(), 'utf-8'))

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'e:', ['env='])
    except getopt.error as e:
        print("Error:", e, file=sys.stderr)
        sys.exit(1)

    env = []
    for o, a in opts:
        if o in ('-e', '--env'):
            env.append(a)
    
    try:
        yql = YQL(env=env)
        if not os.isatty(sys.stdin.fileno()):
            makequery(yql, sys.stdin.read())
            return
        print(InteractiveGreeting.format(url=yql.url))
        uses = []
        query = ''
        while True:
            if query:
                line = input('>>> ')
                query += ' ' + line.strip()
            else:
                line = input('> ')
                if not line:
                    break
                query = line.strip()
            if query.endswith(';'):
                if query[:1] == '\\':
                    client_command(yql, uses, query)
                else:
                    if query.lower().startswith('use '):
                        uses.append(query)
                        print("Using {} additional tables".format(len(uses)))
                    else:
                        makequery(yql, ' '.join(uses) + query)
                        if ResetUses:
                            uses = []
                query = ''
    except EOFError:
        print()
    except KeyboardInterrupt:
        print()
    except:
        print(new_issue())

def makequery(yql, query):
    more = None
    try:
        data = yql(query)
        if data and data['query'] and data['query']['results']:
            if os.isatty(sys.stdout.fileno()):
                more = StringIO()
            else:
                more = sys.stdout
            if query.lower().startswith('desc '):
                desc(data['query']['results']['table'], more)
            else:
                for key in data['query']['results']:
                    print(key + ':', file=more)
                    pprint(data['query']['results'][key], stream=more)
                print("{} results returned".format(data['query']['count']), file=more)
            diagnostics(data['query'].get('diagnostics', {}), file=more)
            if os.isatty(sys.stdout.fileno()):
                pager = Pager()
                pager.write(more.getvalue())
                pager.close()
        else:
            if data and data['query'] and 'diagnostics' in data['query']:
                diagnostics(data['query'].get('diagnostics', {}))
            print("No data")
    except:
        print(new_issue())

def diagnostics(diag, file=None):
    if diag:
        if 'url' in diag:
            # Used external tables.
            if type(diag['url']) == list:
                # Multiple external tables.
                exectimes = ', '.join([ u['execution-time'] for u in diag['url'] ])
            else:
                exectimes = diag['url']['execution-time']
        else:
            exectimes = 'n/a'
        print('Time: Service: {}; User: {}; Exec: {}'.format(
            diag['service-time'],
            diag['user-time'],
            exectimes), file=file)

def desc(data, file=None):
    print("Table: {}".format(data['name']), file=file)
    if 'meta' in data:
        if 'description' in data['meta']:
            print("Description: {}".format(data['meta']['description']), file=file)
        if 'author' in data['meta']:
            print("Author: {}".format(data['meta']['author']), file=file)
        if 'sampleQuery' in data['meta']:
            print("Sample query: {}".format(data['meta']['sampleQuery']), file=file)
    print("Security: {}".format(data['security']), file=file)
    if 'select' in data['request']:
        print("Selects:", file=file)
        if not type(data['request']['select']) == list:
            data['request']['select'] = [ data['request']['select'] ]
        for n, sel in enumerate(data['request']['select']):
            lead = '{:2}. '.format(n + 1)
            if not type(sel['key']) == list:
                sel['key'] = [ sel['key'] ]
            for key in sel['key']:
                print("{}{} {} {}".format(
                    lead,
                    key['type'],
                    key['name'],
                    '(required)' if key.get('required', 'false') == 'true' else ''),
                    file=file)
                lead = '    '

def client_command(yql, uses, cmd):
    assert(cmd[:1] == '\\')
    parts = cmd[1:].rstrip(';').split()
    if parts[0] == 'use':
        print("Using {} additional tables".format(len(uses)))
        for use in uses:
            print(use)
    elif parts[0] == 'env':
        print("Using {} environments".format(len(yql.env)))
        for env in yql.env:
            print(env)
    elif parts[0] == 'help':
        yql_help(parts[1])
    else:
        print(ClientCommandHelp)

if __name__ == '__main__':
    main()