datafusion_app/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 = "huggingface")]
29mod huggingface;
30#[cfg(feature = "s3")]
31mod s3;
32#[cfg(feature = "vortex")]
33mod vortex;
34
35pub use builder::DftSessionStateBuilder;
36
37#[async_trait::async_trait]
38pub trait Extension: Debug {
39    /// Registers this extension with the DataFusion [`SessionStateBuilder`]
40    async fn register(
41        &self,
42        _config: ExecutionConfig,
43        _builder: &mut DftSessionStateBuilder,
44    ) -> Result<()>;
45
46    // Registers this extension after the SessionContext has been created
47    // (this is to match the historic way many extensions were registered)
48    // TODO file a ticket upstream to use the builder pattern
49    fn register_on_ctx(&self, _config: &ExecutionConfig, _ctx: &mut SessionContext) -> Result<()> {
50        Ok(())
51    }
52}
53
54/// Return all extensions currently enabled
55pub fn enabled_extensions() -> Vec<Arc<dyn Extension>> {
56    vec![
57        #[cfg(feature = "s3")]
58        Arc::new(s3::AwsS3Extension::new()),
59        #[cfg(feature = "deltalake")]
60        Arc::new(deltalake::DeltaLakeExtension::new()),
61        #[cfg(feature = "huggingface")]
62        Arc::new(huggingface::HuggingFaceExtension::new()),
63        #[cfg(feature = "vortex")]
64        Arc::new(vortex::VortexExtension::new()),
65    ]
66}
datafusion_app/extensions/mod.rs

datafusion_app/extensions/
mod.rs
