target_csv.py

#!/usr/bin/env python3

import argparse
import io
import os
import sys
import json
import csv
import threading
import http.client
import urllib
import collections
import pkg_resources

import bz2
import gzip
import lzma

from jsonschema.validators import Draft4Validator
import singer

DEFAULT_PARENT_KEY = ''
SEP = '__'

logger = singer.get_logger()

def emit_state(state):
    if state is not None:
        line = json.dumps(state)
        logger.debug('Emitting state {}'.format(line))
        sys.stdout.write("{}\n".format(line))
        sys.stdout.flush()

def flatten(d, parent_key=DEFAULT_PARENT_KEY, sep=SEP):
    items = []
    for k, v in d.items():
        new_key = generate_key(parent_key, sep, k)
        if isinstance(v, collections.MutableMapping):
            items.extend(flatten(v, new_key, sep=sep).items())
        else:
            items.append((new_key, str(v) if type(v) is list else v))
    return dict(items)

def get_headers(schema, parent_key=DEFAULT_PARENT_KEY, sep=SEP):
    headers = []
    for k, v in schema.get('properties', {}).items():
        new_key = generate_key(parent_key, sep, k)
        if v.get('type') == 'object':
            headers = headers + get_headers(v, new_key, sep=sep)
        else:
            headers.append(new_key)
    return list(headers)

def generate_key(parent_key, sep, key):
    return parent_key + sep + key if parent_key else key

def persist_messages(delimiter, quotechar, messages, destination_path, validate_records, compression):
    state = None
    schemas = {}
    key_properties = {}
    validators = {}
    files = {}

    for message in messages:
        try:
            o = singer.parse_message(message).asdict()
        except json.decoder.JSONDecodeError:
            logger.error("Unable to parse:\n{}".format(message))
            raise

        message_type = o['type']

        if message_type in ('RECORD', 'SCHEMA'):
            stream = o['stream']

            filename = stream + '.csv'

            if compression == 'gzip':
                filename += '.gz'
                opener = gzip.open
            elif compression == 'bzip2':
                filename += '.bz2'
                opener = bz2.open
            elif compression == 'xz':
                filename += '.xz'
                opener = lzma.open
            else:
                opener = open

            filename = os.path.expanduser(os.path.join(destination_path, filename))

        if message_type == 'RECORD':
            if stream not in schemas:
                raise Exception("A record for stream {}"
                                "was encountered before a corresponding schema".format(stream))

            if validate_records:
                validators[stream].validate(o['record'])

            csvfile = files.get(filename)
            if not csvfile:
                csvfile = opener(filename, 'wt')
                files[filename] = csvfile
                file_is_empty = True
            else:
                file_is_empty = False

            headers = get_headers(schemas[stream])
            writer = csv.DictWriter(csvfile,
                                    headers,
                                    extrasaction='ignore',
                                    delimiter=delimiter,
                                    quotechar=quotechar)
            if file_is_empty:
                writer.writeheader()

            flattened_record = flatten(o['record'])
            writer.writerow(flattened_record)

            state = None
        elif message_type == 'STATE':
            logger.debug('Setting state to {}'.format(o['value']))
            state = o['value']
        elif message_type == 'SCHEMA':
            schemas[stream] = o['schema']
            validators[stream] = Draft4Validator(o['schema'])
            key_properties[stream] = o['key_properties']

            # Delete previous content of the file
            if os.path.isfile(filename):
                os.remove(filename)
        else:
            logger.warning("Unknown message type {} in message {}"
                            .format(o['type'], o))

    for f in files.values():
        f.close()

    return state


def send_usage_stats():
    try:
        version = pkg_resources.get_distribution('target-csv').version
        conn = http.client.HTTPConnection('collector.singer.io', timeout=10)
        conn.connect()
        params = {
            'e': 'se',
            'aid': 'singer',
            'se_ca': 'target-csv',
            'se_ac': 'open',
            'se_la': version,
        }
        conn.request('GET', '/i?' + urllib.parse.urlencode(params))
        response = conn.getresponse()
        conn.close()
    except:
        logger.debug('Collection request failed')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', help='Config file')
    args = parser.parse_args()

    if args.config:
        with open(args.config) as input_json:
            config = json.load(input_json)
    else:
        config = {}

    if not config.get('disable_collection', False):
        logger.info('Sending version information to singer.io. ' +
                    'To disable sending anonymous usage data, set ' +
                    'the config parameter "disable_collection" to true')
        threading.Thread(target=send_usage_stats).start()

    input_messages = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')
    state = persist_messages(config.get('delimiter', ','),
                             config.get('quotechar', '"'),
                             input_messages,
                             config.get('destination_path', ''),
                             config.get('validate_records', True),
                             config.get('compression', None))

    emit_state(state)
    logger.debug("Exiting normally")


if __name__ == '__main__':
    main()