micromegas/servers/

flight_sql_service_impl.rs

use super::query_audit::{QueryAuditRecord, ScanMetrics, aggregate_scan_metrics};
use super::sqlinfo::{
    SQL_INFO_DATE_TIME_FUNCTIONS, SQL_INFO_NUMERIC_FUNCTIONS, SQL_INFO_SQL_KEYWORDS,
    SQL_INFO_STRING_FUNCTIONS, SQL_INFO_SYSTEM_FUNCTIONS,
};
use anyhow::Result;
use arrow_flight::decode::FlightRecordBatchStream;
use arrow_flight::encode::{DictionaryHandling, FlightDataEncoderBuilder};
use arrow_flight::error::FlightError;
use arrow_flight::sql::DoPutPreparedStatementResult;
use arrow_flight::sql::metadata::{SqlInfoData, SqlInfoDataBuilder};
use arrow_flight::sql::server::PeekableFlightDataStream;
use arrow_flight::sql::{
    ActionBeginSavepointRequest, ActionBeginSavepointResult, ActionBeginTransactionRequest,
    ActionBeginTransactionResult, ActionCancelQueryRequest, ActionCancelQueryResult,
    ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
    ActionCreatePreparedStatementResult, ActionCreatePreparedSubstraitPlanRequest,
    ActionEndSavepointRequest, ActionEndTransactionRequest, Any, CommandGetCatalogs,
    CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
    CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
    CommandGetXdbcTypeInfo, CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
    CommandStatementIngest, CommandStatementQuery, CommandStatementSubstraitPlan,
    CommandStatementUpdate, ProstMessageExt, SqlInfo, TicketStatementQuery,
    server::FlightSqlService,
};
use arrow_flight::{
    Action, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse,
    Ticket, flight_service_server::FlightService,
};
use core::str;
use datafusion::arrow::datatypes::Schema;
use datafusion::arrow::ipc::writer::StreamWriter;
use datafusion::physical_plan::{ExecutionPlan, execute_stream};
use futures::StreamExt;
use futures::{Stream, TryStreamExt};
use micromegas_analytics::lakehouse::lakehouse_context::LakehouseContext;
use micromegas_analytics::lakehouse::partition_cache::QueryPartitionProvider;
use micromegas_analytics::lakehouse::query::make_session_context;
use micromegas_analytics::lakehouse::session_configurator::SessionConfigurator;
use micromegas_analytics::lakehouse::view_factory::ViewFactory;
use micromegas_analytics::replication::bulk_ingest;
use micromegas_analytics::time::TimeRange;
use micromegas_auth::user_attribution::{is_admin, validate_and_resolve_user_attribution_grpc};
use micromegas_tracing::prelude::*;
use once_cell::sync::Lazy;
use prost::Message;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Instant;
use tonic::metadata::MetadataMap;
use tonic::{Request, Response, Status, Streaming};

type FlightDataStream =
    Pin<Box<dyn Stream<Item = Result<arrow_flight::FlightData, Status>> + Send>>;

macro_rules! status {
    ($desc:expr, $err:expr) => {
        Status::internal(format!("{}: {} at {}:{}", $desc, $err, file!(), line!()))
    };
}

macro_rules! api_entry_not_implemented {
    () => {{
        let function_name = micromegas_tracing::__function_name!();
        error!("not implemented: {function_name}");
        Err(Status::unimplemented(format!(
            "{}:{} not implemented: {function_name}",
            file!(),
            line!()
        )))
    }};
}

/// Attribution, per-stage timing, and (once created) the physical plan for
/// one query. Built as soon as attribution is resolved, updated as each
/// setup stage completes, and emitted exactly once as a [`QueryAuditRecord`]
/// under the `flightsql_query_audit` log target — either on an early setup
/// failure, at stream completion/error, or (if the stream is dropped mid-drain)
/// from `CompletionTrackedStream`'s `Drop` impl.
struct QueryAuditState {
    client: String,
    user: String,
    email: String,
    name: Option<String>,
    service_account: bool,
    service_account_name: Option<String>,
    sql: String,
    range_begin: Option<String>,
    range_end: Option<String>,
    limit: Option<u64>,
    context_init_ms: f64,
    planning_ms: f64,
    execution_ms: f64,
    setup_ms: f64,
    request_start: Instant,
    /// `None` until the physical plan is created; set as soon as it is, and
    /// still `None` for a record emitted on a setup failure that happens
    /// before that point.
    plan: Option<Arc<dyn ExecutionPlan>>,
}

