pyinfra/scripts/manage_minio.py
Julius Unverfehrt 55be1c345b Pull request #13: Add storage save response
Merge in RR/pyinfra from add-storage-save-response to master

Squashed commit of the following:

commit cb96b7943c388f020d11e97e7e1f8402efa2dd12
Author: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date:   Mon Feb 28 16:06:29 2022 +0100

    improved config

commit 607e938f16a39215c81c88bb76377bcc90282011
Author: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date:   Mon Feb 28 15:57:21 2022 +0100

    response save on storage logic added

commit 7b66bcaafe32e80a64ad09ef424c28fe1f73cd23
Author: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date:   Mon Feb 28 15:22:19 2022 +0100

    tidy up & black

commit 15a5687bac5bac599d162c86f5cec2cd945c5039
Author: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date:   Mon Feb 28 15:20:50 2022 +0100

    response save on storage logic added WIP
2022-02-28 16:12:57 +01:00

63 lines
1.8 KiB
Python

import argparse
import gzip
import os
from pathlib import Path
from tqdm import tqdm
from pyinfra.config import CONFIG
from pyinfra.storage.storages import get_s3_storage
from pyinfra.utils.file import combine_dossier_id_and_file_id_and_extension
def parse_args():
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(help="sub-command help", dest="command")
parser_add = subparsers.add_parser("add", help="Add file(s) to the MinIO store")
parser_add.add_argument("dossier_id")
add_group = parser_add.add_mutually_exclusive_group(required=True)
add_group.add_argument("--file", "-f")
add_group.add_argument("--directory", "-d")
subparsers.add_parser("purge", help="Delete all files and buckets in the MinIO store")
args = parser.parse_args()
return args
def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
path_gz = combine_dossier_id_and_file_id_and_extension(
dossier_id, Path(path).stem, CONFIG.storage.target_file_extension
)
with open(path, "rb") as f:
data = gzip.compress(f.read())
storage.put_object(bucket_name, path_gz, data)
if __name__ == "__main__":
storage = get_s3_storage()
bucket_name = CONFIG.storage.bucket
if not storage.has_bucket(bucket_name):
storage.make_bucket(bucket_name)
args = parse_args()
if args.command == "add":
if args.file:
add_file_compressed(storage, bucket_name, args.dossier_id, args.file)
elif args.directory:
for fname in tqdm([*os.listdir(args.directory)], desc="Adding files"):
path = Path(args.directory) / fname
add_file_compressed(storage, bucket_name, args.dossier_id, path)
elif args.command == "purge":
storage.clear_bucket(bucket_name)