#!/usr/bin/env python

'''
Usage:  
    bqexport [-p] <project> <dataset> --table <table> --bucket <google_bucket> --format=<fmt> [--delimiter=<delimiter>] [--directory=<directory>]   
    bqexport [-p] <project> <dataset> --table-list-file <tables> --bucket <google_bucket> --format=<fmt> [--delimiter=<delimiter>] [--directory=<directory>]    
         
Options: 
    -p,--preview     : show (but do not execute) export command
'''

import os, sys
import json
from snap import common
import docopt
import sh
from sh import bq  # Google Cloud CLI must already be installed


class EXPORT_FORMAT(object):
    CSV = 'csv'
    JSON = 'json'


def extract_data(source_table_designator, target_designator, export_format, delimiter):
    try:
        if export_format == EXPORT_FORMAT.CSV:
            result = bq.extract('--field_delimiter',
                                delimiter, 
                                '--destination_format',
                                'CSV',
                                source_table_designator,
                                target_designator)
            print('\n### export of "%s" to "%s" complete.\n' % 
                (source_table_designator, target_designator), file=sys.stderr)
        else: # export JSON records
            result = bq.extract('--destination_format',
                                'NEWLINE_DELIMITED_JSON',
                                source_table_designator,
                                target_designator)
            print('\n### export of "%s" to "%s" complete.\n' % 
                (source_table_designator, target_designator), file=sys.stderr)

    except Exception as err:
        print('!!! error exporting table data.', file=sys.stderr)
        print(err, file=sys.stderr)


def main(args):
    export_format = args['--format']
    if export_format == EXPORT_FORMAT.CSV:
        if args.get('--delimiter') is None:
            print('### csv chosen as the export format, but no delimiter specified. Defaulting to comma.', file=sys.stderr)
    elif export_format != EXPORT_FORMAT.JSON:
        print('!!! supported export formats are "csv" and "json".')
        return
    
    tables = []
    if args.get('--table'):
        tables.append(args['<table>'])
    elif args.get('--table-list-file'):
        table_list_file = args['<tables>']
        with open(table_list_file) as f:
            for line in f:
                tables.append(line.lstrip().rstrip())

    project_name = args['<project>']
    dataset = args['<dataset>']    
    bucket = args['<google_bucket>']
    
    delimiter = ','
    if args.get('--delimiter') is not None:
        delimiter = args['--delimiter']

    preview_mode = False
    if args.get('--preview'):
        preview_mode = True
        print('\n### running bqex in preview mode.\n', file=sys.stderr)
    
    if args.get('--directory') is not None: 
        bucket_directory = args['--directory']
    else:
        bucket_directory = ''

    for table_name in tables:
        source_table_designator = '{project}:{dataset}.{table}'.format(project=project_name,
                                                                       dataset=dataset,
                                                                       table=table_name)
        filename = '%s_*.%s' % (table_name, export_format)
        path_string = os.path.join(bucket, bucket_directory, filename)
        target_designator = 'gs://%s' % path_string
        
        if preview_mode:
            if export_format == EXPORT_FORMAT.CSV:
                print(bq.extract.bake('--field_delimiter',
                                        '\'%s\'' % delimiter, 
                                        '--destination_format',
                                        'CSV',
                                        source_table_designator,
                                        target_designator))
            else:
                print(bq.extract.bake('--destination_format',
                                      'NEWLINE_DELIMITED_JSON',
                                      source_table_designator,
                                      target_designator))                    
        else:
            extract_data(source_table_designator, target_designator, export_format, delimiter)
    
    print('\n### exiting.', file=sys.stderr)


if __name__ == '__main__':
    args = docopt.docopt(__doc__)
    main(args)