impl QueryAuditState {
    /// Aggregate the plan's DataFusion metrics (if a physical plan was
    /// created before the failure/completion), assemble the audit record,
    /// and emit it as a single JSON log line.
    ///
    /// Takes `&self` (rather than consuming) so it can be called from a
    /// setup-error `map_err` closure while leaving the state available for
    /// further updates/completion on the success path, and so `Drop` can
    /// call it on an abandoned/cancelled stream without needing to
    /// reconstruct anything.
    fn emit(&self, status: &'static str, error: Option<String>) {
        let scan = match &self.plan {
            Some(plan) => aggregate_scan_metrics(plan.as_ref()),
            None => ScanMetrics {
                output_rows: None,
                bytes_scanned: 0,
            },
        };
        let total_ms = self.request_start.elapsed().as_secs_f64() * 1000.0;
        let record = QueryAuditRecord {
            client: self.client.clone(),
            user: self.user.clone(),
            email: self.email.clone(),
            name: self.name.clone(),
            service_account: self.service_account,
            service_account_name: self.service_account_name.clone(),
            sql: self.sql.clone(),
            range_begin: self.range_begin.clone(),
            range_end: self.range_end.clone(),
            limit: self.limit,
            context_init_ms: self.context_init_ms,
            planning_ms: self.planning_ms,
            execution_ms: self.execution_ms,
            setup_ms: self.setup_ms,
            total_ms,
            status,
            error,
            output_rows: scan.output_rows,
            bytes_scanned: scan.bytes_scanned,
        };
        match serde_json::to_string(&record) {
            Ok(json) => info!(target: "flightsql_query_audit", "{json}"),
            Err(e) => warn!("failed to serialize query audit record: {e}"),
        }
    }
}

/// Stream wrapper that tracks when the stream is fully consumed
struct CompletionTrackedStream<S> {
    inner: S,
    start_time: i64,
    completed: bool,
    audit: Option<QueryAuditState>,
}

impl<S> CompletionTrackedStream<S> {
    fn new(inner: S, start_time: i64, audit: QueryAuditState) -> Self {
        Self {
            inner,
            start_time,
            completed: false,
            audit: Some(audit),
        }
    }
}

impl<S> Drop for CompletionTrackedStream<S> {
    /// If the stream is dropped before yielding `None` or an `Err` (client
    /// disconnect/cancel mid-drain), `poll_next` never ran its completion
    /// arms, so `self.audit` is still `Some(...)`. Emit it here with a
    /// terminal "incomplete" status so cancelled/abandoned queries are still
    /// audited instead of silently vanishing. A no-op for streams that
    /// completed or errored normally, since those arms already `take()` the
    /// audit state.
    fn drop(&mut self) {
        if let Some(state) = self.audit.take() {
            state.emit("incomplete", None);
        }
    }
}

