def main()

in kubernetes/storage/parallelstore-transfer-tool/main.py [0:0]
66 lines of code
8 McCabe index (conditional complexity)

def main():
    parser = argparse.ArgumentParser(
        description="Import or Export data between Parallelstore and GCS",
        epilog="Example usage: python import_data.py --mode import --gcsbucket my-bucket-name --instance my-instance-name --location us-central1-a",
    )
    parser.add_argument(
        "--mode",
        required=True,
        help="Import or Export data to / from a Parallelstore Instance",
    )
    parser.add_argument(
        "--gcsbucket",
        required=False,
        help="specifies the URI to a Cloud Storage bucket, or a path within a bucket, using the format gs://<bucket_name>/<optional_path_inside_bucket>",
    )
    parser.add_argument(
        "--path",
        required=False,
        default="/",
        help="Root directory path to the Parallelstore file system",
    )
    parser.add_argument(
        "--instance", required=False, help="Parallelstore instance name"
    )
    parser.add_argument(
        "--location",
        required=False,
        help="Parallelstore location, must be a supported zone.",
    )
    parser.add_argument("--project-id", required=False, help="Project ID")
    parser.add_argument(
        "--request-id",
        required=False,
        default=None,
        help="allows you to assign a unique ID to this request. If you retry this request using the same request ID, the server will ignore the request if it has already been completed. Must be a valid UUID that is not all zeros.",
    )
    args = parser.parse_args()

    # Get values from environment variables or command-line arguments
    gcs_bucket = os.environ.get("GCS_BUCKET", args.gcsbucket)
    path = os.environ.get("PARALLELSTORE_PATH", args.path)
    instance = os.environ.get("PARALLELSTORE_INSTANCE", args.instance)
    location = os.environ.get("PARALLELSTORE_LOCATION", args.location)
    project_id = os.environ.get("PROJECT_ID", args.project_id)
    request_id = os.environ.get("REQUEST_ID", args.request_id)

    # Check for missing values
    missing_args = []
    for arg_name, arg_value in [
        ("GCS_BUCKET", gcs_bucket),
        ("PARALLELSTORE_INSTANCE", instance),
        ("PARALLELSTORE_LOCATION", location),
        ("project_id", project_id),
    ]:
        if arg_value is None:
            missing_args.append(arg_name)

    if missing_args:
        logger.error(f"Error: Missing required arguments: {', '.join(missing_args)}")
        logger.error(
            "Please provide them through command-line arguments or environment variables."
        )
        exit(1)  # Exit with an error code

    if args.mode == "import":
        import_data_to_parallelstore(
            gcs_bucket, path, instance, location, request_id, project_id
        )
    elif args.mode == "export":
        export_data_to_gcs(gcs_bucket, path, instance, location, request_id, project_id)
    else:
        logger.error("Missing operation mode")