<commit_before>import asyncio
import json
import logging
import os.path
import zlib

import aiobotocore
import botocore
import kinto_http

from buildhub.inventory_to_records import NB_RETRY_REQUEST, csv_to_records
from buildhub.to_kinto import main as to_kinto


REGION_NAME = 'us-east-1'
BUCKET = "net-mozaws-prod-delivery-inventory-us-east-1"
FOLDER = "public/inventories/net-mozaws-prod-delivery-{inventory}/delivery-{inventory}/"
CHUNK_SIZE = 1024 * 256  # 256 KB


logger = logging.getLogger()  # root logger.


async def list_manifest_entries(loop, client, inventory):
    prefix = FOLDER.format(inventory=inventory)
    paginator = client.get_paginator('list_objects')

    manifest_folders = []
    async for result in paginator.paginate(Bucket=BUCKET, Prefix=prefix, Delimiter='/'):
        # Take latest inventory.
        files = list(result.get('CommonPrefixes', []))
        manifest_folders += [f['Prefix'] for f in files]

    # Download latest manifest.json
    last_inventory = sorted(manifest_folders)[-2]  # -1 is data
    logger.info("Latest inventory is {}".format(last_inventory))
    key = last_inventory + 'manifest.json'
    manifest = await client.get_object(Bucket=BUCKET, Key=key)
    async with manifest['Body'] as stream:
        body = await stream.read()
    manifest_content = json.loads(body.decode('utf-8'))
    # Return keys of csv.gz files
    for f in manifest_content['files']:
        yield f['key']


async def download_csv(loop, client, keys_stream, chunk_size=CHUNK_SIZE):
    gzip = zlib.decompressobj(zlib.MAX_WBITS | 16)

    async for key in keys_stream:
        key = "public/" + key
        logger.info("Fetching inventory piece {}".format(key))
        file_csv_gz = await client.get_object(Bucket=BUCKET, Key=key)
        async with file_csv_gz['Body'] as stream:
            while "there are chunks to read":
                gzip_chunk = await stream.read(chunk_size)
                if not gzip_chunk:
                    break
                csv_chunk = gzip.decompress(gzip_chunk)
                yield csv_chunk


async def main(loop, inventory):
    """
    Trigger to populate kinto with the last inventories.
    """
    server_url = os.getenv("SERVER_URL", "http://localhost:8888/v1")
    bucket = os.getenv("BUCKET", "default")
    collection = os.getenv("COLLECTION", "releases")
    kinto_auth = tuple(os.getenv("AUTH", "user:pass").split(":"))

    kinto_client = kinto_http.Client(server_url=server_url, auth=kinto_auth,
                                     bucket=bucket, collection=collection,
                                     retry=NB_RETRY_REQUEST)

    session = aiobotocore.get_session(loop=loop)
    boto_config = botocore.config.Config(signature_version=botocore.UNSIGNED)
    async with session.create_client('s3', region_name=REGION_NAME, config=boto_config) as client:
        keys_stream = list_manifest_entries(loop, client, inventory)
        csv_stream = download_csv(loop, client, keys_stream)
        records_stream = csv_to_records(loop, csv_stream)
        await to_kinto(loop, records_stream, kinto_client, skip_existing=True)

<commit_msg>Do not reuse gzip decompressor instance (fixes #223)<commit_after>import asyncio
import json
import logging
import os.path
import zlib

import aiobotocore
import botocore
import kinto_http

from buildhub.inventory_to_records import NB_RETRY_REQUEST, csv_to_records
from buildhub.to_kinto import main as to_kinto


REGION_NAME = 'us-east-1'
BUCKET = "net-mozaws-prod-delivery-inventory-us-east-1"
FOLDER = "public/inventories/net-mozaws-prod-delivery-{inventory}/delivery-{inventory}/"
CHUNK_SIZE = 1024 * 256  # 256 KB


logger = logging.getLogger()  # root logger.


async def list_manifest_entries(loop, client, inventory):
    prefix = FOLDER.format(inventory=inventory)
    paginator = client.get_paginator('list_objects')

    manifest_folders = []
    async for result in paginator.paginate(Bucket=BUCKET, Prefix=prefix, Delimiter='/'):
        # Take latest inventory.
        files = list(result.get('CommonPrefixes', []))
        manifest_folders += [f['Prefix'] for f in files]

    # Download latest manifest.json
    last_inventory = sorted(manifest_folders)[-2]  # -1 is data
    logger.info("Latest inventory is {}".format(last_inventory))
    key = last_inventory + 'manifest.json'
    manifest = await client.get_object(Bucket=BUCKET, Key=key)
    async with manifest['Body'] as stream:
        body = await stream.read()
    manifest_content = json.loads(body.decode('utf-8'))
    # Return keys of csv.gz files
    for f in manifest_content['files']:
        yield f['key']


async def download_csv(loop, client, keys_stream, chunk_size=CHUNK_SIZE):

    async for key in keys_stream:
        key = "public/" + key
        logger.info("Fetching inventory piece {}".format(key))
        file_csv_gz = await client.get_object(Bucket=BUCKET, Key=key)
        gzip = zlib.decompressobj(zlib.MAX_WBITS | 16)
        async with file_csv_gz['Body'] as stream:
            while "there are chunks to read":
                gzip_chunk = await stream.read(chunk_size)
                if not gzip_chunk:
                    break
                csv_chunk = gzip.decompress(gzip_chunk)
                if not csv_chunk:
                    error_msg = "Empty data (deflating {}B)".format(len(gzip_chunk))
                    raise ValueError(error_msg)
                yield csv_chunk


async def main(loop, inventory):
    """
    Trigger to populate kinto with the last inventories.
    """
    server_url = os.getenv("SERVER_URL", "http://localhost:8888/v1")
    bucket = os.getenv("BUCKET", "default")
    collection = os.getenv("COLLECTION", "releases")
    kinto_auth = tuple(os.getenv("AUTH", "user:pass").split(":"))

    kinto_client = kinto_http.Client(server_url=server_url, auth=kinto_auth,
                                     bucket=bucket, collection=collection,
                                     retry=NB_RETRY_REQUEST)

    session = aiobotocore.get_session(loop=loop)
    boto_config = botocore.config.Config(signature_version=botocore.UNSIGNED)
    async with session.create_client('s3', region_name=REGION_NAME, config=boto_config) as client:
        keys_stream = list_manifest_entries(loop, client, inventory)
        csv_stream = download_csv(loop, client, keys_stream)
        records_stream = csv_to_records(loop, csv_stream)
        await to_kinto(loop, records_stream, kinto_client, skip_existing=True)


{'fraction_non_alphanumeric': 0.10479134466769706, 'fraction_numerical': 0.007573415765069552, 'mean_word_length': 3.251642575558476, 'pattern_counts': {'":': 4, '<': 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': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '231745', 'n_tokens_mistral': 2083, 'n_tokens_neox': 1999, 'n_words': 445}