Remove need to call connect explicitly (StructuredLabs#653)

* Clean up managers/data.py for linter

* Fixed connect idempotency issue - no longer recreating sources upon rerun

* Remove need to call connect explicitly

* Remove connect() from template and tutorial

Author: shivam-singhal

examples/earthquakes/hello.py

@@ -1,17 +1,12 @@
 import pandas as pd
 import plotly.express as px
 
-from preswald import connect, get_df, plotly, slider, table, text
+from preswald import get_df, plotly, slider, table, text
 
 
 # Title
 text("# Earthquake Analytics Dashboard 🌍")
 
-# Load and connect data
-connect()
-
-# ---
-
 # Clickhouse section
 
 query_string = """

examples/fires/hello.py

@@ -1,15 +1,12 @@
 import pandas as pd
 import plotly.express as px
 
-from preswald import connect, get_df, plotly, text
+from preswald import get_df, plotly, text
 
 
 # Display the dashboard title
 text("# Fire Incident Analytics Dashboard 🔥")
 
-# Connect to the data
-connect()
-
 # Load and preprocess the data
 data = get_df("csv")
 data["incident_acres_burned"] = pd.to_numeric(

examples/iris/hello.py

@@ -5,7 +5,6 @@
 
 from preswald import (
     chat,
-    connect,
     # fastplotlib,
     get_df,
     plotly,
@@ -28,7 +27,6 @@
 )
 
 # Load the CSV
-connect()  # Load in all sources, which by default is the iris_csv
 df = get_df("iris_csv")
 
 # 1. Scatter plot - Sepal Length vs Sepal Width

examples/user_event/hello.py

@@ -1,30 +1,30 @@
-from preswald import text, plotly, connect, get_df, table
-import pandas as pd
 import plotly.express as px
 
+from preswald import get_df, plotly, table, text
+
+
 text("# Welcome to Preswald!")
 text("This is your first app. 🎉")
 
 # Load the JSON source defined as "user_events" in preswald.toml
-connect()  # This loads all data sources, including our nested JSON source.
-df = get_df('user_events')
+df = get_df("user_events")
 
 # Create a scatter plot using the flattened data.
 # Assuming the JSON file has been flattened to include "user" and "details.clicks"
 fig = px.scatter(
     df,
-    x='user',
-    y='details.clicks',
-    text='user',
-    title='User Events: Clicks per User',
-    labels={'user': 'User', 'details.clicks': 'Clicks'}
+    x="user",
+    y="details.clicks",
+    text="user",
+    title="User Events: Clicks per User",
+    labels={"user": "User", "details.clicks": "Clicks"},
 )
 
 # Add labels for each point
-fig.update_traces(textposition='top center', marker=dict(size=12, color='lightblue'))
+fig.update_traces(textposition="top center", marker={"size": 12, "color": "lightblue"})
 
 # Style the plot
-fig.update_layout(template='plotly_white')
+fig.update_layout(template="plotly_white")
 
 # Display the plot
 plotly(fig)

preswald/browser/virtual_service.py

@@ -5,7 +5,7 @@
 import asyncio
 import logging
 import sys
-from typing import Any, Dict, Optional
+from typing import Any
 
 from preswald.engine.base_service import BasePreswaldService
 from preswald.engine.utils import RenderBuffer
@@ -37,7 +37,7 @@ def __init__(self, client_id: str):
         )
         console.log(f"[Communication] is browser mode: {self.is_browser_mode}")
 
-    async def send_json(self, data: Dict[str, Any]):
+    async def send_json(self, data: dict[str, Any]):
         """Send JSON data to JavaScript frontend"""
         if not self.is_connected:
             logger.error(f"Cannot send message, connection closed for {self.client_id}")
@@ -185,11 +185,14 @@ def handle_message_from_js(client_id, message_type, data):
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(f"Handling JS message from {client_id}: {message}")
 
-                asyncio.create_task(self.handle_client_message(client_id, message))   # noqa: RUF006
+                asyncio.create_task(self.handle_client_message(client_id, message))  # noqa: RUF006
                 return True
             except Exception:
                 import traceback
-                logger.error("Error in handle_message_from_js: %s", traceback.format_exc())
+
+                logger.error(
+                    "Error in handle_message_from_js: %s", traceback.format_exc()
+                )
                 return False
 
         # Export the function to JavaScript

preswald/engine/base_service.py

@@ -1,8 +1,9 @@
 import logging
 import os
 import time
+from collections.abc import Callable
 from threading import Lock
-from typing import Any, Callable, Dict, Optional
+from typing import Any
 
 from preswald.engine.runner import ScriptRunner
 from preswald.engine.utils import (
@@ -28,19 +29,19 @@ class BasePreswaldService:
     _not_initialized_msg = "Base service not initialized."
 
     def __init__(self):
-        self._component_states: Dict[str, Any] = {}
+        self._component_states: dict[str, Any] = {}
         self._lock = Lock()
 
         # Data management
-        self.data_manager: Optional[DataManager] = None  # set during server creation
+        self.data_manager: DataManager | None = None  # set during server creation
 
         # Initialize service state
-        self._script_path: Optional[str] = None
+        self._script_path: str | None = None
         self._is_shutting_down: bool = False
         self._render_buffer = RenderBuffer()
 
         # Initialize session tracking
-        self.script_runners: Dict[str, ScriptRunner] = {}
+        self.script_runners: dict[str, ScriptRunner] = {}
 
         # Layout management
         self._layout_manager = LayoutManager()
@@ -61,7 +62,7 @@ def initialize(cls, script_path=None):
         return cls._instance
 
     @property
-    def script_path(self) -> Optional[str]:
+    def script_path(self) -> str | None:
         return self._script_path
 
     @script_path.setter
@@ -76,9 +77,10 @@ def script_path(self, path: str):
     def append_component(self, component):
         """Add a component to the layout manager"""
         try:
-
             if isinstance(component, dict):
-                logger.info(f"[APPEND] Appending component: {component.get('id')}, type: {component.get('type')}")
+                logger.info(
+                    f"[APPEND] Appending component: {component.get('id')}, type: {component.get('type')}"
+                )
                 # Clean any NaN values in the component
                 clean_start = time.time()
                 cleaned_component = clean_nan_values(component)
@@ -101,7 +103,9 @@ def append_component(self, component):
                                     )
                         self._layout_manager.add_component(cleaned_component)
                         if logger.isEnabledFor(logging.DEBUG):
-                            logger.debug(f"Added component with state: {cleaned_component}")
+                            logger.debug(
+                                f"Added component with state: {cleaned_component}"
+                            )
                 else:
                     # Components without IDs are added as-is
                     self._layout_manager.add_component(cleaned_component)
@@ -137,7 +141,7 @@ def get_rendered_components(self):
         rows = self._layout_manager.get_layout()
         return {"rows": rows}
 
-    async def handle_client_message(self, client_id: str, message: Dict[str, Any]):
+    async def handle_client_message(self, client_id: str, message: dict[str, Any]):
         """Process incoming messages from clients"""
         start_time = time.time()
         try:
@@ -196,7 +200,7 @@ async def unregister_client(self, client_id: str):
     def _create_send_callback(self, websocket: Any) -> Callable:
         """Create a message sending callback for a specific websocket"""
 
-        async def send_message(msg: Dict[str, Any]):
+        async def send_message(msg: dict[str, Any]):
             if not self._is_shutting_down:
                 try:
                     await websocket.send_json(msg)
@@ -206,7 +210,7 @@ async def send_message(msg: Dict[str, Any]):
         return send_message
 
     async def _broadcast_state_updates(
-        self, states: Dict[str, Any], exclude_client: Optional[str] = None
+        self, states: dict[str, Any], exclude_client: str | None = None
     ):
         """Broadcast state updates to all clients except the sender"""
 
@@ -231,18 +235,15 @@ async def _broadcast_state_updates(
                     except Exception as e:
                         logger.error(f"Error broadcasting to {client_id}: {e}")
 
-    async def _handle_component_update(self, client_id: str, message: Dict[str, Any]):
+    async def _handle_component_update(self, client_id: str, message: dict[str, Any]):
         """Handle component state update messages"""
         states = message.get("states", {})
         if not states:
             await self._send_error(client_id, "Component update missing states")
             raise ValueError("Component update missing states")
 
         # Only rerun if any state actually changed
-        changed_states = {
-            k: v for k, v in states.items()
-            if self.should_render(k, v)
-        }
+        changed_states = {k: v for k, v in states.items() if self.should_render(k, v)}
 
         if not changed_states:
             logger.debug("[STATE] No actual state changes detected. Skipping rerun.")
@@ -260,6 +261,10 @@ async def _handle_component_update(self, client_id: str, message: Dict[str, Any]
         # Broadcast updates to other clients
         await self._broadcast_state_updates(changed_states, exclude_client=client_id)
 
+    def connect_data_manager(self):
+        """Connect the data manager"""
+        self.data_manager.connect()
+
     def _initialize_data_manager(self, script_path: str) -> None:
         script_dir = os.path.dirname(script_path)
         preswald_path = os.path.join(script_dir, "preswald.toml")
@@ -269,7 +274,9 @@ def _initialize_data_manager(self, script_path: str) -> None:
             preswald_path=preswald_path, secrets_path=secrets_path
         )
 
-    async def _register_common_client_setup(self, client_id: str, websocket: Any) -> ScriptRunner:
+    async def _register_common_client_setup(
+        self, client_id: str, websocket: Any
+    ) -> ScriptRunner:
         logger.info(f"Registering client: {client_id}")
 
         self.websocket_connections[client_id] = websocket
@@ -309,7 +316,7 @@ async def _send_initial_states(self, websocket: Any):
         except Exception as e:
             logger.error(f"Error sending initial states: {e}")
 
-    def _update_component_states(self, states: Dict[str, Any]):
+    def _update_component_states(self, states: dict[str, Any]):
         """Update internal state dictionary with cleaned component values."""
         with self._lock:
             logger.debug("[STATE] Updating states")