use tokio::runtime::{Builder, Handle, Runtime};
#[cfg(test)]
use super::session::OrchestratorMssqlOutputWriter;
#[cfg(test)]
use crate::MssqlWorkflowOutputWriter;
use crate::{
DeltaFunnelError, DeltaFunnelSession, LazyTable, MssqlDryRunOutputReport,
MssqlDryRunWorkflowReport, MssqlWriteReport, OutputWritePlan, WriteAllOptions, WriteAllReport,
};
pub struct DeltaFunnelRuntime {
runtime: Runtime,
}
impl DeltaFunnelRuntime {
pub fn new() -> Result<Self, DeltaFunnelError> {
let runtime = Builder::new_multi_thread()
.enable_all()
.thread_name("delta-funnel-runtime")
.build()
.map_err(|error| DeltaFunnelError::Config {
message: format!("failed to create DeltaFunnel runtime: {error}"),
})?;
Ok(Self { runtime })
}
pub fn table_from_sql(
&self,
session: &mut DeltaFunnelSession,
sql: &str,
) -> Result<LazyTable, DeltaFunnelError> {
reject_nested_runtime()?;
self.runtime.block_on(session.table_from_sql(sql))
}
pub fn dry_run_to_mssql(
&self,
session: &DeltaFunnelSession,
request: &OutputWritePlan,
) -> Result<MssqlDryRunOutputReport, DeltaFunnelError> {
reject_nested_runtime()?;
session.dry_run_to_mssql(request)
}
pub fn dry_run_all_to_mssql(
&self,
session: &DeltaFunnelSession,
requests: &[OutputWritePlan],
) -> Result<MssqlDryRunWorkflowReport, DeltaFunnelError> {
reject_nested_runtime()?;
session.dry_run_all_to_mssql_with_tracing(requests)
}
pub fn dry_run_all_to_mssql_with_scan_summary(
&self,
session: &DeltaFunnelSession,
requests: &[OutputWritePlan],
) -> Result<MssqlDryRunWorkflowReport, DeltaFunnelError> {
reject_nested_runtime()?;
self.runtime
.block_on(session.dry_run_all_to_mssql_with_scan_summary_with_tracing(requests))
}
pub fn write_to_mssql(
&self,
session: &DeltaFunnelSession,
request: &OutputWritePlan,
) -> Result<MssqlWriteReport, DeltaFunnelError> {
reject_nested_runtime()?;
self.runtime.block_on(session.write_to_mssql(request))
}
#[cfg(test)]
pub(crate) fn write_to_mssql_with_writer<W>(
&self,
session: &DeltaFunnelSession,
request: &OutputWritePlan,
writer: &mut W,
) -> Result<MssqlWriteReport, DeltaFunnelError>
where
W: OrchestratorMssqlOutputWriter,
{
reject_nested_runtime()?;
self.runtime
.block_on(session.write_to_mssql_with_writer(request, writer))
}
pub fn write_all(
&self,
session: &DeltaFunnelSession,
requests: &[OutputWritePlan],
) -> Result<WriteAllReport, DeltaFunnelError> {
reject_nested_runtime()?;
self.runtime.block_on(session.write_all(requests))
}
pub fn write_all_with_options(
&self,
session: &DeltaFunnelSession,
requests: &[OutputWritePlan],
options: WriteAllOptions,
) -> Result<WriteAllReport, DeltaFunnelError> {
reject_nested_runtime()?;
self.runtime
.block_on(session.write_all_with_options(requests, options))
}
#[cfg(test)]
pub(crate) fn write_all_with_writer<W>(
&self,
session: &DeltaFunnelSession,
requests: &[OutputWritePlan],
writer: W,
) -> Result<WriteAllReport, DeltaFunnelError>
where
W: MssqlWorkflowOutputWriter,
{
reject_nested_runtime()?;
self.runtime
.block_on(session.write_all_with_writer(requests, writer))
}
}
fn reject_nested_runtime() -> Result<(), DeltaFunnelError> {
if Handle::try_current().is_ok() {
return Err(DeltaFunnelError::Config {
message: "DeltaFunnelRuntime blocking methods cannot be called from inside an active Tokio runtime; use DeltaFunnelSession async APIs directly".to_owned(),
});
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use datafusion::arrow::datatypes::SchemaRef;
use futures_util::StreamExt;
use crate::{
DeltaSourceConfig, DryRunScanSummaryMode, LoadMode, MssqlConnectionConfig,
MssqlConnectionSource, MssqlOutputBatchStream, MssqlOutputTarget, MssqlTargetCleanupStatus,
MssqlTargetConfig, MssqlTargetOutputPlan, MssqlTargetTable, MssqlWriteOptions,
ResolvedMssqlTarget, RunMode, SessionOptions, ValidationOptions,
table_formats::RealParquetDeltaTable,
};
fn secret_connection() -> Result<MssqlConnectionConfig, DeltaFunnelError> {
Ok(MssqlConnectionConfig::new(
"server=tcp:sql.example.com;database=warehouse;user=admin;password=secret-token",
)?
.with_display_label("warehouse-primary"))
}
fn output_request(
table: LazyTable,
output_name: &str,
target_table: &str,
) -> Result<OutputWritePlan, DeltaFunnelError> {
let target = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", target_table)?)
.with_load_mode(LoadMode::AppendExisting);
Ok(OutputWritePlan::new(
table,
MssqlOutputTarget::new(output_name, target, RunMode::DryRun),
))
}
fn execute_output_request(
table: LazyTable,
output_name: &str,
target_table: &str,
load_mode: LoadMode,
) -> Result<OutputWritePlan, DeltaFunnelError> {
let target = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", target_table)?)
.with_load_mode(load_mode);
Ok(OutputWritePlan::new(
table,
MssqlOutputTarget::new(output_name, target, RunMode::Execute),
))
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct FakeRuntimeWriteCall {
output_name: String,
target_table: MssqlTargetTable,
connection_source: MssqlConnectionSource,
rows: u64,
batches: u64,
schema_fields: usize,
}
#[derive(Default)]
struct FakeRuntimeWriter {
calls: Vec<FakeRuntimeWriteCall>,
}
#[async_trait]
impl OrchestratorMssqlOutputWriter for FakeRuntimeWriter {
async fn write_output(
&mut self,
output_schema: SchemaRef,
output_plan: MssqlTargetOutputPlan,
resolved_target: ResolvedMssqlTarget,
mut batches: MssqlOutputBatchStream,
_write_options: MssqlWriteOptions,
_validation_options: ValidationOptions,
) -> Result<MssqlWriteReport, DeltaFunnelError> {
let mut rows = 0_u64;
let mut batch_count = 0_u64;
while let Some(batch) = batches.next().await {
let batch = batch?;
rows = rows.saturating_add(u64::try_from(batch.num_rows()).map_err(|_| {
DeltaFunnelError::Config {
message: "fake runtime writer row count overflowed u64".to_owned(),
}
})?);
batch_count = batch_count.saturating_add(1);
}
self.calls.push(FakeRuntimeWriteCall {
output_name: resolved_target.output_name().to_owned(),
target_table: resolved_target.table().clone(),
connection_source: resolved_target.connection_source(),
rows,
batches: batch_count,
schema_fields: output_schema.fields().len(),
});
Ok(MssqlWriteReport::from_output_plan(
&output_plan,
rows,
batch_count,
0,
false,
MssqlTargetCleanupStatus::NotApplicable,
))
}
}
#[async_trait]
impl MssqlWorkflowOutputWriter for FakeRuntimeWriter {
async fn write_output(
&mut self,
output_schema: SchemaRef,
resolved_target: ResolvedMssqlTarget,
schema_options: crate::MssqlSchemaPlanOptions,
batches: MssqlOutputBatchStream,
write_options: MssqlWriteOptions,
validation_options: ValidationOptions,
) -> Result<MssqlWriteReport, DeltaFunnelError> {
let output_plan = crate::plan_mssql_target_for_resolved_output(
output_schema.as_ref(),
&resolved_target,
schema_options,
)?;
OrchestratorMssqlOutputWriter::write_output(
self,
output_schema,
output_plan,
resolved_target,
batches,
write_options,
validation_options,
)
.await
}
}
#[test]
fn runtime_constructs_without_starting_session_work() -> Result<(), DeltaFunnelError> {
let _runtime = DeltaFunnelRuntime::new()?;
Ok(())
}
#[test]
fn runtime_drives_table_sql_and_single_output_dry_run() -> Result<(), DeltaFunnelError> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new().with_default_mssql_connection(secret_connection()?),
)?;
let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
let request = output_request(output, "orders_output", "orders_sink")?;
let report = runtime.dry_run_to_mssql(&session, &request)?;
assert_eq!(report.output_name(), "orders_output");
assert_eq!(report.run_mode(), RunMode::DryRun);
assert!(!report.sql_server_contacted());
assert!(!report.row_production_started());
Ok(())
}
#[test]
fn runtime_drives_multi_output_dry_run() -> Result<(), DeltaFunnelError> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new().with_default_mssql_connection(secret_connection()?),
)?;
let west = runtime.table_from_sql(&mut session, "select 1 as id")?;
let east = runtime.table_from_sql(&mut session, "select 2 as id")?;
let west = output_request(west, "west_output", "west_orders")?;
let east = output_request(east, "east_output", "east_orders")?;
let report = runtime.dry_run_all_to_mssql(&session, &[west, east])?;
assert_eq!(report.len(), 2);
assert_eq!(report.outputs()[0].output_name(), "west_output");
assert_eq!(report.outputs()[1].output_name(), "east_output");
assert!(!report.sql_server_contacted());
assert!(!report.row_production_started());
Ok(())
}
#[test]
fn runtime_drives_dry_run_scan_summary() -> Result<(), Box<dyn std::error::Error>> {
let runtime = DeltaFunnelRuntime::new()?;
let table = RealParquetDeltaTable::new_default("orders")?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new()
.with_default_mssql_connection(secret_connection()?)
.with_validation_options(
ValidationOptions::new()
.with_dry_run_scan_summary_mode(DryRunScanSummaryMode::ExhaustScanMetadata),
),
)?;
let source = session.delta_lake(DeltaSourceConfig::new(
"orders",
table.path().to_string_lossy().to_string(),
))?;
let request = output_request(source, "orders_output", "orders_sink")?;
let report = runtime.dry_run_all_to_mssql_with_scan_summary(&session, &[request])?;
assert_eq!(report.sources().len(), 1);
assert!(report.sources()[0].provider_read_stats().is_some());
assert!(!report.row_production_started());
Ok(())
}
#[test]
fn runtime_preserves_sanitized_session_errors() -> Result<(), DeltaFunnelError> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
let request = output_request(output, "orders_output", "orders_sink")?;
let error = runtime.dry_run_to_mssql(&session, &request);
assert!(matches!(
error,
Err(DeltaFunnelError::MissingMssqlConnection { output_name })
if output_name == "orders_output"
));
Ok(())
}
#[test]
fn runtime_drives_single_output_write_with_injected_writer()
-> Result<(), Box<dyn std::error::Error>> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new().with_default_mssql_connection(secret_connection()?),
)?;
let output =
runtime.table_from_sql(&mut session, "select 1 as id union all select 2 as id")?;
let request = execute_output_request(
output,
"orders_output",
"orders_sink",
LoadMode::CreateAndLoad,
)?;
let mut writer = FakeRuntimeWriter::default();
let report = runtime.write_to_mssql_with_writer(&session, &request, &mut writer)?;
assert_eq!(writer.calls.len(), 1);
let call = writer
.calls
.first()
.ok_or("expected fake runtime writer call")?;
assert_eq!(call.output_name, "orders_output");
assert_eq!(call.target_table.table(), "orders_sink");
assert_eq!(
call.connection_source,
MssqlConnectionSource::ContextDefault
);
assert_eq!(call.rows, 2);
assert!(call.batches >= 1);
assert_eq!(call.schema_fields, 1);
assert_eq!(report.output_name(), "orders_output");
assert_eq!(report.stats().rows_written(), 2);
assert_eq!(report.stats().batches_written(), call.batches);
Ok(())
}
#[test]
fn runtime_drives_multi_output_write_with_injected_writer()
-> Result<(), Box<dyn std::error::Error>> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new().with_default_mssql_connection(secret_connection()?),
)?;
let west =
runtime.table_from_sql(&mut session, "select 1 as id union all select 2 as id")?;
let east = runtime.table_from_sql(&mut session, "select 3 as id")?;
let west =
execute_output_request(west, "west_output", "west_orders", LoadMode::AppendExisting)?;
let east =
execute_output_request(east, "east_output", "east_orders", LoadMode::CreateAndLoad)?;
let writer = FakeRuntimeWriter::default();
let report = runtime.write_all_with_writer(&session, &[west, east], writer)?;
assert_eq!(report.len(), 2);
assert!(report.all_succeeded());
assert_eq!(report.outputs()[0].output_name(), "west_output");
assert_eq!(report.outputs()[1].output_name(), "east_output");
let crate::MssqlOutputWriteStatus::Succeeded(west_report) = &report.outputs()[0] else {
return Err(format!("expected succeeded status, got {:?}", report.outputs()[0]).into());
};
let crate::MssqlOutputWriteStatus::Succeeded(east_report) = &report.outputs()[1] else {
return Err(format!("expected succeeded status, got {:?}", report.outputs()[1]).into());
};
assert_eq!(west_report.stats().rows_written(), 2);
assert_eq!(east_report.stats().rows_written(), 1);
Ok(())
}
#[test]
fn runtime_rejects_blocking_calls_inside_active_tokio_runtime()
-> Result<(), Box<dyn std::error::Error>> {
let runtime = DeltaFunnelRuntime::new()?;
let mut session = DeltaFunnelSession::new(
SessionOptions::new().with_default_mssql_connection(secret_connection()?),
)?;
let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
let request = output_request(output, "orders_output", "orders_sink")?;
let host_runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let error = host_runtime.block_on(async { runtime.dry_run_to_mssql(&session, &request) });
assert!(matches!(
error,
Err(DeltaFunnelError::Config { message })
if message.contains("cannot be called from inside an active Tokio runtime")
&& message.contains("DeltaFunnelSession async APIs")
));
Ok(())
}
#[test]
fn lower_level_modules_do_not_create_hidden_tokio_runtimes() {
let low_level_sources = [
include_str!("../pipeline/batch_handoff.rs"),
include_str!("../query_engine/datafusion/execution.rs"),
include_str!("../query_engine/datafusion/execution/planning_exec.rs"),
include_str!("../sql_server/execution/connection.rs"),
include_str!("../sql_server/execution/sink.rs"),
include_str!("../sql_server/execution/workflow.rs"),
include_str!("../sql_server/execution/write.rs"),
];
for source in low_level_sources {
assert!(!source.contains("tokio::runtime"));
assert!(!source.contains("Runtime::new"));
assert!(!source.contains("block_on"));
}
}
}