algeon-datafusion 0.1.1

use arrow::{
    array::{Array, Int32Array, StringArray, UInt8Array},
    datatypes::Schema,
    ipc::writer::IpcWriteOptions,
};
use arrow_flight::{
    FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
    error::FlightError,
    flight_service_server::FlightService,
    sql::{
        ActionCancelQueryRequest, ActionCancelQueryResult, ActionClosePreparedStatementRequest,
        ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs,
        CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
        CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
        CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate,
        ProstMessageExt, SqlInfo, SqlNullOrdering, SqlOuterJoinsSupportLevel, SqlSupportedCaseSensitivity,
        SqlSupportedGroupBy, SqlSupportedResultSetConcurrency, SqlSupportedResultSetType, SqlSupportedTransaction,
        SqlSupportedUnions, SqlTransactionIsolationLevel, SupportedSqlGrammar, TicketStatementQuery,
        metadata::{GetDbSchemasBuilder, GetTablesBuilder, SqlInfoDataBuilder},
        server::{FlightSqlService, PeekableFlightDataStream},
    },
};
use datafusion::{
    common::TableReference,
    datasource::TableType,
    error::DataFusionError,
    execution::SessionStateBuilder,
    logical_expr::{DdlStatement, LogicalPlan, Statement},
    physical_plan::{SendableRecordBatchStream, execute_stream},
    prelude::SessionContext,
};
use futures::{Stream, StreamExt, TryStreamExt, stream::BoxStream};
use prost::Message;
use query_planning::plan::MlForestModelDigest;
use query_runtime::observability::RuntimeObservationConfig;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
use tonic::{Code, Request, Response, Status, Streaming, metadata::MetadataValue};
use tracing::Instrument;

use super::{
    adapter_error::FailedStatement,
    admission::QueryAdmission,
    auth::{AuthInterceptor, handle_handshake},
    cancel::QueryCancellation,
    coordinator::{
        CoordinatorTicketPayload, QueryCoordinator, QueryCoordinatorConfig, QueryCoordinatorError,
        QueryCoordinatorErrorKind, QueryExecutionLease, QueryPlanningLease, QueryPlanningReportDeliveryTerminal,
        StatementExecutionDomain,
    },
    diagnostics::log_native_planning_report_for_tables,
    error::{adapter_error_to_status, datafusion_error_to_captured_error, datafusion_error_to_status_for_operation},
    observability::{CapturedError, QueryStatsView},
};
use crate::server::{
    adapter_error::{DataFusionErrorExt, Error, ErrorCode, ErrorLogDiagnostics},
    adapter_session::{
        ALGEON_NATIVE_COUNT_ONLY_PHYSICAL_OPTIMIZER_RULE_NAME, ALGEON_NATIVE_PHYSICAL_OPTIMIZER_RULE_NAME,
    },
    cuml_sql::forest_model_digests,
    gpu_coverage::{ExplainGpuOutputMode, coverage_batch, coverage_schema, parse_explain_gpu, validate_query},
    planner::native::{Invocation, InvocationExecution},
};

mod correlation;
mod ddl;
mod execution;
mod frame_limiter;
mod metadata;
mod ownership;
mod prepared;
mod rpc;
mod service;
mod statement;
mod status;
#[cfg(test)]
mod tests;
mod tickets;

use correlation::tracing_correlation_id;
pub(super) use correlation::{attach_query_response_metadata, request_correlation_id};
use ddl::referenced_table_names;
pub(in crate::server) use ddl::reject_drop_object_type_mismatch;
use frame_limiter::encode_flight_data_stream;
use metadata::{
    GetTablesPrefilters, build_table_types_batch, empty_foreign_keys_batch, empty_primary_keys_batch,
    flight_sql_foreign_keys_schema, flight_sql_primary_keys_schema, flight_sql_table_types_schema,
    metadata_flight_info, metadata_record_batch_stream, sql_info_data, table_type_to_flight_sql,
};
use ownership::{Owner, QueryOwnership};
use prepared::{PreparedStatementEntry, PreparedStatementStore};
pub use service::DatafusionAlgeonFlightSqlService;
use statement::verify_forest_model_identity;
pub(in crate::server) use status::verify_update_ddl;
use status::{
    adapter_error_status, adapter_error_to_status_and_capture, arrow_error_to_status_and_capture,
    arrow_error_to_status_for_operation, arrow_flight_error_to_status_and_capture,
    arrow_flight_error_to_status_for_operation, cancellation_error, datafusion_error_to_flight_error,
    datafusion_error_to_status_and_capture, flight_result_schema, flight_sql_transactions_unsupported_error,
    log_statement_status, reject_parameterized_plan, reject_parameterized_statement, schema_to_ipc_bytes,
    server_shutdown_error, status_with_query_id, status_with_query_id_and_correlation_id, verify_flight_info_plan,
    verify_sql_statement_namespace, verify_update_plan,
};
use tickets::{
    CoordinatorExecutionTerminal, CoordinatorReportHandoffTerminal, QueryExecutionGuard, StatementTicketKind,
    UpdateStatementKind,
};

const FLIGHT_SQL_CANCEL_RESULT_CANCELLING: i32 = 2;
const FLIGHT_SQL_CANCEL_RESULT_NOT_CANCELLABLE: i32 = 3;
const GET_ADMISSION_STATS_ACTION: &str = "GetAdmissionStats";
const GET_QUERY_STATS_ACTION: &str = "GetQueryStats";
const GET_QUERY_STATS_SUMMARY_ACTION: &str = "GetQueryStatsSummary";
const GET_QUERY_RUNTIME_ARTIFACT_ACTION: &str = "GetQueryRuntimeArtifact";
const NO_UPDATED_RECORDS_COUNT: i64 = 0;
const SQL_ANSI92_ENTRY_SQL_BIT_INDEX: i32 = 0;
const QUERY_ID_METADATA_KEY: &str = "x-algeon-query-id";
const CORRELATION_ID_METADATA_KEY: &str = "x-algeon-correlation-id";
const CORRELATION_ID_METADATA_ALIASES: &[&str] = &[CORRELATION_ID_METADATA_KEY, "x-correlation-id"];
const MAX_CORRELATION_ID_LEN: usize = 256;