transcoder/Transcoder.py

#
# Copyright 2023 Google LLC
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# pylint: disable=broad-except

import importlib
import logging
import os
import signal
import sys
from datetime import datetime

from transcoder.message import DatacastParser, NoParser
from transcoder.message.ErrorWriter import ErrorWriter, TranscodeStep
from transcoder.message.MessageUtil import get_message_parser, parse_handler_config
from transcoder.output import get_output_manager
from transcoder.source import get_message_source


# pylint: disable=invalid-name
class Transcoder:  # pylint: disable=too-many-instance-attributes
    """ Main entry point for transcodihg sessions, bounded by a schema, source and parser """

    def __init__(self,  # pylint: disable=too-many-arguments),too-many-locals
                 factory: str, schema_file_path: str, source_file_path: str, source_file_encoding: str,
                 source_file_format_type: str, source_file_endian: str, prefix_length: int, skip_lines: int,
                 skip_bytes: int, message_skip_bytes: int, quiet: bool, output_type: str, output_encoding: str,
                 output_path: str, error_output_path: str, destination_project_id: str, destination_dataset_id: str,
                 message_handlers: str, lazy_create_resources: bool, frame_only: bool, stats_only: bool,
                 create_schemas_only: bool, continue_on_error: bool, create_schema_enforcing_topics: bool,
                 sampling_count: int, message_type_inclusions: str, message_type_exclusions: str, fix_header_tags: str,
                 fix_separator: int, base64: bool, base64_urlsafe: bool):

        signal.signal(signal.SIGINT, self.trap)

        self.message_handler_spec = message_handlers
        self.message_handlers = {}
        self.all_message_type_handlers = []
        self.all_handlers = []
        self.handlers_enabled = False
        self.continue_on_error = continue_on_error
        self.error_output_path = error_output_path
        self.output_encoding = output_encoding
        self.frame_only = frame_only
        self.create_schemas_only = create_schemas_only
        self.lazy_create_resources = lazy_create_resources
        self.prefix_length = prefix_length
        self.quiet = quiet
        self.start_time = None
        self.stats_only = stats_only
        self.sampling_count = sampling_count
        self.transcoded_count = 0

        self.output_prefix = os.path.basename(
            os.path.splitext(source_file_path)[0]) if source_file_path else 'stdin'

        self.error_writer = ErrorWriter(prefix=self.output_prefix,
                                        output_path=self.error_output_path)

        self.output_manager = get_output_manager(output_type, self.output_prefix, output_path,
                                                 output_encoding, self.prefix_length, destination_project_id,
                                                 destination_dataset_id, lazy_create_resources,
                                                 create_schema_enforcing_topics)

        # TODO: think about this abstraction some more
        if self.output_manager.supports_data_writing() is False:
            self.create_schemas_only = True
        else:
            self.source = get_message_source(source_file_path, source_file_encoding,
                                             source_file_format_type, source_file_endian,
                                             skip_bytes, skip_lines, message_skip_bytes,
                                             prefix_length, base64, base64_urlsafe)

        self.message_parser: DatacastParser = NoParser() if self.frame_only else get_message_parser(
            factory,
            schema_file_path,
            stats_only,
            message_type_inclusions,
            message_type_exclusions,
            fix_header_tags,
            fix_separator
        )

        self.setup_handlers()

    def transcode(self):
        """Entry point for transcoding session"""
        self.start_time = datetime.now()
        if self.frame_only is False:
            self.process_schemas()

        with self.source:
            for raw_msg in self.source.get_message_iterator():
                if self.frame_only:  # don't parse message
                    self.message_parser.process_message(raw_msg)
                    self.output_manager.write_record(None, raw_msg)
                else:  # parse message
                    if self.stats_only is False:  # output message
                        self.transcode_message(raw_msg)

                if self.transcoded_count == self.sampling_count:
                    break

        self.print_summary()

    def transcode_message(self, raw):
        """ Transcoding steps executed on each source message """
        self.error_writer.set_step(TranscodeStep.PARSE_MESSAGE)
        msg = None
        try:
            msg = self.message_parser.process_message(raw)

            if msg.exception is not None:
                self.handle_exception(raw, msg, msg.exception)

            if msg.ignored is False:  # passed inclusions / exclusions
                self.execute_handlers(msg)
                if msg.ignored is False:  # passed filters
                    self.error_writer.set_step(TranscodeStep.WRITE_OUTPUT_RECORD)
                    self.output_manager.write_record(msg.name, msg.dictionary)
                    self.transcoded_count += 1

        except Exception as ex:
            self.handle_exception(raw, msg, ex)

    def execute_handlers(self, message):
        """ Executes in sequence the message handlers specified for this transcoding instance """
        if self.handlers_enabled is True:  # execute handlers
            self.error_writer.set_step(TranscodeStep.EXECUTE_HANDLERS)
            for handler in self.all_message_type_handlers + self.message_handlers.get(message.type, []):
                self.error_writer.set_step(TranscodeStep.EXECUTE_HANDLER, type(handler).__name__)
                handler.handle(message)

    def setup_handlers(self):
        """Initialize MessageHandler instances to employ at runtime"""

        if self.message_handler_spec is None or self.message_handler_spec == "":
            return

        self.handlers_enabled = True
        handler_strs = self.message_handler_spec.split(',')
        for handler_spec in handler_strs:
            cls_name = None
            config_dict = None
            if handler_spec.find(':') == -1:  # no handler params
                cls_name = handler_spec
            else:
                cls_name = handler_spec.split(':')[0]
                config_dict = parse_handler_config(handler_spec)

            module = importlib.import_module('transcoder.message.handler')
            class_ = getattr(module, cls_name)
            instance = class_(config_dict)
            self.all_handlers.append(instance)

            if instance.supports_all_message_types is True:
                self.all_message_type_handlers.append(instance)
                continue

            supported_msg_types = instance.supported_message_types
            for supported_type in supported_msg_types:
                if supported_type in self.message_handlers:
                    handler_list = self.message_handlers[supported_type]
                    if instance not in handler_list:
                        self.message_handlers[supported_type].append(instance)
                else:
                    self.message_handlers[supported_type] = [instance]

    def print_summary(self):
        """Print summary of the messages that were processed"""
        if logging.getLogger().isEnabledFor(logging.INFO):
            end_time = datetime.now()
            time_diff = end_time - self.start_time
            total_seconds = time_diff.total_seconds()

            if self.create_schemas_only is True:
                logging.info('Run in create_schemas_only mode')

            if self.output_manager.supports_data_writing() is False:
                logging.info('Output manager \'%s\' does not support message writes',
                             self.output_manager.output_type_identifier())

            if self.frame_only is False:

                if self.message_parser.stats_only is True:
                    logging.info('Run in stats_only mode')

                if self.sampling_count is not None:
                    logging.info('Sampled messages: %s', self.sampling_count)

                if self.message_parser.message_type_inclusions is not None:
                    logging.info('Message type inclusions: %s', self.message_parser.message_type_inclusions)
                elif self.message_parser.message_type_exclusions is not None:
                    logging.info('Message type exclusions: %s', self.message_parser.message_type_exclusions)

                if self.create_schemas_only is False:
                    logging.info('Source message count: %s', self.source.record_count)
                    logging.info('Processed message count: %s', self.message_parser.record_count)
                    logging.info('Transcoded message count: %s', self.transcoded_count)
                    logging.info('Processed schema count: %s', self.message_parser.total_schema_count)
                    logging.info('Summary of message counts: %s', self.message_parser.record_type_count)
                    logging.info('Summary of error message counts: %s', self.message_parser.error_record_type_count)
                    logging.info('Message rate: %s per second', round(self.source.record_count / total_seconds, 6))

            else:
                logging.info('Source record count: %s', self.source.record_count)

            logging.info('Total runtime in seconds: %s', round(total_seconds, 6))
            logging.info('Total runtime in minutes: %s', round(total_seconds / 60, 6))

    def process_schemas(self):
        """Process the schema specified at runtime"""
        spec_schemas = self.message_parser.process_schema()
        for schema in spec_schemas:
            if self.output_manager.supports_zero_field_schemas() is False and len(schema.fields) == 0:
                logging.info('Schema "%s" contains no field definitions, skipping schema creation', schema.name)
                continue

            for handler in self.all_handlers:
                if handler.supports_all_message_types is True \
                        or schema.message_id in handler.supported_message_types:
                    handler.append_manufactured_fields(schema)

            self.output_manager.enqueue_schema(schema)

        # Only need to wait if lazy create is off, and you want to force creation before data is read
        if self.lazy_create_resources is False:
            self.output_manager.wait_for_schema_creation()

    def handle_exception(self, raw_record, message, exception):
        """Process exceptions encountered in the message processing runtime"""
        if message is not None:
            self.message_parser.increment_error_summary_count(message.name)
        else:
            self.message_parser.increment_error_summary_count()

        self.error_writer.write_error(raw_record, message, exception)

        if self.continue_on_error is False:
            raise exception

    def trap(self, _signum, _frame):
        """Trap SIGINT to suppress noisy stack traces and show interim summary"""
        print()
        self.print_summary()
        sys.exit(1)