impl<S> Stream for CompletionTrackedStream<S>
where
    S: Stream<Item = Result<arrow_flight::FlightData, Status>> + Unpin + Send,
{
    type Item = Result<arrow_flight::FlightData, Status>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(cx) {
            Poll::Ready(Some(result)) => {
                // Check if this is an error result and log it
                if let Err(ref err) = result {
                    let sql = self.audit.as_ref().map(|state| state.sql.as_str());
                    error!("stream error occurred: {err:?} sql={sql:?}");
                    if !self.completed {
                        let total_duration = now() - self.start_time;
                        imetric!("query_duration_with_error", "ticks", total_duration as u64);
                        imetric!("query_failed", "count", 1);
                        self.completed = true;
                        if let Some(state) = self.audit.take() {
                            state.emit("error", Some(err.to_string()));
                        }
                    }
                }
                Poll::Ready(Some(result))
            }
            Poll::Ready(None) => {
                // Stream completed successfully
                if !self.completed {
                    let total_duration = now() - self.start_time;
                    imetric!("query_duration_total", "ticks", total_duration as u64);
                    imetric!("query_completed_successfully", "count", 1);
                    self.completed = true;
                    if let Some(state) = self.audit.take() {
                        state.emit("ok", None);
                    }
                }
                Poll::Ready(None)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

static INSTANCE_SQL_DATA: Lazy<SqlInfoData> = Lazy::new(|| {
    let mut builder = SqlInfoDataBuilder::new();
    // Server information
    builder.append(SqlInfo::FlightSqlServerName, "Micromegas Flight SQL Server");
    builder.append(SqlInfo::FlightSqlServerVersion, "1");
    // 1.3 comes from https://github.com/apache/arrow/blob/f9324b79bf4fc1ec7e97b32e3cce16e75ef0f5e3/format/Schema.fbs#L24
    builder.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
    builder.append(SqlInfo::SqlKeywords, SQL_INFO_SQL_KEYWORDS);
    builder.append(SqlInfo::SqlNumericFunctions, SQL_INFO_NUMERIC_FUNCTIONS);
    builder.append(SqlInfo::SqlStringFunctions, SQL_INFO_STRING_FUNCTIONS);
    builder.append(SqlInfo::SqlSystemFunctions, SQL_INFO_SYSTEM_FUNCTIONS);
    builder.append(SqlInfo::SqlDatetimeFunctions, SQL_INFO_DATE_TIME_FUNCTIONS);
    builder.build().unwrap()
});

/// Implementation of the Flight SQL service.
#[derive(Clone)]
pub struct FlightSqlServiceImpl {
    lakehouse: Arc<LakehouseContext>,
    part_provider: Arc<dyn QueryPartitionProvider>,
    view_factory: Arc<ViewFactory>,
    session_configurator: Arc<dyn SessionConfigurator>,
}

impl FlightSqlServiceImpl {
    pub fn new(
        lakehouse: Arc<LakehouseContext>,
        part_provider: Arc<dyn QueryPartitionProvider>,
        view_factory: Arc<ViewFactory>,
        session_configurator: Arc<dyn SessionConfigurator>,
    ) -> Self {
        Self {
            lakehouse,
            part_provider,
            view_factory,
            session_configurator,
        }
    }

    fn should_preserve_dictionary(metadata: &MetadataMap) -> bool {
        metadata
            .get("preserve_dictionary")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.eq_ignore_ascii_case("true"))
            .unwrap_or(false)
    }

    #[span_fn]
    async fn execute_query(
        &self,
        ticket_stmt: TicketStatementQuery,
        metadata: &MetadataMap,
    ) -> Result<Response<FlightDataStream>, Status> {
        let begin_request = now();
        let request_start = Instant::now();
        let sql = std::str::from_utf8(&ticket_stmt.statement_handle)
            .map_err(|e| status!("Unable to parse query", e))?;

        let mut begin = metadata.get("query_range_begin");
        if let Some(s) = &begin
            && s.is_empty()
        {
            begin = None;
        }
        let mut end = metadata.get("query_range_end");
        if let Some(s) = &end
            && s.is_empty()
        {
            end = None;
        }
        let query_range = if begin.is_some() && end.is_some() {
            let begin_datetime = chrono::DateTime::parse_from_rfc3339(
                begin
                    .unwrap()
                    .to_str()
                    .map_err(|e| status!("Unable to convert query_range_begin to string", e))?,
            )
            .map_err(|e| status!("Unable to parse query_range_begin as a rfc3339 datetime", e))?;
            let end_datetime = chrono::DateTime::parse_from_rfc3339(
                end.unwrap()
                    .to_str()
                    .map_err(|e| status!("Unable to convert query_range_end to string", e))?,
            )
            .map_err(|e| status!("Unable to parse query_range_end as a rfc3339 datetime", e))?;
            Some(TimeRange::new(begin_datetime.into(), end_datetime.into()))
        } else {
            None
        };

        // Validate and resolve user attribution
        let attr = validate_and_resolve_user_attribution_grpc(metadata).map_err(|e| *e)?;

        let client_type = metadata
            .get("x-client-type")
            .and_then(|v| v.to_str().ok())
            .unwrap_or("unknown");

        let user_name_display = attr.user_name.as_deref().unwrap_or("");

        // Log query with full attribution
        if let Some(service_account_name) = &attr.service_account {
            info!(
                "execute_query range={query_range:?} sql={sql:?} limit={:?} user={} email={} name={user_name_display:?} service_account={service_account_name} client={client_type}",
                metadata.get("limit"),
                attr.user_id,
                attr.user_email
            );
        } else {
            info!(
                "execute_query range={query_range:?} sql={sql:?} limit={:?} user={} email={} name={user_name_display:?} client={client_type}",
                metadata.get("limit"),
                attr.user_id,
                attr.user_email
            );
        }

        // Attribution is resolved from here on, so build the audit state now
        // (durations/limit/plan filled in as they become known) instead of
        // only after the physical plan exists. This lets every subsequent
        // setup failure (session context, planning, limit, physical plan,
        // stream construction) still emit an "error" audit record instead of
        // silently disappearing on an early `?` return.
        let mut audit_state = QueryAuditState {
            client: client_type.to_string(),
            user: attr.user_id.clone(),
            email: attr.user_email.clone(),
            name: attr.user_name.clone(),
            service_account: attr.service_account.is_some(),
            service_account_name: attr.service_account.clone(),
            sql: sql.to_string(),
            range_begin: query_range.as_ref().map(|r| r.begin.to_rfc3339()),
            range_end: query_range.as_ref().map(|r| r.end.to_rfc3339()),
            limit: None,
            context_init_ms: 0.0,
            planning_ms: 0.0,
            execution_ms: 0.0,
            setup_ms: 0.0,
            request_start,
            plan: None,
        };

        // Session context creation phase
        let session_begin = now();
        let session_begin_instant = Instant::now();
        let ctx = make_session_context(
            self.lakehouse.clone(),
            self.part_provider.clone(),
            query_range,
            self.view_factory.clone(),
            self.session_configurator.clone(),
            is_admin(metadata),
        )
        .await
        .map_err(|e| {
            audit_state.emit("error", Some(format!("error in make_session_context: {e}")));
            status!("error in make_session_context", e)
        })?;
        let context_init_duration = now() - session_begin;
        audit_state.context_init_ms = session_begin_instant.elapsed().as_secs_f64() * 1000.0;

        // Query planning phase
        let planning_begin = now();
        let planning_begin_instant = Instant::now();
        let mut df = ctx.sql(sql).await.map_err(|e| {
            audit_state.emit("error", Some(format!("error building dataframe: {e}")));
            status!("error building dataframe", e)
        })?;
        let planning_duration = now() - planning_begin;
        audit_state.planning_ms = planning_begin_instant.elapsed().as_secs_f64() * 1000.0;

        if let Some(limit_str) = metadata.get("limit") {
            let parsed_limit: usize = usize::from_str(limit_str.to_str().map_err(|e| {
                audit_state.emit("error", Some(format!("error converting limit to str: {e}")));
                status!("error converting limit to str", e)
            })?)
            .map_err(|e| {
                audit_state.emit("error", Some(format!("error parsing limit: {e}")));
                status!("error parsing limit", e)
            })?;
            audit_state.limit = Some(parsed_limit as u64);
            df = df.limit(0, Some(parsed_limit)).map_err(|e| {
                audit_state.emit(
                    "error",
                    Some(format!("error building dataframe with limit: {e}")),
                );
                status!("error building dataframe with limit", e)
            })?;
        }

        // Query execution phase: build the physical plan (kept for post-drain metrics)
        // and run it via the free-function `execute_stream`, which is what
        // `DataFrame::execute_stream` does internally, minus dropping the plan.
        let execution_begin = now();
        let execution_begin_instant = Instant::now();
        let schema = Arc::new(df.schema().as_arrow().clone());
        let task_ctx = Arc::new(df.task_ctx());
        let plan = df.create_physical_plan().await.map_err(|e| {
            audit_state.emit("error", Some(format!("error creating physical plan: {e}")));
            status!("error creating physical plan", e)
        })?;
        audit_state.plan = Some(plan.clone());
        let stream = execute_stream(plan, task_ctx)
            .map_err(|e| {
                audit_state.emit("error", Some(format!("Error executing plan: {e:?}")));
                Status::internal(format!("Error executing plan: {e:?}"))
            })?
            .map_err(|e| FlightError::ExternalError(Box::new(e)));
        let builder = if Self::should_preserve_dictionary(metadata) {
            FlightDataEncoderBuilder::new()
                .with_schema(schema.clone())
                .with_dictionary_handling(DictionaryHandling::Resend)
        } else {
            FlightDataEncoderBuilder::new().with_schema(schema.clone())
        };
        let flight_data_stream = builder.build(stream);
        let execution_duration = now() - execution_begin;
        audit_state.execution_ms = execution_begin_instant.elapsed().as_secs_f64() * 1000.0;

        // Calculate total setup time and record detailed metrics
        let total_setup_duration = now() - begin_request;
        audit_state.setup_ms = request_start.elapsed().as_secs_f64() * 1000.0;

        // Record detailed timing metrics
        imetric!(
            "context_init_duration",
            "ticks",
            context_init_duration as u64
        );
        imetric!("query_planning_duration", "ticks", planning_duration as u64);
        imetric!(
            "query_execution_duration",
            "ticks",
            execution_duration as u64
        );
        imetric!("query_setup_duration", "ticks", total_setup_duration as u64);

        // Create instrumented stream that tracks completion
        let instrumented_stream =
            flight_data_stream.map_err(|e| status!("error building data stream", e));
        let completion_tracked_stream =
            CompletionTrackedStream::new(instrumented_stream.boxed(), begin_request, audit_state);
        Ok(Response::new(
            Box::pin(completion_tracked_stream) as FlightDataStream
        ))
    }
}

#[tonic::async_trait]
impl FlightSqlService for FlightSqlServiceImpl {
    type FlightService = FlightSqlServiceImpl;

    async fn do_handshake(
        &self,
        _request: Request<Streaming<HandshakeRequest>>,
    ) -> Result<
        Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
        Status,
    > {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_get_fallback(
        &self,
        request: Request<Ticket>,
        _message: Any,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        let ticket_stmt = TicketStatementQuery::decode(request.get_ref().ticket.clone())
            .map_err(|e| status!("Could not read ticket", e))?;
        self.execute_query(ticket_stmt, request.metadata()).await
    }

    #[span_fn]
    async fn get_flight_info_statement(
        &self,
        query: CommandStatementQuery,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        let begin_request = now();
        info!("get_flight_info_statement {query:?} ");
        let CommandStatementQuery { query, .. } = query;
        let schema = Schema::empty();
        let ticket = TicketStatementQuery {
            statement_handle: query.into(),
        };
        let mut bytes: Vec<u8> = Vec::new();
        if ticket.encode(&mut bytes).is_ok() {
            let info = FlightInfo::new()
                .try_with_schema(&schema)
                .unwrap()
                .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(bytes)));
            let duration = now() - begin_request;
            imetric!("request_duration", "ticks", duration as u64);
            Ok(Response::new(info))
        } else {
            error!("Error encoding ticket");
            Err(Status::internal("Error encoding ticket"))
        }
    }

    async fn get_flight_info_substrait_plan(
        &self,
        _query: CommandStatementSubstraitPlan,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_prepared_statement(
        &self,
        _cmd: CommandPreparedStatementQuery,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_catalogs(
        &self,
        _query: CommandGetCatalogs,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_schemas(
        &self,
        _query: CommandGetDbSchemas,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn get_flight_info_tables(
        &self,
        query: CommandGetTables,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        let begin_request = now();
        info!("get_flight_info_tables");
        let flight_descriptor = request.into_inner();
        let ticket = Ticket {
            ticket: query.as_any().encode_to_vec().into(),
        };
        let endpoint = FlightEndpoint::new().with_ticket(ticket);
        let flight_info = FlightInfo::new()
            .try_with_schema(&query.into_builder().schema())
            .map_err(|e| status!("Unable to encode schema", e))?
            .with_endpoint(endpoint)
            .with_descriptor(flight_descriptor);
        let duration = now() - begin_request;
        imetric!("request_duration", "ticks", duration as u64);
        Ok(tonic::Response::new(flight_info))
    }

    async fn get_flight_info_table_types(
        &self,
        _query: CommandGetTableTypes,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn get_flight_info_sql_info(
        &self,
        query: CommandGetSqlInfo,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        let begin_request = now();
        info!("get_flight_info_sql_info");
        let flight_descriptor = request.into_inner();
        let ticket = Ticket::new(query.as_any().encode_to_vec());
        let endpoint = FlightEndpoint::new().with_ticket(ticket);
        let flight_info = FlightInfo::new()
            .try_with_schema(query.into_builder(&INSTANCE_SQL_DATA).schema().as_ref())
            .map_err(|e| status!("Unable to encode schema", e))?
            .with_endpoint(endpoint)
            .with_descriptor(flight_descriptor);
        let duration = now() - begin_request;
        imetric!("request_duration", "ticks", duration as u64);
        Ok(tonic::Response::new(flight_info))
    }

    async fn get_flight_info_primary_keys(
        &self,
        _query: CommandGetPrimaryKeys,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_exported_keys(
        &self,
        _query: CommandGetExportedKeys,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_imported_keys(
        &self,
        _query: CommandGetImportedKeys,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_cross_reference(
        &self,
        _query: CommandGetCrossReference,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    async fn get_flight_info_xdbc_type_info(
        &self,
        _query: CommandGetXdbcTypeInfo,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_get_statement(
        &self,
        ticket: TicketStatementQuery,
        request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        self.execute_query(ticket, request.metadata()).await
    }

    async fn do_get_prepared_statement(
        &self,
        _query: CommandPreparedStatementQuery,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_catalogs(
        &self,
        _query: CommandGetCatalogs,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_schemas(
        &self,
        _query: CommandGetDbSchemas,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_get_tables(
        &self,
        query: CommandGetTables,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        let begin_request = now();
        info!("do_get_tables {query:?}");
        let mut builder = query.into_builder();
        for view in self.view_factory.get_global_views() {
            let catalog_name = "";
            let schema_name = "";
            builder
                .append(
                    catalog_name,
                    schema_name,
                    &*view.get_view_set_name(),
                    "table",
                    &view.get_file_schema(),
                )
                .map_err(Status::from)?;
        }
        let schema = builder.schema();
        let batch = builder.build();
        let stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            .build(futures::stream::once(async { batch }))
            .map_err(Status::from);
        let duration = now() - begin_request;
        imetric!("request_duration", "ticks", duration as u64);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn do_get_table_types(
        &self,
        _query: CommandGetTableTypes,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_get_sql_info(
        &self,
        query: CommandGetSqlInfo,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        info!("do_get_sql_info");
        let builder = query.into_builder(&INSTANCE_SQL_DATA);
        let schema = builder.schema();
        let batch = builder.build();
        let stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            .build(futures::stream::once(async { batch }))
            .map_err(Status::from);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn do_get_primary_keys(
        &self,
        _query: CommandGetPrimaryKeys,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_exported_keys(
        &self,
        _query: CommandGetExportedKeys,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_imported_keys(
        &self,
        _query: CommandGetImportedKeys,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_cross_reference(
        &self,
        _query: CommandGetCrossReference,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_get_xdbc_type_info(
        &self,
        _query: CommandGetXdbcTypeInfo,
        _request: Request<Ticket>,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        api_entry_not_implemented!()
    }

    async fn do_put_statement_update(
        &self,
        _ticket: CommandStatementUpdate,
        _request: Request<PeekableFlightDataStream>,
    ) -> Result<i64, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_put_statement_ingest(
        &self,
        command: CommandStatementIngest,
        request: Request<PeekableFlightDataStream>,
    ) -> Result<i64, Status> {
        let table_name = command.table;
        info!("do_put_statement_ingest table_name={table_name}");
        let stream = FlightRecordBatchStream::new_from_flight_data(
            request.into_inner().map_err(|e| e.into()),
        );
        bulk_ingest(self.lakehouse.lake().clone(), &table_name, stream)
            .await
            .map_err(|e| {
                let msg = format!("error ingesting into {table_name}: {e:?}");
                error!("{msg}");
                status!(msg, e)
            })
    }

    async fn do_put_substrait_plan(
        &self,
        _ticket: CommandStatementSubstraitPlan,
        _request: Request<PeekableFlightDataStream>,
    ) -> Result<i64, Status> {
        api_entry_not_implemented!()
    }

    async fn do_put_prepared_statement_query(
        &self,
        _query: CommandPreparedStatementQuery,
        _request: Request<PeekableFlightDataStream>,
    ) -> Result<DoPutPreparedStatementResult, Status> {
        api_entry_not_implemented!()
    }

    async fn do_put_prepared_statement_update(
        &self,
        _query: CommandPreparedStatementUpdate,
        _request: Request<PeekableFlightDataStream>,
    ) -> Result<i64, Status> {
        api_entry_not_implemented!()
    }

    #[span_fn]
    async fn do_action_create_prepared_statement(
        &self,
        query: ActionCreatePreparedStatementRequest,
        request: Request<Action>,
    ) -> Result<ActionCreatePreparedStatementResult, Status> {
        info!("do_action_create_prepared_statement query={}", &query.query);

        let ctx = make_session_context(
            self.lakehouse.clone(),
            self.part_provider.clone(),
            None,
            self.view_factory.clone(),
            self.session_configurator.clone(),
            is_admin(request.metadata()),
        )
        .await
        .map_err(|e| status!("error in make_session_context", e))?;

        let df = ctx
            .sql(&query.query)
            .await
            .map_err(|e| status!("error building dataframe", e))?;
        let schema = df.schema().as_arrow();
        let mut schema_buffer = Vec::new();
        let mut writer = StreamWriter::try_new(&mut schema_buffer, schema)
            .map_err(|e| status!("error writing schema to in-memory buffer", e))?;
        writer
            .finish()
            .map_err(|e| status!("error closing arrow ipc stream writer", e))?;
        // here we could serialize the logical plan and return that as the prepared statement, but we would
        // need to register LogicalExtensionCodec for user-defined functions
        // instead, we are sending back the sql as we received it
        let result = ActionCreatePreparedStatementResult {
            prepared_statement_handle: query.query.into(),
            dataset_schema: schema_buffer.into(),
            parameter_schema: "".into(),
        };
        Ok(result)
    }

    async fn do_action_close_prepared_statement(
        &self,
        _query: ActionClosePreparedStatementRequest,
        _request: Request<Action>,
    ) -> Result<(), Status> {
        info!("do_action_close_prepared_statement");
        Ok(())
    }

    async fn do_action_create_prepared_substrait_plan(
        &self,
        _query: ActionCreatePreparedSubstraitPlanRequest,
        _request: Request<Action>,
    ) -> Result<ActionCreatePreparedStatementResult, Status> {
        api_entry_not_implemented!()
    }

    async fn do_action_begin_transaction(
        &self,
        _query: ActionBeginTransactionRequest,
        _request: Request<Action>,
    ) -> Result<ActionBeginTransactionResult, Status> {
        api_entry_not_implemented!()
    }

    async fn do_action_end_transaction(
        &self,
        _query: ActionEndTransactionRequest,
        _request: Request<Action>,
    ) -> Result<(), Status> {
        api_entry_not_implemented!()
    }

    async fn do_action_begin_savepoint(
        &self,
        _query: ActionBeginSavepointRequest,
        _request: Request<Action>,
    ) -> Result<ActionBeginSavepointResult, Status> {
        api_entry_not_implemented!()
    }

    async fn do_action_end_savepoint(
        &self,
        _query: ActionEndSavepointRequest,
        _request: Request<Action>,
    ) -> Result<(), Status> {
        api_entry_not_implemented!()
    }

    async fn do_action_cancel_query(
        &self,
        _query: ActionCancelQueryRequest,
        _request: Request<Action>,
    ) -> Result<ActionCancelQueryResult, Status> {
        api_entry_not_implemented!()
    }

    async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {
        info!("register_sql_info");
    }
}