subprocess-util/btrfs_send_chunks.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
from pathlib import Path
import shlex

from exec_print_transfer import execute_print_transfer_chunks
from common import _get_chunk_file, _get_remote_socket
from transfer_inform import transfer_inform


def main():
    args = parse_args()

    command_parts = (
        ['btrfs', 'send'],
        ['-p', str(args.parent)] if args.parent else [],
        ['--compress-data'] if args.compressed_data else [],
        [str(args.child)]
    )
    command = [x for xs in command_parts for x in xs]

    execute_print_transfer_chunks(
        command=command,
        chunk_file_tmpl=args.chunk_tmpl,
        chunk_transfer_fun=chunk_transfer_fun,
        chunk_transfer_args=(args.ssh_target, args.target_path, args.chunk_tmpl),
        chunk_size=args.chunk_size,
    )


def parse_args():
    parser = argparse.ArgumentParser(prog='btrfs-send-chunks')

    parser.add_argument('-p',
                        help='Parent subvolume; forwarded to btrfs-send.',
                        dest='parent',
                        default=None,
                        type=Path,
                        metavar='PARENT_SUBVOLUME'
                        )

    parser.add_argument('--compressed-data',
                        help='Forwarded to btrfs-send.',
                        dest='compressed_data',
                        action='store_true',
                        default=False,
                        )

    parser.add_argument('--chunk-tmpl',
                        help='During btrfs-send, chunks are saved as "CHUNK_TMPL.CHUNK_NUMBER". '
                             'The default value of CHUNK_TMPL is "CHILD_SUBVOLUME.CHUNK". '
                             'One can change it to e.g. "/tmp/chunk/CHILD_SUBVOLUME".',
                        dest='chunk_tmpl',
                        default=None,
                        type=Path,
                        metavar='CHUNK_TMPL'
                        )

    parser.add_argument('--chunk-size',
                        help='Size in bytes; defaults to 64 MB.',
                        dest='chunk_size',
                        type=int,
                        default=64 * 1024 * 1024,
                        )

    parser.add_argument('child',
                        help='Forwarded to btrfs-send.',
                        type=Path,
                        metavar='CHILD_SUBVOLUME'
                        )

    parser.add_argument('ssh_target',
                        help='Hostname of target computer; as configured in ~/.ssh/config.',
                        metavar='SSH_TARGET'
                        )

    parser.add_argument('target_path',
                        help='Path where the subvolume will be created on the remote ssh target.',
                        type=Path,
                        metavar='TARGET_PATH'
                        )

    args = parser.parse_args()

    if not args.chunk_tmpl:
        child: Path = args.child
        args.chunk_tmpl = child.parent.joinpath(f'{child.name}.CHUNK')

    return args


def chunk_transfer_fun(chunk_file: Path, ct: int, eof: bool,
                       ssh_target: str,
                       target_path: Path,
                       chunk_tmpl: Path):
    target_chunk = _get_chunk_file(chunk_tmpl, ct)
    rsync_cmd = ['rsync', str(chunk_file), f'{ssh_target}:{str(target_chunk)}']

    message = 'EOF' if eof else 'OK'
    target_socket = _get_remote_socket(target_path)
    inform_cmd = ['ssh', ssh_target, f'echo {message} | nc -U {shlex.quote(str(target_socket))}']

    transfer_inform(
        rsync_cmd=rsync_cmd,
        inform_cmd=inform_cmd,
        user_input_file=chunk_file.parent.joinpath(f'{chunk_file.name}.SOCKET'),
    )


if __name__ == '__main__':
    main()