Make video download buffer size configurable

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)

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')

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

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