From f6295f93ce99dd29493ed5c870ba4975535b425d Mon Sep 17 00:00:00 2001
From: Sultan Iman <354868+sultaniman@users.noreply.github.com>
Date: Mon, 22 Apr 2024 21:27:25 +0200
Subject: [PATCH] Check for default schema and schema name in streamlit session
 (#1155)

* Check for default schema and schema name in streamlit session

* Do not show resource state if it is not available

* Fix mypy errors

* Remove the message if there is no schema in state

* Simplify code
---
 dlt/helpers/streamlit_app/blocks/load_info.py | 53 ++++++++++---------
 .../streamlit_app/blocks/resource_state.py    |  9 +---
 dlt/helpers/streamlit_app/pages/dashboard.py  | 14 +++--
 dlt/helpers/streamlit_app/widgets/schema.py   |  2 +-
 .../test_streamlit_show_resources.py          |  2 +-
 5 files changed, 40 insertions(+), 40 deletions(-)

diff --git a/dlt/helpers/streamlit_app/blocks/load_info.py b/dlt/helpers/streamlit_app/blocks/load_info.py
index 9482cb5afa..bfd9d386df 100644
--- a/dlt/helpers/streamlit_app/blocks/load_info.py
+++ b/dlt/helpers/streamlit_app/blocks/load_info.py
@@ -8,33 +8,34 @@
 
 
 def last_load_info(pipeline: dlt.Pipeline) -> None:
-    loads_df = query_data_live(
-        pipeline,
-        f"SELECT load_id, inserted_at FROM {pipeline.default_schema.loads_table_name} WHERE"
-        " status = 0 ORDER BY inserted_at DESC LIMIT 101 ",
-    )
-
-    if loads_df is None:
-        st.error(
-            "Load info is not available",
-            icon="🚨",
+    if pipeline.default_schema_name:
+        loads_df = query_data_live(
+            pipeline,
+            f"SELECT load_id, inserted_at FROM {pipeline.default_schema.loads_table_name} WHERE"
+            " status = 0 ORDER BY inserted_at DESC LIMIT 101 ",
         )
-    else:
-        loads_no = loads_df.shape[0]
-        if loads_df.shape[0] > 0:
-            rel_time = (
-                humanize.naturaldelta(
-                    pendulum.now() - pendulum.from_timestamp(loads_df.iloc[0, 1].timestamp())
-                )
-                + " ago"
+
+        if loads_df is None:
+            st.error(
+                "Load info is not available",
+                icon="🚨",
             )
-            last_load_id = loads_df.iloc[0, 0]
-            if loads_no > 100:
-                loads_no = "> " + str(loads_no)
         else:
-            rel_time = "---"
-            last_load_id = "---"
+            loads_no = loads_df.shape[0]
+            if loads_df.shape[0] > 0:
+                rel_time = (
+                    humanize.naturaldelta(
+                        pendulum.now() - pendulum.from_timestamp(loads_df.iloc[0, 1].timestamp())
+                    )
+                    + " ago"
+                )
+                last_load_id = loads_df.iloc[0, 0]
+                if loads_no > 100:
+                    loads_no = "> " + str(loads_no)
+            else:
+                rel_time = "---"
+                last_load_id = "---"
 
-        stat("Last load time", rel_time, border_left_width=4)
-        stat("Last load id", last_load_id)
-        stat("Total number of loads", loads_no)
+            stat("Last load time", rel_time, border_left_width=4)
+            stat("Last load id", last_load_id)
+            stat("Total number of loads", loads_no)
diff --git a/dlt/helpers/streamlit_app/blocks/resource_state.py b/dlt/helpers/streamlit_app/blocks/resource_state.py
index dabbea4d46..69fae92e81 100644
--- a/dlt/helpers/streamlit_app/blocks/resource_state.py
+++ b/dlt/helpers/streamlit_app/blocks/resource_state.py
@@ -22,13 +22,8 @@ def resource_state_info(
     resource_name: str,
 ) -> None:
     sources_state = pipeline.state.get("sources") or {}
-    schema = sources_state.get(schema_name)
-    if not schema:
-        st.error(f"Schema with name: {schema_name} is not found")
-        return
-
-    resource = schema["resources"].get(resource_name)
-
+    schema = sources_state.get(schema_name, {})
+    resource = schema.get("resources", {}).get(resource_name)
     with st.expander("Resource state", expanded=(resource is None)):
         if not resource:
             st.info(f"{resource_name} is missing resource state")
diff --git a/dlt/helpers/streamlit_app/pages/dashboard.py b/dlt/helpers/streamlit_app/pages/dashboard.py
index 656dd6ecdf..420c9b2021 100644
--- a/dlt/helpers/streamlit_app/pages/dashboard.py
+++ b/dlt/helpers/streamlit_app/pages/dashboard.py
@@ -29,12 +29,16 @@ def write_data_explorer_page(
 
     st.subheader("Schemas and tables", divider="rainbow")
     schema_picker(pipeline)
-    tables = sorted(
-        st.session_state["schema"].data_tables(),
-        key=lambda table: table["name"],
-    )
+    if schema := st.session_state["schema"]:
+        tables = sorted(
+            schema.data_tables(),
+            key=lambda table: table["name"],
+        )
+
+        list_table_hints(pipeline, tables)
+    else:
+        st.warning("No schemas found")
 
-    list_table_hints(pipeline, tables)
     maybe_run_query(
         pipeline,
         show_charts=show_charts,
diff --git a/dlt/helpers/streamlit_app/widgets/schema.py b/dlt/helpers/streamlit_app/widgets/schema.py
index f7883bc45e..c9dd87097e 100644
--- a/dlt/helpers/streamlit_app/widgets/schema.py
+++ b/dlt/helpers/streamlit_app/widgets/schema.py
@@ -16,6 +16,6 @@ def schema_picker(pipeline: dlt.Pipeline) -> None:
         )
         schema = pipeline.schemas.get(selected_schema_name)
 
+    st.session_state["schema"] = schema
     if schema:
         st.subheader(f"Schema: {schema.name}")
-        st.session_state["schema"] = schema
diff --git a/tests/helpers/streamlit_tests/test_streamlit_show_resources.py b/tests/helpers/streamlit_tests/test_streamlit_show_resources.py
index dd807260fe..691af8a9d1 100644
--- a/tests/helpers/streamlit_tests/test_streamlit_show_resources.py
+++ b/tests/helpers/streamlit_tests/test_streamlit_show_resources.py
@@ -13,7 +13,7 @@
 
 import dlt
 
-from streamlit.testing.v1 import AppTest  # type: ignore
+from streamlit.testing.v1 import AppTest  # type: ignore[import-not-found]
 
 from dlt.helpers.streamlit_app.utils import render_with_pipeline
 from dlt.pipeline.exceptions import CannotRestorePipelineException