#![allow(
dead_code,
reason = "the staged planning capabilities the conformance cases drive - a local call context, and the dependency views a mixed plan would read - are reached by no production caller yet"
)]
use std::collections::BTreeSet;
use std::fmt;
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use std::time::Duration;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::catalog::{
CatalogProvider, CatalogProviderList, MemoryCatalogProvider, MemoryCatalogProviderList,
MemorySchemaProvider, Session, TableProvider,
};
use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion::datasource::TableType;
use datafusion::error::DataFusionError;
use datafusion::execution::context::{SessionContext, SessionState};
use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::logical_expr::{Expr, LogicalPlan};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::scalar::ScalarValue;
use polyc_projection::family::{LogicalField, LogicalType, TableSchema, conversation_core};
use polyc_state::command::CommandEnvelope;
use polyc_state::context::CallContext;
use polyc_state::deadline::{Clock, ProductionClock};
use polyc_state::digest::ContentDigest;
use polyc_state::error::{BoundKind, StateError};
use polyc_state::id::{Audience, NamespaceId, OperationFamily, OwnerId, PartitionId, Purpose};
use polyc_state::immutable::Classification;
use polyc_state::journal::{
GetJournalSource, JournalAnchor, JournalDirectoryPage, JournalDirectorySnapshot,
JournalSourceHead, ListJournalDirectorySnapshot, MAX_DIRECTORY_PAGE_PARTITIONS,
ReleaseJournalDirectorySnapshot,
};
use polyc_state::projection::{
FamilyId, ProjectionCatalogError, ProjectionKey, ProjectionManifest, ProjectionResolution,
ResolveManifest,
};
use polyc_state::query_audit::{
BeginOutcome, BeginQueryAudit, CompleteQueryAudit, ExecutionPermit, MAX_SOURCE_PINS,
ProjectionPin, QueryAuditError, QueryCompletion, QueryId, RequesterId, SourcePin,
SourceSnapshot,
};
use polyc_state::receipt::Receipt;
use polyc_state::revision::JournalPosition;
use polyc_state_connect::query_audit::{RemoteCompleteQueryAudit, RemoteExecutionPermit};
use polyc_state_connect::wire::DeclaredCall;
use crate::core_execution::PermitGuardian;
use crate::decode::message_content::tool_calls_schema;
use crate::engine::QueryLimits;
use crate::session::QueryScope;
use crate::statement_gate::{AllowedStatement, check_statement_allowed};
const SHAPE_DOMAIN: &[u8] = b"polychrome.query.conversation-core-shape.v1\0";
const BOUNDS_DOMAIN: &[u8] = b"polychrome.query.conversation-core-bounds.v1\0";
const CORE_PARTITION_PREFIX: &str = "conv-";
const DEFAULT_CORE_RESULT_RELEASE_BYTES: u64 = 32 * 1024 * 1024;
const DEFAULT_CORE_RESPONSE_FRAME_BYTES: u64 = 256 * 1024;
const DEFAULT_CORE_ARTIFACT_FILE_BYTES: u64 = 256 * 1024 * 1024;
const DEFAULT_CORE_ARTIFACT_RANGE_BYTES: u64 = 4 * 1024 * 1024;
const DEFAULT_CORE_SOURCE_DECODE_BYTES: u64 = 512 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) enum CoreTable {
Turns,
Messages,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) enum LegacyTable {
ToolCalls,
}
impl LegacyTable {
const fn name(self) -> &'static str {
match self {
Self::ToolCalls => "tool_calls",
}
}
fn schema(self) -> SchemaRef {
match self {
Self::ToolCalls => tool_calls_schema(),
}
}
fn schema_fingerprint(self) -> ContentDigest {
let schema = self.schema();
let mut bytes = b"polychrome.query.legacy-schema.v1\0".to_vec();
push(&mut bytes, self.name().as_bytes());
for field in schema.fields() {
push(&mut bytes, field.name().as_bytes());
push(&mut bytes, field.data_type().to_string().as_bytes());
bytes.push(u8::from(field.is_nullable()));
}
digest(&bytes)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) enum CoreRealm {
Visible,
Fleet,
}
impl CoreRealm {
const fn from_scope(scope: &QueryScope) -> Self {
match scope {
QueryScope::Fleet => Self::Fleet,
QueryScope::Conversations(_) => Self::Visible,
}
}
}
impl CoreTable {
pub(crate) fn from_name(name: &str) -> Result<Self, CoreResolutionError> {
match name {
"turns" => Ok(Self::Turns),
"messages" => Ok(Self::Messages),
other => Err(CoreResolutionError::UnknownDependency(other.to_owned())),
}
}
pub(crate) const fn name(self) -> &'static str {
match self {
Self::Turns => "turns",
Self::Messages => "messages",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DeclaredTable {
Projected(CoreTable),
Legacy(LegacyTable),
}
impl DeclaredTable {
fn from_name(name: &str) -> Result<Self, CoreResolutionError> {
if name == LegacyTable::ToolCalls.name() {
return Ok(Self::Legacy(LegacyTable::ToolCalls));
}
CoreTable::from_name(name).map(Self::Projected)
}
}
pub(crate) struct CompiledCoreQuery {
normalized_plan: String,
dependencies: Vec<CoreTable>,
legacy_dependencies: Vec<LegacyTable>,
statement: AllowedStatement,
explain_enabled: bool,
plan: LogicalPlan,
base_state: SessionState,
}
impl fmt::Debug for CompiledCoreQuery {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CompiledCoreQuery")
.field("dependencies", &self.dependencies)
.field("legacy_dependencies", &self.legacy_dependencies)
.field("statement", &self.statement)
.field("explain_enabled", &self.explain_enabled)
.finish_non_exhaustive()
}
}
impl CompiledCoreQuery {
pub(crate) fn dependencies(&self) -> &[CoreTable] {
&self.dependencies
}
pub(crate) fn legacy_dependencies(&self) -> &[LegacyTable] {
&self.legacy_dependencies
}
pub(crate) fn into_parts(self) -> CompiledCoreParts {
let Self {
normalized_plan,
dependencies,
legacy_dependencies,
statement,
explain_enabled,
plan,
base_state,
} = self;
CompiledCoreParts {
normalized_plan,
dependencies,
legacy_dependencies,
statement,
explain_enabled,
plan,
base_state,
}
}
}
pub(crate) struct CompiledCoreParts {
pub(crate) normalized_plan: String,
pub(crate) dependencies: Vec<CoreTable>,
pub(crate) legacy_dependencies: Vec<LegacyTable>,
pub(crate) statement: AllowedStatement,
pub(crate) explain_enabled: bool,
pub(crate) plan: LogicalPlan,
pub(crate) base_state: SessionState,
}
#[derive(Debug)]
struct SchemaOnlyTable {
schema: SchemaRef,
scans: Arc<AtomicUsize>,
}
#[async_trait]
impl TableProvider for SchemaOnlyTable {
fn schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
fn table_type(&self) -> TableType {
TableType::Base
}
async fn scan(
&self,
_state: &dyn Session,
_projection: Option<&Vec<usize>>,
_filters: &[Expr],
_limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
self.scans.fetch_add(1, Ordering::SeqCst);
Err(DataFusionError::Plan(
"schema-only tables cannot create a physical scan".to_owned(),
))
}
}
struct ArrowLogicalType(LogicalType);
impl From<ArrowLogicalType> for DataType {
fn from(value: ArrowLogicalType) -> Self {
let ArrowLogicalType(logical) = value;
match logical {
LogicalType::Utf8 => Self::Utf8,
LogicalType::FixedBytes { len } => {
Self::FixedSizeBinary(i32::try_from(len).unwrap_or(i32::MAX))
}
LogicalType::UInt64 => Self::UInt64,
LogicalType::Boolean => Self::Boolean,
}
}
}
pub(crate) fn arrow_schema(table: &TableSchema) -> SchemaRef {
let fields = table
.fields()
.iter()
.map(|field: &LogicalField| {
Field::new(
field.name(),
DataType::from(ArrowLogicalType(field.logical_type())),
field.nullable(),
)
})
.collect::<Vec<_>>();
Arc::new(Schema::new(fields))
}
#[derive(Debug)]
pub(crate) struct CatalogCompiler {
state: SessionState,
scans: Arc<AtomicUsize>,
}
impl CatalogCompiler {
pub(crate) fn new(state: SessionState) -> Self {
Self {
state,
scans: Arc::new(AtomicUsize::new(0)),
}
}
pub(crate) async fn compile(
&self,
sql: &str,
parameters: &[CoreParameter],
allow_explain: bool,
) -> Result<CompiledCoreQuery, CoreResolutionError> {
let statement = check_statement_allowed(sql, allow_explain)
.map_err(|error| CoreResolutionError::Statement(error.to_string()))?;
let catalog_name = self.state.config_options().catalog.default_catalog.clone();
let schema_name = self.state.config_options().catalog.default_schema.clone();
let catalog_list = Arc::new(MemoryCatalogProviderList::new());
let catalog = Arc::new(MemoryCatalogProvider::new());
catalog.register_schema(&schema_name, Arc::new(MemorySchemaProvider::new()))?;
catalog_list.register_catalog(catalog_name, catalog);
let state = SessionStateBuilder::new_from_existing(self.state.clone())
.with_catalog_list(catalog_list)
.build();
let context = SessionContext::new_with_state(state);
for table in conversation_core().tables() {
context.register_table(
table.table().as_str(),
Arc::new(SchemaOnlyTable {
schema: arrow_schema(table),
scans: Arc::clone(&self.scans),
}),
)?;
}
let legacy = LegacyTable::ToolCalls;
context.register_table(
legacy.name(),
Arc::new(SchemaOnlyTable {
schema: legacy.schema(),
scans: Arc::clone(&self.scans),
}),
)?;
let dataframe = context.sql(sql).await?;
let actual = dataframe
.logical_plan()
.get_parameter_names()?
.into_iter()
.collect::<BTreeSet<_>>();
let expected = (1..=parameters.len())
.map(|index| format!("${index}"))
.collect::<BTreeSet<_>>();
if actual != expected {
return Err(CoreResolutionError::ParameterMismatch);
}
let values = parameters
.iter()
.map(BoundCoreParameter)
.map(ScalarValue::from)
.collect::<Vec<_>>();
let dataframe = dataframe.with_param_values(values)?;
let plan = dataframe.logical_plan().clone();
let mut dependencies = BTreeSet::new();
let mut legacy_dependencies = BTreeSet::new();
plan.apply(|node| {
if let LogicalPlan::TableScan(scan) = node {
match DeclaredTable::from_name(scan.table_name.table())
.map_err(|error| DataFusionError::Plan(error.to_string()))?
{
DeclaredTable::Projected(table) => {
dependencies.insert(table);
}
DeclaredTable::Legacy(table) => {
legacy_dependencies.insert(table);
}
}
}
Ok(TreeNodeRecursion::Continue)
})?;
if dependencies.is_empty() && legacy_dependencies.is_empty() {
return Err(CoreResolutionError::NoSourceDependency);
}
let normalized_plan = plan.display_indent().to_string();
Ok(CompiledCoreQuery {
normalized_plan,
dependencies: dependencies.into_iter().collect(),
legacy_dependencies: legacy_dependencies.into_iter().collect(),
statement,
explain_enabled: allow_explain,
plan,
base_state: self.state.clone(),
})
}
#[cfg(test)]
fn physical_scan_count(&self) -> usize {
self.scans.load(Ordering::SeqCst)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CoreConsistency {
Projected,
RequireProjectedThrough(JournalPosition),
}
#[derive(Clone, PartialEq, Eq)]
pub(crate) enum CoreParameter {
Utf8(String),
UInt64(u64),
Boolean(bool),
Null,
}
impl fmt::Debug for CoreParameter {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Utf8(_) => "utf8",
Self::UInt64(_) => "uint64",
Self::Boolean(_) => "boolean",
Self::Null => "null",
})
}
}
struct BoundCoreParameter<'a>(&'a CoreParameter);
impl From<BoundCoreParameter<'_>> for ScalarValue {
fn from(value: BoundCoreParameter<'_>) -> Self {
let BoundCoreParameter(parameter) = value;
match parameter {
CoreParameter::Utf8(value) => Self::Utf8(Some(value.clone())),
CoreParameter::UInt64(value) => Self::UInt64(Some(*value)),
CoreParameter::Boolean(value) => Self::Boolean(Some(*value)),
CoreParameter::Null => Self::Null,
}
}
}
pub(crate) struct CoreQueryRequest {
pub(crate) sql: String,
pub(crate) parameters: Vec<CoreParameter>,
pub(crate) consistency: CoreConsistency,
requested_bounds: CoreRequestedBounds,
}
impl fmt::Debug for CoreQueryRequest {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CoreQueryRequest")
.field("sql_bytes", &self.sql.len())
.field("parameters", &self.parameters.len())
.field("consistency", &self.consistency)
.field("requested_bounds", &self.requested_bounds)
.finish()
}
}
impl CoreQueryRequest {
pub(crate) const fn new(
sql: String,
parameters: Vec<CoreParameter>,
consistency: CoreConsistency,
requested_bounds: CoreRequestedBounds,
) -> Self {
Self {
sql,
parameters,
consistency,
requested_bounds,
}
}
pub(crate) const fn requested_bounds(&self) -> CoreRequestedBounds {
self.requested_bounds
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct CoreRequestedBounds {
timeout: Duration,
rows: u64,
result_release_bytes: u64,
response_frame_bytes: u64,
}
impl CoreRequestedBounds {
pub(crate) const fn from_requested(
timeout: Duration,
rows: u64,
result_release_bytes: u64,
response_frame_bytes: u64,
) -> Self {
Self {
timeout,
rows,
result_release_bytes,
response_frame_bytes,
}
}
#[cfg(test)]
pub(crate) const fn unbounded() -> Self {
Self {
timeout: Duration::MAX,
rows: u64::MAX,
result_release_bytes: u64::MAX,
response_frame_bytes: u64::MAX,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(
clippy::struct_field_names,
reason = "every bound names its byte unit at the deployment boundary"
)]
pub(crate) struct ProjectedCorePolicy {
result_release_bytes: u64,
response_frame_bytes: u64,
manifest_bytes: u64,
artifact_file_bytes: u64,
artifact_range_bytes: u64,
source_decode_bytes: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(
clippy::struct_field_names,
reason = "every bound names its byte unit at the deployment boundary"
)]
pub(crate) struct ProjectedCorePolicyInput {
pub(crate) result_release_bytes: u64,
pub(crate) response_frame_bytes: u64,
pub(crate) manifest_bytes: u64,
pub(crate) artifact_file_bytes: u64,
pub(crate) artifact_range_bytes: u64,
pub(crate) source_decode_bytes: u64,
}
impl TryFrom<ProjectedCorePolicyInput> for ProjectedCorePolicy {
type Error = CoreResolutionError;
fn try_from(value: ProjectedCorePolicyInput) -> Result<Self, Self::Error> {
let ProjectedCorePolicyInput {
result_release_bytes,
response_frame_bytes,
manifest_bytes,
artifact_file_bytes,
artifact_range_bytes,
source_decode_bytes,
} = value;
let policy = Self {
result_release_bytes,
response_frame_bytes,
manifest_bytes,
artifact_file_bytes,
artifact_range_bytes,
source_decode_bytes,
};
policy.validate()?;
Ok(policy)
}
}
impl ProjectedCorePolicy {
const fn validate(self) -> Result<(), CoreResolutionError> {
if self.result_release_bytes == 0
|| self.response_frame_bytes == 0
|| self.manifest_bytes == 0
|| self.artifact_file_bytes == 0
|| self.artifact_range_bytes == 0
|| self.source_decode_bytes == 0
|| self.response_frame_bytes > self.result_release_bytes
|| self.artifact_range_bytes > self.artifact_file_bytes
{
return Err(CoreResolutionError::InvalidBounds);
}
Ok(())
}
}
impl Default for ProjectedCorePolicy {
fn default() -> Self {
Self {
result_release_bytes: DEFAULT_CORE_RESULT_RELEASE_BYTES,
response_frame_bytes: DEFAULT_CORE_RESPONSE_FRAME_BYTES,
manifest_bytes: polyc_state::projection::artifact::MAX_MANIFEST_BYTES,
artifact_file_bytes: DEFAULT_CORE_ARTIFACT_FILE_BYTES,
artifact_range_bytes: DEFAULT_CORE_ARTIFACT_RANGE_BYTES,
source_decode_bytes: DEFAULT_CORE_SOURCE_DECODE_BYTES,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct EffectiveCoreBounds {
timeout: Duration,
rows: u64,
result_release_bytes: u64,
response_frame_bytes: u64,
manifest_bytes: u64,
artifact_file_bytes: u64,
artifact_range_bytes: u64,
source_decode_bytes: u64,
}
impl EffectiveCoreBounds {
fn mint(
limits: &QueryLimits,
policy: ProjectedCorePolicy,
requested: CoreRequestedBounds,
) -> Result<Self, CoreResolutionError> {
let query_row_ceiling = u64::try_from(limits.row_cap).unwrap_or(u64::MAX);
let effective = Self {
timeout: limits.timeout.min(requested.timeout),
rows: query_row_ceiling.min(requested.rows),
result_release_bytes: policy
.result_release_bytes
.min(requested.result_release_bytes),
response_frame_bytes: policy
.response_frame_bytes
.min(requested.response_frame_bytes),
manifest_bytes: policy.manifest_bytes,
artifact_file_bytes: policy.artifact_file_bytes,
artifact_range_bytes: policy.artifact_range_bytes,
source_decode_bytes: policy.source_decode_bytes,
};
if effective.timeout.is_zero()
|| effective.rows == 0
|| effective.result_release_bytes == 0
|| effective.response_frame_bytes == 0
|| effective.manifest_bytes == 0
|| effective.artifact_file_bytes == 0
|| effective.artifact_range_bytes == 0
|| effective.source_decode_bytes == 0
|| effective.artifact_range_bytes > effective.artifact_file_bytes
|| effective.response_frame_bytes > effective.result_release_bytes
{
return Err(CoreResolutionError::InvalidBounds);
}
Ok(effective)
}
fn canonical_bytes(self) -> Vec<u8> {
let mut bytes = BOUNDS_DOMAIN.to_vec();
bytes.extend_from_slice(
&u64::try_from(self.timeout.as_nanos())
.unwrap_or(u64::MAX)
.to_be_bytes(),
);
for value in [
self.rows,
self.result_release_bytes,
self.response_frame_bytes,
self.manifest_bytes,
self.artifact_file_bytes,
self.artifact_range_bytes,
self.source_decode_bytes,
] {
bytes.extend_from_slice(&value.to_be_bytes());
}
bytes
}
pub(crate) const fn timeout(self) -> Duration {
self.timeout
}
pub(crate) const fn rows(self) -> u64 {
self.rows
}
pub(crate) const fn response_frame_bytes(self) -> u64 {
self.response_frame_bytes
}
pub(crate) const fn result_release_bytes(self) -> u64 {
self.result_release_bytes
}
pub(crate) const fn manifest_bytes(self) -> u64 {
self.manifest_bytes
}
pub(crate) const fn artifact_file_bytes(self) -> u64 {
self.artifact_file_bytes
}
pub(crate) const fn artifact_range_bytes(self) -> u64 {
self.artifact_range_bytes
}
pub(crate) const fn source_decode_bytes(self) -> u64 {
self.source_decode_bytes
}
}
pub(crate) struct CoreAuditContext {
query: QueryId,
requester: RequesterId,
bounds: EffectiveCoreBounds,
operation: CoreOperationContext,
}
impl fmt::Debug for CoreAuditContext {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CoreAuditContext")
.field("bounds", &self.bounds)
.finish_non_exhaustive()
}
}
impl CoreAuditContext {
pub(crate) fn from_scoped(
query: QueryId,
requester: RequesterId,
declared: &DeclaredCall,
bounds: EffectiveCoreBounds,
) -> Self {
Self {
query,
requester,
bounds,
operation: CoreOperationContext::from_declared(declared, bounds.timeout),
}
}
}
#[derive(Debug)]
pub(crate) struct CoreOperationContext {
context: CallContext,
clock: ProductionClock,
audience: Audience,
}
impl CoreOperationContext {
fn from_declared(declared: &DeclaredCall, timeout_ceiling: Duration) -> Self {
let clock = ProductionClock::new();
let mut clamped = declared.clone();
clamped.budget = clamped.budget.min(timeout_ceiling);
let context = clamped.origin_relative_context().in_frame(clock.now());
Self {
context,
clock,
audience: polyc_state_connect::state_audience(),
}
}
#[cfg(test)]
#[cfg(test)]
pub(crate) fn for_test(timeout: Duration) -> Self {
Self::from_declared(
&DeclaredCall::live(polyc_state_connect::state_audience(), timeout),
timeout,
)
}
pub(crate) fn check(&self) -> Result<(), CoreResolutionError> {
self.context
.check(self.clock.now(), &core_operation_family())
.map_err(CoreResolutionError::from)
}
pub(crate) fn remaining(&self) -> Result<Duration, CoreResolutionError> {
self.check()?;
Ok(self.context.remaining(self.clock.now()))
}
pub(crate) fn declared(&self) -> Result<DeclaredCall, CoreResolutionError> {
self.check()?;
Ok(DeclaredCall::bounded(
self.audience.clone(),
self.context.remaining(self.clock.now()),
))
}
pub(crate) fn local_context(&self) -> Result<&CallContext, CoreResolutionError> {
self.check()?;
Ok(&self.context)
}
}
#[derive(Debug)]
pub(crate) struct CoreCompletionContext(CoreOperationContext);
impl CoreCompletionContext {
pub(crate) fn server_owned(timeout: Duration) -> Self {
let declared = DeclaredCall::live(polyc_state_connect::state_audience(), timeout);
Self(CoreOperationContext::from_declared(&declared, timeout))
}
pub(crate) fn check(&self) -> Result<(), CoreResolutionError> {
self.0.check()
}
pub(crate) fn remaining(&self) -> Result<Duration, CoreResolutionError> {
self.0.remaining()
}
pub(crate) fn declared(&self) -> Result<DeclaredCall, CoreResolutionError> {
self.0.declared()
}
pub(crate) fn local_context(&self) -> Result<&CallContext, CoreResolutionError> {
self.0.local_context()
}
}
fn core_operation_family() -> OperationFamily {
OperationFamily::new("query.conversation-core.resolve")
}
#[async_trait]
pub(crate) trait CoreMetadataAuthority: Send + Sync {
async fn create_directory_snapshot(
&self,
operation: &CoreOperationContext,
) -> Result<JournalDirectorySnapshot, CoreResolutionError>;
async fn directory_page(
&self,
operation: &CoreOperationContext,
request: ListJournalDirectorySnapshot,
) -> Result<JournalDirectoryPage, CoreResolutionError>;
async fn release_directory_snapshot(
&self,
operation: &CoreOperationContext,
request: ReleaseJournalDirectorySnapshot,
) -> Result<(), CoreResolutionError>;
async fn source_head(
&self,
operation: &CoreOperationContext,
request: GetJournalSource,
) -> Result<Option<JournalSourceHead>, CoreResolutionError>;
async fn resolve_manifest(
&self,
operation: &CoreOperationContext,
request: ResolveManifest,
) -> Result<ProjectionResolution, CoreResolutionError>;
async fn begin_audit(
&self,
operation: &CoreOperationContext,
command: BeginQueryAudit,
) -> Result<BeginOutcome<CoreExecutionPermit>, CoreResolutionError>;
async fn complete_audit(
&self,
operation: &CoreCompletionContext,
command: &CoreCompletionCommand,
) -> Result<Receipt, CoreResolutionError>;
async fn completion_receipt(
&self,
operation: &CoreCompletionContext,
command: &CoreCompletionCommand,
) -> Result<Option<Receipt>, CoreResolutionError>;
}
#[derive(Clone, PartialEq, Eq)]
pub(crate) enum CoreCompletionCommand {
Local(CompleteQueryAudit),
Remote(RemoteCompleteQueryAudit),
}
impl fmt::Debug for CoreCompletionCommand {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Local(_) => "local",
Self::Remote(_) => "remote",
})
}
}
impl CoreCompletionCommand {
pub(crate) const fn metadata(&self) -> &polyc_state::command::CommandMetadata {
match self {
Self::Local(command) => command.metadata(),
Self::Remote(command) => command.metadata(),
}
}
pub(crate) const fn completion(&self) -> &QueryCompletion {
match self {
Self::Local(command) => command.completion(),
Self::Remote(command) => command.completion(),
}
}
}
#[derive(PartialEq, Eq)]
pub(crate) enum CoreExecutionPermit {
Local(ExecutionPermit),
Remote(RemoteExecutionPermit),
}
impl fmt::Debug for CoreExecutionPermit {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Local(_) => "local",
Self::Remote(_) => "remote",
})
}
}
impl From<ExecutionPermit> for CoreExecutionPermit {
fn from(value: ExecutionPermit) -> Self {
Self::Local(value)
}
}
impl From<RemoteExecutionPermit> for CoreExecutionPermit {
fn from(value: RemoteExecutionPermit) -> Self {
Self::Remote(value)
}
}
impl CoreExecutionPermit {
const fn query(&self) -> &QueryId {
match self {
Self::Local(permit) => permit.query(),
Self::Remote(permit) => permit.query(),
}
}
const fn namespace(&self) -> &NamespaceId {
match self {
Self::Local(permit) => permit.namespace(),
Self::Remote(permit) => permit.namespace(),
}
}
pub(crate) const fn source(&self) -> &SourceSnapshot {
match self {
Self::Local(permit) => permit.source(),
Self::Remote(permit) => permit.source(),
}
}
pub(crate) fn into_completion(
self,
completion: QueryCompletion,
) -> Result<CoreCompletionCommand, CoreResolutionError> {
match self {
Self::Local(permit) => {
let digest = digest(&permit.completion_canonical_bytes(&completion));
Ok(CoreCompletionCommand::Local(CompleteQueryAudit::new(
permit,
completion,
digest,
audit_envelope(),
)))
}
Self::Remote(permit) => {
let digest = digest(&permit.completion_canonical_bytes(&completion));
Ok(CoreCompletionCommand::Remote(permit.into_completion(
completion,
digest,
audit_envelope(),
)?))
}
}
}
}
pub(crate) struct CorePlanningAuthority {
namespace: NamespaceId,
projection_owner: OwnerId,
policy: ProjectedCorePolicy,
metadata: Arc<dyn CoreMetadataAuthority>,
}
impl fmt::Debug for CorePlanningAuthority {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CorePlanningAuthority")
.field("namespace", &self.namespace)
.field("projection_owner", &self.projection_owner)
.field("policy", &self.policy)
.finish_non_exhaustive()
}
}
impl CorePlanningAuthority {
pub(crate) fn new(
namespace: NamespaceId,
projection_owner: OwnerId,
policy: ProjectedCorePolicy,
metadata: Arc<dyn CoreMetadataAuthority>,
) -> Result<Self, CoreResolutionError> {
if namespace.is_empty() || projection_owner.is_empty() {
return Err(CoreResolutionError::InvalidComposition);
}
policy.validate()?;
Ok(Self {
namespace,
projection_owner,
policy,
metadata,
})
}
pub(crate) fn effective_bounds(
&self,
limits: &QueryLimits,
requested: CoreRequestedBounds,
) -> Result<EffectiveCoreBounds, CoreResolutionError> {
EffectiveCoreBounds::mint(limits, self.policy, requested)
}
pub(crate) async fn plan(
&self,
compiler: &CatalogCompiler,
limits: &QueryLimits,
scope: &QueryScope,
allow_explain: bool,
audit: CoreAuditContext,
request: CoreQueryRequest,
) -> Result<CorePlanOutcome, CoreResolutionError> {
if let CoreConsistency::RequireProjectedThrough(position) = request.consistency {
return Err(CoreResolutionError::FreshnessUnsupported { position });
}
if audit.bounds != self.effective_bounds(limits, request.requested_bounds)? {
return Err(CoreResolutionError::InvalidBounds);
}
audit.operation.check()?;
let logical = compiler
.compile(&request.sql, &request.parameters, allow_explain)
.await?;
let partitions = self.authorized_partitions(scope, &audit.operation).await?;
let projected = !logical.dependencies.is_empty();
let journal = !logical.legacy_dependencies.is_empty();
validate_source_pin_count(partitions.len(), projected, journal)?;
let sources = self.resolve_sources(&partitions, &audit.operation).await?;
let manifests = if projected {
self.resolve_all(&partitions, &sources, &audit.operation)
.await?
} else {
Vec::new()
};
let realm = CoreRealm::from_scope(scope);
let source = hybrid_source_snapshot(&manifests, &sources, journal)?;
let shape = shape_digest(
&audit,
&request,
&logical,
&partitions,
realm,
&self.namespace,
&self.projection_owner,
&source,
audit.bounds,
);
let placeholder = ContentDigest::from_bytes([0; ContentDigest::LEN]);
let envelope = audit_envelope();
let draft = BeginQueryAudit::new(
audit.query.clone(),
self.namespace.clone(),
audit.requester.clone(),
shape,
source.clone(),
placeholder,
envelope.clone(),
);
let expected_query = audit.query.clone();
let command = BeginQueryAudit::new(
audit.query,
self.namespace.clone(),
audit.requester,
shape,
source.clone(),
digest(&draft.canonical_bytes()),
envelope,
);
audit.operation.check()?;
match self.metadata.begin_audit(&audit.operation, command).await? {
BeginOutcome::Granted(permit) => {
let crossed = permit.query() != &expected_query
|| permit.namespace() != &self.namespace
|| permit.source() != &source;
let guardian = PermitGuardian::new(permit, Arc::clone(&self.metadata));
if crossed {
guardian.abandon();
return Err(CoreResolutionError::CrossedPermit);
}
Ok(CorePlanOutcome::Granted(Box::new(PreparedCoreQuery {
guardian,
compiled: logical,
manifests,
partitions,
scope: scope.clone(),
realm,
metadata: Arc::clone(&self.metadata),
operation: audit.operation,
bounds: audit.bounds,
})))
}
BeginOutcome::AlreadyRecorded(receipt) => Ok(CorePlanOutcome::AlreadyRecorded(receipt)),
}
}
async fn authorized_partitions(
&self,
scope: &QueryScope,
operation: &CoreOperationContext,
) -> Result<Vec<PartitionId>, CoreResolutionError> {
match scope {
QueryScope::Conversations(conversations) => {
if conversations.iter().any(String::is_empty) {
return Err(CoreResolutionError::EmptyConversationIdentity);
}
let mut partitions = conversations
.iter()
.map(|conversation| {
PartitionId::new(format!("{CORE_PARTITION_PREFIX}{conversation}"))
})
.collect::<Vec<_>>();
canonical_partitions(&mut partitions)?;
Ok(partitions)
}
QueryScope::Fleet => self.fleet_partitions(operation).await,
}
}
async fn fleet_partitions(
&self,
operation: &CoreOperationContext,
) -> Result<Vec<PartitionId>, CoreResolutionError> {
operation.check()?;
let snapshot = self.metadata.create_directory_snapshot(operation).await?;
let id = snapshot.id().clone();
let result = self.read_snapshot(&snapshot, operation).await;
let release = self
.metadata
.release_directory_snapshot(operation, ReleaseJournalDirectorySnapshot::new(id.clone()))
.await;
match (result, release) {
(Ok(partitions), Ok(())) => Ok(partitions),
(Err(error), _) | (Ok(_), Err(error)) => Err(error),
}
}
async fn read_snapshot(
&self,
snapshot: &JournalDirectorySnapshot,
operation: &CoreOperationContext,
) -> Result<Vec<PartitionId>, CoreResolutionError> {
let mut partitions = Vec::new();
let mut observed = 0_u64;
let mut after = None;
loop {
let mut request = ListJournalDirectorySnapshot::new(
snapshot.id().clone(),
MAX_DIRECTORY_PAGE_PARTITIONS,
);
if let Some(cursor) = after.take() {
request = request.after(cursor);
}
operation.check()?;
let page = self
.metadata
.directory_page(operation, request.clone())
.await?;
observed = snapshot.validate_page(&request, &page, observed)?;
partitions.extend(
page.partitions()
.iter()
.filter(|partition| is_conversation_partition(partition))
.cloned(),
);
if partitions.len() > MAX_SOURCE_PINS as usize {
return Err(source_bound(partitions.len()));
}
if page.is_truncated() {
after = page.next_after().cloned();
} else {
break;
}
}
canonical_partitions(&mut partitions)?;
Ok(partitions)
}
async fn resolve_all(
&self,
partitions: &[PartitionId],
sources: &[JournalSourceHead],
operation: &CoreOperationContext,
) -> Result<Vec<ProjectionManifest>, CoreResolutionError> {
if partitions.len() != sources.len() {
return Err(CoreResolutionError::SourceVectorMismatch);
}
let family = conversation_core();
let versions = family.versions();
let mut manifests = Vec::with_capacity(partitions.len());
for (partition, source) in partitions.iter().zip(sources) {
if source.source().partition() != partition {
return Err(CoreResolutionError::SourceMismatch(partition.clone()));
}
let key = ProjectionKey::new(FamilyId::new(family.family_str()), partition.clone());
operation.check()?;
let resolution = self
.metadata
.resolve_manifest(
operation,
ResolveManifest::new(
key.clone(),
source.source().clone(),
self.projection_owner.clone(),
),
)
.await?;
let manifest = resolution.current().ok_or_else(|| {
if resolution.is_superseded() {
CoreResolutionError::Superseded(partition.clone())
} else {
CoreResolutionError::MissingProjection(partition.clone())
}
})?;
if manifest.key() != &key
|| manifest.checkpoint().source() != source.source()
|| manifest.schema_version() != versions.schema().get()
|| manifest.fact_version() != versions.fact_model().get()
|| manifest.object_descriptor().owner() != &self.projection_owner
|| manifest.object_descriptor().classification() != Classification::Confidential
{
return Err(CoreResolutionError::IncompatibleDescriptor(
partition.clone(),
));
}
manifest.validate_structure()?;
manifests.push(manifest.clone());
}
manifests.sort_by(|left, right| left.key().cmp(right.key()));
if manifests
.windows(2)
.any(|pair| pair[0].key() >= pair[1].key())
{
return Err(CoreResolutionError::DuplicateDescriptor);
}
Ok(manifests)
}
async fn resolve_sources(
&self,
partitions: &[PartitionId],
operation: &CoreOperationContext,
) -> Result<Vec<JournalSourceHead>, CoreResolutionError> {
let mut sources = Vec::with_capacity(partitions.len());
for partition in partitions {
operation.check()?;
let source = self
.metadata
.source_head(operation, GetJournalSource::new(partition.clone()))
.await?
.ok_or_else(|| CoreResolutionError::MissingSource(partition.clone()))?;
if source.source().partition() != partition {
return Err(CoreResolutionError::SourceMismatch(partition.clone()));
}
sources.push(source);
}
Ok(sources)
}
}
fn canonical_partitions(partitions: &mut [PartitionId]) -> Result<(), CoreResolutionError> {
if partitions.len() > MAX_SOURCE_PINS as usize {
return Err(source_bound(partitions.len()));
}
partitions.sort();
if partitions.iter().any(PartitionId::is_empty)
|| partitions.windows(2).any(|pair| pair[0] >= pair[1])
{
return Err(CoreResolutionError::DuplicatePartition);
}
Ok(())
}
fn is_conversation_partition(partition: &PartitionId) -> bool {
partition
.as_str()
.strip_prefix(CORE_PARTITION_PREFIX)
.is_some_and(|suffix| !suffix.is_empty())
}
fn source_bound(requested: usize) -> CoreResolutionError {
StateError::BoundsExceeded {
bound: BoundKind::CommandRecords,
limit: u64::from(MAX_SOURCE_PINS),
requested: u64::try_from(requested).unwrap_or(u64::MAX),
}
.into()
}
fn hybrid_source_snapshot(
manifests: &[ProjectionManifest],
sources: &[JournalSourceHead],
journal: bool,
) -> Result<SourceSnapshot, CoreResolutionError> {
let mut pins = manifests
.iter()
.cloned()
.map(ProjectionPin::new)
.map(SourcePin::Projected)
.collect::<Vec<_>>();
if journal {
pins.extend(sources.iter().map(|source| {
SourcePin::Journal(JournalAnchor::new(
source.source().clone(),
source.position(),
))
}));
}
SourceSnapshot::try_new(pins).map_err(CoreResolutionError::from)
}
fn validate_source_pin_count(
partitions: usize,
projected: bool,
journal: bool,
) -> Result<(), CoreResolutionError> {
let kinds = usize::from(projected) + usize::from(journal);
let requested = partitions.saturating_mul(kinds);
if requested > MAX_SOURCE_PINS as usize {
return Err(source_bound(requested));
}
Ok(())
}
fn digest(bytes: &[u8]) -> ContentDigest {
ContentDigest::from_bytes(*blake3::hash(bytes).as_bytes())
}
fn audit_envelope() -> CommandEnvelope {
CommandEnvelope::new(
Purpose::new("conversation-core-query"),
Audience::new("state"),
polyc_state::query_audit::command_bounds(),
)
}
fn push(bytes: &mut Vec<u8>, value: &[u8]) {
bytes.extend_from_slice(&u64::try_from(value.len()).unwrap_or(u64::MAX).to_be_bytes());
bytes.extend_from_slice(value);
}
fn push_parameters(bytes: &mut Vec<u8>, parameters: &[CoreParameter]) {
bytes.extend_from_slice(
&u64::try_from(parameters.len())
.unwrap_or(u64::MAX)
.to_be_bytes(),
);
for parameter in parameters {
match parameter {
CoreParameter::Utf8(value) => {
bytes.push(0);
push(bytes, value.as_bytes());
}
CoreParameter::UInt64(value) => {
bytes.push(1);
bytes.extend_from_slice(&value.to_be_bytes());
}
CoreParameter::Boolean(value) => {
bytes.push(2);
bytes.push(u8::from(*value));
}
CoreParameter::Null => bytes.push(3),
}
}
}
#[allow(clippy::too_many_arguments)]
fn shape_digest(
audit: &CoreAuditContext,
request: &CoreQueryRequest,
compiled: &CompiledCoreQuery,
partitions: &[PartitionId],
realm: CoreRealm,
namespace: &NamespaceId,
owner: &OwnerId,
source: &SourceSnapshot,
bounds: EffectiveCoreBounds,
) -> ContentDigest {
let mut bytes = SHAPE_DOMAIN.to_vec();
push(&mut bytes, compiled.normalized_plan.as_bytes());
bytes.push(match compiled.statement {
AllowedStatement::Query => 0,
AllowedStatement::Explain => 1,
});
bytes.push(u8::from(compiled.explain_enabled));
bytes.extend_from_slice(
&u64::try_from(compiled.dependencies.len())
.unwrap_or(u64::MAX)
.to_be_bytes(),
);
for dependency in &compiled.dependencies {
push(&mut bytes, dependency.name().as_bytes());
}
bytes.extend_from_slice(
&u64::try_from(compiled.legacy_dependencies.len())
.unwrap_or(u64::MAX)
.to_be_bytes(),
);
for dependency in &compiled.legacy_dependencies {
push(&mut bytes, dependency.name().as_bytes());
bytes.extend_from_slice(dependency.schema_fingerprint().as_bytes());
}
push(&mut bytes, &conversation_core().fingerprint());
push(&mut bytes, namespace.as_str().as_bytes());
push(&mut bytes, owner.as_str().as_bytes());
push(&mut bytes, audit.query.as_str().as_bytes());
push(&mut bytes, audit.requester.as_str().as_bytes());
push_parameters(&mut bytes, &request.parameters);
bytes.push(match realm {
CoreRealm::Visible => 0,
CoreRealm::Fleet => 1,
});
bytes.extend_from_slice(
&u64::try_from(partitions.len())
.unwrap_or(u64::MAX)
.to_be_bytes(),
);
for partition in partitions {
push(&mut bytes, partition.as_str().as_bytes());
}
match request.consistency {
CoreConsistency::Projected => bytes.push(0),
CoreConsistency::RequireProjectedThrough(position) => {
bytes.push(1);
bytes.extend_from_slice(&position.get().to_be_bytes());
}
}
push(&mut bytes, &source.canonical_bytes());
push(&mut bytes, &bounds.canonical_bytes());
digest(&bytes)
}
#[derive(Debug)]
pub(crate) enum CorePlanOutcome {
Granted(Box<PreparedCoreQuery>),
AlreadyRecorded(Box<Receipt>),
}
pub(crate) struct PreparedCoreQuery {
guardian: PermitGuardian,
compiled: CompiledCoreQuery,
manifests: Vec<ProjectionManifest>,
partitions: Vec<PartitionId>,
scope: QueryScope,
realm: CoreRealm,
metadata: Arc<dyn CoreMetadataAuthority>,
operation: CoreOperationContext,
bounds: EffectiveCoreBounds,
}
impl fmt::Debug for PreparedCoreQuery {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("PreparedCoreQuery")
.field("guardian", &self.guardian)
.field("compiled", &self.compiled)
.field("manifests", &self.manifests.len())
.field("partitions", &self.partitions.len())
.field("realm", &self.realm)
.field("bounds", &self.bounds)
.finish_non_exhaustive()
}
}
impl PreparedCoreQuery {
#[cfg(test)]
pub(crate) fn pause_completion_dispatch(&self) -> crate::core_execution::GuardianDispatchPause {
self.guardian.pause_dispatch()
}
pub(crate) fn into_parts(self) -> PreparedCoreParts {
let Self {
guardian,
compiled,
manifests,
partitions,
scope,
realm,
metadata,
operation,
bounds,
} = self;
PreparedCoreParts {
guardian,
compiled,
manifests,
partitions,
scope,
realm,
metadata,
operation,
bounds,
}
}
}
pub(crate) struct PreparedCoreParts {
pub(crate) guardian: PermitGuardian,
pub(crate) compiled: CompiledCoreQuery,
pub(crate) manifests: Vec<ProjectionManifest>,
pub(crate) partitions: Vec<PartitionId>,
pub(crate) scope: QueryScope,
pub(crate) realm: CoreRealm,
pub(crate) metadata: Arc<dyn CoreMetadataAuthority>,
pub(crate) operation: CoreOperationContext,
pub(crate) bounds: EffectiveCoreBounds,
}
#[derive(thiserror::Error)]
pub(crate) enum CoreResolutionError {
#[error("projected conversation-core planning is not composed")]
Unavailable,
#[error("the query statement was refused")]
Statement(String),
#[error("schema-only planning failed")]
DataFusion(#[from] DataFusionError),
#[error("query has no declared projected or journal dependency")]
NoSourceDependency,
#[error("the query names a table outside the conversation-core catalog")]
UnknownDependency(String),
#[error("the typed parameter vector does not match the SQL placeholders")]
ParameterMismatch,
#[error("core planning composition has an empty namespace or owner")]
InvalidComposition,
#[error("the verified query scope has no durable audit attribution")]
InvalidAttribution,
#[error("the verified conversation identity is empty")]
EmptyConversationIdentity,
#[error("the requested projection freshness is not implemented")]
FreshnessUnsupported { position: JournalPosition },
#[error("projected query bounds are empty, crossed, or exceed their parent posture")]
InvalidBounds,
#[error("audit authority returned a permit for another query, tenant, or source")]
CrossedPermit,
#[error("the durable query completion disagrees with the presented command")]
CompletionReceiptMismatch,
#[error("the current source is absent for an authorized partition")]
MissingSource(PartitionId),
#[error("the current source vector does not match the authorized partitions")]
SourceVectorMismatch,
#[error("a current source response named another partition")]
SourceMismatch(PartitionId),
#[error("the current projection is absent for an authorized partition")]
MissingProjection(PartitionId),
#[error("the current projection belongs to a recreated source")]
Superseded(PartitionId),
#[error("the projection descriptor is incompatible with this build")]
IncompatibleDescriptor(PartitionId),
#[error("authority scope contains a duplicate or empty partition")]
DuplicatePartition,
#[error("resolved descriptors contain a duplicate key")]
DuplicateDescriptor,
#[error("State metadata refused descriptor planning")]
State(#[from] StateError),
#[error("State projection catalog refused descriptor planning")]
Projection(#[from] ProjectionCatalogError),
#[error("State query audit refused descriptor planning")]
Audit(#[from] QueryAuditError),
}
impl fmt::Debug for CoreResolutionError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Unavailable => "Unavailable",
Self::Statement(_) => "Statement",
Self::DataFusion(_) => "DataFusion",
Self::NoSourceDependency => "NoSourceDependency",
Self::UnknownDependency(_) => "UnknownDependency",
Self::ParameterMismatch => "ParameterMismatch",
Self::InvalidComposition => "InvalidComposition",
Self::InvalidAttribution => "InvalidAttribution",
Self::EmptyConversationIdentity => "EmptyConversationIdentity",
Self::FreshnessUnsupported { .. } => "FreshnessUnsupported",
Self::InvalidBounds => "InvalidBounds",
Self::CrossedPermit => "CrossedPermit",
Self::CompletionReceiptMismatch => "CompletionReceiptMismatch",
Self::MissingSource(_) => "MissingSource",
Self::SourceVectorMismatch => "SourceVectorMismatch",
Self::SourceMismatch(_) => "SourceMismatch",
Self::MissingProjection(_) => "MissingProjection",
Self::Superseded(_) => "Superseded",
Self::IncompatibleDescriptor(_) => "IncompatibleDescriptor",
Self::DuplicatePartition => "DuplicatePartition",
Self::DuplicateDescriptor => "DuplicateDescriptor",
Self::State(_) => "State",
Self::Projection(_) => "Projection",
Self::Audit(_) => "Audit",
})
}
}
#[cfg(test)]
mod tests;