<commit_before>import gzip
import json
import os
import tempfile
import re
from exporters.progress_callback import BotoDownloadProgress
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_long
from exporters.exceptions import ConfigurationError


def get_bucket(bucket, aws_access_key_id, aws_secret_access_key, **kwargs):
    import boto
    connection = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    return connection.get_bucket(bucket)


class S3BucketKeysFetcher(object):
    def __init__(self, reader_options):
        self.source_bucket = get_bucket(**reader_options)
        self.pattern = reader_options.get('pattern', None)
        single_prefix = reader_options.get('prefix', '')
        self.prefix_pointer = reader_options.get('prefix_pointer', '')
        if single_prefix and self.prefix_pointer:
            raise ConfigurationError("prefix and prefix_pointer options cannot be used together")
        self.prefixes = [single_prefix]
        if self.prefix_pointer:
            self.prefixes = self._fetch_prefixes_from_pointer(self.prefix_pointer)

    @retry_long
    def _download_pointer(self, prefix_pointer):
        return self.source_bucket.get_key(prefix_pointer).get_contents_as_string()

    def _fetch_prefixes_from_pointer(self, prefix_pointer):
        return self._download_pointer(prefix_pointer).split('\n')

    def _get_keys_from_bucket(self):
        keys = []
        for prefix in self.prefixes:
            for key in self.source_bucket.list(prefix=prefix):
                if self.pattern:
                    if self._should_add_key(key, prefix):
                        keys.append(key.name)
                else:
                    keys.append(key.name)
        return keys

    def _should_add_key(self, key, prefix):
        return bool(re.findall(self.pattern, key.name))

    def pending_keys(self):
        return self._get_keys_from_bucket()


