delta_funnel/orchestrator/session/registry/
source.rs

1use std::fmt;
2
3use datafusion::arrow::datatypes::SchemaRef;
4
5use crate::{
6    DeltaFunnelError, DeltaProtocolReport, DeltaSourceConfig, DeltaTableProviderConfig,
7    PhaseTimingReport, RegisteredDeltaSource, register_delta_sources_with_scan_execution_options,
8    report::PhaseTimer,
9    table_formats::{load_delta_source_with_tracing, preflight_delta_protocol_with_tracing},
10};
11
12use super::super::{DeltaFunnelSession, LazyTable};
13
14const SOURCE_LOADING_PHASE: &str = "source_loading";
15const PROTOCOL_PREFLIGHT_PHASE: &str = "protocol_preflight";
16const DATAFUSION_REGISTRATION_PHASE: &str = "datafusion_registration";
17
18/// Registered Delta source tracked by a query-load session.
19#[derive(Clone, PartialEq, Eq)]
20pub struct RegisteredSessionSource {
21    table: LazyTable,
22    source_uri: String,
23    snapshot_version: u64,
24    schema: SchemaRef,
25    protocol: DeltaProtocolReport,
26    phase_timings: Vec<PhaseTimingReport>,
27}
28
29impl RegisteredSessionSource {
30    pub(super) fn from_registered(
31        table: LazyTable,
32        registered: RegisteredDeltaSource,
33        phase_timings: Vec<PhaseTimingReport>,
34    ) -> Self {
35        Self {
36            table,
37            source_uri: registered.table_uri,
38            snapshot_version: registered.snapshot_version,
39            schema: registered.schema,
40            protocol: registered.protocol,
41            phase_timings,
42        }
43    }
44
45    /// Returns the lazy table handle for this registered source.
46    #[must_use]
47    pub const fn table(&self) -> &LazyTable {
48        &self.table
49    }
50
51    /// Returns the DataFusion table name for this source.
52    #[must_use]
53    pub fn name(&self) -> &str {
54        self.table.name()
55    }
56
57    /// Returns the sanitized Delta source URI or display summary.
58    #[must_use]
59    pub fn source_uri(&self) -> &str {
60        &self.source_uri
61    }
62
63    /// Returns the resolved Delta snapshot version.
64    #[must_use]
65    pub const fn snapshot_version(&self) -> u64 {
66        self.snapshot_version
67    }
68
69    /// Returns the logical Arrow schema exposed to DataFusion.
70    #[must_use]
71    pub fn schema(&self) -> &SchemaRef {
72        &self.schema
73    }
74
75    /// Returns the sanitized protocol report captured before registration.
76    #[must_use]
77    pub const fn protocol(&self) -> &DeltaProtocolReport {
78        &self.protocol
79    }
80
81    /// Returns durable phase timings captured while registering this source.
82    #[must_use]
83    pub fn phase_timings(&self) -> &[PhaseTimingReport] {
84        &self.phase_timings
85    }
86}
87
88impl fmt::Debug for RegisteredSessionSource {
89    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
90        formatter
91            .debug_struct("RegisteredSessionSource")
92            .field("table", &self.table)
93            .field("source_uri", &self.source_uri)
94            .field("snapshot_version", &self.snapshot_version)
95            .field("schema", &self.schema)
96            .field("protocol", &self.protocol)
97            .field("phase_timings", &self.phase_timings)
98            .finish()
99    }
100}
101
102impl DeltaFunnelSession {
103    /// Registers one Delta source and returns its lazy table handle.
104    ///
105    /// The method performs source setup only: Delta snapshot metadata loading,
106    /// protocol preflight, and DataFusion table registration. It does not scan
107    /// data files for row production, parse user SQL, contact SQL Server, or
108    /// execute an output action.
109    ///
110    /// # Errors
111    ///
112    /// Returns the first Delta source loading, protocol preflight, duplicate
113    /// alias, schema conversion, or DataFusion registration error. Session
114    /// source state is updated only after the DataFusion registration succeeds.
115    pub fn delta_lake(&mut self, source: DeltaSourceConfig) -> Result<LazyTable, DeltaFunnelError> {
116        self.reject_registered_alias_name(&source.name)?;
117        let mut phase_timings = Vec::new();
118
119        let source_timer = PhaseTimer::start(SOURCE_LOADING_PHASE);
120        let planned = match load_delta_source_with_tracing(source) {
121            Ok(planned) => {
122                phase_timings.push(source_timer.completed());
123                planned
124            }
125            Err(error) => {
126                phase_timings.push(source_timer.failed());
127                return Err(error);
128            }
129        };
130
131        let preflight_timer = PhaseTimer::start(PROTOCOL_PREFLIGHT_PHASE);
132        let preflight = match preflight_delta_protocol_with_tracing(&planned) {
133            Ok(preflight) => {
134                phase_timings.push(preflight_timer.completed());
135                preflight
136            }
137            Err(error) => {
138                phase_timings.push(preflight_timer.failed());
139                return Err(error);
140            }
141        };
142
143        let registration_timer = PhaseTimer::start(DATAFUSION_REGISTRATION_PHASE);
144        let configs = vec![DeltaTableProviderConfig {
145            source: planned,
146            protocol: preflight,
147            scan_target_partitions: None,
148        }];
149        let registered = match register_delta_sources_with_scan_execution_options(
150            &self.context,
151            configs,
152            self.options.provider_scan_options(),
153        ) {
154            Ok(registered) => {
155                phase_timings.push(registration_timer.completed());
156                registered
157            }
158            Err(error) => {
159                phase_timings.push(registration_timer.failed());
160                return Err(error);
161            }
162        };
163        let registered =
164            registered
165                .sources
166                .into_iter()
167                .next()
168                .ok_or_else(|| DeltaFunnelError::Config {
169                    message: "Delta source registration returned no registered source".to_owned(),
170                })?;
171        let table = self.allocate_delta_source_table(registered.name.clone());
172        let session_source =
173            RegisteredSessionSource::from_registered(table.clone(), registered, phase_timings);
174        self.sources.push(session_source);
175        Ok(table)
176    }
177
178    fn allocate_delta_source_table(&mut self, name: String) -> LazyTable {
179        let id = self.next_table_id;
180        self.next_table_id = self.next_table_id.saturating_add(1);
181        LazyTable::delta_source(id, name)
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::{DATAFUSION_REGISTRATION_PHASE, PROTOCOL_PREFLIGHT_PHASE, SOURCE_LOADING_PHASE};
188
189    use crate::{
190        DeltaFunnelError, DeltaProviderReaderBackend, DeltaProviderScanExecutionOptions,
191        DeltaSourceConfig, DeltaStorageOptions, QueryOptions,
192    };
193
194    use super::super::super::{
195        DeltaFunnelSession, LazyTableKind, SessionOptions, SourceUsageStatus,
196        test_support::DeltaLogTable,
197    };
198
199    const UNSUPPORTED_PROTOCOL_JSON: &str =
200        r#"{"protocol":{"minReaderVersion":99,"minWriterVersion":2}}"#;
201
202    #[test]
203    fn delta_lake_registers_source_and_returns_lazy_table() -> Result<(), Box<dyn std::error::Error>>
204    {
205        let table = DeltaLogTable::new("orders")?;
206        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
207
208        let lazy = session.delta_lake(DeltaSourceConfig::new("orders", table.uri()))?;
209
210        assert_eq!(lazy.id(), 0);
211        assert_eq!(lazy.kind(), LazyTableKind::DeltaSource);
212        assert_eq!(lazy.name(), "orders");
213        assert_eq!(session.next_table_id(), 1);
214        assert_eq!(session.sources().len(), 1);
215        let registered = session
216            .registered_source("ORDERS")
217            .ok_or("expected registered source")?;
218        assert_eq!(registered.table(), &lazy);
219        assert_eq!(registered.name(), "orders");
220        assert!(registered.source_uri().starts_with("file://"));
221        assert_eq!(registered.snapshot_version(), 1);
222        assert_eq!(registered.protocol().source_name, "orders");
223        assert_eq!(registered.schema().fields().len(), 2);
224        let source_reports = session.source_reports();
225        assert_eq!(source_reports.len(), 1);
226        let report = &source_reports[0];
227        assert_eq!(report.source_name(), "orders");
228        assert_eq!(report.source_uri(), registered.source_uri());
229        assert_eq!(report.snapshot_version(), 1);
230        assert_eq!(report.protocol().source_name, "orders");
231        assert_eq!(report.scheduling().query_target_partitions(), None);
232        assert_eq!(
233            report.scheduling().reader_backend(),
234            DeltaProviderReaderBackend::NativeAsync
235        );
236        assert_eq!(
237            report.scheduling().max_concurrent_file_reads_per_scan(),
238            None
239        );
240        assert_eq!(report.file_count(), crate::FileCount::unavailable());
241        assert_eq!(
242            report.file_count_reason(),
243            Some(crate::ReportReasonCode::CostAvoidance)
244        );
245        assert!(!report.scan_metadata_exhausted());
246        assert_eq!(report.usage_status(), SourceUsageStatus::Unknown);
247        assert!(report.used_by_output_names().is_empty());
248        assert!(report.provider_read_stats().is_none());
249        assert_eq!(
250            report.provider_stats_reason(),
251            Some(crate::ReportReasonCode::NotExecuted)
252        );
253        let phase_timings = report.phase_timings();
254        assert_eq!(
255            phase_timings
256                .iter()
257                .map(crate::PhaseTimingReport::phase_name)
258                .collect::<Vec<_>>(),
259            vec![
260                SOURCE_LOADING_PHASE,
261                PROTOCOL_PREFLIGHT_PHASE,
262                DATAFUSION_REGISTRATION_PHASE
263            ]
264        );
265        assert!(
266            phase_timings
267                .iter()
268                .all(|timing| timing.status().is_completed())
269        );
270        assert!(
271            phase_timings
272                .iter()
273                .all(|timing| timing.elapsed_micros().is_some())
274        );
275
276        Ok(())
277    }
278
279    #[test]
280    fn delta_lake_registers_multiple_distinct_sources() -> Result<(), Box<dyn std::error::Error>> {
281        let orders = DeltaLogTable::new("orders")?;
282        let customers = DeltaLogTable::new("customers")?;
283        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
284
285        let orders = session.delta_lake(DeltaSourceConfig::new("orders", orders.uri()))?;
286        let customers = session.delta_lake(DeltaSourceConfig::new("customers", customers.uri()))?;
287
288        assert_eq!(orders.id(), 0);
289        assert_eq!(customers.id(), 1);
290        assert_eq!(session.sources().len(), 2);
291        assert!(session.registered_source("orders").is_some());
292        assert!(session.registered_source("customers").is_some());
293        Ok(())
294    }
295
296    #[test]
297    fn duplicate_source_alias_fails_before_loading_second_source()
298    -> Result<(), Box<dyn std::error::Error>> {
299        let table = DeltaLogTable::new("orders")?;
300        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
301        session.delta_lake(DeltaSourceConfig::new("orders", table.uri()))?;
302
303        let error = session.delta_lake(DeltaSourceConfig::new("ORDERS", ""));
304
305        assert!(matches!(
306            error,
307            Err(DeltaFunnelError::DuplicateSourceName { name }) if name == "ORDERS"
308        ));
309        assert_eq!(session.sources().len(), 1);
310        assert_eq!(session.next_table_id(), 1);
311        Ok(())
312    }
313
314    #[test]
315    fn invalid_source_alias_fails_before_registration() -> Result<(), DeltaFunnelError> {
316        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
317
318        let error = session.delta_lake(DeltaSourceConfig::new("select", ""));
319
320        assert!(matches!(
321            error,
322            Err(DeltaFunnelError::InvalidSourceName { name, .. }) if name == "select"
323        ));
324        assert!(session.sources().is_empty());
325        assert_eq!(session.next_table_id(), 0);
326        Ok(())
327    }
328
329    #[test]
330    fn protocol_preflight_failure_does_not_register_source()
331    -> Result<(), Box<dyn std::error::Error>> {
332        let table = DeltaLogTable::new_with_protocol("unsupported", UNSUPPORTED_PROTOCOL_JSON)?;
333        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
334
335        let error = session.delta_lake(DeltaSourceConfig::new("unsupported", table.uri()));
336
337        let display = format!("{}", error.as_ref().err().ok_or("expected error")?);
338        assert!(display.contains("unsupported"));
339        assert!(display.contains("unsupported Delta minReaderVersion"));
340        assert!(matches!(
341            error,
342            Err(DeltaFunnelError::DeltaProtocolCompatibility { .. })
343        ));
344        assert!(session.sources().is_empty());
345        assert_eq!(session.next_table_id(), 0);
346        Ok(())
347    }
348
349    #[test]
350    fn protocol_preflight_failure_redacts_secret_uri_parts()
351    -> Result<(), Box<dyn std::error::Error>> {
352        let table = DeltaLogTable::new_with_protocol("unsupported", UNSUPPORTED_PROTOCOL_JSON)?;
353        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
354
355        let error = session
356            .delta_lake(DeltaSourceConfig::new(
357                "unsupported",
358                table.file_uri_with_secret_parts()?,
359            ))
360            .map(|_| ())
361            .map_err(|error| error.to_string());
362
363        assert!(
364            matches!(error, Err(display) if display.contains("unsupported")
365            && display.contains("unsupported Delta minReaderVersion")
366            && !display.contains("super-secret")
367            && !display.contains("debug-secret")
368            && !display.contains("token"))
369        );
370        assert!(session.sources().is_empty());
371        assert_eq!(session.next_table_id(), 0);
372        Ok(())
373    }
374
375    #[tokio::test]
376    async fn protocol_preflight_failure_does_not_leak_datafusion_alias()
377    -> Result<(), Box<dyn std::error::Error>> {
378        let table = DeltaLogTable::new_with_protocol("unsupported", UNSUPPORTED_PROTOCOL_JSON)?;
379        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
380
381        let error = session.delta_lake(DeltaSourceConfig::new("unsupported", table.uri()));
382
383        assert!(matches!(
384            error,
385            Err(DeltaFunnelError::DeltaProtocolCompatibility { .. })
386        ));
387        assert!(session.context().table("unsupported").await.is_err());
388        assert!(session.sources().is_empty());
389        assert_eq!(session.next_table_id(), 0);
390        Ok(())
391    }
392
393    #[tokio::test]
394    async fn registered_source_sql_analysis_does_not_read_data_files_for_rows()
395    -> Result<(), Box<dyn std::error::Error>> {
396        let table = DeltaLogTable::new("orders")?;
397        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
398
399        session.delta_lake(DeltaSourceConfig::new("orders", table.uri()))?;
400        let dataframe = session
401            .context()
402            .sql("select id, customer_name from orders")
403            .await?;
404        let schema = dataframe.schema();
405
406        assert_eq!(schema.fields().len(), 2);
407        assert_eq!(schema.field(0).name(), "id");
408        assert_eq!(schema.field(1).name(), "customer_name");
409        assert_eq!(session.sources().len(), 1);
410        Ok(())
411    }
412
413    #[test]
414    fn source_debug_does_not_expose_storage_option_values() -> Result<(), Box<dyn std::error::Error>>
415    {
416        let table = DeltaLogTable::new("storage-options")?;
417        let mut storage_options = DeltaStorageOptions::new();
418        storage_options.insert(
419            "AWS_SECRET_ACCESS_KEY".to_owned(),
420            "super-secret".to_owned(),
421        );
422        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
423
424        session.delta_lake(
425            DeltaSourceConfig::new("orders", table.uri()).with_storage_options(storage_options),
426        )?;
427
428        let debug = format!("{session:?}");
429        assert!(debug.contains("orders"));
430        assert!(!debug.contains("super-secret"));
431        assert!(!debug.contains("AWS_SECRET_ACCESS_KEY"));
432        let report_debug = format!("{:?}", session.source_reports());
433        assert!(report_debug.contains("orders"));
434        assert!(!report_debug.contains("super-secret"));
435        assert!(!report_debug.contains("AWS_SECRET_ACCESS_KEY"));
436        Ok(())
437    }
438
439    #[test]
440    fn source_registration_honors_configured_provider_options()
441    -> Result<(), Box<dyn std::error::Error>> {
442        let table = DeltaLogTable::new("configured-provider")?;
443        let provider_scan_options = DeltaProviderScanExecutionOptions::try_new_with_reader_backend(
444            DeltaProviderReaderBackend::OfficialKernel,
445            2,
446            1,
447        )?
448        .with_output_buffer_capacity_per_partition(3)?;
449        let mut session = DeltaFunnelSession::new(
450            SessionOptions::new()
451                .with_query_options(QueryOptions {
452                    target_partitions: Some(4),
453                    output_batch_size: None,
454                })
455                .with_provider_scan_options(provider_scan_options),
456        )?;
457
458        session.delta_lake(DeltaSourceConfig::new("orders", table.uri()))?;
459
460        assert_eq!(session.sources().len(), 1);
461        assert!(session.registered_source("orders").is_some());
462        let reports = session.source_reports();
463        assert_eq!(reports.len(), 1);
464        let scheduling = reports[0].scheduling();
465        assert_eq!(scheduling.query_target_partitions(), Some(4));
466        assert_eq!(
467            scheduling.reader_backend(),
468            DeltaProviderReaderBackend::OfficialKernel
469        );
470        assert_eq!(scheduling.max_concurrent_file_reads_per_scan(), Some(2));
471        assert_eq!(scheduling.max_concurrent_file_reads_per_partition(), 1);
472        assert_eq!(scheduling.output_buffer_capacity_per_partition(), 3);
473        assert_eq!(
474            scheduling.native_async_prefetch_file_count_per_partition(),
475            0
476        );
477        Ok(())
478    }
479}
delta_funnel/orchestrator/session/registry/source.rs

delta_funnel/orchestrator/session/registry/
source.rs
