Remove SSM parameter query for project and venue and define from envi…

…ronment variable
unity-sds · Dec 16, 2024 · 6bdd735 · 6bdd735
1 parent 71429ed
commit 6bdd735
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 10 deletions.
diff --git a/airflow/dags/cwl_dag_modular.py b/airflow/dags/cwl_dag_modular.py
@@ -172,14 +172,8 @@ def select_stage_out(ti):
     """Retrieve stage out input parameters from SSM parameter store."""
     ssm_client = boto3.client("ssm", region_name="us-west-2")
 
-    project = ssm_client.get_parameter(Name=unity_sps_utils.SPS_PROJECT_PARAM, WithDecryption=True)[
-        "Parameter"
-    ]["Value"]
-
-    venue = ssm_client.get_parameter(Name=unity_sps_utils.SPS_VENUE_PARAM, WithDecryption=True)["Parameter"][
-        "Value"
-    ]
-
+    project = os.environ["AIRFLOW_VAR_UNITY_PROJECT"]
+    venue = os.environ["AIRFLOW_VAR_UNITY_VENUE"]
     staging_bucket = ssm_client.get_parameter(Name=unity_sps_utils.DS_S3_BUCKET_PARAM, WithDecryption=True)[
         "Parameter"
     ]["Value"]

diff --git a/airflow/plugins/unity_sps_utils.py b/airflow/plugins/unity_sps_utils.py
@@ -14,8 +14,6 @@
 NODE_POOL_DEFAULT = "airflow-kubernetes-pod-operator"
 NODE_POOL_HIGH_WORKLOAD = "airflow-kubernetes-pod-operator-high-workload"
 
-SPS_PROJECT_PARAM = "/unity/unity-nikki-1/dev/project-name"
-SPS_VENUE_PARAM = "/unity/unity-nikki-1/dev/venue-name"
 DS_S3_BUCKET_PARAM = "/unity/unity-nikki-1/dev/ds/staging/s3/bucket-name"