datafusion_dft/extensions/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! This module has code to register DataFusion extensions
19
20use crate::config::ExecutionConfig;
21use datafusion::common::Result;
22use datafusion::prelude::SessionContext;
23use std::{fmt::Debug, sync::Arc};
24
25mod builder;
26#[cfg(feature = "deltalake")]
27mod deltalake;
28#[cfg(feature = "hudi")]
29mod hudi;
30#[cfg(feature = "huggingface")]
31mod huggingface;
32#[cfg(feature = "iceberg")]
33mod iceberg;
34#[cfg(feature = "s3")]
35mod s3;
36
37pub use builder::DftSessionStateBuilder;
38
39#[async_trait::async_trait]
40pub trait Extension: Debug {
41    /// Registers this extension with the DataFusion [`SessionStateBuilder`]
42    async fn register(
43        &self,
44        _config: ExecutionConfig,
45        _builder: &mut DftSessionStateBuilder,
46    ) -> Result<()>;
47
48    // Registers this extension after the SessionContext has been created
49    // (this is to match the historic way many extensions were registered)
50    // TODO file a ticket upstream to use the builder pattern
51    fn register_on_ctx(&self, _config: &ExecutionConfig, _ctx: &mut SessionContext) -> Result<()> {
52        Ok(())
53    }
54}
55
56/// Return all extensions currently enabled
57pub fn enabled_extensions() -> Vec<Arc<dyn Extension>> {
58    vec![
59        #[cfg(feature = "s3")]
60        Arc::new(s3::AwsS3Extension::new()),
61        #[cfg(feature = "deltalake")]
62        Arc::new(deltalake::DeltaLakeExtension::new()),
63        #[cfg(feature = "hudi")]
64        Arc::new(hudi::HudiExtension::new()),
65        #[cfg(feature = "iceberg")]
66        Arc::new(iceberg::IcebergExtension::new()),
67        #[cfg(feature = "huggingface")]
68        Arc::new(huggingface::HuggingFaceExtension::new()),
69    ]
70}
datafusion_dft/extensions/mod.rs

datafusion_dft/extensions/
mod.rs
