Add unit specifier logic for chunk-size option

fixes: pulp#260
gerrod3 · Aug 26, 2021 · a69a184 · a69a184
1 parent a0f44a7
commit a69a184
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 12 deletions.
diff --git a/CHANGES/260.feature b/CHANGES/260.feature
@@ -0,0 +1 @@
+Chunked artifact and content uploads now allow unit specifier in ``--chunk-size`` option
diff --git a/pulpcore/cli/common/generic.py b/pulpcore/cli/common/generic.py
@@ -234,6 +234,22 @@ def _callback(
     return _callback
 
 
+# based on https://stackoverflow.com/a/42865957/2002471
+units = {"B": 1, "KB": 10 ** 3, "MB": 10 ** 6, "GB": 10 ** 9, "TB": 10 ** 12}
+
+
+def parse_size_callback(ctx: click.Context, param: click.Parameter, value: str) -> int:
+    size = value.strip().upper()
+    if not size or not re.match(r"^\s*([0-9]*)\s*([KMGT]?B)?\s*$", size):
+        raise click.ClickException("Please pass in a valid size of form: [0-9] [K/M/G]B")
+    size = re.sub(r"([KMGT]?B)", r" \1", size)
+    chunk = [string.strip() for string in size.split()]
+    if len(chunk) == 1:
+        chunk.append("B")
+    number, unit = chunk
+    return int(float(number) * units[unit])
+
+
 ##############################################################################
 # Decorator common options
 
@@ -420,6 +436,13 @@ def _multi_option_callback(
     callback=load_json_callback,
 )
 
+chunk_size_option = pulp_option(
+    "--chunk-size",
+    help=_("Chunk size to break up {entity} into. Defaults to 1MB"),
+    default="1MB",
+    callback=parse_size_callback,
+)
+
 pulp_created_gte_option = pulp_option(
     "--created-after",
     "pulp_created__gte",

diff --git a/pulpcore/cli/core/artifact.py b/pulpcore/cli/core/artifact.py
@@ -4,7 +4,7 @@
 import click
 
 from pulpcore.cli.common.context import PulpContext, pass_entity_context, pass_pulp_context
-from pulpcore.cli.common.generic import href_option, list_command, show_command
+from pulpcore.cli.common.generic import chunk_size_option, href_option, list_command, show_command
 from pulpcore.cli.core.context import PulpArtifactContext
 
 _ = gettext.gettext
@@ -26,9 +26,7 @@ def artifact(ctx: click.Context, pulp_ctx: PulpContext) -> None:
 
 @artifact.command()
 @click.option("--file", type=click.File("rb"), required=True)
-@click.option(
-    "--chunk-size", default=1000000, type=int, help=_("Chunk size in bytes (default is 1 MB)")
-)
+@chunk_size_option
 @pass_entity_context
 @pass_pulp_context
 def upload(

diff --git a/pulpcore/cli/file/content.py b/pulpcore/cli/file/content.py
@@ -9,7 +9,13 @@
     pass_entity_context,
     pass_pulp_context,
 )
-from pulpcore.cli.common.generic import create_command, href_option, list_command, show_command
+from pulpcore.cli.common.generic import (
+    chunk_size_option,
+    create_command,
+    href_option,
+    list_command,
+    show_command,
+)
 from pulpcore.cli.core.context import PulpArtifactContext
 from pulpcore.cli.file.context import PulpFileContentContext
 
@@ -91,9 +97,7 @@ def content(ctx: click.Context, pulp_ctx: PulpContext, content_type: str) -> Non
 @content.command()
 @click.option("--relative-path", required=True)
 @click.option("--file", type=click.File("rb"), required=True)
-@click.option(
-    "--chunk-size", default=1000000, type=int, help=_("Chunk size in bytes (default is 1 MB)")
-)
+@chunk_size_option
 @pass_entity_context
 @pass_pulp_context
 def upload(

diff --git a/pulpcore/cli/python/content.py b/pulpcore/cli/python/content.py
@@ -9,7 +9,13 @@
     pass_entity_context,
     pass_pulp_context,
 )
-from pulpcore.cli.common.generic import create_command, href_option, list_command, show_command
+from pulpcore.cli.common.generic import (
+    chunk_size_option,
+    create_command,
+    href_option,
+    list_command,
+    show_command,
+)
 from pulpcore.cli.core.context import PulpArtifactContext
 from pulpcore.cli.python.context import PulpPythonContentContext
 
@@ -68,9 +74,7 @@ def content(ctx: click.Context, pulp_ctx: PulpContext, content_type: str) -> Non
 @content.command()
 @click.option("--relative-path", required=True, help=_("Exact name of file"))
 @click.option("--file", type=click.File("rb"), required=True, help=_("Path to file"))
-@click.option(
-    "--chunk-size", default=1000000, type=int, help=_("Chunk size in bytes (default is 1 MB)")
-)
+@chunk_size_option
 @pass_entity_context
 @pass_pulp_context
 def upload(

diff --git a/tests/scripts/pulpcore/test_artifact.sh b/tests/scripts/pulpcore/test_artifact.sh
@@ -9,11 +9,18 @@ cleanup() {
 trap cleanup EXIT
 
 dd if=/dev/urandom of=test.txt bs=2MiB count=1
+dd if=/dev/urandom of=test2.txt bs=10KiB count=1
 sha256=$(sha256sum test.txt | cut -d' ' -f1)
+sha2256=$(sha256sum test2.txt | cut -d' ' -f1)
 
 expect_succ pulp artifact upload --file test.txt
 expect_succ pulp artifact list --sha256 "$sha256"
 test "$(echo "$OUTPUT" | jq -r length)" -eq "1"
 
+expect_succ pulp artifact upload --file test2.txt --chunk-size 1KB
+expect_fail pulp artifact upload --file test2.txt --chunk-size 1KKB
+expect_succ pulp artifact list --sha256 "$sha2256"
+test "$(echo "$OUTPUT" | jq -r length)" -eq "1"
+
 # attempt to reupload the file
 expect_succ pulp artifact upload --file test.txt
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Chunked artifact and content uploads now allow unit specifier in ``--chunk-size`` option