From 6bdd7355913c218c1574e3abfcd47a421405110a Mon Sep 17 00:00:00 2001
From: Nikki <17799906+nikki-t@users.noreply.github.com>
Date: Mon, 16 Dec 2024 17:14:01 -0500
Subject: [PATCH] Remove SSM parameter query for project and venue and define
 from environment variable

---
 airflow/dags/cwl_dag_modular.py    | 10 ++--------
 airflow/plugins/unity_sps_utils.py |  2 --
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/airflow/dags/cwl_dag_modular.py b/airflow/dags/cwl_dag_modular.py
index 8d9898fc..4938fdb3 100644
--- 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
index e4123712..914c9dad 100644
--- 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"