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;