Skip to content

[SPARK-52428] Harden streaming and session error handling #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rafafrdz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[SPARK-52428] Harden streaming and session error handling #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions crates/connect/src/session.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,8 @@ impl Default for SparkSessionBuilder {

impl SparkSessionBuilder {
fn new(connection: &str) -> Self {
let channel_builder = ChannelBuilder::create(connection).unwrap();
let channel_builder = ChannelBuilder::create(connection)
.expect("Failed to parse the Spark Connect connection string");

Self {
channel_builder,
Expand Down Expand Up @@ -207,9 +208,17 @@ impl SparkSession {
.execute_command_and_fetch(plan)
.await?;

let relation = resp.sql_command_result.to_owned().unwrap().relation;

let logical_plan = LogicalPlanBuilder::new(relation.unwrap());
let relation = resp
.sql_command_result
.to_owned()
.ok_or(SparkError::AnalysisException(
"SQL command result is empty".to_string(),
))?
.relation;

let logical_plan = LogicalPlanBuilder::new(relation.ok_or(
SparkError::AnalysisException("SQL relation result is empty".to_string()),
)?);

Ok(DataFrame::new(self.session(), logical_plan))
}
Expand Down
36 changes: 25 additions & 11 deletions crates/connect/src/streaming/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,14 @@ impl DataStreamWriter {
format: self.format.unwrap_or("".to_string()),
options: self.write_options,
partitioning_column_names: self.partition_by,
output_mode: self.output_mode.unwrap().as_str_name().to_string(),
output_mode: self
.output_mode
.as_ref()
.ok_or(SparkError::AnalysisException(
"output_mode is not set".to_string(),
))?
.as_str_name()
.to_string(),
query_name: self.query_name.unwrap_or("".to_string()),
foreach_batch: None,
foreach_writer: None,
Expand All @@ -252,10 +259,11 @@ impl DataStreamWriter {
.await?
.write_stream_operation_start_result;

Ok(StreamingQuery::new(
self.dataframe.spark_session,
operation_start_resp.unwrap(),
))
let operation_start_resp = operation_start_resp.ok_or(SparkError::AnalysisException(
"write stream operation start result is empty".to_string(),
))?;

StreamingQuery::new(self.dataframe.spark_session, operation_start_resp)
}

/// Start a streaming job to save the contents of the [StreamingQuery] to a data source.
Expand Down Expand Up @@ -293,14 +301,16 @@ impl StreamingQuery {
pub fn new(
spark_session: Box<SparkSession>,
write_stream: spark::WriteStreamOperationStartResult,
) -> Self {
let query_instance = write_stream.query_id.unwrap();
) -> Result<Self, SparkError> {
let query_instance = write_stream.query_id.ok_or(SparkError::AnalysisException(
"streaming query id is empty".to_string(),
))?;

Self {
Ok(Self {
spark_session,
query_instance,
name: Some(write_stream.name),
}
})
}

fn streaming_query_cmd() -> spark::StreamingQueryCommand {
Expand Down Expand Up @@ -565,7 +575,9 @@ impl StreamingQuery {
}

fn to_json_object(val: Vec<String>) -> Result<serde_json::Value, SparkError> {
let val = &val.first().unwrap();
let val = val.first().ok_or(SparkError::AnalysisException(
"empty progress response".to_string(),
))?;
Ok(serde_json::from_str::<serde_json::Value>(val)?)
}

Expand Down Expand Up @@ -636,7 +648,9 @@ impl StreamingQueryManager {
for stream in active_result.active_queries {
let query = StreamingQuery {
spark_session: self.spark_session.clone(),
query_instance: stream.id.clone().unwrap(),
query_instance: stream.id.clone().ok_or(SparkError::AnalysisException(
"streaming query instance id is empty".to_string(),
))?,
name: stream.name,
};

Expand Down