mangleframes 0.3.6__tar.gz → 0.3.7__tar.gz

{mangleframes-0.3.6 → mangleframes-0.3.7}/Cargo.lock

@@ -1486,6 +1486,7 @@ dependencies = [
  "spark-connect",
  "thiserror",
  "tokio",
+ "tonic",
  "tower-http",
  "tracing",
  "tracing-subscriber",

{mangleframes-0.3.6 → mangleframes-0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mangleframes
-Version: 0.3.6
+Version: 0.3.7
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Rust
 Classifier: License :: OSI Approved :: MIT License

{mangleframes-0.3.6 → mangleframes-0.3.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "mangleframes"
-version = "0.3.6"
+version = "0.3.7"
 description = "PySpark DataFrame viewer with modern web UI"
 requires-python = ">=3.12"
 license = { text = "MIT" }

{mangleframes-0.3.6 → mangleframes-0.3.7}/python/mangleframes/__init__.py

@@ -32,7 +32,7 @@ from .session import SparkSession, get_proxy_port, get_spark_session
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame
 
-__version__ = "0.3.6"
+__version__ = "0.3.7"
 
 # Import alert classes for convenience (optional dependency)
 try:

{mangleframes-0.3.6 → mangleframes-0.3.7}/viewer/Cargo.toml

@@ -29,6 +29,7 @@ chrono = "0.4"
 base64 = "0.22"
 reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
 spark-connect = { path = "../spark-connect" }
+tonic = { version = "0.12", default-features = false }
 
 [profile.release]
 lto = true

{mangleframes-0.3.6 → mangleframes-0.3.7}/viewer/src/spark_client.rs

@@ -15,10 +15,23 @@ pub struct SqlResponse {
     pub execution_ms: u64,
 }
 
+/// Stored credentials/URL used to reconnect after session expiry.
+enum ConnectionConfig {
+    Direct {
+        host: String,
+        token: String,
+        cluster_id: Option<String>,
+    },
+    Proxy {
+        url: String,
+    },
+}
+
 /// Wrapper around SparkConnectClient for Databricks integration.
 pub struct DatabricksClient {
     client: Arc<RwLock<Option<SparkConnectClient>>>,
     connected: Arc<RwLock<bool>>,
+    config: Arc<RwLock<Option<ConnectionConfig>>>,
 }
 
 impl DatabricksClient {
@@ -27,6 +40,7 @@ impl DatabricksClient {
         Self {
             client: Arc::new(RwLock::new(None)),
             connected: Arc::new(RwLock::new(false)),
+            config: Arc::new(RwLock::new(None)),
         }
     }
 
@@ -46,6 +60,11 @@ impl DatabricksClient {
 
         *self.client.write().await = Some(client);
         *self.connected.write().await = true;
+        *self.config.write().await = Some(ConnectionConfig::Direct {
+            host: host.to_string(),
+            token: token.to_string(),
+            cluster_id: cluster_id.map(str::to_string),
+        });
 
         info!("Databricks {} connection established", mode);
         Ok(())
@@ -64,6 +83,9 @@ impl DatabricksClient {
 
         *self.client.write().await = Some(client);
         *self.connected.write().await = true;
+        *self.config.write().await = Some(ConnectionConfig::Proxy {
+            url: proxy_url.to_string(),
+        });
 
         info!("Connected to proxy at {}", proxy_url);
         Ok(())
@@ -79,6 +101,50 @@ impl DatabricksClient {
         &self,
         query: &str,
         limit: usize,
+    ) -> Result<SqlResponse, SparkConnectError> {
+        match self.execute_sql_inner(query, limit).await {
+            Err(ref e) if is_session_timeout(e) => {
+                self.reconnect().await?;
+                self.execute_sql_inner(query, limit).await
+            }
+            other => other,
+        }
+    }
+
+    /// Execute SQL with reattachable execution for large result sets (>10K rows).
+    /// Uses ReattachExecute RPC to ensure complete results.
+    pub async fn execute_sql_reattachable(
+        &self,
+        query: &str,
+        limit: usize,
+    ) -> Result<SqlResponse, SparkConnectError> {
+        match self.execute_sql_reattachable_inner(query, limit).await {
+            Err(ref e) if is_session_timeout(e) => {
+                self.reconnect().await?;
+                self.execute_sql_reattachable_inner(query, limit).await
+            }
+            other => other,
+        }
+    }
+
+    /// Register Arrow batches as a temporary view in Spark.
+    pub async fn create_temp_view(
+        &self,
+        view_name: &str,
+        batches: &[RecordBatch],
+    ) -> Result<(), SparkConnectError> {
+        let guard = self.client.read().await;
+        let client = guard
+            .as_ref()
+            .ok_or_else(|| SparkConnectError::Config("Not connected".to_string()))?;
+
+        client.create_temp_view(view_name, batches).await
+    }
+
+    async fn execute_sql_inner(
+        &self,
+        query: &str,
+        limit: usize,
     ) -> Result<SqlResponse, SparkConnectError> {
         let guard = self.client.read().await;
         let client = guard
@@ -103,9 +169,7 @@ impl DatabricksClient {
         })
     }
 
-    /// Execute SQL with reattachable execution for large result sets (>10K rows).
-    /// Uses ReattachExecute RPC to ensure complete results.
-    pub async fn execute_sql_reattachable(
+    async fn execute_sql_reattachable_inner(
         &self,
         query: &str,
         limit: usize,
@@ -133,18 +197,26 @@ impl DatabricksClient {
         })
     }
 
-    /// Register Arrow batches as a temporary view in Spark.
-    pub async fn create_temp_view(
-        &self,
-        view_name: &str,
-        batches: &[RecordBatch],
-    ) -> Result<(), SparkConnectError> {
-        let guard = self.client.read().await;
-        let client = guard
+    async fn reconnect(&self) -> Result<(), SparkConnectError> {
+        info!("Spark Connect session expired due to inactivity, reconnecting...");
+        let config_guard = self.config.read().await;
+        let config = config_guard
             .as_ref()
-            .ok_or_else(|| SparkConnectError::Config("Not connected".to_string()))?;
+            .ok_or_else(|| SparkConnectError::Config("No connection config stored".to_string()))?;
 
-        client.create_temp_view(view_name, batches).await
+        let new_client = match config {
+            ConnectionConfig::Direct { host, token, cluster_id } => {
+                SparkConnectClient::connect(host, token, cluster_id.as_deref()).await?
+            }
+            ConnectionConfig::Proxy { url } => {
+                SparkConnectClient::connect_via_proxy(url).await?
+            }
+        };
+
+        drop(config_guard);
+        *self.client.write().await = Some(new_client);
+        info!("Spark Connect session reinitialized successfully");
+        Ok(())
     }
 }
 
@@ -154,6 +226,15 @@ impl Default for DatabricksClient {
     }
 }
 
+fn is_session_timeout(err: &SparkConnectError) -> bool {
+    if let SparkConnectError::Grpc(status) = err {
+        status.code() == tonic::Code::FailedPrecondition
+            && status.message().contains("INACTIVITY_TIMEOUT")
+    } else {
+        false
+    }
+}
+
 /// Try to create and connect a DatabricksClient from environment variables.
 /// Uses serverless if DATABRICKS_CLUSTER_ID is not set.
 pub async fn try_connect_from_env() -> Option<Arc<DatabricksClient>> {
@@ -171,3 +252,34 @@ pub async fn try_connect_from_env() -> Option<Arc<DatabricksClient>> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn is_session_timeout_detects_inactivity_timeout() {
+        let status = tonic::Status::failed_precondition(
+            "Session expired: reason=INACTIVITY_TIMEOUT after 10 minutes",
+        );
+        assert!(is_session_timeout(&SparkConnectError::Grpc(status)));
+    }
+
+    #[test]
+    fn is_session_timeout_ignores_other_precondition() {
+        let status = tonic::Status::failed_precondition("Some other precondition failure");
+        assert!(!is_session_timeout(&SparkConnectError::Grpc(status)));
+    }
+
+    #[test]
+    fn is_session_timeout_ignores_non_grpc_errors() {
+        let err = SparkConnectError::Config("Not connected".to_string());
+        assert!(!is_session_timeout(&err));
+    }
+
+    #[test]
+    fn is_session_timeout_ignores_wrong_grpc_code() {
+        let status = tonic::Status::unavailable("INACTIVITY_TIMEOUT");
+        assert!(!is_session_timeout(&SparkConnectError::Grpc(status)));
+    }
+}