def _configure_partition_parser()

in data_validation/cli_tools.py [0:0]

• 67 lines of code
• 8 McCabe index (conditional complexity)

def _configure_partition_parser(subparsers):
    """Configure arguments to generate partitioned config files."""
    partition_parser = subparsers.add_parser(
        "generate-table-partitions",
        help=("Generate table partitions and store validation config files"),
    )
    optional_arguments = partition_parser.add_argument_group("optional arguments")
    required_arguments = partition_parser.add_argument_group("required arguments")

    _configure_row_parser(
        partition_parser,
        optional_arguments,
        required_arguments,
        is_generate_partitions=True,
    )
    optional_arguments.add_argument(
        "--parts-per-file",
        "-ppf",
        type=_check_positive,
        default=1,
        help="Number of partitions to be validated in a single yaml file.",
    )

    required_arguments.add_argument(
        "--config-dir",
        "-cdir",
        required=True,
        help="Directory path to store YAML config files. "
        "GCS: Provide a full gs:// path of the target directory. "
        "Eg: `gs://<BUCKET>/partitons_dir`. "
        "Local: Provide a relative path of the target directory. "
        "Eg: `partitions_dir`",
    )
    required_arguments.add_argument(
        "--partition-num",
        "-pn",
        required=True,
        help="Number of partitions into which the table should be split",
        type=_check_gt_one,
    )
    # User can provide tables or custom queries, but not both
    # However, Argparse does not support adding an argument_group to an argument_group or adding a
    # mutually_exclusive_group or argument_group to a mutually_exclusive_group since version 3.11.
    # We are only ensuring leaf level mutual exclusivity here and will need to check higher level
    # mutual exclusivity in the code - i.e. a) when --tables-list is present, there can be no custom
    # query parameters and b) when custom query parameters are specified, both source and target must be
    # specified.
    optional_arguments.add_argument(
        "--tables-list",
        "-tbls",
        help=(
            "Comma separated tables list in the form "
            "'schema.table=target_schema.target_table'"
        ),
    )

    source_mutually_exclusive = optional_arguments.add_mutually_exclusive_group()
    source_mutually_exclusive.add_argument(
        "--source-query-file",
        "-sqf",
        help="File containing the source sql query",
    )
    source_mutually_exclusive.add_argument(
        "--source-query",
        "-sq",
        help="Source sql query",
    )

    # Group for mutually exclusive target query arguments. Either must be supplied
    target_mutually_exclusive = optional_arguments.add_mutually_exclusive_group()
    target_mutually_exclusive.add_argument(
        "--target-query-file",
        "-tqf",
        help="File containing the target sql query",
    )
    target_mutually_exclusive.add_argument(
        "--target-query",
        "-tq",
        help="Target sql query",
    )