<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division


import math
import time
import uuid
import logging

from apps.core.models import LayerImage, Layer
from apps.workers.image_validator import ImageValidator
from apps.workers.sqs_manager import SQSManager
from apps.workers.thumbnail import make_thumbs_for_layer
from apps.workers.emr import create_cluster
import apps.core.enums as enums
import apps.workers.status_updates as status_updates
from apps.workers.copy_images import s3_copy

log = logging.getLogger(__name__)

TIMEOUT_SECONDS = (60 * 60)  # 60 Minutes.

JOB_COPY_IMAGE = 'copy_image'
JOB_VALIDATE = [
    'ObjectCreated:CompleteMultipartUpload',
    'ObjectCreated:Put',
    'ObjectCreated:Post',
    'ObjectCreated:Copy'
]
JOB_THUMBNAIL = 'thumbnail'
JOB_HANDOFF = 'emr_handoff'
JOB_TIMEOUT = 'timeout'

MAX_WAIT = 20  # Seconds.
TIMEOUT_DELAY = 60

# Error messages
ERROR_MESSAGE_THUMBNAIL_FAILED = 'Thumbnail generation failed.'
ERROR_MESSAGE_IMAGE_TRANSFER = 'Transferring image failed.'
ERROR_MESSAGE_JOB_TIMEOUT = 'Layer could not be processed. ' + \
                            'The job timed out after ' + \
                            str(math.ceil(TIMEOUT_SECONDS / 60)) + ' minutes.'


