target_tools.py

import http.client
import io
import json
import pkg_resources
import sys
import threading

import singer
from singer import utils, metadata, metrics

from target_postgres import json_schema
from target_postgres.exceptions import TargetError
from target_postgres.singer_stream import BufferedSingerStream
from target_postgres.stream_tracker import StreamTracker

LOGGER = singer.get_logger()


def main(target):
    """
    Given a target, stream stdin input as a text stream.
    :param target: object which implements `write_batch` and `activate_version`
    :return: None
    """
    config = utils.parse_args([]).config
    input_stream = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')
    stream_to_target(input_stream, target, config=config)

    return None


def stream_to_target(stream, target, config={}):
    """
    Persist `stream` to `target` with optional `config`.
    :param stream: iterator which represents a Singer data stream
    :param target: object which implements `write_batch` and `activate_version`
    :param config: [optional] configuration for buffers etc.
    :return: None
    """

    state_support = config.get('state_support', True)
    state_tracker = StreamTracker(target, state_support)
    _run_sql_hook('before_run_sql', config, target)

    line_stats = {} # dict of <line_type: count>

    try:
        if not config.get('disable_collection', False):
            _async_send_usage_stats()

        invalid_records_detect = config.get('invalid_records_detect')
        invalid_records_threshold = config.get('invalid_records_threshold')
        max_batch_rows = config.get('max_batch_rows', 200000)
        max_batch_size = config.get('max_batch_size', 104857600)  # 100MB
        batch_detection_threshold = config.get('batch_detection_threshold', max(max_batch_rows / 40, 50))

        LOGGER.info('Streaming to target with the following configuration: {}'.format(
            {
                'state_support': state_support,
                'invalid_records_detect': invalid_records_detect,
                'invalid_records_threshold': invalid_records_threshold,
                'max_batch_rows': max_batch_rows,
                'max_batch_size': max_batch_size,
                'batch_detection_threshold': batch_detection_threshold
            }))

        line_count = 0
        for line in stream:
            _line_handler(line_stats,
                          state_tracker,
                          target,
                          invalid_records_detect,
                          invalid_records_threshold,
                          max_batch_rows,
                          max_batch_size,
                          line
                          )

            line_count += 1

            if line_count % batch_detection_threshold == 0:
                LOGGER.debug('Attempting to flush streams at `line_count` {}, with records distribution of: {}'.format(
                    line_count,
                    _records_distribution(line_count, line_stats)
                ))
                state_tracker.flush_streams()

        LOGGER.debug('Forcing flush of streams. Input depleted.')

        state_tracker.flush_streams(force=True)
        _run_sql_hook('after_run_sql', config, target)

        return None

    except Exception as e:
        LOGGER.critical(e)
        raise e
    finally:
        _report_invalid_records(state_tracker.streams)


def _records_distribution(count, stats):
    return { (k, '{:.2%} ({})'.format(
        v / count,
        v
    )) for k, v in stats.items() }


def _report_invalid_records(streams):
    for stream_buffer in streams.values():
        if stream_buffer.peek_invalid_records():
            LOGGER.warning("Invalid records detected for stream {}: {}".format(
                stream_buffer.stream,
                stream_buffer.peek_invalid_records()
            ))


def _line_handler(line_stats, state_tracker, target, invalid_records_detect, invalid_records_threshold, max_batch_rows,
                  max_batch_size, line):
    try:
        line_data = json.loads(line)
    except json.decoder.JSONDecodeError:
        LOGGER.error("Unable to parse JSON: {}".format(line))
        raise

    if 'type' not in line_data:
        raise TargetError('`type` is a required key: {}'.format(line))

    line_type = line_data['type']

    if line_type == 'SCHEMA':
        if 'stream' not in line_data:
            raise TargetError('`stream` is a required key: {}'.format(line))

        stream = line_data['stream']

        if 'schema' not in line_data:
            raise TargetError('`schema` is a required key: {}'.format(line))

        schema = line_data['schema']

        schema_validation_errors = json_schema.validation_errors(schema)
        if schema_validation_errors:
            raise TargetError('`schema` is an invalid JSON Schema instance: {}'.format(line), *schema_validation_errors)

        if 'key_properties' in line_data:
            key_properties = line_data['key_properties']
        else:
            key_properties = None

        if stream not in state_tracker.streams:
            buffered_stream = BufferedSingerStream(stream,
                                                   schema,
                                                   key_properties,
                                                   invalid_records_detect=invalid_records_detect,
                                                   invalid_records_threshold=invalid_records_threshold,
                                                   max_rows=max_batch_rows,
                                                   max_buffer_size=max_batch_size)

            state_tracker.register_stream(stream, buffered_stream)
        else:
            state_tracker.streams[stream].update_schema(schema, key_properties)
    elif line_type == 'RECORD':
        if 'stream' not in line_data:
            raise TargetError('`stream` is a required key: {}'.format(line))

        state_tracker.handle_record_message(line_data['stream'], line_data)
    elif line_type == 'ACTIVATE_VERSION':
        if 'stream' not in line_data:
            raise TargetError('`stream` is a required key: {}'.format(line))
        if 'version' not in line_data:
            raise TargetError('`version` is a required key: {}'.format(line))
        if line_data['stream'] not in state_tracker.streams:
            raise TargetError('A ACTIVATE_VERSION for stream {} was encountered before a corresponding schema'
                              .format(line_data['stream']))

        stream_buffer = state_tracker.streams[line_data['stream']]
        state_tracker.flush_stream(line_data['stream'])
        target.activate_version(stream_buffer, line_data['version'])
    elif line_type == 'STATE':
        state_tracker.handle_state_message(line_data)
    else:
        raise TargetError('Unknown message type {} in message {}'.format(
            line_type,
            line))

    line_stats[line_type] = line_stats.get(line_type, 0) + 1


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


def _async_send_usage_stats():
    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()


def _run_sql_hook(hook_name, config, target):
    if hook_name in config:
        with target.conn.cursor() as cur:
            cur.execute(config[hook_name])
            LOGGER.debug('{} SQL executed'.format(hook_name))