Extract pipeline attaching logic

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)

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)

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,