From af707fb9a0df6c49e3b6c4cb0a8a3a27798cf885 Mon Sep 17 00:00:00 2001
From: Sultan Iman <sultan.imanhodjaev@gmail.com>
Date: Mon, 18 Mar 2024 16:44:35 +0100
Subject: [PATCH] Extract pipeline attaching logic

---
 dlt/helpers/streamlit_app/dashboard.py       | 7 ++-----
 dlt/helpers/streamlit_app/pages/load_info.py | 7 +++----
 dlt/helpers/streamlit_app/utils.py           | 9 +++++++++
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/dlt/helpers/streamlit_app/dashboard.py b/dlt/helpers/streamlit_app/dashboard.py
index c73b4d2d28..58268b6d2d 100644
--- a/dlt/helpers/streamlit_app/dashboard.py
+++ b/dlt/helpers/streamlit_app/dashboard.py
@@ -7,6 +7,7 @@
 from dlt.helpers.streamlit_app.blocks.query import maybe_run_query
 from dlt.helpers.streamlit_app.blocks.table_hints import list_table_hints
 from dlt.helpers.streamlit_app.menu import menu
+from dlt.helpers.streamlit_app.utils import attach_to_pipeline
 from dlt.helpers.streamlit_app.widgets import schema_picker
 from dlt.pipeline import Pipeline
 
@@ -46,11 +47,7 @@ def write_data_explorer_page(
 
 
 def display(pipeline_name: str) -> None:
-    pipeline = dlt.attach(pipeline_name)
-    if pipelines_dir := os.getenv("DLT_PIPELINES_DIR"):
-        pipeline.pipelines_dir = pipelines_dir
-
-    st.session_state["pipeline_name"] = pipeline_name
+    pipeline = attach_to_pipeline(pipeline_name)
 
     with st.sidebar:
         menu(pipeline)
diff --git a/dlt/helpers/streamlit_app/pages/load_info.py b/dlt/helpers/streamlit_app/pages/load_info.py
index 3208cf83af..50d3d1d25e 100644
--- a/dlt/helpers/streamlit_app/pages/load_info.py
+++ b/dlt/helpers/streamlit_app/pages/load_info.py
@@ -8,7 +8,7 @@
 from dlt.common.libs.pandas import pandas as pd
 from dlt.helpers.streamlit_app.menu import menu
 from dlt.helpers.streamlit_app.widgets import stat, pipeline_summary
-from dlt.helpers.streamlit_app.utils import cache_data
+from dlt.helpers.streamlit_app.utils import attach_to_pipeline, cache_data
 from dlt.pipeline import Pipeline
 from dlt.pipeline.exceptions import CannotRestorePipelineException, SqlClientNotAvailable
 from dlt.pipeline.state_sync import load_pipeline_state_from_destination
@@ -127,12 +127,11 @@ def show_state_versions(pipeline: dlt.Pipeline) -> None:
 
 
 def show() -> None:
+    pipeline_name = st.session_state.get("pipeline_name")
     if not st.session_state.get("pipeline_name"):
         st.switch_page("dashboard.py")
 
-    pipeline = dlt.attach(st.session_state["pipeline_name"])
-    if pipelines_dir := os.getenv("DLT_PIPELINES_DIR"):
-        pipeline.pipelines_dir = pipelines_dir
+    pipeline = attach_to_pipeline(pipeline_name)
 
     st.subheader("Load info", divider="rainbow")
     write_load_status_page(pipeline)
diff --git a/dlt/helpers/streamlit_app/utils.py b/dlt/helpers/streamlit_app/utils.py
index 9f9e82c815..83d9e8ff11 100644
--- a/dlt/helpers/streamlit_app/utils.py
+++ b/dlt/helpers/streamlit_app/utils.py
@@ -1,3 +1,5 @@
+import os
+
 from pathlib import Path
 from typing import Optional
 
@@ -15,6 +17,13 @@
     cache_data = st.experimental_memo
 
 
+def attach_to_pipeline(pipeline_name: str) -> dlt.Pipeline:
+    st.session_state["pipeline_name"] = pipeline_name
+    pipelines_dir = os.getenv("DLT_PIPELINES_DIR")
+    pipeline = dlt.attach(pipeline_name, pipelines_dir=pipelines_dir)
+    return pipeline
+
+
 # FIXME: make something to DRY the code
 def query_data(
     pipeline: dlt.Pipeline,