cnn-classification-service/cnn_classification_service/main.py

#!/usr/bin/env python3
import logging
import os
import sys
import pika
import json

from sentry_sdk.integrations.logging import LoggingIntegration
import sentry_sdk

from magic_doer import run_everything


def message_callback(ch, method, properties, body):
    msg = json.loads(body.decode('utf-8'))
    results = run_everything(msg)

    # TODO: Ez azért elég gettó, de legalább csatlakozik
    connection = pika.BlockingConnection(pika.connection.URLParameters(os.environ['PIKA_URL']))
    channel = connection.channel()
    channel.exchange_declare(exchange=os.environ['PIKA_OUTPUT_EXCHANGE'], exchange_type='fanout')
    channel.basic_publish(exchange=os.environ['PIKA_OUTPUT_EXCHANGE'], routing_key='classification-result',
                          body=json.dumps(results).encode("utf-8"))


def main():
    logging.basicConfig(stream=sys.stdout, format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
                        level=logging.DEBUG if '--debug' in sys.argv else logging.INFO)

    SENTRY_DSN = os.environ.get("SENTRY_DSN")
    if SENTRY_DSN:
        sentry_logging = LoggingIntegration(
            level=logging.DEBUG,  # Capture info and above as breadcrumbs
            event_level=logging.ERROR  # Send errors as events
        )
        sentry_sdk.init(
            dsn=SENTRY_DSN,
            integrations=[sentry_logging],
            send_default_pii=True,
            release=os.environ.get('RELEASE_ID', 'test'),
            environment=os.environ.get('RELEASEMODE', 'dev')
        )

    logging.info("Connecting to MQ service...")
    connection = pika.BlockingConnection(pika.connection.URLParameters(os.environ['PIKA_URL']))
    channel = connection.channel()
    channel.exchange_declare(exchange=os.environ['PIKA_INPUT_EXCHANGE'], exchange_type='fanout')

    queue_declare_result = channel.queue_declare(queue='', exclusive=True)
    queue_name = queue_declare_result.method.queue

    channel.queue_bind(exchange=os.environ['PIKA_INPUT_EXCHANGE'], queue=queue_name)
    channel.basic_consume(queue=queue_name, on_message_callback=message_callback, auto_ack=True)

    logging.info("Connection complete! Listening to messages...")
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        logging.info("SIGINT Received! Stopping stuff...")
        channel.stop_consuming()


if __name__ == '__main__':
    main()