<commit_before>#!/usr/bin/env python
"""
Repackage a USGS Collection-1 tar for faster read access.

They arrive as a *.tar.gz with inner uncompressed tiffs, which Josh's tests have found to be too slow to read.

We compress the inner tiffs and store them in an uncompressed tar. This allows random reads within the files.
We also append a checksum file at the end of the tar.
"""
import copy
import io
import socket
import stat
import sys
import tarfile
import tempfile
import traceback
from contextlib import suppress
from functools import partial
from itertools import chain
from pathlib import Path
from subprocess import call
from typing import List, Iterable, Tuple, Callable, IO, Dict

import click
import numpy
import rasterio
import structlog
from structlog.processors import (
    StackInfoRenderer,
    TimeStamper,
    format_exc_info,
    JSONRenderer,
)

from eodatasets3.ui import PathPath
from eodatasets3.verify import PackageChecksum

_PREDICTOR_TABLE = {
    "int8": 2,
    "uint8": 2,
    "int16": 2,
    "uint16": 2,
    "int32": 2,
    "uint32": 2,
    "int64": 2,
    "uint64": 2,
    "float32": 3,
    "float64": 3,
}

# The info of a file, and a method to open the file for reading.
ReadableMember = Tuple[tarfile.TarInfo, Callable[[], IO]]

_LOG = structlog.get_logger()


class RecompressFailure(Exception):
    pass


def _create_tarinfo(path: Path, name=None) -> tarfile.TarInfo:
    """
    Create a TarInfo ("tar member") based on the given filesystem path.

    (these contain the information of a file, such as permissions, when writing to a tar file)

    This code is based on TarFile.gettarinfo(), but doesn't need an existing tar file.
    """
    # We're avoiding convenience methods like `path.is_file()`, to minimise repeated `stat()` calls on lustre.
    s = path.stat()
    info = tarfile.TarInfo(name or path.name)

    if stat.S_ISREG(s.st_mode):
        info.size = s.st_size
        info.type = tarfile.REGTYPE
    elif stat.S_ISDIR(s.st_mode):
        info.type = tarfile.DIRTYPE
        info.size = 0
    else:
        raise NotImplementedError(
            f"Only regular files and directories are supported for extracted datasets. "
            f"({path.name} in {path.absolute().parent})"
        )

    info.mode = s.st_mode
    info.uid = s.st_uid
    info.gid = s.st_gid
    info.mtime = s.st_mtime

    if tarfile.pwd:
        try:
            info.uname = tarfile.pwd.getpwuid(info.uid)[0]
        except KeyError:
            pass
    if tarfile.grp:
        try:
            info.gname = tarfile.grp.getgrgid(info.gid)[0]
        except KeyError:
            pass
    return info


def _tar_members(in_tar: tarfile.TarFile) -> Iterable[ReadableMember]:
    """Get readable files (members) from a tar"""
    members: List[tarfile.TarInfo] = in_tar.getmembers()

    for member in members:
        # We return a lambda/callable so that the file isn't opened until it's needed.
        yield member, partial(in_tar.extractfile, member)


def _folder_members(path: Path, base_path: Path = None) -> Iterable[ReadableMember]:
    """
    Get readable files (presented as tar members) from a directory.
    """
    if not base_path:
        base_path = path

    # Note that the members in input USGS tars are sorted alphabetically.
    # We'll sort our own inputs to match.
    # (The primary practical benefit is predictable outputs in tests)

    for item in sorted(path.iterdir()):
        member = _create_tarinfo(item, name=str(item.relative_to(base_path)))
        if member.type == tarfile.DIRTYPE:
            yield member, None
            yield from _folder_members(item, base_path=path)
        else:
            # We return a lambda/callable so that the file isn't opened until it's needed.
            yield member, partial(item.open, "rb")


def repackage_tar(
    input_path: Path,
    input_files: Iterable[ReadableMember],
    output_tar_path: Path,
    clean_inputs: bool,
    **compress_args,
) -> bool:
    log = _LOG.bind(
        name=output_tar_path.stem,
        in_path=str(input_path.absolute()),
        out_path=str(output_tar_path.absolute()),
    )

    if output_tar_path.exists():
        log.info("skip.exists")
        return True

    try:
        members = list(input_files)

        # Add the MTL file to the beginning of the output tar, so it can be accessed faster.
        # This slows down this repackage a little, as we're seeking/decompressing the input stream an extra time.
        _reorder_tar_members(members, input_path.name)

        _create_tar_with_files(input_path, members, output_tar_path, **compress_args)

        log.info(
            "complete",
            in_size=sum(m.size for m, _ in members),
            in_count=len(members),
            # The user/group give us a hint as to whether this was repackaged outside of USGS.
            in_users=list({(member.uname, member.gname) for member, _ in members}),
            out_size=output_tar_path.stat().st_size,
        )

        result_exists = output_tar_path.exists()
        if not result_exists:
            # This should never happen, so it's an exception.
            raise RuntimeError(f"No output after a success? Expected {output_tar_path}")

        if clean_inputs:
            # Be extra sure the new file was flushed to disk before we clean the original.
            call("sync")
            log.info("input.cleanup")<commit_msg>Don't call sync on gadi<commit_after>#!/usr/bin/env python
