cli_tools.py

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Data Validation CLI tool is intended to help to build and execute
data validation runs with ease.

The Data Validator can be called either using:
    data-validation -h
    python -m data_validation -h

ex.
Step 1) Store Connection to be used in validation
data-validation connections add -c my_bq_conn BigQuery --project-id pso-project

Step 2) Run Validation using supplied connections
data-validation validate column -sc my_bq_conn -tc my_bq_conn \
-tbls bigquery-public-data.new_york_citibike.citibike_trips,bigquery-public-data.new_york_citibike.citibike_stations \
--sum '*' --count '*'

python -m data_validation validate column -sc my_bq_conn -tc my_bq_conn \
-tbls bigquery-public-data.new_york_citibike.citibike_trips \
--grouped-columns starttime \
--sum tripduration --count tripduration

data-validation validate column \
-sc my_bq_conn -tc my_bq_conn \
-tbls bigquery-public-data.new_york_citibike.citibike_trips,bigquery-public-data.new_york_citibike.citibike_stations \
--sum tripduration,start_station_name --count tripduration,start_station_name \
-rh pso-project.pso_data_validator.results \
-c ex_yaml.yaml

command:
data-validation
"""

import argparse
import copy
import csv
import json
import logging
import sys
import uuid
import os
import math
from typing import Dict, List, Optional, TYPE_CHECKING
from yaml import Dumper, Loader, dump, load

from data_validation import (
    clients,
    consts,
    exceptions,
    find_tables,
    state_manager,
    gcs_helper,
    util,
)
from data_validation.validation_builder import list_to_sublists

if TYPE_CHECKING:
    from argparse import Namespace


CONNECTION_SOURCE_FIELDS = {
    consts.SOURCE_TYPE_BIGQUERY: [
        ["project_id", "GCP Project to use for BigQuery"],
        ["google_service_account_key_path", "(Optional) GCP SA Key Path"],
        [
            "api_endpoint",
            '(Optional) GCP BigQuery API endpoint (e.g. "https://mybq.p.googleapis.com")',
        ],
    ],
    consts.SOURCE_TYPE_TERADATA: [
        ["host", "Desired Teradata host"],
        ["port", "Teradata port to connect on"],
        ["user_name", "User used to connect"],
        ["password", "Password for supplied user"],
        ["logmech", "(Optional) Log on mechanism"],
        ["use_no_lock_tables", "Use an access lock for queries (defaults to False)"],
        ["json_params", "(Optional) Additional teradatasql JSON string parameters"],
    ],
    consts.SOURCE_TYPE_ORACLE: [
        ["host", "Desired Oracle host"],
        ["port", "Oracle port to connect on"],
        ["user", "User used to connect"],
        ["password", "Password for supplied user"],
        ["database", "Database to connect to"],
        ["url", "Oracle SQLAlchemy connection URL"],
    ],
    consts.SOURCE_TYPE_MSSQL: [
        ["host", "Desired SQL Server host (default localhost)"],
        ["port", "SQL Server port to connect on (default 1433)"],
        ["user", "User used to connect"],
        ["password", "Password for supplied user"],
        ["database", "Database to connect to (default master)"],
        ["query", "Connection query parameters"],
        ["url", "SQL Server SQLAlchemy connection URL"],
    ],
    consts.SOURCE_TYPE_MYSQL: [
        ["host", "Desired MySQL host (default localhost)"],
        ["port", "MySQL port to connect on (default 3306)"],
        ["user", "User used to connect"],
        ["password", "Password for supplied user"],
        ["database", "Database to connect to (default master)"],
    ],
    consts.SOURCE_TYPE_SNOWFLAKE: [
        ["user", "Username to connect to"],
        ["password", "Password for authentication of user"],
        ["account", "Snowflake account to connect to"],
        ["database", "Database in snowflake to connect to"],
        ["connect_args", "(Optional) Additional connection arg mapping"],
    ],
    consts.SOURCE_TYPE_POSTGRES: [
        ["host", "Desired PostgreSQL host."],
        ["port", "PostgreSQL port to connect on (e.g. 5432)"],
        ["user", "Username to connect to"],
        ["password", "Password for authentication of user"],
        ["database", "Database in PostgreSQL to connect to (default postgres)"],
    ],
    consts.SOURCE_TYPE_REDSHIFT: [
        ["host", "Desired Redshift host."],
        ["port", "Redshift port to connect on (e.g. 5439)"],
        ["user", "Username to connect to"],
        ["password", "Password for authentication of user"],
        ["database", "Database in Redshift to connect to"],
    ],
    consts.SOURCE_TYPE_SPANNER: [
        ["project_id", "GCP Project to use for Spanner"],
        ["instance_id", "ID of Spanner instance to connect to"],
        ["database_id", "ID of Spanner database (schema) to connect to"],
        ["google_service_account_key_path", "(Optional) GCP SA Key Path"],
        [
            "api_endpoint",
            '(Optional) GCP Spanner API endpoint (e.g. "https://mycs.p.googleapis.com")',
        ],
    ],
    consts.SOURCE_TYPE_FILESYSTEM: [
        ["table_name", "Table name to use as reference for file data"],
        ["file_path", "The local, s3, or GCS file path to the data"],
        ["file_type", "The file type of the file. 'csv', 'orc', 'parquet' or 'json'"],
    ],
    consts.SOURCE_TYPE_IMPALA: [
        ["host", "Desired Impala host"],
        ["port", "Desired Impala port (10000 if not provided)"],
        ["database", "Desired Impala database (default if not provided)"],
        ["auth_mechanism", "Desired Impala auth mechanism (PLAIN if not provided)"],
        [
            "kerberos_service_name",
            "Desired Kerberos service name ('impala' if not provided)",
        ],
        ["use_ssl", "Use SSL when connecting to HiveServer2 (default is False)"],
        [
            "timeout",
            "Connection timeout in seconds when communicating with HiveServer2 (default is 45)",
        ],
        [
            "ca_cert",
            "Local path to 3rd party CA certificate or copy of server certificate for self-signed certificates. If SSL is enabled, but this argument is None, then certificate validation is skipped.",
        ],
        ["user", "LDAP user to authenticate"],
        ["password", "LDAP password to authenticate"],
        [
            "pool_size",
            "Size of the connection pool. Typically this is not necessary to configure. (default is 8)",
        ],
        ["hdfs_client", "An existing HDFS client"],
        ["use_http_transport", "Boolean if HTTP proxy is provided (default is False)"],
        ["http_path", "URL path of HTTP proxy"],
    ],
    consts.SOURCE_TYPE_DB2: [
        ["host", "Desired DB2 host"],
        ["port", "Desired DB2 port (50000 if not provided)"],
        ["user", "Username to connect to"],
        ["password", "Password for authentication of user"],
        ["database", "Database in DB2 to connect to"],
        ["url", "URL link in DB2 to connect to"],
        ["driver", "Driver link in DB2 to connect to (default ibm_db_sa)"],
    ],
}

VALIDATE_HELP_TEXT = "Run a validation and optionally store to config"
VALIDATE_COLUMN_HELP_TEXT = "Run a column validation"
VALIDATE_ROW_HELP_TEXT = "Run a row validation"
VALIDATE_SCHEMA_HELP_TEXT = "Run a schema validation"
VALIDATE_CUSTOM_QUERY_HELP_TEXT = "Run a custom query validation"


def _check_custom_query_args(parser: argparse.ArgumentParser, parsed_args: "Namespace"):
    # This is where we make additional checks if the arguments provided are what we expect
    # For example, only one of -tbls and custom query options can be provided
    if hasattr(parsed_args, "tables_list") and hasattr(
        parsed_args, "source_query"
    ):  # New Format
        if (
            parsed_args.tables_list
        ):  # Tables_list is not None - so source and target queries all must be None
            if (
                parsed_args.source_query_file
                or parsed_args.source_query
                or parsed_args.target_query_file
                or parsed_args.target_query
            ):
                parser.error(
                    f"{parsed_args.command}: when --tables-list/-tbls is specified, --source-query-file/-sqf, --source-query/-sq, --target-query-file/-tqf and --target-query/-tq must not be specified"
                )
            else:
                return
        elif (parsed_args.source_query_file or parsed_args.source_query) and (
            parsed_args.target_query_file or parsed_args.target_query
        ):
            return
        else:
            parser.error(
                f"{parsed_args.command}: Must specify both source (--source-query-file/-sqf or --source-query/-sq) and target (--target-query-file/-tqf or --target-query/-tq) - when --tables-list/-tbls is not specified"
            )
    else:
        return  # old format - only one of them is present


def get_parsed_args() -> "Namespace":
    """Return ArgParser with configured CLI arguments."""
    parser = configure_arg_parser()
    args = ["--help"] if len(sys.argv) == 1 else None
    parsed_args = parser.parse_args(args)
    _check_custom_query_args(parser, parsed_args)
    return parsed_args


def configure_arg_parser():
    """Extract Args for Run."""
    parser = argparse.ArgumentParser(
        usage=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")
    parser.add_argument(
        "--log-level",
        "-ll",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Log Level to be assigned. This will print logs with level same or above",
    )

    subparsers = parser.add_subparsers(dest="command")
    _configure_validate_parser(subparsers)
    _configure_validation_config_parser(subparsers)
    _configure_connection_parser(subparsers)
    _configure_find_tables(subparsers)
    _configure_raw_query(subparsers)
    _configure_beta_parser(subparsers)
    _configure_partition_parser(subparsers)
    return parser


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",
    )


def _configure_beta_parser(subparsers):
    """Configure beta commands for the parser."""
    connection_parser = subparsers.add_parser(
        "beta", help="Run a Beta command for new utilities and features."
    )
    beta_subparsers = connection_parser.add_subparsers(dest="beta_cmd")

    _configure_validate_parser(beta_subparsers)
    _configure_deploy(beta_subparsers)


def _configure_deploy(subparsers):
    """Configure arguments for deploying as a service."""
    subparsers.add_parser(
        "deploy", help="Deploy Data Validation as a Service (w/ Flask)"
    )


def _configure_find_tables(subparsers):
    """Configure arguments for text search table matching."""
    find_tables_parser = subparsers.add_parser(
        "find-tables", help="Build tables list using approx string matching."
    )
    find_tables_parser.add_argument(
        "--source-conn", "-sc", help="Source connection name."
    )
    find_tables_parser.add_argument(
        "--target-conn", "-tc", help="Target connection name."
    )
    find_tables_parser.add_argument(
        "--allowed-schemas", "-as", help="List of source schemas to match."
    )
    find_tables_parser.add_argument(
        "--include-views",
        "-iv",
        default=False,
        action="store_true",
        help="Include views in results.",
    )
    find_tables_parser.add_argument(
        "--score-cutoff",
        "-score",
        type=float,
        help="The minimum distance score allowed to match tables (0 to 1).",
    )


def _configure_raw_query(subparsers):
    """Configure arguments for text search table matching."""
    query_parser = subparsers.add_parser(
        "query", help="Run an adhoc query against the supplied connection"
    )
    query_parser.add_argument("--conn", "-c", help="Connection name to query")
    query_parser.add_argument("--query", "-q", help="Raw query to execute")
    query_parser.add_argument(
        "--format",
        "-f",
        dest="output_format",
        choices=consts.RAW_QUERY_FORMAT_TYPES,
        default=consts.FORMAT_TYPE_PYTHON,
        help=f"Format for query output (default: {consts.FORMAT_TYPE_PYTHON})",
    )


def _configure_validation_config_parser(subparsers):
    """Configure arguments to run a data validation YAML config file."""
    validation_config_parser = subparsers.add_parser(
        "configs", help="Run validations stored in a YAML config file"
    )
    configs_subparsers = validation_config_parser.add_subparsers(
        dest="validation_config_cmd"
    )
    list_parser = configs_subparsers.add_parser(
        "list", help="List your validation configs"
    )
    list_parser.add_argument(
        "--config-dir",
        "-cdir",
        help="Directory path from which to list validation YAML configs.",
    )

    run_parser = configs_subparsers.add_parser(
        "run", help="Run your validation configs"
    )
    run_parser.add_argument(
        "--dry-run",
        "-dr",
        action="store_true",
        help="Prints source and target SQL to stdout in lieu of performing a validation.",
    )
    run_parser.add_argument(
        "--config-file",
        "-c",
        help="YAML Config File path to be used for building or running validations.",
    )
    run_parser.add_argument(
        "--config-dir",
        "-cdir",
        help="Directory path containing YAML Config Files to be used for running validations.",
    )
    run_parser.add_argument(
        "--kube-completions",
        "-kc",
        action="store_true",
        help="When validating multiple table partitions generated by generate-table-partitions, using DVT in Kubernetes in index completion mode use this flag so that all the validations are completed",
    )

    get_parser = configs_subparsers.add_parser(
        "get", help="Get and print a validation config"
    )
    get_parser.add_argument(
        "--config-file",
        "-c",
        help="YAML Config File Path to be used for building or running validations.",
    )


def _configure_connection_parser(subparsers):
    """Configure the Parser for Connection Management."""

    connection_parser = subparsers.add_parser(
        "connections", help="Manage & Store connections to your Databases"
    )
    connect_subparsers = connection_parser.add_subparsers(dest="connect_cmd")
    _ = connect_subparsers.add_parser("list", help="List your connections")
    add_parser = connect_subparsers.add_parser("add", help="Store a new connection")
    add_parser.add_argument(
        "--connection-name", "-c", help="Name of connection used as reference"
    )
    add_parser.add_argument(
        "--secret-manager-type",
        "-sm",
        default=None,
        help="Secret manager type to store credentials by default will be None ",
    )
    add_parser.add_argument(
        "--secret-manager-project-id",
        "-sm-prj-id",
        default=None,
        help="Project ID for the secret manager that stores the credentials",
    )
    _configure_database_specific_parsers(add_parser)

    delete_parser = connect_subparsers.add_parser(
        "delete", help="Delete an existing connection"
    )
    delete_parser.add_argument(
        "--connection-name", "-c", required=True, help="Name of connection to delete"
    )
    describe_parser = connect_subparsers.add_parser(
        "describe", help="Describe an existing connection"
    )
    describe_parser.add_argument(
        "--connection-name", "-c", required=True, help="Name of connection to describe"
    )
    describe_parser.add_argument(
        "--format",
        "-f",
        dest="output_format",
        choices=["json", "yaml"],
        default="yaml",
        help="Output format for the configuration (default: yaml)",
    )


def _configure_database_specific_parsers(parser):
    """Configure a separate subparser for each supported DB."""
    subparsers = parser.add_subparsers(dest="connect_type")

    raw_parser = subparsers.add_parser(
        "Raw", help="Supply Raw JSON config for a connection"
    )
    raw_parser.add_argument("--json", "-j", help="Json string config")

    for database in CONNECTION_SOURCE_FIELDS:
        article = "an" if database[0].lower() in "aeiou" else "a"
        db_parser = subparsers.add_parser(
            database, help=f"Store {article} {database} connection"
        )

        for field_obj in CONNECTION_SOURCE_FIELDS[database]:
            arg_field = "--" + field_obj[0].replace("_", "-")
            help_txt = field_obj[1]
            db_parser.add_argument(arg_field, help=help_txt)


def _configure_validate_parser(subparsers):
    """Configure arguments to run validations."""
    validate_parser = subparsers.add_parser("validate", help=VALIDATE_HELP_TEXT)

    validate_parser.add_argument(
        "--dry-run",
        "-dr",
        action="store_true",
        help="Prints source and target SQL to stdout in lieu of performing a validation.",
    )

    validate_subparsers = validate_parser.add_subparsers(dest="validate_cmd")

    column_parser = validate_subparsers.add_parser(
        "column", help=VALIDATE_COLUMN_HELP_TEXT
    )
    _configure_column_parser(column_parser)

    row_parser = validate_subparsers.add_parser("row", help=VALIDATE_ROW_HELP_TEXT)
    optional_arguments = row_parser.add_argument_group("optional arguments")
    required_arguments = row_parser.add_argument_group("required arguments")
    _configure_row_parser(row_parser, optional_arguments, required_arguments)

    schema_parser = validate_subparsers.add_parser(
        "schema", help=VALIDATE_SCHEMA_HELP_TEXT
    )
    _configure_schema_parser(schema_parser)

    custom_query_parser = validate_subparsers.add_parser(
        "custom-query", help=VALIDATE_CUSTOM_QUERY_HELP_TEXT
    )
    _configure_custom_query_parser(custom_query_parser)


def _configure_row_parser(
    parser,
    optional_arguments,
    required_arguments,
    is_generate_partitions=False,
    is_custom_query=False,
):
    """Configure arguments to run row level validations."""
    # Group optional arguments
    optional_arguments.add_argument(
        "--primary-keys",
        "-pk",
        help=(
            "Comma separated list of primary key columns 'col_a,col_b', "
            "when not specified the value will be inferred from the source or target table if available"
        ),
    )

    optional_arguments.add_argument(
        "--threshold",
        "-th",
        type=threshold_float,
        default=0.0,
        help="Float max threshold for percent difference",
    )
    optional_arguments.add_argument(
        "--exclude-columns",
        "-ec",
        action="store_true",
        help="Flag to indicate the list of columns should be excluded from hash or concat instead of included.",
    )
    optional_arguments.add_argument(
        "--filters",
        "-filters",
        type=get_filters,
        default=[],
        help="Filters in the format source_filter:target_filter",
    )
    optional_arguments.add_argument(
        "--trim-string-pks",
        "-tsp",
        action="store_true",
        help=(
            "Trims string based primary key values, intended for use when one engine uses "
            "padded string semantics (e.g. CHAR(n)) and the other does not (e.g. VARCHAR(n))."
        ),
    )
    optional_arguments.add_argument(
        "--case-insensitive-match",
        "-cim",
        action="store_true",
        help=(
            "Performs a case insensitive match by adding an UPPER() before comparison."
        ),
    )
    optional_arguments.add_argument(
        "--max-concat-columns",
        "-mcc",
        type=int,
        help=(
            "The maximum number of columns accepted by a --hash or --concat validation. When there are "
            "more columns than this the validation will implicitly be split into multiple validations. "
            "This option has engine specific defaults."
        ),
    )
    # Generate-table-partitions and custom-query does not support random row
    if not (is_generate_partitions or is_custom_query):
        optional_arguments.add_argument(
            "--use-random-row",
            "-rr",
            action="store_true",
            help="Finds a set of random rows of the first primary key supplied.",
        )
        optional_arguments.add_argument(
            "--random-row-batch-size",
            "-rbs",
            help="Row batch size used for random row filters (default 10,000).",
        )
        # Generate table partitions follows a new argument spec where either the table names or queries can be provided, but not both.
        # that is specified in configure_partition_parser. If we use the same spec for row and column validation, the custom query commands
        # may get subsumed by validate and validate commands by specifying tables name or queries. Until this -tbls will be
        # a required argument for validate row, validate column and validate schema.
        required_arguments.add_argument(
            "--tables-list",
            "-tbls",
            default=None,
            required=True,
            help="Comma separated tables list in the form 'schema.table=target_schema.target_table'",
        )

    # Group for mutually exclusive required arguments. Either must be supplied
    mutually_exclusive_arguments = required_arguments.add_mutually_exclusive_group(
        required=True
    )
    mutually_exclusive_arguments.add_argument(
        "--hash",
        "-hash",
        help=(
            "Comma separated list of columns for hash 'col_a,col_b' or * for "
            "all columns"
        ),
    )
    mutually_exclusive_arguments.add_argument(
        "--concat",
        "-concat",
        help=(
            "Comma separated list of columns for concat 'col_a,col_b' or * "
            "for all columns"
        ),
    )

    mutually_exclusive_arguments.add_argument(
        "--comparison-fields",
        "-comp-fields",
        help=(
            "Individual columns to compare. If comparing a calculated field use "
            "the column alias."
        ),
    )
    _add_common_arguments(
        optional_arguments,
        required_arguments,
        is_generate_partitions=is_generate_partitions,
    )


def _configure_column_parser(column_parser):
    """Configure arguments to run column level validations."""
    # Group optional arguments
    optional_arguments = column_parser.add_argument_group("optional arguments")
    optional_arguments.add_argument(
        "--count",
        "-count",
        help="Comma separated list of columns for count 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--sum",
        "-sum",
        help="Comma separated list of columns for sum 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--avg",
        "-avg",
        help="Comma separated list of columns for avg 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--min",
        "-min",
        help="Comma separated list of columns for min 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--max",
        "-max",
        help="Comma separated list of columns for max 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--bit_xor",
        "-bit_xor",
        help="Comma separated list of columns for hashing a concatenate 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--std",
        "-std",
        help="Comma separated list of columns for standard deviation 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--grouped-columns",
        "-gc",
        help="Comma separated list of columns to use in GroupBy 'col_a,col_b'",
    )
    optional_arguments.add_argument(
        "--exclude-columns",
        "-ec",
        action="store_true",
        help="Flag to indicate the list of columns should be excluded from validation and not included.",
    )
    optional_arguments.add_argument(
        "--threshold",
        "-th",
        type=threshold_float,
        default=0.0,
        help="Float max threshold for percent difference",
    )
    optional_arguments.add_argument(
        "--filters",
        "-filters",
        type=get_filters,
        default=[],
        help="Filters in the format source_filter:target_filter",
    )

    optional_arguments.add_argument(
        "--wildcard-include-string-len",
        "-wis",
        action="store_true",
        help="Include string fields for wildcard aggregations.",
    )
    optional_arguments.add_argument(
        "--wildcard-include-timestamp",
        "-wit",
        action="store_true",
        help="Include timestamp/date fields for wildcard aggregations.",
    )
    optional_arguments.add_argument(
        "--cast-to-bigint",
        "-ctb",
        action="store_true",
        help="Cast any int32 fields to int64 for large aggregations.",
    )

    # Group required arguments
    required_arguments = column_parser.add_argument_group("required arguments")
    required_arguments.add_argument(
        "--tables-list",
        "-tbls",
        default=None,
        required=True,
        help="Comma separated tables list in the form 'schema.table=target_schema.target_table'. Or shorthand schema.* for all tables.",
    )
    _add_common_arguments(optional_arguments, required_arguments)


def _configure_schema_parser(schema_parser):
    """Configure arguments to run schema level validations."""

    # Group optional arguments
    optional_arguments = schema_parser.add_argument_group("optional arguments")
    optional_arguments.add_argument(
        "--exclusion-columns",
        "-ec",
        help="Comma separated list of columns 'col_a,col_b' to be excluded from the schema validation",
    )
    optional_arguments.add_argument(
        "--allow-list",
        "-al",
        help="Comma separated list of datatype mappings due to incompatible datatypes in source and target. e.g.: decimal(12,2):decimal(38,9),!string:string,decimal(10-18,0):int64",
    )
    optional_arguments.add_argument(
        "--allow-list-file",
        "-alf",
        help="YAML file containing default --allow-list mappings. Can be used in conjunction with --allow-list. e.g.: samples/allow_list/oracle_to_bigquery.yaml or gs://dvt-allow-list-files/oracle_to_bigquery.yaml. See example files in samples/allow_list/",
    )

    # Group required arguments
    required_arguments = schema_parser.add_argument_group("required arguments")
    required_arguments.add_argument(
        "--tables-list",
        "-tbls",
        default=None,
        required=True,
        help="Comma separated tables list in the form 'schema.table=target_schema.target_table'",
    )
    _add_common_arguments(optional_arguments, required_arguments)


def _configure_custom_query_parser(custom_query_parser):
    """Configure arguments to run custom-query validations."""

    custom_query_subparsers = custom_query_parser.add_subparsers(
        dest="custom_query_type"
    )

    # Add arguments for custom-query row parser
    custom_query_row_parser = custom_query_subparsers.add_parser(
        "row", help="Run a custom query row validation"
    )
    _configure_custom_query_row_parser(custom_query_row_parser)

    # Add arguments for custom-query column parser
    custom_query_column_parser = custom_query_subparsers.add_parser(
        "column", help="Run a custom query column validation"
    )
    _configure_custom_query_column_parser(custom_query_column_parser)


def _configure_custom_query_row_parser(custom_query_row_parser):
    optional_arguments = custom_query_row_parser.add_argument_group(
        "optional arguments"
    )
    required_arguments = custom_query_row_parser.add_argument_group(
        "required arguments"
    )

    _configure_row_parser(
        custom_query_row_parser,
        optional_arguments,
        required_arguments,
        is_custom_query=True,
    )

    # Group for mutually exclusive source query arguments. Either must be supplied
    source_mutually_exclusive = required_arguments.add_mutually_exclusive_group(
        required=True
    )
    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 = required_arguments.add_mutually_exclusive_group(
        required=True
    )
    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",
    )


def _configure_custom_query_column_parser(custom_query_column_parser):
    # Group optional arguments
    optional_arguments = custom_query_column_parser.add_argument_group(
        "optional arguments"
    )
    optional_arguments.add_argument(
        "--count",
        "-count",
        help="Comma separated list of columns for count 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--sum",
        "-sum",
        help="Comma separated list of columns for sum 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--avg",
        "-avg",
        help="Comma separated list of columns for avg 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--min",
        "-min",
        help="Comma separated list of columns for min 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--max",
        "-max",
        help="Comma separated list of columns for max 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--bit_xor",
        "-bit_xor",
        help="Comma separated list of columns for hashing a concatenate 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--std",
        "-std",
        help="Comma separated list of columns for standard deviation 'col_a,col_b' or * for all columns",
    )
    optional_arguments.add_argument(
        "--exclude-columns",
        "-ec",
        action="store_true",
        help="Flag to indicate the list of columns should be excluded from validation and not included.",
    )
    optional_arguments.add_argument(
        "--wildcard-include-string-len",
        "-wis",
        action="store_true",
        help="Include string fields for wildcard aggregations.",
    )
    optional_arguments.add_argument(
        "--wildcard-include-timestamp",
        "-wit",
        action="store_true",
        help="Include timestamp/date fields for wildcard aggregations.",
    )
    optional_arguments.add_argument(
        "--cast-to-bigint",
        "-ctb",
        action="store_true",
        help="Cast any int32 fields to int64 for large aggregations.",
    )
    optional_arguments.add_argument(
        "--filters",
        "-filters",
        type=get_filters,
        default=[],
        help="Filters in the format source_filter:target_filter",
    )
    optional_arguments.add_argument(
        "--threshold",
        "-th",
        type=threshold_float,
        default=0.0,
        help="Float max threshold for percent difference",
    )

    # Group required arguments
    required_arguments = custom_query_column_parser.add_argument_group(
        "required arguments"
    )

    # Group for mutually exclusive source query arguments. Either must be supplied
    source_mutually_exclusive = required_arguments.add_mutually_exclusive_group(
        required=True
    )
    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 = required_arguments.add_mutually_exclusive_group(
        required=True
    )
    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",
    )

    _add_common_arguments(optional_arguments, required_arguments)