#1112 use templating for conn details!

CityofToronto · Dec 20, 2024 · 015b8eb · 015b8eb
1 parent 8645bdb
commit 015b8eb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 18 deletions.
diff --git a/dags/pull_here.py b/dags/pull_here.py
@@ -81,14 +81,14 @@ def get_download_link(request_id: str, access_token: str):
     request_id =  get_request_id(access_token)
     download_url = get_download_link(request_id, access_token)
 
-    @task.bash(append_env=True)
-    def load_data_run(download_url)->str:
-        conn = BaseHook.get_connection("here_bot")
-        os.environ['HOST'] = conn.host
-        os.environ['LOGIN'] = conn.login
-        os.environ['PW'] = conn.password
-        os.environ['DOWNLOAD_URL'] = download_url
-        return '''curl $DOWNLOAD_URL | gunzip | PGPASSWORD=$PW psql -h $HOST -U $LOGIN -d bigdata -c "\\COPY here.ta_view FROM STDIN WITH (FORMAT csv, HEADER TRUE);" '''
+    @task.bash(env = {
+        'HOST': '{{ conn.here_bot.host }}',
+        'LOGIN': '{{ conn.here_bot.login }}',
+        'PGPASSWORD': '{{ conn.here_bot.password }}',
+        'DOWNLOAD_URL': download_url
+    })
+    def load_data_run()->str:
+        return '''curl $DOWNLOAD_URL | gunzip | psql -h $HOST -U $LOGIN -d bigdata -c "\\COPY here.ta_view FROM STDIN WITH (FORMAT csv, HEADER TRUE);" '''
 
     # Create a task group for triggering the DAGs
     @task_group(group_id='trigger_dags_tasks')
@@ -105,6 +105,6 @@ def trigger_dags(**kwargs):
             )
             trigger_operators.append(trigger_operator)
 
-    load_data_run(download_url) >> trigger_dags()
+    load_data_run() >> trigger_dags()
 
 pull_here()
diff --git a/dags/pull_here_path.py b/dags/pull_here_path.py
@@ -76,15 +76,15 @@ def get_download_link(request_id: str, access_token: str):
     request_id =  get_request_id(access_token)
     download_url = get_download_link(request_id, access_token)
 
-    @task.bash(append_env = True)
-    def load_data(download_url)->str:
-        conn = BaseHook.get_connection("here_bot")
-        os.environ['HOST'] = conn.host
-        os.environ['LOGIN'] = conn.login
-        os.environ['PW'] = conn.password
-        os.environ['DOWNLOAD_URL'] = download_url
-        return '''curl $DOWNLOAD_URL | gunzip | PGPASSWORD=$PW psql -h $HOST -U $LOGIN -d bigdata -c "\\COPY here.ta_path_view FROM STDIN WITH (FORMAT csv, HEADER TRUE);" '''
+    @task.bash(env = {
+        'HOST': '{{ conn.here_bot.host }}',
+        'LOGIN': '{{ conn.here_bot.login }}',
+        'PGPASSWORD': '{{ conn.here_bot.password }}',
+        'DOWNLOAD_URL': download_url
+    })
+    def load_data()->str:
+        return '''curl $DOWNLOAD_URL | gunzip | psql -h $HOST -U $LOGIN -d bigdata -c "\\COPY here.ta_path_view FROM STDIN WITH (FORMAT csv, HEADER TRUE);" '''
 
-    load_data(download_url)
+    load_data()
 
 pull_here_path()