1use tokio::runtime::{Builder, Handle, Runtime};
9
10#[cfg(test)]
11use super::session::OrchestratorMssqlOutputWriter;
12#[cfg(test)]
13use crate::MssqlWorkflowOutputWriter;
14use crate::{
15 DeltaFunnelError, DeltaFunnelSession, LazyTable, MssqlDryRunOutputReport,
16 MssqlDryRunWorkflowReport, MssqlWriteReport, OutputWritePlan, WriteAllOptions, WriteAllReport,
17};
18
19pub struct DeltaFunnelRuntime {
30 runtime: Runtime,
31}
32
33impl DeltaFunnelRuntime {
34 pub fn new() -> Result<Self, DeltaFunnelError> {
40 let runtime = Builder::new_multi_thread()
41 .enable_all()
42 .thread_name("delta-funnel-runtime")
43 .build()
44 .map_err(|error| DeltaFunnelError::Config {
45 message: format!("failed to create DeltaFunnel runtime: {error}"),
46 })?;
47
48 Ok(Self { runtime })
49 }
50
51 pub fn table_from_sql(
57 &self,
58 session: &mut DeltaFunnelSession,
59 sql: &str,
60 ) -> Result<LazyTable, DeltaFunnelError> {
61 reject_nested_runtime()?;
62 self.runtime.block_on(session.table_from_sql(sql))
63 }
64
65 pub fn dry_run_to_mssql(
71 &self,
72 session: &DeltaFunnelSession,
73 request: &OutputWritePlan,
74 ) -> Result<MssqlDryRunOutputReport, DeltaFunnelError> {
75 reject_nested_runtime()?;
76 session.dry_run_to_mssql(request)
77 }
78
79 pub fn dry_run_all_to_mssql(
85 &self,
86 session: &DeltaFunnelSession,
87 requests: &[OutputWritePlan],
88 ) -> Result<MssqlDryRunWorkflowReport, DeltaFunnelError> {
89 reject_nested_runtime()?;
90 session.dry_run_all_to_mssql_with_tracing(requests)
91 }
92
93 pub fn dry_run_all_to_mssql_with_scan_summary(
100 &self,
101 session: &DeltaFunnelSession,
102 requests: &[OutputWritePlan],
103 ) -> Result<MssqlDryRunWorkflowReport, DeltaFunnelError> {
104 reject_nested_runtime()?;
105 self.runtime
106 .block_on(session.dry_run_all_to_mssql_with_scan_summary_with_tracing(requests))
107 }
108
109 pub fn write_to_mssql(
115 &self,
116 session: &DeltaFunnelSession,
117 request: &OutputWritePlan,
118 ) -> Result<MssqlWriteReport, DeltaFunnelError> {
119 reject_nested_runtime()?;
120 self.runtime.block_on(session.write_to_mssql(request))
121 }
122
123 #[cfg(test)]
124 pub(crate) fn write_to_mssql_with_writer<W>(
125 &self,
126 session: &DeltaFunnelSession,
127 request: &OutputWritePlan,
128 writer: &mut W,
129 ) -> Result<MssqlWriteReport, DeltaFunnelError>
130 where
131 W: OrchestratorMssqlOutputWriter,
132 {
133 reject_nested_runtime()?;
134 self.runtime
135 .block_on(session.write_to_mssql_with_writer(request, writer))
136 }
137
138 pub fn write_all(
144 &self,
145 session: &DeltaFunnelSession,
146 requests: &[OutputWritePlan],
147 ) -> Result<WriteAllReport, DeltaFunnelError> {
148 reject_nested_runtime()?;
149 self.runtime.block_on(session.write_all(requests))
150 }
151
152 pub fn write_all_with_options(
158 &self,
159 session: &DeltaFunnelSession,
160 requests: &[OutputWritePlan],
161 options: WriteAllOptions,
162 ) -> Result<WriteAllReport, DeltaFunnelError> {
163 reject_nested_runtime()?;
164 self.runtime
165 .block_on(session.write_all_with_options(requests, options))
166 }
167
168 #[cfg(test)]
169 pub(crate) fn write_all_with_writer<W>(
170 &self,
171 session: &DeltaFunnelSession,
172 requests: &[OutputWritePlan],
173 writer: W,
174 ) -> Result<WriteAllReport, DeltaFunnelError>
175 where
176 W: MssqlWorkflowOutputWriter,
177 {
178 reject_nested_runtime()?;
179 self.runtime
180 .block_on(session.write_all_with_writer(requests, writer))
181 }
182}
183
184fn reject_nested_runtime() -> Result<(), DeltaFunnelError> {
185 if Handle::try_current().is_ok() {
186 return Err(DeltaFunnelError::Config {
187 message: "DeltaFunnelRuntime blocking methods cannot be called from inside an active Tokio runtime; use DeltaFunnelSession async APIs directly".to_owned(),
188 });
189 }
190
191 Ok(())
192}
193
194#[cfg(test)]
195mod tests {
196 use super::*;
197 use async_trait::async_trait;
198 use datafusion::arrow::datatypes::SchemaRef;
199 use futures_util::StreamExt;
200
201 use crate::{
202 DeltaSourceConfig, DryRunScanSummaryMode, LoadMode, MssqlConnectionConfig,
203 MssqlConnectionSource, MssqlOutputBatchStream, MssqlOutputTarget, MssqlTargetCleanupStatus,
204 MssqlTargetConfig, MssqlTargetOutputPlan, MssqlTargetTable, MssqlWriteOptions,
205 ResolvedMssqlTarget, RunMode, SessionOptions, ValidationOptions,
206 table_formats::RealParquetDeltaTable,
207 };
208
209 fn secret_connection() -> Result<MssqlConnectionConfig, DeltaFunnelError> {
210 Ok(MssqlConnectionConfig::new(
211 "server=tcp:sql.example.com;database=warehouse;user=admin;password=secret-token",
212 )?
213 .with_display_label("warehouse-primary"))
214 }
215
216 fn output_request(
217 table: LazyTable,
218 output_name: &str,
219 target_table: &str,
220 ) -> Result<OutputWritePlan, DeltaFunnelError> {
221 let target = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", target_table)?)
222 .with_load_mode(LoadMode::AppendExisting);
223
224 Ok(OutputWritePlan::new(
225 table,
226 MssqlOutputTarget::new(output_name, target, RunMode::DryRun),
227 ))
228 }
229
230 fn execute_output_request(
231 table: LazyTable,
232 output_name: &str,
233 target_table: &str,
234 load_mode: LoadMode,
235 ) -> Result<OutputWritePlan, DeltaFunnelError> {
236 let target = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", target_table)?)
237 .with_load_mode(load_mode);
238
239 Ok(OutputWritePlan::new(
240 table,
241 MssqlOutputTarget::new(output_name, target, RunMode::Execute),
242 ))
243 }
244
245 #[derive(Debug, Clone, PartialEq, Eq)]
246 struct FakeRuntimeWriteCall {
247 output_name: String,
248 target_table: MssqlTargetTable,
249 connection_source: MssqlConnectionSource,
250 rows: u64,
251 batches: u64,
252 schema_fields: usize,
253 }
254
255 #[derive(Default)]
256 struct FakeRuntimeWriter {
257 calls: Vec<FakeRuntimeWriteCall>,
258 }
259
260 #[async_trait]
261 impl OrchestratorMssqlOutputWriter for FakeRuntimeWriter {
262 async fn write_output(
263 &mut self,
264 output_schema: SchemaRef,
265 output_plan: MssqlTargetOutputPlan,
266 resolved_target: ResolvedMssqlTarget,
267 mut batches: MssqlOutputBatchStream,
268 _write_options: MssqlWriteOptions,
269 _validation_options: ValidationOptions,
270 ) -> Result<MssqlWriteReport, DeltaFunnelError> {
271 let mut rows = 0_u64;
272 let mut batch_count = 0_u64;
273
274 while let Some(batch) = batches.next().await {
275 let batch = batch?;
276 rows = rows.saturating_add(u64::try_from(batch.num_rows()).map_err(|_| {
277 DeltaFunnelError::Config {
278 message: "fake runtime writer row count overflowed u64".to_owned(),
279 }
280 })?);
281 batch_count = batch_count.saturating_add(1);
282 }
283
284 self.calls.push(FakeRuntimeWriteCall {
285 output_name: resolved_target.output_name().to_owned(),
286 target_table: resolved_target.table().clone(),
287 connection_source: resolved_target.connection_source(),
288 rows,
289 batches: batch_count,
290 schema_fields: output_schema.fields().len(),
291 });
292
293 Ok(MssqlWriteReport::from_output_plan(
294 &output_plan,
295 rows,
296 batch_count,
297 0,
298 false,
299 MssqlTargetCleanupStatus::NotApplicable,
300 ))
301 }
302 }
303
304 #[async_trait]
305 impl MssqlWorkflowOutputWriter for FakeRuntimeWriter {
306 async fn write_output(
307 &mut self,
308 output_schema: SchemaRef,
309 resolved_target: ResolvedMssqlTarget,
310 schema_options: crate::MssqlSchemaPlanOptions,
311 batches: MssqlOutputBatchStream,
312 write_options: MssqlWriteOptions,
313 validation_options: ValidationOptions,
314 ) -> Result<MssqlWriteReport, DeltaFunnelError> {
315 let output_plan = crate::plan_mssql_target_for_resolved_output(
316 output_schema.as_ref(),
317 &resolved_target,
318 schema_options,
319 )?;
320
321 OrchestratorMssqlOutputWriter::write_output(
322 self,
323 output_schema,
324 output_plan,
325 resolved_target,
326 batches,
327 write_options,
328 validation_options,
329 )
330 .await
331 }
332 }
333
334 #[test]
335 fn runtime_constructs_without_starting_session_work() -> Result<(), DeltaFunnelError> {
336 let _runtime = DeltaFunnelRuntime::new()?;
337
338 Ok(())
339 }
340
341 #[test]
342 fn runtime_drives_table_sql_and_single_output_dry_run() -> Result<(), DeltaFunnelError> {
343 let runtime = DeltaFunnelRuntime::new()?;
344 let mut session = DeltaFunnelSession::new(
345 SessionOptions::new().with_default_mssql_connection(secret_connection()?),
346 )?;
347 let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
348 let request = output_request(output, "orders_output", "orders_sink")?;
349
350 let report = runtime.dry_run_to_mssql(&session, &request)?;
351
352 assert_eq!(report.output_name(), "orders_output");
353 assert_eq!(report.run_mode(), RunMode::DryRun);
354 assert!(!report.sql_server_contacted());
355 assert!(!report.row_production_started());
356 Ok(())
357 }
358
359 #[test]
360 fn runtime_drives_multi_output_dry_run() -> Result<(), DeltaFunnelError> {
361 let runtime = DeltaFunnelRuntime::new()?;
362 let mut session = DeltaFunnelSession::new(
363 SessionOptions::new().with_default_mssql_connection(secret_connection()?),
364 )?;
365 let west = runtime.table_from_sql(&mut session, "select 1 as id")?;
366 let east = runtime.table_from_sql(&mut session, "select 2 as id")?;
367 let west = output_request(west, "west_output", "west_orders")?;
368 let east = output_request(east, "east_output", "east_orders")?;
369
370 let report = runtime.dry_run_all_to_mssql(&session, &[west, east])?;
371
372 assert_eq!(report.len(), 2);
373 assert_eq!(report.outputs()[0].output_name(), "west_output");
374 assert_eq!(report.outputs()[1].output_name(), "east_output");
375 assert!(!report.sql_server_contacted());
376 assert!(!report.row_production_started());
377 Ok(())
378 }
379
380 #[test]
381 fn runtime_drives_dry_run_scan_summary() -> Result<(), Box<dyn std::error::Error>> {
382 let runtime = DeltaFunnelRuntime::new()?;
383 let table = RealParquetDeltaTable::new_default("orders")?;
384 let mut session = DeltaFunnelSession::new(
385 SessionOptions::new()
386 .with_default_mssql_connection(secret_connection()?)
387 .with_validation_options(
388 ValidationOptions::new()
389 .with_dry_run_scan_summary_mode(DryRunScanSummaryMode::ExhaustScanMetadata),
390 ),
391 )?;
392 let source = session.delta_lake(DeltaSourceConfig::new(
393 "orders",
394 table.path().to_string_lossy().to_string(),
395 ))?;
396 let request = output_request(source, "orders_output", "orders_sink")?;
397
398 let report = runtime.dry_run_all_to_mssql_with_scan_summary(&session, &[request])?;
399
400 assert_eq!(report.sources().len(), 1);
401 assert!(report.sources()[0].provider_read_stats().is_some());
402 assert!(!report.row_production_started());
403 Ok(())
404 }
405
406 #[test]
407 fn runtime_preserves_sanitized_session_errors() -> Result<(), DeltaFunnelError> {
408 let runtime = DeltaFunnelRuntime::new()?;
409 let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
410 let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
411 let request = output_request(output, "orders_output", "orders_sink")?;
412
413 let error = runtime.dry_run_to_mssql(&session, &request);
414
415 assert!(matches!(
416 error,
417 Err(DeltaFunnelError::MissingMssqlConnection { output_name })
418 if output_name == "orders_output"
419 ));
420 Ok(())
421 }
422
423 #[test]
424 fn runtime_drives_single_output_write_with_injected_writer()
425 -> Result<(), Box<dyn std::error::Error>> {
426 let runtime = DeltaFunnelRuntime::new()?;
427 let mut session = DeltaFunnelSession::new(
428 SessionOptions::new().with_default_mssql_connection(secret_connection()?),
429 )?;
430 let output =
431 runtime.table_from_sql(&mut session, "select 1 as id union all select 2 as id")?;
432 let request = execute_output_request(
433 output,
434 "orders_output",
435 "orders_sink",
436 LoadMode::CreateAndLoad,
437 )?;
438 let mut writer = FakeRuntimeWriter::default();
439
440 let report = runtime.write_to_mssql_with_writer(&session, &request, &mut writer)?;
441
442 assert_eq!(writer.calls.len(), 1);
443 let call = writer
444 .calls
445 .first()
446 .ok_or("expected fake runtime writer call")?;
447 assert_eq!(call.output_name, "orders_output");
448 assert_eq!(call.target_table.table(), "orders_sink");
449 assert_eq!(
450 call.connection_source,
451 MssqlConnectionSource::ContextDefault
452 );
453 assert_eq!(call.rows, 2);
454 assert!(call.batches >= 1);
455 assert_eq!(call.schema_fields, 1);
456 assert_eq!(report.output_name(), "orders_output");
457 assert_eq!(report.stats().rows_written(), 2);
458 assert_eq!(report.stats().batches_written(), call.batches);
459 Ok(())
460 }
461
462 #[test]
463 fn runtime_drives_multi_output_write_with_injected_writer()
464 -> Result<(), Box<dyn std::error::Error>> {
465 let runtime = DeltaFunnelRuntime::new()?;
466 let mut session = DeltaFunnelSession::new(
467 SessionOptions::new().with_default_mssql_connection(secret_connection()?),
468 )?;
469 let west =
470 runtime.table_from_sql(&mut session, "select 1 as id union all select 2 as id")?;
471 let east = runtime.table_from_sql(&mut session, "select 3 as id")?;
472 let west =
473 execute_output_request(west, "west_output", "west_orders", LoadMode::AppendExisting)?;
474 let east =
475 execute_output_request(east, "east_output", "east_orders", LoadMode::CreateAndLoad)?;
476 let writer = FakeRuntimeWriter::default();
477
478 let report = runtime.write_all_with_writer(&session, &[west, east], writer)?;
479
480 assert_eq!(report.len(), 2);
481 assert!(report.all_succeeded());
482 assert_eq!(report.outputs()[0].output_name(), "west_output");
483 assert_eq!(report.outputs()[1].output_name(), "east_output");
484 let crate::MssqlOutputWriteStatus::Succeeded(west_report) = &report.outputs()[0] else {
485 return Err(format!("expected succeeded status, got {:?}", report.outputs()[0]).into());
486 };
487 let crate::MssqlOutputWriteStatus::Succeeded(east_report) = &report.outputs()[1] else {
488 return Err(format!("expected succeeded status, got {:?}", report.outputs()[1]).into());
489 };
490 assert_eq!(west_report.stats().rows_written(), 2);
491 assert_eq!(east_report.stats().rows_written(), 1);
492 Ok(())
493 }
494
495 #[test]
496 fn runtime_rejects_blocking_calls_inside_active_tokio_runtime()
497 -> Result<(), Box<dyn std::error::Error>> {
498 let runtime = DeltaFunnelRuntime::new()?;
499 let mut session = DeltaFunnelSession::new(
500 SessionOptions::new().with_default_mssql_connection(secret_connection()?),
501 )?;
502 let output = runtime.table_from_sql(&mut session, "select 1 as id")?;
503 let request = output_request(output, "orders_output", "orders_sink")?;
504 let host_runtime = tokio::runtime::Builder::new_current_thread()
505 .enable_all()
506 .build()?;
507
508 let error = host_runtime.block_on(async { runtime.dry_run_to_mssql(&session, &request) });
509
510 assert!(matches!(
511 error,
512 Err(DeltaFunnelError::Config { message })
513 if message.contains("cannot be called from inside an active Tokio runtime")
514 && message.contains("DeltaFunnelSession async APIs")
515 ));
516 Ok(())
517 }
518
519 #[test]
520 fn lower_level_modules_do_not_create_hidden_tokio_runtimes() {
521 let low_level_sources = [
522 include_str!("../pipeline/batch_handoff.rs"),
523 include_str!("../query_engine/datafusion/execution.rs"),
524 include_str!("../query_engine/datafusion/execution/planning_exec.rs"),
525 include_str!("../sql_server/execution/connection.rs"),
526 include_str!("../sql_server/execution/sink.rs"),
527 include_str!("../sql_server/execution/workflow.rs"),
528 include_str!("../sql_server/execution/write.rs"),
529 ];
530
531 for source in low_level_sources {
532 assert!(!source.contains("tokio::runtime"));
533 assert!(!source.contains("Runtime::new"));
534 assert!(!source.contains("block_on"));
535 }
536 }
537}