spark_connect_core/
error.rs

1//! Error types mirroring PySpark's error-class framework.
2//!
3//! PySpark raises typed exceptions carrying an ``errorClass`` and
4//! ``messageParameters`` (see ``pyspark/errors/exceptions``). We preserve the
5//! error class so the PyO3 layer can re-raise the exact Python exception type
6//! with the same class/parameters, giving byte-compatible error messages.
7//!
8//! The error-conditions registry is embedded at build time from
9//! `error-conditions.json`, allowing us to render error messages with parameter
10//! substitution matching Python's `PySparkException.getMessage()` exactly.
11
12use prost::Message;
13use std::collections::BTreeMap;
14use std::collections::HashMap;
15use std::fmt;
16use tonic::Status;
17
18// Embed error-conditions.json at compile time for offline message rendering.
19const ERROR_CONDITIONS_JSON: &str = include_str!("../error-conditions.json");
20
21// ============================================================================
22// Query context types mirroring PySpark's QueryContext
23// ============================================================================
24
25/// The type of query context (SQL or DataFrame).
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum QueryContextType {
28    SQL,
29    DataFrame,
30}
31
32impl fmt::Display for QueryContextType {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            QueryContextType::SQL => write!(f, "SQL"),
36            QueryContextType::DataFrame => write!(f, "DataFrame"),
37        }
38    }
39}
40
41/// Query context describing where an error occurred.
42#[derive(Debug, Clone)]
43pub struct QueryContext {
44    pub context_type: QueryContextType,
45    pub object_type: String,
46    pub object_name: String,
47    pub start_index: i32,
48    pub stop_index: i32,
49    pub fragment: String,
50    pub call_site: String,
51    pub summary: String,
52}
53
54impl QueryContext {
55    /// Create a new QueryContext.
56    pub fn new(
57        context_type: QueryContextType,
58        object_type: String,
59        object_name: String,
60        start_index: i32,
61        stop_index: i32,
62        fragment: String,
63        call_site: String,
64        summary: String,
65    ) -> Self {
66        Self {
67            context_type,
68            object_type,
69            object_name,
70            start_index,
71            stop_index,
72            fragment,
73            call_site,
74            summary,
75        }
76    }
77
78    /// Get the context type.
79    pub fn context_type(&self) -> QueryContextType {
80        self.context_type
81    }
82
83    /// Get the object type (e.g., "VIEW", or empty string for main query).
84    pub fn object_type(&self) -> &str {
85        &self.object_type
86    }
87
88    /// Get the object name (e.g., view name, or empty string for main query).
89    pub fn object_name(&self) -> &str {
90        &self.object_name
91    }
92
93    /// Get the starting index in the query text (0-based).
94    pub fn start_index(&self) -> i32 {
95        self.start_index
96    }
97
98    /// Get the stopping index in the query text (0-based).
99    pub fn stop_index(&self) -> i32 {
100        self.stop_index
101    }
102
103    /// Get the corresponding fragment of the query.
104    pub fn fragment(&self) -> &str {
105        &self.fragment
106    }
107
108    /// Get the user code (call site of the API) that caused the exception.
109    pub fn call_site(&self) -> &str {
110        &self.call_site
111    }
112
113    /// Get a summary of the exception cause.
114    pub fn summary(&self) -> &str {
115        &self.summary
116    }
117}
118
119// ============================================================================
120// Minimal prost message definitions for google.rpc types (needed for error parsing)
121// ============================================================================
122
123/// Minimal google.rpc.Status for decoding error details.
124#[derive(Clone, PartialEq, Message)]
125struct RpcStatus {
126    #[prost(int32, tag = "1")]
127    code: i32,
128    #[prost(string, tag = "2")]
129    message: String,
130    #[prost(message, repeated, tag = "3")]
131    details: Vec<RpcAny>,
132}
133
134/// Minimal google.rpc.Any for decoding packed types.
135#[derive(Clone, PartialEq, Message)]
136struct RpcAny {
137    #[prost(string, tag = "1")]
138    type_url: String,
139    #[prost(bytes, tag = "2")]
140    value: Vec<u8>,
141}
142
143/// Minimal google.rpc.ErrorInfo for structured error details.
144#[derive(Clone, PartialEq, Message)]
145struct RpcErrorInfo {
146    #[prost(string, tag = "1")]
147    reason: String,
148    #[prost(string, tag = "2")]
149    domain: String,
150    #[prost(map = "string, string", tag = "3")]
151    metadata: HashMap<String, String>,
152}
153
154/// A structured error carrying a PySpark error class and its message parameters.
155#[derive(Debug, Clone)]
156pub struct SparkError {
157    /// Which Python exception type to raise at the PyO3 boundary.
158    pub kind: SparkErrorKind,
159    /// PySpark error class (e.g. "INVALID_CONNECT_URL"), or empty for plain messages.
160    pub error_class: String,
161    /// Message parameters keyed by name, as in ``messageParameters``.
162    pub params: BTreeMap<String, String>,
163    /// A pre-rendered message, used when no error class applies.
164    pub message: String,
165    /// SQL state code, if known (from the error-conditions registry).
166    pub sql_state: Option<String>,
167    /// Query contexts where the error occurred.
168    pub contexts: Vec<QueryContext>,
169    /// Server-side stack trace, if available.
170    pub server_stacktrace: Option<String>,
171    /// Raw gRPC status code, if this error originated from a gRPC Status.
172    pub grpc_code: Option<i32>,
173    /// Raw `grpc-status-details-bin` (a serialized `google.rpc.Status`), if present.
174    /// Lets the PyO3 layer reconstruct the exact typed pyspark exception.
175    pub grpc_details: Option<Vec<u8>>,
176}
177
178/// The PySpark exception hierarchy leaf we map to.
179///
180/// These variants correspond to the Python exception classes in
181/// `pyspark/errors/exceptions/base.py` and `pyspark/errors/exceptions/connect.py`.
182#[derive(Debug, Clone, Copy, PartialEq, Eq)]
183pub enum SparkErrorKind {
184    /// ``pyspark.errors.PySparkValueError``
185    ValueError,
186    /// ``pyspark.errors.PySparkTypeError``
187    TypeError,
188    /// ``pyspark.errors.PySparkIndexError``
189    IndexError,
190    /// ``pyspark.errors.PySparkAttributeError``
191    AttributeError,
192    /// ``pyspark.errors.PySparkKeyError``
193    KeyError,
194    /// ``pyspark.errors.PySparkRuntimeError``
195    RuntimeError,
196    /// ``pyspark.errors.PySparkNotImplementedError``
197    NotImplementedError,
198    /// ``pyspark.errors.PySparkAssertionError``
199    AssertionError,
200    /// ``pyspark.errors.PySparkPicklingError``
201    PicklingError,
202    /// ``pyspark.errors.PySparkImportError``
203    ImportError,
204    /// ``pyspark.errors.exceptions.connect.SparkConnectException`` and subclasses
205    Connect,
206    /// ``pyspark.errors.exceptions.connect.SparkConnectGrpcException``
207    ConnectGrpc,
208    /// ``pyspark.errors.exceptions.base.AnalysisException``
209    Analysis,
210    /// ``pyspark.errors.exceptions.base.SessionNotSameException``
211    SessionNotSame,
212    /// ``pyspark.errors.exceptions.base.TempTableAlreadyExistsException``
213    TempTableAlreadyExists,
214    /// ``pyspark.errors.exceptions.base.ParseException``
215    Parse,
216    /// ``pyspark.errors.exceptions.base.IllegalArgumentException``
217    IllegalArgument,
218    /// ``pyspark.errors.exceptions.base.ArithmeticException``
219    Arithmetic,
220    /// ``pyspark.errors.exceptions.base.UnsupportedOperationException``
221    UnsupportedOperation,
222    /// ``pyspark.errors.exceptions.base.ArrayIndexOutOfBoundsException``
223    ArrayIndexOutOfBounds,
224    /// ``pyspark.errors.exceptions.base.DateTimeException``
225    DateTime,
226    /// ``pyspark.errors.exceptions.base.NumberFormatException``
227    NumberFormat,
228    /// ``pyspark.errors.exceptions.base.StreamingQueryException``
229    StreamingQuery,
230    /// ``pyspark.errors.exceptions.base.StreamingPythonRunnerInitializationException``
231    StreamingPythonRunnerInitialization,
232    /// ``pyspark.errors.exceptions.base.QueryExecutionException``
233    QueryExecution,
234    /// ``pyspark.errors.exceptions.base.PythonException``
235    Python,
236    /// ``pyspark.errors.exceptions.base.SparkRuntimeException``
237    SparkRuntime,
238    /// ``pyspark.errors.exceptions.base.SparkUpgradeException``
239    SparkUpgrade,
240    /// ``pyspark.errors.exceptions.base.SparkNoSuchElementException``
241    SparkNoSuchElement,
242    /// ``pyspark.errors.exceptions.base.UnknownException``
243    Unknown,
244    /// ``pyspark.errors.exceptions.connect.InvalidPlanInput``
245    InvalidPlanInput,
246    /// ``pyspark.errors.exceptions.connect.PickleException``
247    PickleException,
248}
249
250impl SparkError {
251    /// Create a ValueError-kind error with an error class and parameters.
252    pub fn value(error_class: &str, params: &[(&str, &str)]) -> Self {
253        Self::classed(SparkErrorKind::ValueError, error_class, params)
254    }
255
256    /// Create a Connect-kind error with a plain message (no error class).
257    pub fn connect_msg(message: impl Into<String>) -> Self {
258        Self {
259            kind: SparkErrorKind::Connect,
260            error_class: String::new(),
261            params: BTreeMap::new(),
262            message: message.into(),
263            sql_state: None,
264            contexts: Vec::new(),
265            server_stacktrace: None,
266            grpc_code: None,
267            grpc_details: None,
268        }
269    }
270
271    /// Create a ValueError-kind error with a plain message (no error class).
272    ///
273    /// For invalid client-side configuration (e.g. a malformed connection string)
274    /// that has no registered error class. Surfaces as `PySparkValueError` on the
275    /// Python side, matching how the reference client reports connection-string
276    /// problems, rather than the generic runtime kind.
277    pub fn value_msg(message: impl Into<String>) -> Self {
278        Self {
279            kind: SparkErrorKind::ValueError,
280            error_class: String::new(),
281            params: BTreeMap::new(),
282            message: message.into(),
283            sql_state: None,
284            contexts: Vec::new(),
285            server_stacktrace: None,
286            grpc_code: None,
287            grpc_details: None,
288        }
289    }
290
291    /// Create an error of a specific kind with an error class and parameters.
292    pub fn classed(kind: SparkErrorKind, error_class: &str, params: &[(&str, &str)]) -> Self {
293        Self {
294            kind,
295            error_class: error_class.to_string(),
296            params: params
297                .iter()
298                .map(|(k, v)| (k.to_string(), v.to_string()))
299                .collect(),
300            message: String::new(),
301            sql_state: None,
302            contexts: Vec::new(),
303            server_stacktrace: None,
304            grpc_code: None,
305            grpc_details: None,
306        }
307    }
308
309    /// Get the rendered message for this error, performing template substitution.
310    ///
311    /// If an error_class is set, looks up the message template in error-conditions.json,
312    /// substitutes `<param_name>` placeholders with values from `params`, and returns
313    /// `[ERROR_CLASS] rendered_message`. Otherwise, returns the pre-set message.
314    ///
315    /// Mirrors `PySparkException.getMessage()` from `pyspark/errors/exceptions/base.py`.
316    pub fn message(&self) -> String {
317        if self.error_class.is_empty() {
318            return self.message.clone();
319        }
320
321        // Try to render from error-conditions registry
322        if let Ok(rendered) = render_message(&self.error_class, &self.params) {
323            format!("[{}] {}", self.error_class, rendered)
324        } else {
325            // Fallback: format as-is if error class not found
326            if self.message.is_empty() {
327                format!("[{}]", self.error_class)
328            } else {
329                format!("[{}] {}", self.error_class, self.message)
330            }
331        }
332    }
333
334    /// Get the SQL state code for this error, if known.
335    pub fn sql_state(&self) -> Option<String> {
336        if self.sql_state.is_some() {
337            self.sql_state.clone()
338        } else if !self.error_class.is_empty() {
339            get_sql_state(&self.error_class).map(|s| s.to_string())
340        } else {
341            None
342        }
343    }
344
345    /// Convert a gRPC Status to a SparkError, extracting error class and parameters from metadata.
346    ///
347    /// Mirrors `convert_exception` from `pyspark.errors.exceptions.connect`.
348    pub fn from_grpc_status(status: Status) -> Self {
349        // Extract error metadata from trailer if present
350        let message = status.message().to_string();
351        let code = status.code() as i32;
352
353        // Try to parse details from the Status (details is a byte slice)
354        let details = status.details().to_vec();
355        let mut err = if details.is_empty() {
356            Self::connect_msg(format!("[{}] {}", status.code(), message))
357        } else {
358            // Decode grpc-status-details-bin as a google.rpc.Status with Any-packed ErrorInfo.
359            match parse_error_info_from_details(&details) {
360                Some((error_class, params, sql_state, server_stacktrace)) => Self {
361                    kind: classify_error_kind_with_classes(&error_class, &params),
362                    error_class,
363                    params,
364                    message,
365                    sql_state,
366                    contexts: Vec::new(),
367                    server_stacktrace,
368                    grpc_code: None,
369                    grpc_details: None,
370                },
371                None => Self::connect_msg(format!("[{}] {}", status.code(), message)),
372            }
373        };
374        // Preserve the raw status so the PyO3 layer can rebuild the exact pyspark exception.
375        err.grpc_code = Some(code);
376        if !details.is_empty() {
377            err.grpc_details = Some(details);
378        }
379        err
380    }
381
382    /// Get the error condition (error class).
383    ///
384    /// Mirrors `PySparkException.getCondition()` from PySpark.
385    pub fn get_condition(&self) -> Option<String> {
386        if self.error_class.is_empty() {
387            None
388        } else {
389            Some(self.error_class.clone())
390        }
391    }
392
393    /// Get the error class (deprecated, use `get_condition` instead).
394    ///
395    /// Mirrors `PySparkException.getErrorClass()` from PySpark (deprecated).
396    pub fn get_error_class(&self) -> Option<String> {
397        self.get_condition()
398    }
399
400    /// Get the message parameters.
401    ///
402    /// Mirrors `PySparkException.getMessageParameters()` from PySpark.
403    pub fn get_message_parameters(&self) -> Option<BTreeMap<String, String>> {
404        if self.params.is_empty() && self.error_class.is_empty() {
405            None
406        } else {
407            Some(self.params.clone())
408        }
409    }
410
411    /// Get the SQL state code, if known.
412    ///
413    /// Mirrors `PySparkException.getSqlState()` from PySpark.
414    pub fn get_sql_state(&self) -> Option<String> {
415        if self.sql_state.is_some() {
416            self.sql_state.clone()
417        } else if !self.error_class.is_empty() {
418            get_sql_state(&self.error_class).map(|s| s.to_string())
419        } else {
420            None
421        }
422    }
423
424    /// Get the rendered error message.
425    ///
426    /// Mirrors `PySparkException.getMessage()` from PySpark.
427    pub fn get_message(&self) -> String {
428        self.message()
429    }
430
431    /// Get the query contexts where this error occurred.
432    ///
433    /// Mirrors `PySparkException.getQueryContext()` from PySpark.
434    pub fn get_query_context(&self) -> Vec<QueryContext> {
435        self.contexts.clone()
436    }
437
438    /// Get the server-side stack trace, if available.
439    pub fn get_stacktrace(&self) -> Option<String> {
440        self.server_stacktrace.clone()
441    }
442
443    /// Raw gRPC status code, if this error came from a gRPC Status.
444    pub fn grpc_code(&self) -> Option<i32> {
445        self.grpc_code
446    }
447
448    /// Raw `grpc-status-details-bin` bytes (serialized `google.rpc.Status`), if present.
449    pub fn grpc_details(&self) -> Option<&[u8]> {
450        self.grpc_details.as_deref()
451    }
452}
453
454impl std::fmt::Display for SparkError {
455    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
456        write!(f, "{}", self.message())
457    }
458}
459
460impl std::error::Error for SparkError {}
461
462pub type Result<T> = std::result::Result<T, SparkError>;
463
464// ============================================================================
465// gRPC error detail parsing
466// ============================================================================
467
468/// Parse ErrorInfo from gRPC status details bytes.
469///
470/// Expects the details to be a protobuf-encoded google.rpc.Status containing
471/// Any-packed google.rpc.ErrorInfo messages. Returns (error_class, params, sql_state, stacktrace) if found.
472fn parse_error_info_from_details(
473    details: &[u8],
474) -> Option<(
475    String,
476    BTreeMap<String, String>,
477    Option<String>,
478    Option<String>,
479)> {
480    // Decode the outer google.rpc.Status from details
481    let rpc_status = RpcStatus::decode(details).ok()?;
482
483    // Look for an ErrorInfo in the details list
484    for any_detail in &rpc_status.details {
485        // google.rpc.ErrorInfo type URL
486        if any_detail.type_url == "type.googleapis.com/google.rpc.ErrorInfo" {
487            // Decode the Any.value as google.rpc.ErrorInfo
488            if let Ok(error_info) = RpcErrorInfo::decode(&any_detail.value[..]) {
489                // Extract errorClass and other metadata
490                let error_class = error_info
491                    .metadata
492                    .get("errorClass")
493                    .cloned()
494                    .unwrap_or_else(|| error_info.reason.clone());
495
496                let sql_state = error_info.metadata.get("sqlState").cloned();
497                let stacktrace = error_info.metadata.get("stackTrace").cloned();
498
499                // Convert HashMap to BTreeMap
500                let params: BTreeMap<String, String> = error_info
501                    .metadata
502                    .iter()
503                    .map(|(k, v)| (k.clone(), v.clone()))
504                    .collect();
505
506                // Return the parsed error info
507                return Some((error_class, params, sql_state, stacktrace));
508            }
509        }
510    }
511
512    None
513}
514
515/// Classify a server error into a `SparkErrorKind`, preferring the JVM exception-class
516/// hierarchy (the `classes` metadata, e.g. `["org.apache.spark.sql.AnalysisException",…]`)
517/// over the error-condition name. Spark's analysis errors are named `UNRESOLVED_COLUMN`,
518/// `TABLE_OR_VIEW_NOT_FOUND`, … — none start with "ANALYSIS" — so classifying by the
519/// condition string alone mislabels them as generic runtime errors and breaks
520/// `except AnalysisException:`.
521fn classify_error_kind_with_classes(
522    error_class: &str,
523    params: &BTreeMap<String, String>,
524) -> SparkErrorKind {
525    if let Some(classes) = params.get("classes") {
526        // Most-specific first; the string holds the whole hierarchy.
527        let checks: &[(&str, SparkErrorKind)] = &[
528            ("AnalysisException", SparkErrorKind::Analysis),
529            ("ParseException", SparkErrorKind::Parse),
530            ("StreamingQueryException", SparkErrorKind::StreamingQuery),
531            ("SparkUpgradeException", SparkErrorKind::SparkUpgrade),
532            ("NumberFormatException", SparkErrorKind::NumberFormat),
533            (
534                "ArrayIndexOutOfBoundsException",
535                SparkErrorKind::ArrayIndexOutOfBounds,
536            ),
537            ("DateTimeException", SparkErrorKind::DateTime),
538            ("ArithmeticException", SparkErrorKind::Arithmetic),
539            (
540                "UnsupportedOperationException",
541                SparkErrorKind::UnsupportedOperation,
542            ),
543            ("IllegalArgumentException", SparkErrorKind::IllegalArgument),
544            ("NoSuchElementException", SparkErrorKind::SparkNoSuchElement),
545            ("PythonException", SparkErrorKind::Python),
546            ("SparkRuntimeException", SparkErrorKind::SparkRuntime),
547        ];
548        for (needle, kind) in checks {
549            if classes.contains(needle) {
550                return *kind;
551            }
552        }
553    }
554    classify_error_kind(error_class)
555}
556
557/// Classify error_class string into a SparkErrorKind.
558///
559/// Maps Spark error classes to the corresponding Python exception type.
560fn classify_error_kind(error_class: &str) -> SparkErrorKind {
561    if error_class.starts_with("ANALYSIS") {
562        SparkErrorKind::Analysis
563    } else if error_class.starts_with("PARSE") {
564        SparkErrorKind::Parse
565    } else if error_class.starts_with("ILLEGAL") {
566        SparkErrorKind::IllegalArgument
567    } else if error_class.starts_with("UNSUPPORTED") {
568        SparkErrorKind::UnsupportedOperation
569    } else if error_class.starts_with("ARITHMETIC") {
570        SparkErrorKind::Arithmetic
571    } else if error_class.contains("NOT_IMPLEMENTED") {
572        SparkErrorKind::NotImplementedError
573    } else if error_class.contains("RUNTIME") {
574        SparkErrorKind::RuntimeError
575    } else if error_class.contains("VALUE") {
576        SparkErrorKind::ValueError
577    } else {
578        match error_class {
579            "INVALID_CONNECT_URL" => SparkErrorKind::RuntimeError,
580            "SYNTAX_ERROR" => SparkErrorKind::Parse,
581            "INVALID_PLAN_INPUT" => SparkErrorKind::InvalidPlanInput,
582            "PICKLE_ERROR" => SparkErrorKind::PickleException,
583            "RESPONSE_ALREADY_RECEIVED" | "INVALID_HANDLE" => SparkErrorKind::ConnectGrpc,
584            _ => SparkErrorKind::RuntimeError, // Default fallback
585        }
586    }
587}
588
589// ============================================================================
590// Error conditions registry parsing and message rendering
591// ============================================================================
592
593/// Render an error message by substituting parameters into the template.
594///
595/// Looks up the error class in the error-conditions registry, finds the message
596/// template, and substitutes `<param_name>` placeholders with values from the
597/// params map. Returns the rendered message without the error class prefix.
598fn render_message(
599    error_class: &str,
600    params: &BTreeMap<String, String>,
601) -> std::result::Result<String, String> {
602    let template = get_message_template(error_class)?;
603
604    // Find all placeholders in the template
605    let mut result = template.clone();
606
607    // Use regex to find and replace all <name> placeholders
608    // Pattern: <([a-zA-Z0-9_-]+)>
609    let pattern = regex::Regex::new(r"<([a-zA-Z0-9_\-]+)>").map_err(|e| e.to_string())?;
610
611    for cap in pattern.captures_iter(&template) {
612        if let Some(param_name) = cap.get(1) {
613            let name = param_name.as_str();
614            if let Some(value) = params.get(name) {
615                let placeholder = format!("<{}>", name);
616                result = result.replace(&placeholder, value);
617            }
618        }
619    }
620
621    Ok(result)
622}
623
624/// Get the message template for an error class from the registry.
625fn get_message_template(error_class: &str) -> std::result::Result<String, String> {
626    // Parse error_class which may be "MAIN_CLASS" or "MAIN_CLASS.SUB_CLASS"
627    let parts: Vec<&str> = error_class.split('.').collect();
628
629    let json_obj = parse_error_conditions_json()
630        .map_err(|e| format!("Failed to parse error-conditions.json: {}", e))?;
631
632    match parts.len() {
633        1 => {
634            let main_class = parts[0];
635            if let Some(entry) = json_obj.get(main_class) {
636                if let Some(msg_array) = entry.get("message") {
637                    if let Some(msg_list) = msg_array.as_array() {
638                        let message_parts: Vec<String> = msg_list
639                            .iter()
640                            .filter_map(|m| m.as_str().map(|s| s.to_string()))
641                            .collect();
642                        return Ok(message_parts.join("\n"));
643                    }
644                }
645            }
646            Err(format!("Error class not found: {}", main_class))
647        }
648        2 => {
649            let main_class = parts[0];
650            let sub_class = parts[1];
651            if let Some(entry) = json_obj.get(main_class) {
652                // Get main message
653                let mut message = String::new();
654                if let Some(msg_array) = entry.get("message") {
655                    if let Some(msg_list) = msg_array.as_array() {
656                        let message_parts: Vec<String> = msg_list
657                            .iter()
658                            .filter_map(|m| m.as_str().map(|s| s.to_string()))
659                            .collect();
660                        message = message_parts.join("\n");
661                    }
662                }
663
664                // Get sub-class message if it exists
665                if let Some(subclasses) = entry.get("sub_class") {
666                    if let Some(sub) = subclasses.get(sub_class) {
667                        if let Some(sub_msg_array) = sub.get("message") {
668                            if let Some(sub_msg_list) = sub_msg_array.as_array() {
669                                let sub_message_parts: Vec<String> = sub_msg_list
670                                    .iter()
671                                    .filter_map(|m| m.as_str().map(|s| s.to_string()))
672                                    .collect();
673                                if !message.is_empty() {
674                                    message.push(' ');
675                                }
676                                message.push_str(&sub_message_parts.join("\n"));
677                            }
678                        }
679                    }
680                }
681
682                if !message.is_empty() {
683                    Ok(message)
684                } else {
685                    Err(format!("No message found for error class: {}", error_class))
686                }
687            } else {
688                Err(format!("Main error class not found: {}", main_class))
689            }
690        }
691        _ => Err(format!("Invalid error class format: {}", error_class)),
692    }
693}
694
695/// Get the SQL state code for an error class from the registry.
696fn get_sql_state(error_class: &str) -> Option<String> {
697    let parts: Vec<&str> = error_class.split('.').collect();
698
699    let json_obj = parse_error_conditions_json().ok()?;
700
701    match parts.len() {
702        1 => {
703            let main_class = parts[0];
704            json_obj
705                .get(main_class)
706                .and_then(|entry| entry.get("sqlState"))
707                .and_then(|state| state.as_str())
708                .map(|s| s.to_string())
709        }
710        2 => {
711            let main_class = parts[0];
712            let sub_class = parts[1];
713            json_obj
714                .get(main_class)
715                .and_then(|entry| entry.get("sub_class"))
716                .and_then(|subclasses| subclasses.get(sub_class))
717                .and_then(|sub| sub.get("sqlState"))
718                .and_then(|state| state.as_str())
719                .map(|s| s.to_string())
720        }
721        _ => None,
722    }
723}
724
725/// Parse the embedded error-conditions.json file.
726fn parse_error_conditions_json(
727) -> std::result::Result<serde_json::Map<String, serde_json::Value>, String> {
728    let json: serde_json::Value = serde_json::from_str(ERROR_CONDITIONS_JSON)
729        .map_err(|e| format!("Failed to parse JSON: {}", e))?;
730    match json {
731        serde_json::Value::Object(map) => Ok(map),
732        _ => Err("Expected JSON object at root".to_string()),
733    }
734}
735
736#[cfg(test)]
737mod tests {
738    use super::*;
739
740    #[test]
741    fn test_render_simple_message() {
742        let mut params = BTreeMap::new();
743        params.insert("arg_name".to_string(), "x".to_string());
744
745        let result = render_message("CANNOT_BE_NONE", &params).expect("should render");
746        assert_eq!(result, "Argument `x` cannot be None.");
747    }
748
749    #[test]
750    fn test_error_message_format() {
751        let mut params = BTreeMap::new();
752        params.insert("arg_name".to_string(), "x".to_string());
753
754        let err = SparkError::classed(
755            SparkErrorKind::ValueError,
756            "CANNOT_BE_NONE",
757            &[("arg_name", "x")],
758        );
759
760        let msg = err.message();
761        assert_eq!(msg, "[CANNOT_BE_NONE] Argument `x` cannot be None.");
762    }
763
764    #[test]
765    fn test_error_display() {
766        let err = SparkError::classed(
767            SparkErrorKind::ValueError,
768            "CANNOT_BE_NONE",
769            &[("arg_name", "test_arg")],
770        );
771
772        let display = format!("{}", err);
773        assert_eq!(
774            display,
775            "[CANNOT_BE_NONE] Argument `test_arg` cannot be None."
776        );
777    }
778
779    #[test]
780    fn test_connect_msg() {
781        let err = SparkError::connect_msg("Connection failed");
782        let msg = err.message();
783        assert_eq!(msg, "Connection failed");
784    }
785
786    #[test]
787    fn test_get_sql_state_known() {
788        let sql_state = get_sql_state("ATTRIBUTE_NOT_SUPPORTED");
789        assert_eq!(sql_state, Some("0A000".to_string()));
790    }
791
792    #[test]
793    fn test_get_sql_state_unknown() {
794        let sql_state = get_sql_state("UNKNOWN_ERROR_CLASS_XXXXX");
795        assert_eq!(sql_state, None);
796    }
797
798    #[test]
799    fn test_multiple_params() {
800        let result = render_message(
801            "INVALID_CONNECT_URL",
802            &[(
803                "detail".to_string(),
804                "The URL must start with 'sc://'".to_string(),
805            )]
806            .iter()
807            .cloned()
808            .collect(),
809        );
810        assert!(result.is_ok());
811    }
812
813    #[test]
814    fn test_full_error_message_with_params() {
815        // Test: CANNOT_BE_NONE
816        let err = SparkError::classed(
817            SparkErrorKind::ValueError,
818            "CANNOT_BE_NONE",
819            &[("arg_name", "my_arg")],
820        );
821        assert_eq!(
822            err.message(),
823            "[CANNOT_BE_NONE] Argument `my_arg` cannot be None."
824        );
825
826        // Test: ARGUMENT_REQUIRED
827        let err = SparkError::classed(
828            SparkErrorKind::ValueError,
829            "ARGUMENT_REQUIRED",
830            &[("arg_name", "foo"), ("condition", "x > 0")],
831        );
832        assert_eq!(
833            err.message(),
834            "[ARGUMENT_REQUIRED] Argument `foo` is required when x > 0."
835        );
836
837        // Test: INVALID_CONNECT_URL
838        let err = SparkError::classed(
839            SparkErrorKind::RuntimeError,
840            "INVALID_CONNECT_URL",
841            &[("detail", "The URL must start with sc://")],
842        );
843        assert_eq!(
844            err.message(),
845            "[INVALID_CONNECT_URL] Invalid URL for Spark Connect: The URL must start with sc://"
846        );
847
848        // Test: ATTRIBUTE_NOT_CALLABLE
849        let err = SparkError::classed(
850            SparkErrorKind::ValueError,
851            "ATTRIBUTE_NOT_CALLABLE",
852            &[("attr_name", "compute"), ("obj_name", "MyClass")],
853        );
854        assert_eq!(
855            err.message(),
856            "[ATTRIBUTE_NOT_CALLABLE] Attribute `compute` in provided object `MyClass` is not callable."
857        );
858    }
859
860    #[test]
861    fn test_error_without_error_class() {
862        let mut err = SparkError::connect_msg("Custom error message");
863        assert_eq!(err.message(), "Custom error message");
864
865        // Also test when message field is set but error_class is empty
866        err.message = "Another message".to_string();
867        assert_eq!(err.message(), "Another message");
868    }
869
870    #[test]
871    fn test_error_kind_variants() {
872        // Ensure all error kinds are defined and accessible
873        let _ = SparkErrorKind::ValueError;
874        let _ = SparkErrorKind::TypeError;
875        let _ = SparkErrorKind::IndexError;
876        let _ = SparkErrorKind::AttributeError;
877        let _ = SparkErrorKind::KeyError;
878        let _ = SparkErrorKind::RuntimeError;
879        let _ = SparkErrorKind::NotImplementedError;
880        let _ = SparkErrorKind::AssertionError;
881        let _ = SparkErrorKind::PicklingError;
882        let _ = SparkErrorKind::ImportError;
883        let _ = SparkErrorKind::Connect;
884        let _ = SparkErrorKind::ConnectGrpc;
885        let _ = SparkErrorKind::Analysis;
886        let _ = SparkErrorKind::Parse;
887        let _ = SparkErrorKind::Python;
888        let _ = SparkErrorKind::Unknown;
889    }
890
891    #[test]
892    fn test_message_rendering_parity_with_python() {
893        // These test cases are verified to match PySpark's str(e) output exactly.
894        // Run these in Python to verify:
895        //   from pyspark.errors import PySparkValueError, PySparkRuntimeError
896        //   e = PySparkValueError(errorClass='CANNOT_BE_NONE', messageParameters={'arg_name':'x'})
897        //   print(str(e))  # Should be: [CANNOT_BE_NONE] Argument `x` cannot be None.
898
899        // Test 1: CANNOT_BE_NONE
900        let err = SparkError::classed(
901            SparkErrorKind::ValueError,
902            "CANNOT_BE_NONE",
903            &[("arg_name", "x")],
904        );
905        assert_eq!(
906            err.message(),
907            "[CANNOT_BE_NONE] Argument `x` cannot be None."
908        );
909
910        // Test 2: ARGUMENT_REQUIRED
911        let err = SparkError::classed(
912            SparkErrorKind::ValueError,
913            "ARGUMENT_REQUIRED",
914            &[("arg_name", "foo"), ("condition", "x > 0")],
915        );
916        assert_eq!(
917            err.message(),
918            "[ARGUMENT_REQUIRED] Argument `foo` is required when x > 0."
919        );
920
921        // Test 3: INVALID_CONNECT_URL
922        let err = SparkError::classed(
923            SparkErrorKind::RuntimeError,
924            "INVALID_CONNECT_URL",
925            &[("detail", "test")],
926        );
927        assert_eq!(
928            err.message(),
929            "[INVALID_CONNECT_URL] Invalid URL for Spark Connect: test"
930        );
931
932        // Test 4: ATTRIBUTE_NOT_CALLABLE
933        let err = SparkError::classed(
934            SparkErrorKind::ValueError,
935            "ATTRIBUTE_NOT_CALLABLE",
936            &[("attr_name", "compute"), ("obj_name", "MyClass")],
937        );
938        assert_eq!(
939            err.message(),
940            "[ATTRIBUTE_NOT_CALLABLE] Attribute `compute` in provided object `MyClass` is not callable."
941        );
942    }
943
944    #[test]
945    fn test_parse_error_info_from_grpc_status() {
946        // Create a synthetic google.rpc.ErrorInfo
947        let mut error_metadata = HashMap::new();
948        error_metadata.insert("errorClass".to_string(), "ANALYSIS_ERROR".to_string());
949        error_metadata.insert("sqlState".to_string(), "42601".to_string());
950        error_metadata.insert(
951            "messageParameters".to_string(),
952            r#"{"message":"test error"}"#.to_string(),
953        );
954        error_metadata.insert(
955            "stackTrace".to_string(),
956            "at org.apache.spark.sql....".to_string(),
957        );
958
959        let error_info = RpcErrorInfo {
960            reason: "ANALYSIS_ERROR".to_string(),
961            domain: "org.apache.spark.connect".to_string(),
962            metadata: error_metadata,
963        };
964
965        // Pack into Any
966        let error_info_bytes = error_info.encode_to_vec();
967        let any = RpcAny {
968            type_url: "type.googleapis.com/google.rpc.ErrorInfo".to_string(),
969            value: error_info_bytes,
970        };
971
972        // Create outer Status
973        let rpc_status = RpcStatus {
974            code: 0,
975            message: "Analysis error".to_string(),
976            details: vec![any],
977        };
978
979        let status_bytes = rpc_status.encode_to_vec();
980
981        // Parse it back
982        let result = parse_error_info_from_details(&status_bytes);
983        assert!(result.is_some());
984
985        let (error_class, params, sql_state, stacktrace) = result.unwrap();
986        assert_eq!(error_class, "ANALYSIS_ERROR");
987        assert_eq!(sql_state, Some("42601".to_string()));
988        assert_eq!(stacktrace, Some("at org.apache.spark.sql....".to_string()));
989        assert!(params.contains_key("errorClass"));
990    }
991
992    #[test]
993    fn test_classify_error_kind() {
994        assert_eq!(
995            classify_error_kind("ANALYSIS_ERROR"),
996            SparkErrorKind::Analysis
997        );
998        assert_eq!(classify_error_kind("PARSE_ERROR"), SparkErrorKind::Parse);
999        assert_eq!(
1000            classify_error_kind("ILLEGAL_ARGUMENT"),
1001            SparkErrorKind::IllegalArgument
1002        );
1003        assert_eq!(
1004            classify_error_kind("ARITHMETIC_ERROR"),
1005            SparkErrorKind::Arithmetic
1006        );
1007        assert_eq!(
1008            classify_error_kind("UNSUPPORTED_OPERATION"),
1009            SparkErrorKind::UnsupportedOperation
1010        );
1011        assert_eq!(
1012            classify_error_kind("INVALID_PLAN_INPUT"),
1013            SparkErrorKind::InvalidPlanInput
1014        );
1015        assert_eq!(
1016            classify_error_kind("RESPONSE_ALREADY_RECEIVED"),
1017            SparkErrorKind::ConnectGrpc
1018        );
1019    }
1020
1021    #[test]
1022    fn test_accessor_methods_get_condition() {
1023        // Test with error class
1024        let err = SparkError::classed(
1025            SparkErrorKind::ValueError,
1026            "CANNOT_BE_NONE",
1027            &[("arg_name", "test")],
1028        );
1029        assert_eq!(err.get_condition(), Some("CANNOT_BE_NONE".to_string()));
1030        assert_eq!(err.get_error_class(), Some("CANNOT_BE_NONE".to_string()));
1031
1032        // Test without error class
1033        let err = SparkError::connect_msg("Plain message");
1034        assert_eq!(err.get_condition(), None);
1035        assert_eq!(err.get_error_class(), None);
1036    }
1037
1038    #[test]
1039    fn test_accessor_methods_get_message_parameters() {
1040        // Test with parameters
1041        let err = SparkError::classed(
1042            SparkErrorKind::ValueError,
1043            "CANNOT_BE_NONE",
1044            &[("arg_name", "x"), ("other", "y")],
1045        );
1046        let params = err.get_message_parameters();
1047        assert!(params.is_some());
1048        let params = params.unwrap();
1049        assert_eq!(params.get("arg_name"), Some(&"x".to_string()));
1050        assert_eq!(params.get("other"), Some(&"y".to_string()));
1051
1052        // Test without parameters
1053        let err = SparkError::connect_msg("Plain message");
1054        assert_eq!(err.get_message_parameters(), None);
1055    }
1056
1057    #[test]
1058    fn test_accessor_methods_get_sql_state() {
1059        // Test with known SQL state
1060        let err = SparkError::classed(
1061            SparkErrorKind::ValueError,
1062            "ATTRIBUTE_NOT_SUPPORTED",
1063            &[("attr_name", "test")],
1064        );
1065        let sql_state = err.get_sql_state();
1066        assert_eq!(sql_state, Some("0A000".to_string()));
1067
1068        // Test without SQL state
1069        let err = SparkError::classed(
1070            SparkErrorKind::ValueError,
1071            "CANNOT_BE_NONE",
1072            &[("arg_name", "x")],
1073        );
1074        let sql_state = err.get_sql_state();
1075        assert_eq!(sql_state, None);
1076    }
1077
1078    #[test]
1079    fn test_accessor_methods_get_message() {
1080        let err = SparkError::classed(
1081            SparkErrorKind::ValueError,
1082            "CANNOT_BE_NONE",
1083            &[("arg_name", "x")],
1084        );
1085        assert_eq!(
1086            err.get_message(),
1087            "[CANNOT_BE_NONE] Argument `x` cannot be None."
1088        );
1089
1090        // Test without error class
1091        let err = SparkError::connect_msg("Custom message");
1092        assert_eq!(err.get_message(), "Custom message");
1093    }
1094
1095    #[test]
1096    fn test_accessor_methods_get_query_context() {
1097        let mut err = SparkError::connect_msg("Test error");
1098        assert!(err.get_query_context().is_empty());
1099
1100        // Add a query context
1101        let ctx = QueryContext::new(
1102            QueryContextType::SQL,
1103            "VIEW".to_string(),
1104            "my_view".to_string(),
1105            0,
1106            10,
1107            "SELECT * FROM".to_string(),
1108            "file.py:10".to_string(),
1109            "In DataFrame operation".to_string(),
1110        );
1111        err.contexts.push(ctx);
1112        let contexts = err.get_query_context();
1113        assert_eq!(contexts.len(), 1);
1114        assert_eq!(contexts[0].context_type(), QueryContextType::SQL);
1115        assert_eq!(contexts[0].object_type(), "VIEW");
1116        assert_eq!(contexts[0].object_name(), "my_view");
1117    }
1118
1119    #[test]
1120    fn test_query_context_creation() {
1121        let ctx = QueryContext::new(
1122            QueryContextType::DataFrame,
1123            "".to_string(),
1124            "".to_string(),
1125            5,
1126            15,
1127            "df.select(col('x'))".to_string(),
1128            "test.py:42".to_string(),
1129            "Selecting column x".to_string(),
1130        );
1131        assert_eq!(ctx.context_type(), QueryContextType::DataFrame);
1132        assert_eq!(ctx.start_index(), 5);
1133        assert_eq!(ctx.stop_index(), 15);
1134        assert_eq!(ctx.fragment(), "df.select(col('x'))");
1135        assert_eq!(ctx.call_site(), "test.py:42");
1136        assert_eq!(ctx.summary(), "Selecting column x");
1137    }
1138
1139    #[test]
1140    fn test_stacktrace_accessor() {
1141        let mut err = SparkError::connect_msg("Error with stack trace");
1142        assert_eq!(err.get_stacktrace(), None);
1143
1144        err.server_stacktrace = Some("at org.apache.spark.sql...".to_string());
1145        assert_eq!(
1146            err.get_stacktrace(),
1147            Some("at org.apache.spark.sql...".to_string())
1148        );
1149    }
1150
1151    #[test]
1152    fn test_parity_cannot_be_none() {
1153        // Verify that message rendering matches PySpark exactly
1154        // Expected in PySpark: [CANNOT_BE_NONE] Argument `x` cannot be None.
1155        let err = SparkError::classed(
1156            SparkErrorKind::ValueError,
1157            "CANNOT_BE_NONE",
1158            &[("arg_name", "x")],
1159        );
1160        assert_eq!(
1161            err.message(),
1162            "[CANNOT_BE_NONE] Argument `x` cannot be None."
1163        );
1164        assert_eq!(err.get_condition(), Some("CANNOT_BE_NONE".to_string()));
1165        let params = err.get_message_parameters().unwrap();
1166        assert_eq!(params.get("arg_name"), Some(&"x".to_string()));
1167    }
1168
1169    #[test]
1170    fn test_parity_argument_required() {
1171        // Expected in PySpark: [ARGUMENT_REQUIRED] Argument `foo` is required when x > 0.
1172        let err = SparkError::classed(
1173            SparkErrorKind::ValueError,
1174            "ARGUMENT_REQUIRED",
1175            &[("arg_name", "foo"), ("condition", "x > 0")],
1176        );
1177        assert_eq!(
1178            err.message(),
1179            "[ARGUMENT_REQUIRED] Argument `foo` is required when x > 0."
1180        );
1181    }
1182
1183    #[test]
1184    fn test_parity_invalid_connect_url() {
1185        // Expected in PySpark: [INVALID_CONNECT_URL] Invalid URL for Spark Connect: <detail>
1186        let err = SparkError::classed(
1187            SparkErrorKind::RuntimeError,
1188            "INVALID_CONNECT_URL",
1189            &[("detail", "must start with sc://")],
1190        );
1191        assert_eq!(
1192            err.message(),
1193            "[INVALID_CONNECT_URL] Invalid URL for Spark Connect: must start with sc://"
1194        );
1195    }
1196
1197    #[test]
1198    fn test_parity_attribute_not_callable() {
1199        // Expected in PySpark: [ATTRIBUTE_NOT_CALLABLE] Attribute `compute` in provided object `MyClass` is not callable.
1200        let err = SparkError::classed(
1201            SparkErrorKind::ValueError,
1202            "ATTRIBUTE_NOT_CALLABLE",
1203            &[("attr_name", "compute"), ("obj_name", "MyClass")],
1204        );
1205        assert_eq!(
1206            err.message(),
1207            "[ATTRIBUTE_NOT_CALLABLE] Attribute `compute` in provided object `MyClass` is not callable."
1208        );
1209    }
1210
1211    #[test]
1212    fn test_analysis_exception_mapping() {
1213        let err = SparkError::classed(
1214            SparkErrorKind::Analysis,
1215            "ANALYSIS_ERROR",
1216            &[("message", "Table not found")],
1217        );
1218        assert_eq!(err.kind, SparkErrorKind::Analysis);
1219        assert_eq!(err.get_condition(), Some("ANALYSIS_ERROR".to_string()));
1220    }
1221
1222    #[test]
1223    fn test_illegal_argument_exception_mapping() {
1224        let err = SparkError::classed(SparkErrorKind::IllegalArgument, "ILLEGAL_ARGUMENT", &[]);
1225        assert_eq!(err.kind, SparkErrorKind::IllegalArgument);
1226        assert_eq!(err.get_condition(), Some("ILLEGAL_ARGUMENT".to_string()));
1227    }
1228
1229    #[test]
1230    fn test_arithmetic_exception_mapping() {
1231        let err = SparkError::classed(
1232            SparkErrorKind::Arithmetic,
1233            "ARITHMETIC_ERROR",
1234            &[("message", "Division by zero")],
1235        );
1236        assert_eq!(err.kind, SparkErrorKind::Arithmetic);
1237    }
1238
1239    #[test]
1240    fn test_unsupported_operation_exception_mapping() {
1241        let err = SparkError::classed(
1242            SparkErrorKind::UnsupportedOperation,
1243            "UNSUPPORTED_OPERATION",
1244            &[],
1245        );
1246        assert_eq!(err.kind, SparkErrorKind::UnsupportedOperation);
1247    }
1248
1249    #[test]
1250    fn test_query_execution_exception_mapping() {
1251        let err = SparkError::classed(
1252            SparkErrorKind::QueryExecution,
1253            "QUERY_EXECUTION_ERROR",
1254            &[("message", "Stage failed")],
1255        );
1256        assert_eq!(err.kind, SparkErrorKind::QueryExecution);
1257    }
1258
1259    #[test]
1260    fn test_streaming_query_exception_mapping() {
1261        let err = SparkError::classed(SparkErrorKind::StreamingQuery, "STREAMING_QUERY_ERROR", &[]);
1262        assert_eq!(err.kind, SparkErrorKind::StreamingQuery);
1263    }
1264
1265    #[test]
1266    fn test_python_exception_mapping() {
1267        let err = SparkError::classed(
1268            SparkErrorKind::Python,
1269            "PYTHON_ERROR",
1270            &[("message", "Python worker failed")],
1271        );
1272        assert_eq!(err.kind, SparkErrorKind::Python);
1273    }
1274
1275    #[test]
1276    fn test_spark_runtime_exception_mapping() {
1277        let err = SparkError::classed(SparkErrorKind::SparkRuntime, "SPARK_RUNTIME_ERROR", &[]);
1278        assert_eq!(err.kind, SparkErrorKind::SparkRuntime);
1279    }
1280
1281    #[test]
1282    fn test_connect_grpc_exception_mapping() {
1283        let err = SparkError::classed(
1284            SparkErrorKind::ConnectGrpc,
1285            "RESPONSE_ALREADY_RECEIVED",
1286            &[],
1287        );
1288        assert_eq!(err.kind, SparkErrorKind::ConnectGrpc);
1289    }
1290
1291    #[test]
1292    fn test_invalid_plan_input_exception_mapping() {
1293        let err = SparkError::classed(
1294            SparkErrorKind::InvalidPlanInput,
1295            "INVALID_PLAN_INPUT",
1296            &[("message", "Invalid plan")],
1297        );
1298        assert_eq!(err.kind, SparkErrorKind::InvalidPlanInput);
1299    }
1300}
spark_connect_core/error.rs

spark_connect_core/
error.rs
