use crate::{
db::{
DbSession, Query, QueryError, QueryTracePlan, TraceExecutionFamily,
access::summarize_executable_access_plan,
executor::{
EntityAuthority, ExecutionFamily, initial_read_plan_requires_materialized_sort,
},
query::admission::{
QueryAdmissionPolicy, QueryAdmissionSummary, QueryMaterializationSummary,
},
query::builder::{AggregateExplain, ProjectionExplain},
query::explain::{
ExplainAggregateTerminalPlan, ExplainExecutionNodeDescriptor, ExplainPlan,
},
query::plan::{AccessPlannedQuery, QueryMode, VisibleIndexes},
session::query::{QueryPlanCacheAttribution, query_plan_cache_reuse_event},
},
traits::{CanisterKind, EntityKind, EntityValue},
};
const fn trace_execution_family_from_executor(family: ExecutionFamily) -> TraceExecutionFamily {
match family {
ExecutionFamily::PrimaryKey => TraceExecutionFamily::PrimaryKey,
ExecutionFamily::Ordered => TraceExecutionFamily::Ordered,
ExecutionFamily::Grouped => TraceExecutionFamily::Grouped,
}
}
impl<C: CanisterKind> DbSession<C> {
fn cached_finalized_explain_plan<E>(
&self,
query: &Query<E>,
visible_indexes: &VisibleIndexes<'_>,
) -> Result<
(
AccessPlannedQuery,
EntityAuthority,
QueryPlanCacheAttribution,
),
QueryError,
>
where
E: EntityKind<Canister = C>,
{
let (prepared_plan, cache_attribution) =
self.cached_shared_query_plan_for_entity::<E>(query)?;
let mut plan = prepared_plan.logical_plan().clone();
let authority = prepared_plan.authority();
let schema_info = authority
.accepted_schema_info()
.ok_or_else(QueryError::invariant)?;
plan.finalize_access_choice_for_model_with_semantic_indexes_and_schema(
query.structural().model(),
visible_indexes.accepted_semantic_index_contracts(),
schema_info,
);
Ok((plan, authority, cache_attribution))
}
pub(in crate::db) fn explain_query_with_visible_indexes<E>(
&self,
query: &Query<E>,
) -> Result<ExplainPlan, QueryError>
where
E: EntityKind<Canister = C>,
{
self.with_query_visible_indexes(query, |query, visible_indexes| {
let (plan, _, _) = self.cached_finalized_explain_plan::<E>(query, visible_indexes)?;
Ok(plan.explain())
})
}
pub(in crate::db) fn query_plan_hash_hex_with_visible_indexes<E>(
&self,
query: &Query<E>,
) -> Result<String, QueryError>
where
E: EntityKind<Canister = C>,
{
let (prepared_plan, _) = self.cached_shared_query_plan_for_entity::<E>(query)?;
Ok(prepared_plan.plan_hash_hex())
}
#[doc(hidden)]
pub fn ensure_default_query_read_admission<E>(
&self,
query: &Query<E>,
) -> Result<QueryAdmissionSummary, QueryError>
where
E: EntityKind<Canister = C>,
{
self.ensure_query_read_admission_policy(
query,
&QueryAdmissionPolicy::default_bounded_read(),
)
}
pub(in crate::db) fn evaluate_query_read_admission_policy<E>(
&self,
query: &Query<E>,
policy: &QueryAdmissionPolicy,
) -> Result<QueryAdmissionSummary, QueryError>
where
E: EntityKind<Canister = C>,
{
let (prepared_plan, _) = self.cached_shared_query_plan_for_entity::<E>(query)?;
let mut summary =
QueryAdmissionSummary::from_plan(policy.lane(), prepared_plan.logical_plan());
if initial_read_plan_requires_materialized_sort(&prepared_plan)
.map_err(QueryError::execute)?
{
let returned_row_bound = summary.returned_row_bound();
let returned_row_bound_kind = summary.returned_row_bound_kind();
summary = summary.with_materialization(QueryMaterializationSummary::sort(
returned_row_bound,
returned_row_bound_kind,
));
}
Ok(policy.evaluate(summary))
}
pub(in crate::db) fn ensure_query_read_admission_policy<E>(
&self,
query: &Query<E>,
policy: &QueryAdmissionPolicy,
) -> Result<QueryAdmissionSummary, QueryError>
where
E: EntityKind<Canister = C>,
{
let admission = self.evaluate_query_read_admission_policy(query, policy)?;
if let Some(rejection) = admission.rejection() {
Err(QueryError::from(rejection.code()))
} else {
Ok(admission)
}
}
pub(in crate::db) fn explain_query_execution_with_visible_indexes<E>(
&self,
query: &Query<E>,
) -> Result<ExplainExecutionNodeDescriptor, QueryError>
where
E: EntityValue + EntityKind<Canister = C>,
{
self.with_query_visible_indexes(query, |query, visible_indexes| {
let (plan, authority, _) =
self.cached_finalized_explain_plan::<E>(query, visible_indexes)?;
query
.structural()
.explain_execution_descriptor_from_plan_with_authority(&plan, &authority)
})
}
pub(in crate::db) fn explain_query_execution_verbose_with_visible_indexes<E>(
&self,
query: &Query<E>,
) -> Result<String, QueryError>
where
E: EntityValue + EntityKind<Canister = C>,
{
self.with_query_visible_indexes(query, |query, visible_indexes| {
let (plan, authority, cache_attribution) =
self.cached_finalized_explain_plan::<E>(query, visible_indexes)?;
query
.structural()
.finalized_execution_diagnostics_from_plan_with_authority_and_descriptor_mutator(
&plan,
&authority,
Some(query_plan_cache_reuse_event(cache_attribution)),
|_| {},
)
.map(|diagnostics| diagnostics.render_text_verbose())
})
}
pub(in crate::db) fn explain_query_execution_json_with_visible_indexes<E>(
&self,
query: &Query<E>,
) -> Result<String, QueryError>
where
E: EntityValue + EntityKind<Canister = C>,
{
self.with_query_visible_indexes(query, |query, visible_indexes| {
let (plan, authority, cache_attribution) =
self.cached_finalized_explain_plan::<E>(query, visible_indexes)?;
query
.structural()
.finalized_execution_diagnostics_from_plan_with_authority_and_descriptor_mutator(
&plan,
&authority,
Some(query_plan_cache_reuse_event(cache_attribution)),
|_| {},
)
.map(|diagnostics| diagnostics.render_json_canonical())
})
}
pub(in crate::db) fn explain_query_prepared_aggregate_terminal_with_visible_indexes<E, S>(
&self,
query: &Query<E>,
strategy: &S,
) -> Result<ExplainAggregateTerminalPlan, QueryError>
where
E: EntityKind<Canister = C>,
S: AggregateExplain,
{
let (plan, _) = self.cached_shared_query_plan_for_entity::<E>(query)?;
plan.explain_prepared_aggregate_terminal(strategy)
}
pub(in crate::db) fn explain_query_bytes_by_with_visible_indexes<E>(
&self,
query: &Query<E>,
target_field: &str,
) -> Result<ExplainExecutionNodeDescriptor, QueryError>
where
E: EntityKind<Canister = C>,
{
let (plan, _) = self.cached_shared_query_plan_for_entity::<E>(query)?;
plan.explain_bytes_by_terminal(target_field)
}
pub(in crate::db) fn explain_query_prepared_projection_terminal_with_visible_indexes<E, S>(
&self,
query: &Query<E>,
strategy: &S,
) -> Result<ExplainExecutionNodeDescriptor, QueryError>
where
E: EntityKind<Canister = C>,
S: ProjectionExplain,
{
let (plan, _) = self.cached_shared_query_plan_for_entity::<E>(query)?;
plan.explain_prepared_projection_terminal(strategy)
}
pub fn trace_query<E>(&self, query: &Query<E>) -> Result<QueryTracePlan, QueryError>
where
E: EntityKind<Canister = C>,
{
let (prepared_plan, cache_attribution) =
self.cached_shared_query_plan_for_entity::<E>(query)?;
let logical_plan = prepared_plan.logical_plan();
let explain = logical_plan.explain();
let plan_hash = prepared_plan.plan_hash_hex();
let executable_access = prepared_plan.access().executable_contract();
let access_strategy = summarize_executable_access_plan(&executable_access);
let execution_family = match prepared_plan.mode() {
QueryMode::Load(_) => Some(trace_execution_family_from_executor(
prepared_plan
.execution_family()
.map_err(QueryError::execute)?,
)),
QueryMode::Delete(_) => None,
};
let reuse = query_plan_cache_reuse_event(cache_attribution);
Ok(QueryTracePlan::new(
plan_hash,
access_strategy,
execution_family,
reuse,
explain,
))
}
}