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-hub · Apr 22, 2024 · f6295f9 · f6295f9
1 parent 10d9e20
commit f6295f9
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 40 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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