From e92ed5920f18eca8c6f9087b078a23dd17094e4a Mon Sep 17 00:00:00 2001
From: Keyur Shah <keyurs@google.com>
Date: Mon, 26 Feb 2024 13:20:54 -0800
Subject: [PATCH] Make docker image runnable via a config file. (#284)

---
 build/simple/run.sh        | 37 ++++++++++++++++++++++++-------------
 tools/simple/README.md     | 10 +++++++---
 tools/simple/env.list      |  2 ++
 tools/simple/run_simple.sh | 22 +++++++++++++++-------
 4 files changed, 48 insertions(+), 23 deletions(-)

diff --git a/build/simple/run.sh b/build/simple/run.sh
index 409ff1f0..26a8032c 100755
--- a/build/simple/run.sh
+++ b/build/simple/run.sh
@@ -13,24 +13,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-if [[ $GCS_INPUT_DIR != "" ]]; then
-    export INPUT_DIR=$GCS_INPUT_DIR
-else
-    export INPUT_DIR=/input/
-fi
-
 if [[ $GCS_OUTPUT_DIR != "" ]]; then
     export OUTPUT_DIR=$GCS_OUTPUT_DIR
 else
     export OUTPUT_DIR=/output/
 fi
 
-echo "DC_API_KEY=$DC_API_KEY"
-echo "INPUT_DIR=$INPUT_DIR"
-echo "OUTPUT_DIR=$INPUT_DIR"
+if [[ $CONFIG_FILE != "" ]]; then
+    echo "DC_API_KEY=$DC_API_KEY"
+    echo "CONFIG_FILE=$CONFIG_FILE"
+    echo "OUTPUT_DIR=$OUTPUT_DIR"
 
-python3 -m stats.main \
-    --mode=maindc \
-    --input_dir=$INPUT_DIR \
-    --output_dir=$OUTPUT_DIR
+    python3 -m stats.main \
+        --mode=maindc \
+        --config_file=$CONFIG_FILE \
+        --output_dir=$OUTPUT_DIR
+else
+    if [[ $GCS_INPUT_DIR != "" ]]; then
+        export INPUT_DIR=$GCS_INPUT_DIR
+    else
+        export INPUT_DIR=/input/
+    fi
+
+    echo "DC_API_KEY=$DC_API_KEY"
+    echo "INPUT_DIR=$INPUT_DIR"
+    echo "OUTPUT_DIR=$OUTPUT_DIR"
+
+    python3 -m stats.main \
+        --mode=maindc \
+        --input_dir=$INPUT_DIR \
+        --output_dir=$OUTPUT_DIR
+fi
 
diff --git a/tools/simple/README.md b/tools/simple/README.md
index f07a3e1e..e50942c5 100644
--- a/tools/simple/README.md
+++ b/tools/simple/README.md
@@ -3,6 +3,10 @@
 `run_simple.sh` is an example script illustrating how the simple import docker image
 can be run locally.
 
-In the `env.list` file, specify the `DC_API_KEY`.
-If using local input and output dirs, specify `INPUT_DIR` and `OUTPUT_DIR` respectively.
-If using GCS input and output dirs, specify `GCS_INPUT_DIR` and `GCS_OUTPUT_DIR` respectively.
+* In the `env.list` file, specify the `DC_API_KEY`.
+* If using GCS:
+  * Run `gcloud beta auth application-default login`.
+  * Specify your cloud project via the `GOOGLE_CLOUD_PROJECT` variable.
+* If using config file and output dirs, specify `CONFIG_FILE` and `GCS_OUTPUT_DIR` (or `OUTPUT_DIR`) respectively.
+* If using local input and output dirs, specify `INPUT_DIR` and `OUTPUT_DIR` respectively.
+* If using GCS input and output dirs, specify `GCS_INPUT_DIR` and `GCS_OUTPUT_DIR` respectively.
diff --git a/tools/simple/env.list b/tools/simple/env.list
index fe2f5e87..c988ec04 100644
--- a/tools/simple/env.list
+++ b/tools/simple/env.list
@@ -3,3 +3,5 @@ INPUT_DIR=
 OUTPUT_DIR=
 GCS_INPUT_DIR=
 GCS_OUTPUT_DIR=
+CONFIG_FILE=
+GOOGLE_CLOUD_PROJECT=
diff --git a/tools/simple/run_simple.sh b/tools/simple/run_simple.sh
index 3486f3a5..8087989d 100755
--- a/tools/simple/run_simple.sh
+++ b/tools/simple/run_simple.sh
@@ -17,10 +17,18 @@ set -o allexport
 source env.list
 set +o allexport
 
-docker run -it \
-    --env-file env.list \
-    -v $INPUT_DIR:/input \
-    -v $OUTPUT_DIR:/output \
-    -e GOOGLE_APPLICATION_CREDENTIALS=/gcp/creds.json \
-    -v $HOME/.config/gcloud/application_default_credentials.json:/gcp/creds.json:ro \
-    gcr.io/datcom-ci/datacommons-simple:latest
\ No newline at end of file
+if [[ $CONFIG_FILE != "" ]]; then
+    docker run -it \
+        --env-file env.list \
+        -e GOOGLE_APPLICATION_CREDENTIALS=/gcp/creds.json \
+        -v $HOME/.config/gcloud/application_default_credentials.json:/gcp/creds.json:ro \
+        -e CONFIG_FILE=/config.json \
+        -v $CONFIG_FILE:/config.json:ro \
+        gcr.io/datcom-ci/datacommons-simple:latest
+else
+    docker run -it \
+        --env-file env.list \
+        -e GOOGLE_APPLICATION_CREDENTIALS=/gcp/creds.json \
+        -v $HOME/.config/gcloud/application_default_credentials.json:/gcp/creds.json:ro \
+        gcr.io/datcom-ci/datacommons-simple:latest
+fi