datafusion_dft/
execution.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
18pub use datafusion_app::{collect_plan_io_stats, ExecutionStats};
19
20use color_eyre::Result;
21use datafusion::prelude::*;
22#[cfg(feature = "flightsql")]
23use datafusion_app::flightsql::{FlightSQLClient, FlightSQLContext};
24use datafusion_app::{local::ExecutionContext, ExecOptions, ExecResult};
25
26/// Provides all core execution functionality for execution queries from either a local
27/// `SessionContext` or a remote `FlightSQL` service
28#[derive(Clone, Debug)]
29pub struct AppExecution {
30    local: ExecutionContext,
31    #[cfg(feature = "flightsql")]
32    flightsql: FlightSQLContext,
33}
34
35impl AppExecution {
36    pub fn new(local: ExecutionContext) -> Self {
37        Self {
38            local,
39            #[cfg(feature = "flightsql")]
40            flightsql: FlightSQLContext::default(),
41        }
42    }
43
44    pub fn execution_ctx(&self) -> &ExecutionContext {
45        &self.local
46    }
47
48    pub fn session_ctx(&self) -> &SessionContext {
49        self.local.session_ctx()
50    }
51
52    #[cfg(feature = "flightsql")]
53    pub fn flightsql_client(&self) -> &FlightSQLClient {
54        self.flightsql.client()
55    }
56
57    #[cfg(feature = "flightsql")]
58    pub fn flightsql_ctx(&self) -> &FlightSQLContext {
59        &self.flightsql
60    }
61
62    #[cfg(feature = "flightsql")]
63    pub fn with_flightsql_ctx(&mut self, flightsql_ctx: FlightSQLContext) {
64        self.flightsql = flightsql_ctx;
65    }
66
67    pub async fn execute_sql_with_opts(&self, sql: &str, opts: ExecOptions) -> Result<ExecResult> {
68        #[cfg(feature = "flightsql")]
69        if opts.flightsql {
70            return self
71                .flightsql
72                .execute_sql_with_opts(sql, opts)
73                .await
74                .map_err(|e| e.into());
75        }
76
77        // If flightsql is not enabled or `opts.flightsql` is false, fall back to local:
78        self.local
79            .execute_sql_with_opts(sql, opts)
80            .await
81            .map_err(|e| e.into())
82    }
83}
datafusion_dft/execution.rs

datafusion_dft/
execution.rs
