From 171b503aabf4fd9cf20892a8f21a39d36fde036e Mon Sep 17 00:00:00 2001
From: Pete Walsh <epwalsh10@gmail.com>
Date: Wed, 15 Jan 2025 15:22:19 -0800
Subject: [PATCH] Add `--dataset-secret` option (#110)

---
 CHANGELOG.md           |  4 ++++
 gantry/commands/run.py | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 15d1eee..a152a94 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added `--dataset-secret` option for mounting secrets to files as Beaker datasets.
+
 ## [v1.11.3](https://github.com/allenai/beaker-gantry/releases/tag/v1.11.3) - 2025-01-09
 
 ## [v1.11.0](https://github.com/allenai/beaker-gantry/releases/tag/v1.11.0) - 2025-01-09
diff --git a/gantry/commands/run.py b/gantry/commands/run.py
index dce9147..caae1af 100644
--- a/gantry/commands/run.py
+++ b/gantry/commands/run.py
@@ -158,6 +158,13 @@
     Should be in the form '{NAME}={SECRET_NAME}'.""",
     multiple=True,
 )
+@click.option(
+    "--dataset-secret",
+    type=str,
+    help="""Mount a Beaker secret to a file as a dataset.
+    Should be in the form '{SECRET_NAME}:{MOUNT_PATH}'.""",
+    multiple=True,
+)
 @click.option(
     "--nfs / --no-nfs",
     default=None,
@@ -297,6 +304,7 @@ def run(
     venv: Optional[str] = None,
     env: Optional[Tuple[str, ...]] = None,
     env_secret: Optional[Tuple[str, ...]] = None,
+    dataset_secret: Optional[Tuple[str, ...]] = None,
     timeout: int = 0,
     nfs: Optional[bool] = None,
     show_logs: bool = True,
@@ -419,6 +427,14 @@ def run(
             raise ValueError(f"Invalid --env-secret option: '{e}'")
         env_secrets.append((env_secret_name, secret))
 
+    dataset_secrets = []
+    for ds in dataset_secret or []:
+        try:
+            secret, mount_path = ds.split(":", 1)
+        except ValueError:
+            raise ValueError(f"Invalid --dataset-secret option: '{ds}'")
+        dataset_secrets.append((secret, mount_path))
+
     mounts = []
     for m in mount or []:
         try:
@@ -473,6 +489,7 @@ def run(
         datasets=datasets_to_use,
         env=env_vars,
         env_secrets=env_secrets,
+        dataset_secrets=dataset_secrets,
         priority=priority,
         install=install,
         no_python=no_python,
@@ -610,6 +627,7 @@ def build_experiment_spec(
     datasets: Optional[List[Tuple[str, Optional[str], str]]] = None,
     env: Optional[List[Tuple[str, str]]] = None,
     env_secrets: Optional[List[Tuple[str, str]]] = None,
+    dataset_secrets: Optional[List[Tuple[str, str]]] = None,
     priority: Optional[Union[str, Priority]] = None,
     install: Optional[str] = None,
     no_python: bool = False,
@@ -728,6 +746,9 @@ def build_experiment_spec(
         for dataset_id, sub_path, path in datasets:
             task_spec = task_spec.with_dataset(path, beaker=dataset_id, sub_path=sub_path)
 
+    for secret, mount_path in dataset_secrets or []:
+        task_spec = task_spec.with_dataset(mount_path, secret=secret)
+
     if mounts:
         for source, target in mounts:
             task_spec = task_spec.with_dataset(target, host_path=source)