class QueueProcessor(object):
    """
    Encapsulates behavior for connecting to an SQS queue and moving image data
    through the processing pipeline by adding new messages to the same queue.
    """
    def __init__(self):
        self.queue = SQSManager()

    def start_polling(self):
        while True:
            message = self.queue.get_message()
            if not message:
                continue

            record = message['payload']

            # Flatten S3 event.
            if 'Records' in record and record['Records'][0]:
                record = record['Records'][0]

            delete_message = self.handle_message(record)
            if delete_message:
                self.queue.remove_message(message)

    def handle_message(self, record):
        # Parse record fields.
        if self.is_test_event(record):
            return True

        try:
            job_type = record['eventName']
            event_source = record['eventSource']
        except KeyError:
            log.exception('Missing fields in message %s', record)
            return False

        log.info('Executing job %s', job_type)

        if job_type == JOB_THUMBNAIL:
            return self.thumbnail(record)
        elif job_type == JOB_HANDOFF:
            return self.emr_hand_off(record)
        elif job_type == JOB_TIMEOUT:
            return self.check_timeout(record)
        elif job_type in JOB_VALIDATE:
            if event_source == 'aws:s3':
                return self.validate_image(record)
        elif job_type == JOB_COPY_IMAGE:
            return self.copy_image(record)

        return False

    def is_test_event(self, record):
        return 'Event' in record and record['Event'] == 's3:TestEvent'

    def validate_image(self, record):
        """
        Use Gdal to verify an image is properly formatted and can be further
        processed.
        record -- attribute data from SQS.
        """
        try:
            key = record['s3']['object']['key']
        except KeyError:
            return False

        s3_uuid = extract_uuid_from_aws_key(key)
        try:
            uuid.UUID(s3_uuid)
        except ValueError:
            return False

        # Ignore thumbnails.
        if not LayerImage.objects.filter(s3_uuid=s3_uuid).exists():
            return True

        byte_range = '0-1000000'  # Get first Mb(ish) of bytes.
        # Image verification.
        validator = ImageValidator(key, byte_range)
        if not validator.image_format_is_supported():
            log.info('Image format not supported %s', key)
            return status_updates.mark_image_invalid(
                s3_uuid, validator.get_error())
        elif not validator.image_has_min_bands(3):
            log.info('Not enough bands %s', key)
            return status_updates.mark_image_invalid(
                s3_uuid, validator.get_error())
        else:
            updated = status_updates.mark_image_valid(s3_uuid)
            if updated:
                log.info('Image validated %s', key)
                layer_id = status_updates.get_layer_id_from_uuid(s3_uuid)
                layer_images = LayerImage.objects.filter(layer_id=layer_id)
                valid_images = LayerImage.objects.filter(
                    layer_id=layer_id,
                    status=enums.STATUS_VALID)

                all_valid = len(layer_images) == len(valid_images)
                some_images = len(layer_images) > 0

                log.info('%d/%d images validated for layer %d',
                         len(valid_images),
                         len(layer_images),
                         layer_id)

                if all_valid and some_images:
                    log.info('Layer %d is valid', layer_id)
                    status_updates.update_layer_status(
                        layer_id,
                        enums.STATUS_VALID)

                    data = {'layer_id': layer_id}
                    log.info('Queue thumbnail job')
                    self.queue.add_message(JOB_THUMBNAIL, data)

                    # Add a message to the queue to watch for timeouts.
                    data = {
                        'timeout': time.time() + TIMEOUT_SECONDS,
                        'layer_id': layer_id
                    }
                    log.info('Queue timeout job')
                    self.queue.add_message(JOB_TIMEOUT, data, TIMEOUT_DELAY)
            return updated

    def thumbnail(self, record):
        """
        Generate thumbnails for all images.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
        except KeyError:
            return False

        try:
            int(layer_id)
        except ValueError:
            return False

        layer_id = data['layer_id']
        log.info('Generating thumbnails for layer %d...', layer_id)

        if make_thumbs_for_layer(layer_id):
            data = {'layer_id': layer_id}
            log.info('Queue handoff job')
            self.queue.add_message(JOB_HANDOFF, data)
            return True
        else:
            log.info('Failed to thumbnail')
            status_updates.update_layer_status(layer_id,
                                               enums.STATUS_FAILED,
                                               ERROR_MESSAGE_THUMBNAIL_FAILED)
            return False

    def emr_hand_off(self, record):
        """
        Passes layer imgages to EMR to begin creating custom rasters.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
        except KeyError:
            return False

        try:
            int(layer_id)
        except ValueError:
            return False

        layer_images = LayerImage.objects.filter(layer_id=layer_id)
        valid_images = LayerImage.objects.filter(
            layer_id=layer_id,
            status=enums.STATUS_THUMBNAILED)
        all_valid = len(layer_images) == len(valid_images)
        some_images = len(layer_images) > 0
        ready_to_process = some_images and all_valid

        log.info('%d/%d images thumbnailed for layer %d',
                 len(valid_images),
                 len(layer_images),
                 layer_id)

        if ready_to_process:
            layer = Layer.objects.get(id=layer_id)
            status_updates.update_layer_status(layer.id,
                                               enums.STATUS_PROCESSING)
            log.info('Launching EMR cluster...')
            create_cluster(layer)
            return True
        else:
            return some_images

    def check_timeout(self, record):
        """
        Check an EMR job to see if it has completed before the timeout has
        been reached.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
            timeout = data['timeout']
        except KeyError:
            # A bad message showed up. Leave it alone. Eventually it'll end
            # in the dead letter queue.
            return False

        try:
            int(layer_id)
            float(timeout)
        except ValueError:
            return False

        if time.time() > timeout:
            try:
                layer = Layer.objects.get(id=layer_id)
            except Layer.DoesNotExist:
                return False

            if layer is not None and layer.status != enums.STATUS_COMPLETED:
                return status_updates \
                    .update_layer_status(layer_id,
                                         enums.STATUS_FAILED,
                                         ERROR_MESSAGE_JOB_TIMEOUT)
        else:
            # Requeue the timeout message.
            self.queue.add_message(JOB_TIMEOUT, data, TIMEOUT_DELAY)

        return True

    def copy_image(self, record):
        """
        Copy an image into the S3 bucket from an external source.
        data -- attribute data from SQS.
        """
        log.info('Copying S3 image...')

        try:
            data = record['data']
            image_id = data['image_id']
        except KeyError:
            return False

        try:
            int(image_id)
        except ValueError:
            return False

        try:<commit_msg>Cast ids to int<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division


import math
import time
import uuid
import logging

from apps.core.models import LayerImage, Layer
from apps.workers.image_validator import ImageValidator
from apps.workers.sqs_manager import SQSManager
from apps.workers.thumbnail import make_thumbs_for_layer
from apps.workers.emr import create_cluster
import apps.core.enums as enums
import apps.workers.status_updates as status_updates
from apps.workers.copy_images import s3_copy

log = logging.getLogger(__name__)

TIMEOUT_SECONDS = (60 * 60)  # 60 Minutes.

JOB_COPY_IMAGE = 'copy_image'
JOB_VALIDATE = [
    'ObjectCreated:CompleteMultipartUpload',
    'ObjectCreated:Put',
    'ObjectCreated:Post',
    'ObjectCreated:Copy'
]
JOB_THUMBNAIL = 'thumbnail'
JOB_HANDOFF = 'emr_handoff'
JOB_TIMEOUT = 'timeout'

MAX_WAIT = 20  # Seconds.
TIMEOUT_DELAY = 60

# Error messages
ERROR_MESSAGE_THUMBNAIL_FAILED = 'Thumbnail generation failed.'
ERROR_MESSAGE_IMAGE_TRANSFER = 'Transferring image failed.'
ERROR_MESSAGE_JOB_TIMEOUT = 'Layer could not be processed. ' + \
                            'The job timed out after ' + \
                            str(math.ceil(TIMEOUT_SECONDS / 60)) + ' minutes.'


class QueueProcessor(object):
    """
    Encapsulates behavior for connecting to an SQS queue and moving image data
    through the processing pipeline by adding new messages to the same queue.
    """
    def __init__(self):
        self.queue = SQSManager()

    def start_polling(self):
        while True:
            message = self.queue.get_message()
            if not message:
                continue

            record = message['payload']

            # Flatten S3 event.
            if 'Records' in record and record['Records'][0]:
                record = record['Records'][0]

            delete_message = self.handle_message(record)
            if delete_message:
                self.queue.remove_message(message)

    def handle_message(self, record):
        # Parse record fields.
        if self.is_test_event(record):
            return True

        try:
            job_type = record['eventName']
            event_source = record['eventSource']
        except KeyError:
            log.exception('Missing fields in message %s', record)
            return False

        log.info('Executing job %s', job_type)

        if job_type == JOB_THUMBNAIL:
            return self.thumbnail(record)
        elif job_type == JOB_HANDOFF:
            return self.emr_hand_off(record)
        elif job_type == JOB_TIMEOUT:
            return self.check_timeout(record)
        elif job_type in JOB_VALIDATE:
            if event_source == 'aws:s3':
                return self.validate_image(record)
        elif job_type == JOB_COPY_IMAGE:
            return self.copy_image(record)

        return False

    def is_test_event(self, record):
        return 'Event' in record and record['Event'] == 's3:TestEvent'

    def validate_image(self, record):
        """
        Use Gdal to verify an image is properly formatted and can be further
        processed.
        record -- attribute data from SQS.
        """
        try:
            key = record['s3']['object']['key']
        except KeyError:
            return False

        s3_uuid = extract_uuid_from_aws_key(key)
        try:
            uuid.UUID(s3_uuid)
        except ValueError:
            return False

        # Ignore thumbnails.
        if not LayerImage.objects.filter(s3_uuid=s3_uuid).exists():
            return True

        byte_range = '0-1000000'  # Get first Mb(ish) of bytes.
        # Image verification.
        validator = ImageValidator(key, byte_range)
        if not validator.image_format_is_supported():
            log.info('Image format not supported %s', key)
            return status_updates.mark_image_invalid(
                s3_uuid, validator.get_error())
        elif not validator.image_has_min_bands(3):
            log.info('Not enough bands %s', key)
            return status_updates.mark_image_invalid(
                s3_uuid, validator.get_error())
        else:
            updated = status_updates.mark_image_valid(s3_uuid)
            if updated:
                log.info('Image validated %s', key)
                layer_id = status_updates.get_layer_id_from_uuid(s3_uuid)
                layer_images = LayerImage.objects.filter(layer_id=layer_id)
                # TODO: Filter `layer_images` instead of extra query.
                valid_images = LayerImage.objects.filter(
                    layer_id=layer_id,
                    status=enums.STATUS_VALID)

                all_valid = len(layer_images) == len(valid_images)
                some_images = len(layer_images) > 0

                log.info('%d/%d images validated for layer %d',
                         len(valid_images),
                         len(layer_images),
                         layer_id)

                if all_valid and some_images:
                    log.info('Layer %d is valid', layer_id)
                    status_updates.update_layer_status(
                        layer_id,
                        enums.STATUS_VALID)

                    data = {'layer_id': layer_id}
                    log.info('Queue thumbnail job')
                    self.queue.add_message(JOB_THUMBNAIL, data)

                    # Add a message to the queue to watch for timeouts.
                    data = {
                        'timeout': time.time() + TIMEOUT_SECONDS,
                        'layer_id': layer_id
                    }
                    log.info('Queue timeout job')
                    self.queue.add_message(JOB_TIMEOUT, data, TIMEOUT_DELAY)
            return updated

    def thumbnail(self, record):
        """
        Generate thumbnails for all images.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
        except KeyError:
            return False

        try:
            layer_id = int(layer_id)
        except ValueError:
            return False

        log.info('Generating thumbnails for layer %d...', layer_id)

        if make_thumbs_for_layer(layer_id):
            data = {'layer_id': layer_id}
            log.info('Queue handoff job')
            self.queue.add_message(JOB_HANDOFF, data)
            return True
        else:
            log.info('Failed to thumbnail')
            status_updates.update_layer_status(layer_id,
                                               enums.STATUS_FAILED,
                                               ERROR_MESSAGE_THUMBNAIL_FAILED)
            return False

    def emr_hand_off(self, record):
        """
        Passes layer imgages to EMR to begin creating custom rasters.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
        except KeyError:
            return False

        try:
            layer_id = int(layer_id)
        except ValueError:
            return False

        layer_images = LayerImage.objects.filter(layer_id=layer_id)
        # TODO: Filter `layer_images` instead of extra query.
        valid_images = LayerImage.objects.filter(
            layer_id=layer_id,
            status=enums.STATUS_THUMBNAILED)
        all_valid = len(layer_images) == len(valid_images)
        some_images = len(layer_images) > 0
        ready_to_process = some_images and all_valid

        log.info('%d/%d images thumbnailed for layer %d',
                 len(valid_images),
                 len(layer_images),
                 layer_id)

        if ready_to_process:
            layer = Layer.objects.get(id=layer_id)
            status_updates.update_layer_status(layer.id,
                                               enums.STATUS_PROCESSING)
            log.info('Launching EMR cluster...')
            create_cluster(layer)
            return True
        else:
            return some_images

    def check_timeout(self, record):
        """
        Check an EMR job to see if it has completed before the timeout has
        been reached.
        record -- attribute data from SQS.
        """
        try:
            data = record['data']
            layer_id = data['layer_id']
            timeout = data['timeout']
        except KeyError:
            # A bad message showed up. Leave it alone. Eventually it'll end
            # in the dead letter queue.
            return False

        try:
            int(layer_id)
            float(timeout)
        except ValueError:
            return False

        if time.time() > timeout:
            try:
                layer = Layer.objects.get(id=layer_id)
            except Layer.DoesNotExist:
                return False

            if layer is not None and layer.status != enums.STATUS_COMPLETED:
                return status_updates \
                    .update_layer_status(layer_id,
                                         enums.STATUS_FAILED,
                                         ERROR_MESSAGE_JOB_TIMEOUT)
        else:
            # Requeue the timeout message.
            self.queue.add_message(JOB_TIMEOUT, data, TIMEOUT_DELAY)

        return True

    def copy_image(self, record):
        """
        Copy an image into the S3 bucket from an external source.
        data -- attribute data from SQS.
        """
        log.info('Copying S3 image...')

        try:
            data = record['data']
            image_id = data['image_id']
        except KeyError:
            return False

        try:
            image_id = int(image_id)
        except ValueError:
            return False

        try:
{'fraction_non_alphanumeric': 0.0773060931163698, 'fraction_numerical': 0.0042092295056721934, 'mean_word_length': 1.6979642708766098, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7533445', 'n_tokens_mistral': 5359, 'n_tokens_neox': 5153, 'n_words': 1315}