fix: Update to support up to 10K partitions (#1139)

* Updated to support 10K partitions * Updated after testing * Fixing syntax issue * Update data_validation/cli_tools.py Co-authored-by: Neha Nene <[email protected]> * Update data_validation/cli_tools.py Co-authored-by: Neha Nene <[email protected]> * Update data_validation/cli_tools.py Co-authored-by: Neha Nene <[email protected]> * Simpler check for value of partition number * Correct check for value of partition number --------- Co-authored-by: Neha Nene <[email protected]>
GoogleCloudPlatform · May 22, 2024 · 210c352 · 210c352
1 parent 5137cf5
commit 210c352
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/data_validation/cli_tools.py b/data_validation/cli_tools.py
@@ -954,6 +954,17 @@ def _add_common_arguments(optional_arguments, required_arguments):
     )
 
 
+def _check_no_partitions(value: str) -> int:
+    """Check that number of partitions is between [2-10,000]
+    Using function to validate rather than choices as error message prints all choices."""
+    if value.isdigit() and 2 <= int(value) <= 10000:
+        return int(value)
+    else:
+        raise argparse.ArgumentTypeError(
+            f"{value} is not valid for number of partitions, use number in range 2 to 10000"
+        )
+
+
 def _add_common_partition_arguments(optional_arguments, required_arguments):
     """Add all arguments common to get-partition command"""
 
@@ -978,10 +989,9 @@ def _add_common_partition_arguments(optional_arguments, required_arguments):
         "--partition-num",
         "-pn",
         required=True,
-        help="Number of partitions/config files to generate",
-        type=int,
-        choices=range(1, 1001),
-        metavar="[1-1000]",
+        help="Number of partitions/config files to generate, a number from 2 to 10,000",
+        type=_check_no_partitions,
+        metavar="[2-10000]",
     )
 
     # Optional arguments