From 8c60d868b3fe518865d7cfcf865bb35e69112f24 Mon Sep 17 00:00:00 2001
From: Dan King <dking@broadinstitute.org>
Date: Wed, 30 Nov 2022 20:52:23 -0500
Subject: [PATCH] [hailctl] allow passing arguments to hailctl batch submit
 scripts

---
 build.yaml                                  | 19 +++++++++++++++++++
 hail/python/hailtop/hailctl/batch/submit.py |  7 +++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/build.yaml b/build.yaml
index b04f721a0ace..38ebc9d54945 100644
--- a/build.yaml
+++ b/build.yaml
@@ -2671,6 +2671,25 @@ steps:
       else
           exit 1;
       fi
+
+      cat >hail_with_args.py <<EOF
+      import hail as hl
+      import sys
+
+      with open('foo/baz.txt') as f:
+          print(f.read())
+
+      hl.init(app_name='test-hailctl-batch-submit-query')
+      assert hl.utils.range_table(sys.argv[1])._force_count() == 100
+      EOF
+
+      BATCH_ID=$(hailctl batch submit --name=test-hailctl-batch-submit --files=foo -o json hail_with_args.py 100 | jq '.id')
+      STATUS=$(hailctl batch wait -o json $BATCH_ID | jq -jr '.state')
+      if [ "$STATUS" == "success" ]; then
+          exit 0;
+      else
+          exit 1;
+      fi
     secrets:
       - name: worker-deploy-config
         namespace:
diff --git a/hail/python/hailtop/hailctl/batch/submit.py b/hail/python/hailtop/hailctl/batch/submit.py
index 40309d59c40f..7b435a7d4932 100644
--- a/hail/python/hailtop/hailctl/batch/submit.py
+++ b/hail/python/hailtop/hailctl/batch/submit.py
@@ -1,6 +1,7 @@
 import asyncio
 import orjson
 import os
+from shlex import quote as shq
 
 import hailtop.batch as hb
 import hailtop.batch_client.client as bc
@@ -13,13 +14,14 @@
 
 
 def init_parser(parser):
-    parser.add_argument('script', type=str, help='Path to script')
     parser.add_argument('--name', type=str, default='', help='Batch name')
     parser.add_argument('--image-name', type=str, required=False,
                         help='Name for Docker image. Defaults to hailgenetics/hail')
     parser.add_argument('--files', nargs='+', action='append', default=[],
                         help='Comma-separated list of files or directories to add to the working directory of job')
     parser.add_argument('-o', type=str, default='text', choices=['text', 'json'])
+    parser.add_argument('script', type=str, help='Path to script')
+    parser.add_argument('arguments', nargs='*', help='Arguments to script')
 
 
 async def async_main(args):
@@ -58,7 +60,8 @@ def cloud_prefix(path):
     j.env('HAIL_QUERY_BACKEND', 'batch')
 
     command = 'python3' if script.endswith('.py') else 'bash'
-    j.command(f'{command} {script_file}')
+    script_arguments = " ".join(shq(x) for x in args.arguments)
+    j.command(f'{command} {script_file} {script_file}')
     batch_handle: bc.Batch = b.run(wait=False, disable_progress_bar=quiet)  # type: ignore
 
     if args.o == 'text':