From bba96e9d862b211fa25bd2b4ac0e059ca866b9af Mon Sep 17 00:00:00 2001 From: Sebastian Goscik Date: Thu, 8 Dec 2022 00:15:11 +0000 Subject: [PATCH] Make video download buffer size configurable --- unifi_protect_backup/cli.py | 8 ++++++++ unifi_protect_backup/downloader.py | 4 ++-- unifi_protect_backup/unifi_protect_backup.py | 5 ++++- unifi_protect_backup/utils.py | 17 ++++++++++++++++- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/unifi_protect_backup/cli.py b/unifi_protect_backup/cli.py index e2aebe7..d265fb2 100644 --- a/unifi_protect_backup/cli.py +++ b/unifi_protect_backup/cli.py @@ -7,6 +7,7 @@ from aiorun import run from unifi_protect_backup import __version__ from unifi_protect_backup.unifi_protect_backup import UnifiProtectBackup +from unifi_protect_backup.utils import human_readable_to_float DETECTION_TYPES = ["motion", "person", "vehicle", "ring"] @@ -117,6 +118,13 @@ all warnings, and websocket data envvar='COLOR_LOGGING', help="Set if you want to use color in logging output", ) +@click.option( + '--download-buffer-size', + default='512MiB', + envvar='DOWNLOAD_BUFFER_SIZE', + help='How big the download buffer should be (you can use suffixes like "B", "KiB", "MiB", "GiB")', + callback=lambda ctx, param, value: human_readable_to_float(value), +) def main(**kwargs): """A Python based tool for backing up Unifi Protect event clips as they occur.""" event_listener = UnifiProtectBackup(**kwargs) diff --git a/unifi_protect_backup/downloader.py b/unifi_protect_backup/downloader.py index ed74f6e..cb486f6 100644 --- a/unifi_protect_backup/downloader.py +++ b/unifi_protect_backup/downloader.py @@ -37,10 +37,10 @@ async def get_video_length(video: bytes) -> float: class VideoDownloader: """Downloads event video clips from Unifi Protect""" - def __init__(self, protect: ProtectApiClient, download_queue: asyncio.Queue, buffer_size: int = 256): + def __init__(self, protect: ProtectApiClient, download_queue: asyncio.Queue, buffer_size: int = 256 * 1024 * 1024): self._protect: ProtectApiClient = protect self._download_queue: asyncio.Queue = download_queue - self.video_queue = VideoQueue(buffer_size * 1024 * 1024) + self.video_queue = VideoQueue(buffer_size) # Check if `ffprobe` is available ffprobe = shutil.which('ffprobe') diff --git a/unifi_protect_backup/unifi_protect_backup.py b/unifi_protect_backup/unifi_protect_backup.py index 4da2100..0c9877f 100644 --- a/unifi_protect_backup/unifi_protect_backup.py +++ b/unifi_protect_backup/unifi_protect_backup.py @@ -198,6 +198,7 @@ class UnifiProtectBackup: ignore_cameras: List[str], file_structure_format: str, verbose: int, + download_buffer_size: int, sqlite_path: str = "events.sqlite", color_logging=False, port: int = 443, @@ -244,6 +245,7 @@ class UnifiProtectBackup: logger.debug(f" {detection_types=}") logger.debug(f" {file_structure_format=}") logger.debug(f" {sqlite_path=}") + logger.debug(f" {download_buffer_size=}") self.rclone_destination = rclone_destination self.retention = parse_rclone_retention(retention) @@ -271,6 +273,7 @@ class UnifiProtectBackup: self._has_ffprobe = False self._sqlite_path = sqlite_path self._db = None + self._download_buffer_size = download_buffer_size async def start(self): """Bootstrap the backup process and kick off the main loop. @@ -313,7 +316,7 @@ class UnifiProtectBackup: # Create downloader task # This will download video files to its buffer - downloader = VideoDownloader(self._protect, event_queue) # TODO: Make buffer size configurable + downloader = VideoDownloader(self._protect, event_queue, buffer_size=self._download_buffer_size) tasks.append(asyncio.create_task(downloader.start())) # Create upload task diff --git a/unifi_protect_backup/utils.py b/unifi_protect_backup/utils.py index 0171c5c..b8135f5 100644 --- a/unifi_protect_backup/utils.py +++ b/unifi_protect_backup/utils.py @@ -8,6 +8,8 @@ from pyunifiprotect import ProtectApiClient logger = logging.getLogger(__name__) +_suffixes = ["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"] + def human_readable_size(num: float): """Turns a number into a human readable number with ISO/IEC 80000 binary prefixes. @@ -17,13 +19,26 @@ def human_readable_size(num: float): Args: num (int): The number to be converted into human readable format """ - for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"]: + for unit in _suffixes: if abs(num) < 1024.0: return f"{num:3.1f}{unit}" num /= 1024.0 raise ValueError("`num` too large, ran out of prefixes") +def human_readable_to_float(num: str): + pattern = r"([\d.]+)(" + "|".join(_suffixes) + ")" + print(pattern) + result = re.match(pattern, num) + if result is None: + raise ValueError(f"Value '{num}' is not a valid ISO/IEC 80000 binary value") + + value = float(result[1]) + suffix = result[2] + multiplier = 1024 ** _suffixes.index(suffix) + return value * multiplier + + async def get_camera_name(protect: ProtectApiClient, id: str): """ Returns the name for the camera with the given ID