"""
Repackage a USGS Collection-1 tar for faster read access.

They arrive as a *.tar.gz with inner uncompressed tiffs, which Josh's tests have found to be too slow to read.

We compress the inner tiffs and store them in an uncompressed tar. This allows random reads within the files.
We also append a checksum file at the end of the tar.
"""
import copy
import io
import socket
import stat
import sys
import tarfile
import tempfile
import traceback
from contextlib import suppress
from functools import partial
from itertools import chain
from pathlib import Path
from typing import List, Iterable, Tuple, Callable, IO, Dict

import click
import numpy
import rasterio
import structlog
from structlog.processors import (
    StackInfoRenderer,
    TimeStamper,
    format_exc_info,
    JSONRenderer,
)

from eodatasets3.ui import PathPath
from eodatasets3.verify import PackageChecksum

_PREDICTOR_TABLE = {
    "int8": 2,
    "uint8": 2,
    "int16": 2,
    "uint16": 2,
    "int32": 2,
    "uint32": 2,
    "int64": 2,
    "uint64": 2,
    "float32": 3,
    "float64": 3,
}

# The info of a file, and a method to open the file for reading.
ReadableMember = Tuple[tarfile.TarInfo, Callable[[], IO]]

_LOG = structlog.get_logger()


class RecompressFailure(Exception):
    pass


def _create_tarinfo(path: Path, name=None) -> tarfile.TarInfo:
    """
    Create a TarInfo ("tar member") based on the given filesystem path.

    (these contain the information of a file, such as permissions, when writing to a tar file)

    This code is based on TarFile.gettarinfo(), but doesn't need an existing tar file.
    """
    # We're avoiding convenience methods like `path.is_file()`, to minimise repeated `stat()` calls on lustre.
    s = path.stat()
    info = tarfile.TarInfo(name or path.name)

    if stat.S_ISREG(s.st_mode):
        info.size = s.st_size
        info.type = tarfile.REGTYPE
    elif stat.S_ISDIR(s.st_mode):
        info.type = tarfile.DIRTYPE
        info.size = 0
    else:
        raise NotImplementedError(
            f"Only regular files and directories are supported for extracted datasets. "
            f"({path.name} in {path.absolute().parent})"
        )

    info.mode = s.st_mode
    info.uid = s.st_uid
    info.gid = s.st_gid
    info.mtime = s.st_mtime

    if tarfile.pwd:
        try:
            info.uname = tarfile.pwd.getpwuid(info.uid)[0]
        except KeyError:
            pass
    if tarfile.grp:
        try:
            info.gname = tarfile.grp.getgrgid(info.gid)[0]
        except KeyError:
            pass
    return info


def _tar_members(in_tar: tarfile.TarFile) -> Iterable[ReadableMember]:
    """Get readable files (members) from a tar"""
    members: List[tarfile.TarInfo] = in_tar.getmembers()

    for member in members:
        # We return a lambda/callable so that the file isn't opened until it's needed.
        yield member, partial(in_tar.extractfile, member)


def _folder_members(path: Path, base_path: Path = None) -> Iterable[ReadableMember]:
    """
    Get readable files (presented as tar members) from a directory.
    """
    if not base_path:
        base_path = path

    # Note that the members in input USGS tars are sorted alphabetically.
    # We'll sort our own inputs to match.
    # (The primary practical benefit is predictable outputs in tests)

    for item in sorted(path.iterdir()):
        member = _create_tarinfo(item, name=str(item.relative_to(base_path)))
        if member.type == tarfile.DIRTYPE:
            yield member, None
            yield from _folder_members(item, base_path=path)
        else:
            # We return a lambda/callable so that the file isn't opened until it's needed.
            yield member, partial(item.open, "rb")


def repackage_tar(
    input_path: Path,
    input_files: Iterable[ReadableMember],
    output_tar_path: Path,
    clean_inputs: bool,
    **compress_args,
) -> bool:
    log = _LOG.bind(
        name=output_tar_path.stem,
        in_path=str(input_path.absolute()),
        out_path=str(output_tar_path.absolute()),
    )

    if output_tar_path.exists():
        log.info("skip.exists")
        return True

    try:
        members = list(input_files)

        # Add the MTL file to the beginning of the output tar, so it can be accessed faster.
        # This slows down this repackage a little, as we're seeking/decompressing the input stream an extra time.
        _reorder_tar_members(members, input_path.name)

        _create_tar_with_files(input_path, members, output_tar_path, **compress_args)

        log.info(
            "complete",
            in_size=sum(m.size for m, _ in members),
            in_count=len(members),
            # The user/group give us a hint as to whether this was repackaged outside of USGS.
            in_users=list({(member.uname, member.gname) for member, _ in members}),
            out_size=output_tar_path.stat().st_size,
        )

        result_exists = output_tar_path.exists()
        if not result_exists:
            # This should never happen, so it's an exception.
            raise RuntimeError(f"No output after a success? Expected {output_tar_path}")

        if clean_inputs:
            log.info("input.cleanup")
{'fraction_non_alphanumeric': 0.08848314606741572, 'fraction_numerical': 0.006367041198501872, 'mean_word_length': 2.7437784787942516, 'pattern_counts': {'":': 20, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2866225', 'n_tokens_mistral': 3342, 'n_tokens_neox': 3215, 'n_words': 1073}