class S3Reader(BaseReader):
    """
    Reads items from s3 keys with a common prefix.

        - batch_size (int)
            Number of items to be returned in each batch

        - bucket (str)
            Name of the bucket to retrieve items from.

        - aws_access_key_id (str)
            Public access key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - prefix (str)
            Prefix of s3 keys to be read.

        - prefix_pointer (str)
            Prefix pointing to the last version of dataset.

        - pattern (str)
            Regex pattern that keys should meet.

    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': int, 'default': 10000},
        'bucket': {'type': basestring},
        'aws_access_key_id': {'type': basestring, 'env_fallback': 'EXPORTERS_S3READER_AWS_KEY'},
        'aws_secret_access_key': {'type': basestring, 'env_fallback': 'EXPORTERS_S3READER_AWS_SECRET'},
        'prefix': {'type': basestring, 'default': ''},
        'prefix_pointer': {'type': basestring, 'default': None},
        'pattern': {'type': basestring, 'default': None}
    }

    def __init__(self, options):
        import boto

        super(S3Reader, self).__init__(options)
        self.batch_size = self.read_option('batch_size')
        self.connection = boto.connect_s3(self.read_option('aws_access_key_id'),
                                          self.read_option('aws_secret_access_key'))
        self.bucket = self.connection.get_bucket(self.read_option('bucket'))
        single_prefix = self.read_option('prefix')
        self.prefix_pointer = self.read_option('prefix_pointer')
        self.pattern = self.read_option('pattern')

        if single_prefix and self.prefix_pointer:
            raise ConfigurationError("prefix and prefix_pointer options cannot be used together")

        self.keys_fetcher = S3BucketKeysFetcher(options['options'])
        self.keys = self.keys_fetcher.pending_keys()
        self.read_keys = []
        self.current_key = None
        self.last_line = 0
        self.logger.info('S3Reader has been initiated')
        self.tmp_folder = tempfile.mkdtemp()

    @retry_long
    def get_key(self, file_path, progress):
        self.logger.info('Downloading key: %s' % self.current_key)
        self.bucket.get_key(self.current_key).get_contents_to_filename(file_path, cb=progress)

    def get_next_batch(self):
        file_path = '{}/ds_dump.gz'.format(self.tmp_folder)
        if not self.current_key:
            progress = BotoDownloadProgress(self.logger)
            self.current_key = self.keys[0]<commit_msg>S3Reader: log before trying to connect<commit_after>import gzip
import json
import os
import tempfile
import re
from exporters.progress_callback import BotoDownloadProgress
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_long
from exporters.exceptions import ConfigurationError


def get_bucket(bucket, aws_access_key_id, aws_secret_access_key, **kwargs):
    import boto
    connection = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    return connection.get_bucket(bucket)


class S3BucketKeysFetcher(object):
    def __init__(self, reader_options):
        self.source_bucket = get_bucket(**reader_options)
        self.pattern = reader_options.get('pattern', None)
        single_prefix = reader_options.get('prefix', '')
        self.prefix_pointer = reader_options.get('prefix_pointer', '')
        if single_prefix and self.prefix_pointer:
            raise ConfigurationError("prefix and prefix_pointer options cannot be used together")
        self.prefixes = [single_prefix]
        if self.prefix_pointer:
            self.prefixes = self._fetch_prefixes_from_pointer(self.prefix_pointer)

    @retry_long
    def _download_pointer(self, prefix_pointer):
        return self.source_bucket.get_key(prefix_pointer).get_contents_as_string()

    def _fetch_prefixes_from_pointer(self, prefix_pointer):
        return self._download_pointer(prefix_pointer).split('\n')

    def _get_keys_from_bucket(self):
        keys = []
        for prefix in self.prefixes:
            for key in self.source_bucket.list(prefix=prefix):
                if self.pattern:
                    if self._should_add_key(key, prefix):
                        keys.append(key.name)
                else:
                    keys.append(key.name)
        return keys

    def _should_add_key(self, key, prefix):
        return bool(re.findall(self.pattern, key.name))

    def pending_keys(self):
        return self._get_keys_from_bucket()


class S3Reader(BaseReader):
    """
    Reads items from s3 keys with a common prefix.

        - batch_size (int)
            Number of items to be returned in each batch

        - bucket (str)
            Name of the bucket to retrieve items from.

        - aws_access_key_id (str)
            Public access key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - prefix (str)
            Prefix of s3 keys to be read.

        - prefix_pointer (str)
            Prefix pointing to the last version of dataset.

        - pattern (str)
            Regex pattern that keys should meet.

    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': int, 'default': 10000},
        'bucket': {'type': basestring},
        'aws_access_key_id': {'type': basestring, 'env_fallback': 'EXPORTERS_S3READER_AWS_KEY'},
        'aws_secret_access_key': {'type': basestring, 'env_fallback': 'EXPORTERS_S3READER_AWS_SECRET'},
        'prefix': {'type': basestring, 'default': ''},
        'prefix_pointer': {'type': basestring, 'default': None},
        'pattern': {'type': basestring, 'default': None}
    }

    def __init__(self, options):
        import boto

        super(S3Reader, self).__init__(options)
        self.batch_size = self.read_option('batch_size')
        self.connection = boto.connect_s3(self.read_option('aws_access_key_id'),
                                          self.read_option('aws_secret_access_key'))
        bucket_name = self.read_option('bucket')
        self.logger.info('Initiating S3Reader with bucket: %s' % bucket_name)

        self.bucket = self.connection.get_bucket(bucket_name)
        single_prefix = self.read_option('prefix')
        self.prefix_pointer = self.read_option('prefix_pointer')
        self.pattern = self.read_option('pattern')

        if single_prefix and self.prefix_pointer:
            raise ConfigurationError("prefix and prefix_pointer options cannot be used together")

        self.keys_fetcher = S3BucketKeysFetcher(options['options'])
        self.keys = self.keys_fetcher.pending_keys()
        self.read_keys = []
        self.current_key = None
        self.last_line = 0
        self.logger.info('S3Reader has been initiated')
        self.tmp_folder = tempfile.mkdtemp()

    @retry_long
    def get_key(self, file_path, progress):
        self.logger.info('Downloading key: %s' % self.current_key)
        self.bucket.get_key(self.current_key).get_contents_to_filename(file_path, cb=progress)

    def get_next_batch(self):
        file_path = '{}/ds_dump.gz'.format(self.tmp_folder)
        if not self.current_key:
            progress = BotoDownloadProgress(self.logger)
            self.current_key = self.keys[0]
{'fraction_non_alphanumeric': 0.09434554973821989, 'fraction_numerical': 0.004397905759162304, 'mean_word_length': 2.8342031312725813, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '901858', 'n_tokens_mistral': 2806, 'n_tokens_neox': 2741, 'n_words': 628}