1use std::fmt::Write as _;
7use std::num::{NonZeroU32, NonZeroU64};
8use std::ops::Bound;
9
10use crate::{
11 db::{
12 access::IndexBranchSetOrderedSuffix,
13 query::plan::{
14 AccessPlanProjection, AccessPlannedQuery, GroupPlan, PrimaryKeyInputResourceSummary,
15 QueryMode, ResidualFilterShape, ScalarPlan, project_access_plan,
16 },
17 },
18 value::Value,
19};
20use icydb_diagnostic_code::{
21 Diagnostic, DiagnosticCode, DiagnosticDetail, ErrorCode, ErrorOrigin, QueryReadAdmissionCode,
22};
23
24pub(in crate::db::query) const DEFAULT_BOUNDED_READ_MAX_ROWS: u32 = 100;
25pub(in crate::db::query) const DEFAULT_BOUNDED_READ_RESPONSE_BYTES: u32 = 128 * 1024;
26const DEFAULT_BOUNDED_READ_MAX_GROUPS: u32 = 100;
27const DEFAULT_BOUNDED_READ_MAX_GROUP_BYTES: u32 = 64 * 1024;
28const DEFAULT_BOUNDED_READ_MAX_DISTINCT_ENTRIES: u32 = 1024;
29const DEFAULT_BOUNDED_READ_MAX_PRIMARY_KEY_INPUT_TERMS: u32 = 1024;
30const DEFAULT_BOUNDED_READ_MAX_PRIMARY_KEY_INPUT_BYTES: u32 = 64 * 1024;
31
32const fn non_zero_default(value: u32) -> NonZeroU32 {
33 match NonZeroU32::new(value) {
34 Some(value) => value,
35 None => NonZeroU32::MIN,
36 }
37}
38
39#[derive(Clone, Copy, Debug, Eq, PartialEq)]
41pub enum QueryAdmissionLane {
42 PublicRead,
44 AdminAdHoc,
46 DiagnosticExplain,
48 DevTest,
50}
51
52impl QueryAdmissionLane {
53 #[must_use]
55 pub const fn as_str(self) -> &'static str {
56 match self {
57 Self::PublicRead => "public_read",
58 Self::AdminAdHoc => "admin_ad_hoc",
59 Self::DiagnosticExplain => "diagnostic_explain",
60 Self::DevTest => "dev_test",
61 }
62 }
63
64 #[must_use]
66 pub const fn executes_rows(self) -> bool {
67 !matches!(self, Self::DiagnosticExplain)
68 }
69}
70
71#[derive(Clone, Copy, Debug, Eq, PartialEq)]
73pub enum QueryBoundKind {
74 Exact,
76 ConservativeUpperBound,
78 EnforcedRuntimeCap,
80 EstimateOnly,
82 Unavailable,
84}
85
86impl QueryBoundKind {
87 #[must_use]
89 pub const fn as_str(self) -> &'static str {
90 match self {
91 Self::Exact => "exact",
92 Self::ConservativeUpperBound => "conservative_upper_bound",
93 Self::EnforcedRuntimeCap => "enforced_runtime_cap",
94 Self::EstimateOnly => "estimate_only",
95 Self::Unavailable => "unavailable",
96 }
97 }
98
99 #[must_use]
101 pub const fn admits_public_read(self) -> bool {
102 matches!(
103 self,
104 Self::Exact | Self::ConservativeUpperBound | Self::EnforcedRuntimeCap
105 )
106 }
107}
108
109#[derive(Clone, Copy, Debug, Eq, PartialEq)]
111pub enum QueryAdmissionDecision {
112 Admitted,
114 Rejected,
116}
117
118impl QueryAdmissionDecision {
119 #[must_use]
121 pub const fn as_str(self) -> &'static str {
122 match self {
123 Self::Admitted => "admitted",
124 Self::Rejected => "rejected",
125 }
126 }
127
128 #[must_use]
130 pub const fn is_admitted(self) -> bool {
131 matches!(self, Self::Admitted)
132 }
133}
134
135#[derive(Clone, Copy, Debug, Eq, PartialEq)]
137pub enum QueryAdmissionAccessKind {
138 Unknown,
140 ByKey,
142 ByKeys,
144 KeyRange,
146 IndexPrefix,
148 IndexMultiLookup,
150 IndexBranchSet,
152 IndexRange,
154 FullScan,
156 Union,
158 Intersection,
160}
161
162impl QueryAdmissionAccessKind {
163 #[must_use]
165 pub const fn as_str(self) -> &'static str {
166 match self {
167 Self::Unknown => "unknown",
168 Self::ByKey => "by_key",
169 Self::ByKeys => "by_keys",
170 Self::KeyRange => "key_range",
171 Self::IndexPrefix => "index_prefix",
172 Self::IndexMultiLookup => "index_multi_lookup",
173 Self::IndexBranchSet => "index_branch_set",
174 Self::IndexRange => "index_range",
175 Self::FullScan => "full_scan",
176 Self::Union => "union",
177 Self::Intersection => "intersection",
178 }
179 }
180
181 #[must_use]
183 pub const fn is_secondary_index(self) -> bool {
184 matches!(
185 self,
186 Self::IndexPrefix | Self::IndexMultiLookup | Self::IndexBranchSet | Self::IndexRange
187 )
188 }
189
190 #[must_use]
192 pub const fn is_full_scan(self) -> bool {
193 matches!(self, Self::FullScan)
194 }
195}
196
197#[derive(Clone, Copy, Debug, Eq, PartialEq)]
199pub enum QueryAdmissionPlanShape {
200 ScalarRead,
202 GroupedAggregate,
204 Delete,
206}
207
208impl QueryAdmissionPlanShape {
209 #[must_use]
211 pub const fn as_str(self) -> &'static str {
212 match self {
213 Self::ScalarRead => "scalar_read",
214 Self::GroupedAggregate => "grouped_aggregate",
215 Self::Delete => "delete",
216 }
217 }
218}
219
220#[derive(Clone, Copy, Debug, Eq, PartialEq)]
222pub enum QueryAdmissionResidualFilter {
223 Absent,
225 Predicate,
227 Expression,
229 ExpressionAndPredicate,
231}
232
233impl QueryAdmissionResidualFilter {
234 #[must_use]
236 pub const fn as_str(self) -> &'static str {
237 match self {
238 Self::Absent => "none",
239 Self::Predicate => "predicate",
240 Self::Expression => "expression",
241 Self::ExpressionAndPredicate => "expression_and_predicate",
242 }
243 }
244
245 #[must_use]
247 pub const fn is_absent(self) -> bool {
248 matches!(self, Self::Absent)
249 }
250}
251
252#[derive(Clone, Copy, Debug, Eq, PartialEq)]
254pub enum QueryAdmissionOrdering {
255 None,
257 Requested,
259 Resolved,
261}
262
263impl QueryAdmissionOrdering {
264 #[must_use]
266 pub const fn as_str(self) -> &'static str {
267 match self {
268 Self::None => "none",
269 Self::Requested => "requested",
270 Self::Resolved => "resolved",
271 }
272 }
273}
274
275#[derive(Clone, Copy, Debug, Eq, PartialEq)]
277pub struct QueryAdmissionGroupedSummary {
278 group_field_count: u32,
279 aggregate_count: u32,
280 distinct_aggregate_count: u32,
281 max_groups: u64,
282 max_group_bytes: u64,
283 having_filter: bool,
284}
285
286impl QueryAdmissionGroupedSummary {
287 #[must_use]
289 pub const fn new(
290 group_field_count: u32,
291 aggregate_count: u32,
292 distinct_aggregate_count: u32,
293 max_groups: u64,
294 max_group_bytes: u64,
295 having_filter: bool,
296 ) -> Self {
297 Self {
298 group_field_count,
299 aggregate_count,
300 distinct_aggregate_count,
301 max_groups,
302 max_group_bytes,
303 having_filter,
304 }
305 }
306
307 #[must_use]
309 pub const fn group_field_count(self) -> u32 {
310 self.group_field_count
311 }
312
313 #[must_use]
315 pub const fn aggregate_count(self) -> u32 {
316 self.aggregate_count
317 }
318
319 #[must_use]
321 pub const fn distinct_aggregate_count(self) -> u32 {
322 self.distinct_aggregate_count
323 }
324
325 #[must_use]
327 pub const fn max_groups(self) -> u64 {
328 self.max_groups
329 }
330
331 #[must_use]
333 pub const fn max_group_bytes(self) -> u64 {
334 self.max_group_bytes
335 }
336
337 #[must_use]
339 pub const fn has_having_filter(self) -> bool {
340 self.having_filter
341 }
342}
343
344#[derive(Clone, Copy, Debug, Eq, PartialEq)]
346pub(in crate::db) struct GroupedAdmissionPolicy {
347 groups: Option<NonZeroU32>,
348 group_bytes: Option<NonZeroU32>,
349 distinct_entries: Option<NonZeroU32>,
350}
351
352impl GroupedAdmissionPolicy {
353 #[must_use]
355 pub(in crate::db) const fn disabled() -> Self {
356 Self {
357 groups: None,
358 group_bytes: None,
359 distinct_entries: None,
360 }
361 }
362
363 #[must_use]
365 pub(in crate::db) const fn bounded(
366 max_groups: NonZeroU32,
367 max_group_bytes: NonZeroU32,
368 max_distinct_entries: Option<NonZeroU32>,
369 ) -> Self {
370 Self {
371 groups: Some(max_groups),
372 group_bytes: Some(max_group_bytes),
373 distinct_entries: max_distinct_entries,
374 }
375 }
376
377 #[must_use]
383 pub(in crate::db) const fn default_bounded_read() -> Self {
384 Self::bounded(
385 non_zero_default(DEFAULT_BOUNDED_READ_MAX_GROUPS),
386 non_zero_default(DEFAULT_BOUNDED_READ_MAX_GROUP_BYTES),
387 Some(non_zero_default(DEFAULT_BOUNDED_READ_MAX_DISTINCT_ENTRIES)),
388 )
389 }
390
391 #[must_use]
393 pub(in crate::db) const fn max_groups(&self) -> Option<NonZeroU32> {
394 self.groups
395 }
396
397 #[must_use]
399 pub(in crate::db) const fn max_group_bytes(&self) -> Option<NonZeroU32> {
400 self.group_bytes
401 }
402
403 #[must_use]
405 pub(in crate::db) const fn max_distinct_entries(&self) -> Option<NonZeroU32> {
406 self.distinct_entries
407 }
408
409 #[must_use]
411 #[cfg(test)]
412 pub(in crate::db) const fn has_hard_limits(&self) -> bool {
413 self.groups.is_some() && self.group_bytes.is_some()
414 }
415
416 #[must_use]
418 #[cfg(all(test, feature = "sql"))]
419 pub(in crate::db) const fn execution_config(
420 &self,
421 ) -> Option<crate::db::query::plan::GroupedExecutionConfig> {
422 match (self.groups, self.group_bytes) {
423 (Some(groups), Some(group_bytes)) => Some(
424 crate::db::query::plan::GroupedExecutionConfig::with_hard_limits(
425 groups.get() as u64,
426 group_bytes.get() as u64,
427 ),
428 ),
429 _ => None,
430 }
431 }
432}
433
434#[derive(Clone, Copy, Debug, Eq, PartialEq)]
435enum LimitRequirement {
436 Required,
437 Optional,
438}
439
440#[derive(Clone, Copy, Debug, Eq, PartialEq)]
441enum IndexRequirement {
442 Required,
443 Optional,
444}
445
446#[derive(Clone, Copy, Debug, Eq, PartialEq)]
447enum FullScanPolicy {
448 Allow,
449 Reject,
450}
451
452#[derive(Clone, Copy, Debug, Eq, PartialEq)]
453enum MaterializedSortPolicy {
454 Allow,
455 Reject,
456}
457
458#[derive(Clone, Copy, Debug, Eq, PartialEq)]
459enum OffsetPolicy {
460 Allow,
461 RejectNonZero,
462}
463
464#[derive(Clone, Debug, Eq, PartialEq)]
466pub(in crate::db) struct QueryAdmissionPolicy {
467 lane: QueryAdmissionLane,
468 limit_requirement: LimitRequirement,
469 max_returned_rows: Option<NonZeroU32>,
470 max_scanned_rows: Option<NonZeroU64>,
471 max_response_bytes: Option<NonZeroU32>,
472 max_primary_key_input_terms: Option<NonZeroU32>,
473 max_primary_key_input_bytes: Option<NonZeroU32>,
474 index_requirement: IndexRequirement,
475 offset_policy: OffsetPolicy,
476 full_scan_policy: FullScanPolicy,
477 materialized_sort_policy: MaterializedSortPolicy,
478 max_materialized_rows: Option<NonZeroU32>,
479 max_projection_columns: Option<NonZeroU32>,
480 grouped: GroupedAdmissionPolicy,
481}
482
483impl QueryAdmissionPolicy {
484 #[must_use]
486 pub(in crate::db) const fn public_read(
487 max_returned_rows: NonZeroU32,
488 max_response_bytes: NonZeroU32,
489 ) -> Self {
490 Self {
491 lane: QueryAdmissionLane::PublicRead,
492 limit_requirement: LimitRequirement::Required,
493 max_returned_rows: Some(max_returned_rows),
494 max_scanned_rows: None,
495 max_response_bytes: Some(max_response_bytes),
496 max_primary_key_input_terms: Some(non_zero_default(
497 DEFAULT_BOUNDED_READ_MAX_PRIMARY_KEY_INPUT_TERMS,
498 )),
499 max_primary_key_input_bytes: Some(non_zero_default(
500 DEFAULT_BOUNDED_READ_MAX_PRIMARY_KEY_INPUT_BYTES,
501 )),
502 index_requirement: IndexRequirement::Required,
503 offset_policy: OffsetPolicy::RejectNonZero,
504 full_scan_policy: FullScanPolicy::Reject,
505 materialized_sort_policy: MaterializedSortPolicy::Reject,
506 max_materialized_rows: None,
507 max_projection_columns: None,
508 grouped: GroupedAdmissionPolicy::disabled(),
509 }
510 }
511
512 #[must_use]
519 pub(in crate::db) const fn default_bounded_read() -> Self {
520 Self::public_read(
521 non_zero_default(DEFAULT_BOUNDED_READ_MAX_ROWS),
522 non_zero_default(DEFAULT_BOUNDED_READ_RESPONSE_BYTES),
523 )
524 .with_grouped_policy(GroupedAdmissionPolicy::default_bounded_read())
525 }
526
527 #[must_use]
532 pub(in crate::db) const fn with_grouped_policy(
533 mut self,
534 grouped: GroupedAdmissionPolicy,
535 ) -> Self {
536 self.grouped = grouped;
537 self
538 }
539
540 #[must_use]
542 #[cfg(test)]
543 pub(in crate::db) const fn admin_ad_hoc(
544 max_returned_rows: NonZeroU32,
545 max_scanned_rows: NonZeroU64,
546 max_response_bytes: NonZeroU32,
547 ) -> Self {
548 Self {
549 lane: QueryAdmissionLane::AdminAdHoc,
550 limit_requirement: LimitRequirement::Optional,
551 max_returned_rows: Some(max_returned_rows),
552 max_scanned_rows: Some(max_scanned_rows),
553 max_response_bytes: Some(max_response_bytes),
554 max_primary_key_input_terms: None,
555 max_primary_key_input_bytes: None,
556 index_requirement: IndexRequirement::Optional,
557 offset_policy: OffsetPolicy::Allow,
558 full_scan_policy: FullScanPolicy::Allow,
559 materialized_sort_policy: MaterializedSortPolicy::Allow,
560 max_materialized_rows: Some(max_returned_rows),
561 max_projection_columns: None,
562 grouped: GroupedAdmissionPolicy::disabled(),
563 }
564 }
565
566 #[must_use]
568 pub(in crate::db) const fn diagnostic_explain() -> Self {
569 Self {
570 lane: QueryAdmissionLane::DiagnosticExplain,
571 limit_requirement: LimitRequirement::Optional,
572 max_returned_rows: None,
573 max_scanned_rows: None,
574 max_response_bytes: None,
575 max_primary_key_input_terms: None,
576 max_primary_key_input_bytes: None,
577 index_requirement: IndexRequirement::Optional,
578 offset_policy: OffsetPolicy::Allow,
579 full_scan_policy: FullScanPolicy::Allow,
580 materialized_sort_policy: MaterializedSortPolicy::Allow,
581 max_materialized_rows: None,
582 max_projection_columns: None,
583 grouped: GroupedAdmissionPolicy::disabled(),
584 }
585 }
586
587 #[must_use]
589 pub(in crate::db) const fn lane(&self) -> QueryAdmissionLane {
590 self.lane
591 }
592
593 #[must_use]
595 pub(in crate::db) const fn require_limit(&self) -> bool {
596 matches!(self.limit_requirement, LimitRequirement::Required)
597 }
598
599 #[must_use]
601 #[cfg(test)]
602 pub(in crate::db) const fn max_returned_rows(&self) -> Option<NonZeroU32> {
603 self.max_returned_rows
604 }
605
606 #[must_use]
608 #[cfg(test)]
609 pub(in crate::db) const fn max_scanned_rows(&self) -> Option<NonZeroU64> {
610 self.max_scanned_rows
611 }
612
613 #[must_use]
615 #[cfg(test)]
616 pub(in crate::db) const fn max_response_bytes(&self) -> Option<NonZeroU32> {
617 self.max_response_bytes
618 }
619
620 #[must_use]
622 pub(in crate::db) const fn require_index(&self) -> bool {
623 matches!(self.index_requirement, IndexRequirement::Required)
624 }
625
626 #[must_use]
628 pub(in crate::db) const fn reject_non_zero_offset(&self) -> bool {
629 matches!(self.offset_policy, OffsetPolicy::RejectNonZero)
630 }
631
632 #[must_use]
634 pub(in crate::db) const fn allow_full_scan(&self) -> bool {
635 matches!(self.full_scan_policy, FullScanPolicy::Allow)
636 }
637
638 #[must_use]
640 pub(in crate::db) const fn allow_materialized_sort(&self) -> bool {
641 matches!(self.materialized_sort_policy, MaterializedSortPolicy::Allow)
642 }
643
644 #[must_use]
646 #[cfg(test)]
647 pub(in crate::db) const fn max_materialized_rows(&self) -> Option<NonZeroU32> {
648 self.max_materialized_rows
649 }
650
651 #[must_use]
653 #[cfg(test)]
654 pub(in crate::db) const fn grouped(&self) -> GroupedAdmissionPolicy {
655 self.grouped
656 }
657
658 #[must_use]
660 #[cfg(test)]
661 pub(in crate::db) const fn public_caps_are_finite(&self) -> bool {
662 !matches!(self.lane, QueryAdmissionLane::PublicRead)
663 || (self.max_returned_rows.is_some() && self.max_response_bytes.is_some())
664 }
665
666 #[must_use]
668 #[cfg(test)]
669 pub(in crate::db) const fn with_primary_key_input_caps(
670 mut self,
671 max_terms: NonZeroU32,
672 max_bytes: NonZeroU32,
673 ) -> Self {
674 self.max_primary_key_input_terms = Some(max_terms);
675 self.max_primary_key_input_bytes = Some(max_bytes);
676 self
677 }
678
679 #[must_use]
681 pub(in crate::db) fn evaluate(
682 &self,
683 mut summary: QueryAdmissionSummary,
684 ) -> QueryAdmissionSummary {
685 summary.lane = self.lane;
686
687 match self.rejection_for_summary(&summary) {
688 Some(rejection) => summary.reject(rejection),
689 None => summary.admit(),
690 }
691 }
692
693 fn rejection_for_summary(
694 &self,
695 summary: &QueryAdmissionSummary,
696 ) -> Option<QueryAdmissionRejection> {
697 if !self.lane.executes_rows() {
698 return Some(QueryAdmissionRejection::DiagnosticLaneDoesNotExecute);
699 }
700
701 if matches!(summary.plan_shape(), QueryAdmissionPlanShape::Delete) {
702 return Some(QueryAdmissionRejection::UnsupportedStatementForQueryLane);
703 }
704
705 if let Some(rejection) = self.grouped_rejection(summary) {
706 return Some(rejection);
707 }
708
709 if !self.allow_full_scan() && summary.selected_access().is_full_scan() {
710 return Some(QueryAdmissionRejection::UnboundedFullScanRejected);
711 }
712
713 if self.require_index()
714 && !access_satisfies_index_requirement(summary.selected_access(), summary.scan_bound())
715 {
716 return Some(QueryAdmissionRejection::PublicQueryRequiresIndex);
717 }
718
719 if self.require_limit()
720 && summary.limit().is_none()
721 && summary.grouped().is_none()
722 && !summary.returned_row_bound_kind().admits_public_read()
723 {
724 return Some(QueryAdmissionRejection::PublicQueryRequiresLimit);
725 }
726
727 if self.reject_non_zero_offset() && summary.offset().unwrap_or_default() != 0 {
728 return Some(QueryAdmissionRejection::PublicQueryOffsetRejected);
729 }
730
731 if let Some(rejection) = self.returned_row_bound_rejection(summary) {
732 return Some(rejection);
733 }
734
735 if let Some(rejection) = self.scan_bound_rejection(summary) {
736 return Some(rejection);
737 }
738
739 if let Some(rejection) = self.primary_key_input_rejection(summary) {
740 return Some(rejection);
741 }
742
743 self.materialization_rejection(summary)
744 }
745
746 fn grouped_rejection(
747 &self,
748 summary: &QueryAdmissionSummary,
749 ) -> Option<QueryAdmissionRejection> {
750 let grouped = summary.grouped()?;
751 let Some(max_groups) = self.grouped.max_groups() else {
752 return Some(QueryAdmissionRejection::GroupedQueryRequiresLimits);
753 };
754 let Some(max_group_bytes) = self.grouped.max_group_bytes() else {
755 return Some(QueryAdmissionRejection::GroupedQueryRequiresLimits);
756 };
757
758 if grouped.max_groups() == u64::MAX || grouped.max_group_bytes() == u64::MAX {
759 return Some(QueryAdmissionRejection::GroupedQueryRequiresLimits);
760 }
761
762 if grouped.max_groups() > u64::from(max_groups.get())
763 || grouped.max_group_bytes() > u64::from(max_group_bytes.get())
764 {
765 return Some(QueryAdmissionRejection::GroupedQueryExceedsBudget);
766 }
767
768 if grouped.distinct_aggregate_count() > 0 && self.grouped.max_distinct_entries().is_none() {
769 return Some(QueryAdmissionRejection::GroupedQueryRequiresLimits);
770 }
771
772 None
773 }
774
775 fn returned_row_bound_rejection(
776 &self,
777 summary: &QueryAdmissionSummary,
778 ) -> Option<QueryAdmissionRejection> {
779 let max_returned_rows = self.max_returned_rows?;
780
781 if matches!(
782 summary.returned_row_bound_kind(),
783 QueryBoundKind::EstimateOnly
784 ) {
785 return Some(QueryAdmissionRejection::EstimatedOnlyBoundRejected);
786 }
787
788 if !summary.returned_row_bound_kind().admits_public_read() {
789 return Some(QueryAdmissionRejection::ScanBoundUnavailable);
790 }
791
792 let Some(returned_row_bound) = summary.returned_row_bound() else {
793 return Some(QueryAdmissionRejection::ScanBoundUnavailable);
794 };
795
796 if returned_row_bound > max_returned_rows.get() {
797 return Some(QueryAdmissionRejection::ReturnedRowBoundExceedsPolicy);
798 }
799
800 None
801 }
802
803 fn scan_bound_rejection(
804 &self,
805 summary: &QueryAdmissionSummary,
806 ) -> Option<QueryAdmissionRejection> {
807 let max_scanned_rows = self.max_scanned_rows?;
808
809 if matches!(summary.scan_bound_kind(), QueryBoundKind::EstimateOnly) {
810 return Some(QueryAdmissionRejection::EstimatedOnlyBoundRejected);
811 }
812
813 if !summary.scan_bound_kind().admits_public_read() {
814 return Some(QueryAdmissionRejection::ScanBoundUnavailable);
815 }
816
817 let Some(scan_bound) = summary.scan_bound() else {
818 return Some(QueryAdmissionRejection::ScanBoundUnavailable);
819 };
820
821 if scan_bound > max_scanned_rows.get() {
822 return Some(QueryAdmissionRejection::ScanBoundExceedsPolicy);
823 }
824
825 None
826 }
827
828 const fn primary_key_input_rejection(
829 &self,
830 summary: &QueryAdmissionSummary,
831 ) -> Option<QueryAdmissionRejection> {
832 if let (Some(bound), Some(max)) = (
833 summary.primary_key_input_terms(),
834 self.max_primary_key_input_terms,
835 ) && bound > max.get()
836 {
837 return Some(QueryAdmissionRejection::PrimaryKeyInputExceedsPolicy);
838 }
839
840 if let (Some(bound), Some(max)) = (
841 summary.primary_key_input_payload_bytes(),
842 self.max_primary_key_input_bytes,
843 ) && bound > max.get()
844 {
845 return Some(QueryAdmissionRejection::PrimaryKeyInputExceedsPolicy);
846 }
847
848 None
849 }
850
851 fn materialization_rejection(
852 &self,
853 summary: &QueryAdmissionSummary,
854 ) -> Option<QueryAdmissionRejection> {
855 if !self.allow_materialized_sort()
856 && summary.materialization().materialized_sort()
857 && !primary_key_materialized_sort_has_exact_candidate_bound(summary)
858 {
859 return Some(QueryAdmissionRejection::SortRequiresMaterialization);
860 }
861
862 let max_materialized_rows = self.max_materialized_rows?;
863 let materialized_rows = summary.materialization().materialized_rows()?;
864
865 if materialized_rows > max_materialized_rows.get() {
866 Some(QueryAdmissionRejection::MaterializationExceedsBudget)
867 } else {
868 None
869 }
870 }
871}
872
873fn primary_key_materialized_sort_has_exact_candidate_bound(
874 summary: &QueryAdmissionSummary,
875) -> bool {
876 if !matches!(
877 summary.selected_access(),
878 QueryAdmissionAccessKind::ByKey | QueryAdmissionAccessKind::ByKeys
879 ) {
880 return false;
881 }
882 if !matches!(summary.scan_bound_kind(), QueryBoundKind::Exact) {
883 return false;
884 }
885 if !summary
886 .materialization()
887 .row_bound_kind()
888 .admits_public_read()
889 {
890 return false;
891 }
892
893 match (
894 summary.scan_bound(),
895 summary.materialization().materialized_rows(),
896 ) {
897 (Some(scan_bound), Some(materialized_rows)) => u64::from(materialized_rows) == scan_bound,
898 _ => false,
899 }
900}
901
902#[derive(Clone, Copy, Debug, Eq, PartialEq)]
904pub struct QueryMaterializationSummary {
905 materialized_sort: bool,
906 materialized_rows: Option<u32>,
907 row_bound_kind: QueryBoundKind,
908}
909
910impl QueryMaterializationSummary {
911 #[must_use]
913 pub const fn none() -> Self {
914 Self {
915 materialized_sort: false,
916 materialized_rows: None,
917 row_bound_kind: QueryBoundKind::Unavailable,
918 }
919 }
920
921 #[must_use]
923 pub const fn sort(materialized_rows: Option<u32>, row_bound_kind: QueryBoundKind) -> Self {
924 Self {
925 materialized_sort: true,
926 materialized_rows,
927 row_bound_kind,
928 }
929 }
930
931 #[must_use]
933 pub const fn materialized_sort(&self) -> bool {
934 self.materialized_sort
935 }
936
937 #[must_use]
939 pub const fn materialized_rows(&self) -> Option<u32> {
940 self.materialized_rows
941 }
942
943 #[must_use]
945 pub const fn row_bound_kind(&self) -> QueryBoundKind {
946 self.row_bound_kind
947 }
948}
949
950#[derive(Clone, Copy, Debug, Eq, PartialEq)]
952pub enum QueryAdmissionRejection {
953 PublicQueryRequiresLimit,
955 PublicQueryRequiresIndex,
957 UnboundedFullScanRejected,
959 ScanBoundUnavailable,
961 ScanBoundExceedsPolicy,
963 EstimatedOnlyBoundRejected,
965 SortRequiresMaterialization,
967 MaterializationExceedsBudget,
969 ProjectionResponseMayExceedLimit,
971 GroupedQueryRequiresLimits,
973 GroupedQueryExceedsBudget,
975 DiagnosticLaneDoesNotExecute,
977 IntrospectionDisabledForLane,
979 UnsupportedStatementForQueryLane,
981 PublicQueryOffsetRejected,
983 ReturnedRowBoundExceedsPolicy,
985 PrimaryKeyInputExceedsPolicy,
987}
988
989impl QueryAdmissionRejection {
990 #[must_use]
992 pub const fn as_str(self) -> &'static str {
993 match self {
994 Self::PublicQueryRequiresLimit => "public_query_requires_limit",
995 Self::PublicQueryRequiresIndex => "public_query_requires_index",
996 Self::UnboundedFullScanRejected => "unbounded_full_scan_rejected",
997 Self::ScanBoundUnavailable => "scan_bound_unavailable",
998 Self::ScanBoundExceedsPolicy => "scan_bound_exceeds_policy",
999 Self::EstimatedOnlyBoundRejected => "estimated_only_bound_rejected",
1000 Self::SortRequiresMaterialization => "sort_requires_materialization",
1001 Self::MaterializationExceedsBudget => "materialization_exceeds_budget",
1002 Self::ProjectionResponseMayExceedLimit => "projection_response_may_exceed_limit",
1003 Self::GroupedQueryRequiresLimits => "grouped_query_requires_limits",
1004 Self::GroupedQueryExceedsBudget => "grouped_query_exceeds_budget",
1005 Self::DiagnosticLaneDoesNotExecute => "diagnostic_lane_does_not_execute",
1006 Self::IntrospectionDisabledForLane => "introspection_disabled_for_lane",
1007 Self::UnsupportedStatementForQueryLane => "unsupported_statement_for_query_lane",
1008 Self::PublicQueryOffsetRejected => "public_query_offset_rejected",
1009 Self::ReturnedRowBoundExceedsPolicy => "returned_row_bound_exceeds_policy",
1010 Self::PrimaryKeyInputExceedsPolicy => "primary_key_input_exceeds_policy",
1011 }
1012 }
1013
1014 #[must_use]
1016 pub const fn code(self) -> QueryReadAdmissionCode {
1017 match self {
1018 Self::PublicQueryRequiresLimit => QueryReadAdmissionCode::PublicQueryRequiresLimit,
1019 Self::PublicQueryRequiresIndex => QueryReadAdmissionCode::PublicQueryRequiresIndex,
1020 Self::UnboundedFullScanRejected => QueryReadAdmissionCode::UnboundedFullScanRejected,
1021 Self::ScanBoundUnavailable => QueryReadAdmissionCode::ScanBoundUnavailable,
1022 Self::ScanBoundExceedsPolicy => QueryReadAdmissionCode::ScanBoundExceedsPolicy,
1023 Self::EstimatedOnlyBoundRejected => QueryReadAdmissionCode::EstimatedOnlyBoundRejected,
1024 Self::SortRequiresMaterialization => {
1025 QueryReadAdmissionCode::SortRequiresMaterialization
1026 }
1027 Self::MaterializationExceedsBudget => {
1028 QueryReadAdmissionCode::MaterializationExceedsBudget
1029 }
1030 Self::ProjectionResponseMayExceedLimit => {
1031 QueryReadAdmissionCode::ProjectionResponseMayExceedLimit
1032 }
1033 Self::GroupedQueryRequiresLimits => QueryReadAdmissionCode::GroupedQueryRequiresLimits,
1034 Self::GroupedQueryExceedsBudget => QueryReadAdmissionCode::GroupedQueryExceedsBudget,
1035 Self::DiagnosticLaneDoesNotExecute => {
1036 QueryReadAdmissionCode::DiagnosticLaneDoesNotExecute
1037 }
1038 Self::IntrospectionDisabledForLane => {
1039 QueryReadAdmissionCode::IntrospectionDisabledForLane
1040 }
1041 Self::UnsupportedStatementForQueryLane => {
1042 QueryReadAdmissionCode::UnsupportedStatementForQueryLane
1043 }
1044 Self::PublicQueryOffsetRejected => QueryReadAdmissionCode::PublicQueryOffsetRejected,
1045 Self::ReturnedRowBoundExceedsPolicy => {
1046 QueryReadAdmissionCode::ReturnedRowBoundExceedsPolicy
1047 }
1048 Self::PrimaryKeyInputExceedsPolicy => {
1049 QueryReadAdmissionCode::PrimaryKeyInputExceedsPolicy
1050 }
1051 }
1052 }
1053
1054 #[must_use]
1056 pub const fn diagnostic(self) -> Diagnostic {
1057 Diagnostic::new(
1058 DiagnosticCode::QueryReadAdmission,
1059 ErrorOrigin::Query,
1060 Some(DiagnosticDetail::QueryReadAdmission {
1061 reason: self.code(),
1062 }),
1063 )
1064 }
1065
1066 #[must_use]
1068 pub const fn error_code(self) -> ErrorCode {
1069 self.diagnostic().error_code()
1070 }
1071}
1072
1073#[derive(Clone, Debug, Eq, PartialEq)]
1075pub struct QueryAdmissionSummary {
1076 lane: QueryAdmissionLane,
1077 decision: QueryAdmissionDecision,
1078 plan_shape: QueryAdmissionPlanShape,
1079 selected_access: QueryAdmissionAccessKind,
1080 selected_index: Option<String>,
1081 limit: Option<u32>,
1082 offset: Option<u32>,
1083 scan_bound: Option<u64>,
1084 scan_bound_kind: QueryBoundKind,
1085 returned_row_bound: Option<u32>,
1086 returned_row_bound_kind: QueryBoundKind,
1087 response_byte_bound: Option<u32>,
1088 response_byte_bound_kind: QueryBoundKind,
1089 primary_key_input_terms: Option<u32>,
1090 primary_key_input_payload_bytes: Option<u32>,
1091 residual_filter: QueryAdmissionResidualFilter,
1092 ordering: QueryAdmissionOrdering,
1093 grouped: Option<QueryAdmissionGroupedSummary>,
1094 materialization: QueryMaterializationSummary,
1095 rejection: Option<QueryAdmissionRejection>,
1096}
1097
1098impl QueryAdmissionSummary {
1099 #[must_use]
1101 pub const fn admitted(
1102 lane: QueryAdmissionLane,
1103 selected_access: QueryAdmissionAccessKind,
1104 ) -> Self {
1105 Self {
1106 lane,
1107 decision: QueryAdmissionDecision::Admitted,
1108 plan_shape: QueryAdmissionPlanShape::ScalarRead,
1109 selected_access,
1110 selected_index: None,
1111 limit: None,
1112 offset: None,
1113 scan_bound: None,
1114 scan_bound_kind: QueryBoundKind::Unavailable,
1115 returned_row_bound: None,
1116 returned_row_bound_kind: QueryBoundKind::Unavailable,
1117 response_byte_bound: None,
1118 response_byte_bound_kind: QueryBoundKind::Unavailable,
1119 primary_key_input_terms: None,
1120 primary_key_input_payload_bytes: None,
1121 residual_filter: QueryAdmissionResidualFilter::Absent,
1122 ordering: QueryAdmissionOrdering::None,
1123 grouped: None,
1124 materialization: QueryMaterializationSummary::none(),
1125 rejection: None,
1126 }
1127 }
1128
1129 #[must_use]
1131 pub const fn rejected(
1132 lane: QueryAdmissionLane,
1133 selected_access: QueryAdmissionAccessKind,
1134 rejection: QueryAdmissionRejection,
1135 ) -> Self {
1136 Self {
1137 lane,
1138 decision: QueryAdmissionDecision::Rejected,
1139 plan_shape: QueryAdmissionPlanShape::ScalarRead,
1140 selected_access,
1141 selected_index: None,
1142 limit: None,
1143 offset: None,
1144 scan_bound: None,
1145 scan_bound_kind: QueryBoundKind::Unavailable,
1146 returned_row_bound: None,
1147 returned_row_bound_kind: QueryBoundKind::Unavailable,
1148 response_byte_bound: None,
1149 response_byte_bound_kind: QueryBoundKind::Unavailable,
1150 primary_key_input_terms: None,
1151 primary_key_input_payload_bytes: None,
1152 residual_filter: QueryAdmissionResidualFilter::Absent,
1153 ordering: QueryAdmissionOrdering::None,
1154 grouped: None,
1155 materialization: QueryMaterializationSummary::none(),
1156 rejection: Some(rejection),
1157 }
1158 }
1159
1160 #[must_use]
1162 pub(in crate::db) fn from_plan(lane: QueryAdmissionLane, plan: &AccessPlannedQuery) -> Self {
1163 let access = summarize_access_plan(plan);
1164 let grouped = plan.grouped_plan().map(summarize_grouped_plan);
1165 let (limit, offset) = scalar_limit_and_offset(plan.scalar_plan());
1166 let (mut returned_row_bound, mut returned_row_bound_kind) =
1167 returned_row_bound_from_plan(limit, grouped);
1168 if returned_row_bound.is_none() && grouped.is_none() {
1169 (returned_row_bound, returned_row_bound_kind) =
1170 returned_row_bound_from_exact_access(&access);
1171 }
1172 let primary_key_input_resource = plan.access_choice().primary_key_input_resource();
1173 let scan_bound_kind = access.scan_bound_kind();
1174 Self {
1175 lane,
1176 decision: QueryAdmissionDecision::Admitted,
1177 plan_shape: plan_shape(plan),
1178 selected_access: access.kind,
1179 selected_index: access.selected_index,
1180 limit,
1181 offset: Some(offset),
1182 scan_bound: access.exact_scan_bound,
1183 scan_bound_kind,
1184 returned_row_bound,
1185 returned_row_bound_kind,
1186 response_byte_bound: None,
1187 response_byte_bound_kind: QueryBoundKind::Unavailable,
1188 primary_key_input_terms: primary_key_input_resource
1189 .map(PrimaryKeyInputResourceSummary::raw_term_count),
1190 primary_key_input_payload_bytes: primary_key_input_resource
1191 .map(PrimaryKeyInputResourceSummary::estimated_payload_bytes),
1192 residual_filter: admission_residual_filter(plan.residual_filter_shape()),
1193 ordering: admission_ordering(plan),
1194 grouped,
1195 materialization: QueryMaterializationSummary::none(),
1196 rejection: None,
1197 }
1198 }
1199
1200 const fn admit(mut self) -> Self {
1201 self.decision = QueryAdmissionDecision::Admitted;
1202 self.rejection = None;
1203 self
1204 }
1205
1206 const fn reject(mut self, rejection: QueryAdmissionRejection) -> Self {
1207 self.decision = QueryAdmissionDecision::Rejected;
1208 self.rejection = Some(rejection);
1209 self
1210 }
1211
1212 #[must_use]
1214 pub const fn lane(&self) -> QueryAdmissionLane {
1215 self.lane
1216 }
1217
1218 #[must_use]
1220 pub const fn decision(&self) -> QueryAdmissionDecision {
1221 self.decision
1222 }
1223
1224 #[must_use]
1226 pub const fn plan_shape(&self) -> QueryAdmissionPlanShape {
1227 self.plan_shape
1228 }
1229
1230 #[must_use]
1232 pub const fn selected_access(&self) -> QueryAdmissionAccessKind {
1233 self.selected_access
1234 }
1235
1236 #[must_use]
1238 pub fn selected_index(&self) -> Option<&str> {
1239 self.selected_index.as_deref()
1240 }
1241
1242 #[must_use]
1244 pub const fn limit(&self) -> Option<u32> {
1245 self.limit
1246 }
1247
1248 #[must_use]
1250 pub const fn offset(&self) -> Option<u32> {
1251 self.offset
1252 }
1253
1254 #[must_use]
1256 pub const fn scan_bound(&self) -> Option<u64> {
1257 self.scan_bound
1258 }
1259
1260 #[must_use]
1262 pub const fn scan_bound_kind(&self) -> QueryBoundKind {
1263 self.scan_bound_kind
1264 }
1265
1266 #[must_use]
1268 pub const fn returned_row_bound(&self) -> Option<u32> {
1269 self.returned_row_bound
1270 }
1271
1272 #[must_use]
1274 pub const fn returned_row_bound_kind(&self) -> QueryBoundKind {
1275 self.returned_row_bound_kind
1276 }
1277
1278 #[must_use]
1280 pub const fn response_byte_bound(&self) -> Option<u32> {
1281 self.response_byte_bound
1282 }
1283
1284 #[must_use]
1286 pub const fn response_byte_bound_kind(&self) -> QueryBoundKind {
1287 self.response_byte_bound_kind
1288 }
1289
1290 #[must_use]
1293 pub const fn primary_key_input_terms(&self) -> Option<u32> {
1294 self.primary_key_input_terms
1295 }
1296
1297 #[must_use]
1299 pub const fn primary_key_input_payload_bytes(&self) -> Option<u32> {
1300 self.primary_key_input_payload_bytes
1301 }
1302
1303 #[must_use]
1305 pub const fn residual_filter(&self) -> QueryAdmissionResidualFilter {
1306 self.residual_filter
1307 }
1308
1309 #[must_use]
1311 pub const fn ordering(&self) -> QueryAdmissionOrdering {
1312 self.ordering
1313 }
1314
1315 #[must_use]
1317 pub const fn grouped(&self) -> Option<QueryAdmissionGroupedSummary> {
1318 self.grouped
1319 }
1320
1321 #[must_use]
1323 pub const fn materialization(&self) -> QueryMaterializationSummary {
1324 self.materialization
1325 }
1326
1327 #[must_use]
1329 #[cfg_attr(not(feature = "sql"), allow(dead_code))]
1330 pub(in crate::db) const fn with_materialization(
1331 mut self,
1332 materialization: QueryMaterializationSummary,
1333 ) -> Self {
1334 self.materialization = materialization;
1335 self
1336 }
1337
1338 #[must_use]
1340 pub const fn rejection(&self) -> Option<QueryAdmissionRejection> {
1341 self.rejection
1342 }
1343
1344 #[must_use]
1346 pub(in crate::db) fn render_text_block(&self) -> String {
1347 let mut out = String::from("admission:");
1348 push_text_field(&mut out, "lane", self.lane().as_str());
1349 push_text_field(&mut out, "decision", self.decision().as_str());
1350 push_text_field(
1351 &mut out,
1352 "reason",
1353 self.rejection()
1354 .map_or("none", QueryAdmissionRejection::as_str),
1355 );
1356 push_text_field(&mut out, "plan_shape", self.plan_shape().as_str());
1357 push_text_field(&mut out, "selected_access", self.selected_access().as_str());
1358 push_text_field(
1359 &mut out,
1360 "selected_index",
1361 self.selected_index().unwrap_or("none"),
1362 );
1363 push_text_option_u32(&mut out, "limit", self.limit());
1364 push_text_option_u32(&mut out, "offset", self.offset());
1365 push_text_option_u64(&mut out, "scan_bound", self.scan_bound());
1366 push_text_field(&mut out, "scan_bound_kind", self.scan_bound_kind().as_str());
1367 push_text_option_u32(&mut out, "returned_row_bound", self.returned_row_bound());
1368 push_text_field(
1369 &mut out,
1370 "returned_row_bound_kind",
1371 self.returned_row_bound_kind().as_str(),
1372 );
1373 push_text_option_u32(&mut out, "response_byte_bound", self.response_byte_bound());
1374 push_text_field(
1375 &mut out,
1376 "response_byte_bound_kind",
1377 self.response_byte_bound_kind().as_str(),
1378 );
1379 push_text_option_u32(
1380 &mut out,
1381 "primary_key_input_terms",
1382 self.primary_key_input_terms(),
1383 );
1384 push_text_option_u32(
1385 &mut out,
1386 "primary_key_input_payload_bytes",
1387 self.primary_key_input_payload_bytes(),
1388 );
1389 push_text_field(&mut out, "residual_filter", self.residual_filter().as_str());
1390 push_text_field(&mut out, "ordering", self.ordering().as_str());
1391 push_text_bool(
1392 &mut out,
1393 "materialized_sort",
1394 self.materialization().materialized_sort(),
1395 );
1396 push_text_option_u32(
1397 &mut out,
1398 "materialized_rows",
1399 self.materialization().materialized_rows(),
1400 );
1401 push_text_field(
1402 &mut out,
1403 "materialized_row_bound_kind",
1404 self.materialization().row_bound_kind().as_str(),
1405 );
1406
1407 if let Some(grouped) = self.grouped() {
1408 push_text_bool(&mut out, "grouped", true);
1409 push_text_u64(
1410 &mut out,
1411 "group_field_count",
1412 u64::from(grouped.group_field_count()),
1413 );
1414 push_text_u64(
1415 &mut out,
1416 "aggregate_count",
1417 u64::from(grouped.aggregate_count()),
1418 );
1419 push_text_u64(
1420 &mut out,
1421 "distinct_aggregate_count",
1422 u64::from(grouped.distinct_aggregate_count()),
1423 );
1424 push_text_u64(&mut out, "max_groups", grouped.max_groups());
1425 push_text_u64(&mut out, "max_group_bytes", grouped.max_group_bytes());
1426 push_text_bool(&mut out, "having_filter", grouped.has_having_filter());
1427 } else {
1428 push_text_bool(&mut out, "grouped", false);
1429 }
1430
1431 out
1432 }
1433}
1434
1435fn push_text_field(out: &mut String, key: &str, value: &str) {
1436 out.push('\n');
1437 out.push_str(" ");
1438 out.push_str(key);
1439 out.push('=');
1440 out.push_str(value);
1441}
1442
1443fn push_text_bool(out: &mut String, key: &str, value: bool) {
1444 push_text_field(out, key, if value { "true" } else { "false" });
1445}
1446
1447fn push_text_u64(out: &mut String, key: &str, value: u64) {
1448 out.push('\n');
1449 out.push_str(" ");
1450 out.push_str(key);
1451 out.push('=');
1452 let _ = write!(out, "{value}");
1453}
1454
1455fn push_text_option_u32(out: &mut String, key: &str, value: Option<u32>) {
1456 match value {
1457 Some(value) => push_text_u64(out, key, u64::from(value)),
1458 None => push_text_field(out, key, "none"),
1459 }
1460}
1461
1462fn push_text_option_u64(out: &mut String, key: &str, value: Option<u64>) {
1463 match value {
1464 Some(value) => push_text_u64(out, key, value),
1465 None => push_text_field(out, key, "none"),
1466 }
1467}
1468
1469const _: fn(QueryAdmissionLane, &AccessPlannedQuery) -> QueryAdmissionSummary =
1471 QueryAdmissionSummary::from_plan;
1472
1473const fn access_satisfies_index_requirement(
1474 kind: QueryAdmissionAccessKind,
1475 scan_bound: Option<u64>,
1476) -> bool {
1477 kind.is_secondary_index()
1478 || matches!(
1479 (kind, scan_bound),
1480 (
1481 QueryAdmissionAccessKind::ByKey | QueryAdmissionAccessKind::ByKeys,
1482 Some(_)
1483 )
1484 )
1485}
1486
1487struct AdmissionAccessProjection;
1488
1489#[derive(Clone, Debug, Eq, PartialEq)]
1490struct AdmissionAccessSummary {
1491 kind: QueryAdmissionAccessKind,
1492 selected_index: Option<String>,
1493 exact_scan_bound: Option<u64>,
1494}
1495
1496impl AdmissionAccessSummary {
1497 const fn non_index(kind: QueryAdmissionAccessKind, exact_scan_bound: Option<u64>) -> Self {
1498 Self {
1499 kind,
1500 selected_index: None,
1501 exact_scan_bound,
1502 }
1503 }
1504
1505 fn secondary_index(kind: QueryAdmissionAccessKind, index_name: &str) -> Self {
1506 Self {
1507 kind,
1508 selected_index: Some(index_name.to_string()),
1509 exact_scan_bound: None,
1510 }
1511 }
1512
1513 const fn composite(kind: QueryAdmissionAccessKind) -> Self {
1514 Self {
1515 kind,
1516 selected_index: None,
1517 exact_scan_bound: None,
1518 }
1519 }
1520
1521 const fn scan_bound_kind(&self) -> QueryBoundKind {
1522 if self.exact_scan_bound.is_some() {
1523 QueryBoundKind::Exact
1524 } else {
1525 QueryBoundKind::Unavailable
1526 }
1527 }
1528}
1529
1530impl AccessPlanProjection<Value> for AdmissionAccessProjection {
1531 type Output = AdmissionAccessSummary;
1532
1533 fn by_key(&mut self, _key: &Value) -> Self::Output {
1534 AdmissionAccessSummary::non_index(QueryAdmissionAccessKind::ByKey, Some(1))
1535 }
1536
1537 fn by_keys(&mut self, keys: &[Value]) -> Self::Output {
1538 AdmissionAccessSummary::non_index(
1539 QueryAdmissionAccessKind::ByKeys,
1540 Some(u64::try_from(keys.len()).unwrap_or(u64::MAX)),
1541 )
1542 }
1543
1544 fn key_range(&mut self, _start: &Value, _end: &Value) -> Self::Output {
1545 AdmissionAccessSummary::non_index(QueryAdmissionAccessKind::KeyRange, None)
1546 }
1547
1548 fn index_prefix(
1549 &mut self,
1550 index_name: &str,
1551 _index_fields: &[String],
1552 _prefix_len: usize,
1553 _values: &[Value],
1554 ) -> Self::Output {
1555 AdmissionAccessSummary::secondary_index(QueryAdmissionAccessKind::IndexPrefix, index_name)
1556 }
1557
1558 fn index_multi_lookup(
1559 &mut self,
1560 index_name: &str,
1561 _index_fields: &[String],
1562 _values: &[Value],
1563 ) -> Self::Output {
1564 AdmissionAccessSummary::secondary_index(
1565 QueryAdmissionAccessKind::IndexMultiLookup,
1566 index_name,
1567 )
1568 }
1569
1570 fn index_branch_set(
1571 &mut self,
1572 index_name: &str,
1573 _index_fields: &[String],
1574 _fixed_values: &[Value],
1575 _branch_values: &[Value],
1576 _ordered_suffix: IndexBranchSetOrderedSuffix,
1577 ) -> Self::Output {
1578 AdmissionAccessSummary::secondary_index(
1579 QueryAdmissionAccessKind::IndexBranchSet,
1580 index_name,
1581 )
1582 }
1583
1584 fn index_range(
1585 &mut self,
1586 index_name: &str,
1587 _index_fields: &[String],
1588 _prefix_len: usize,
1589 _prefix: &[Value],
1590 _lower: &Bound<Value>,
1591 _upper: &Bound<Value>,
1592 ) -> Self::Output {
1593 AdmissionAccessSummary::secondary_index(QueryAdmissionAccessKind::IndexRange, index_name)
1594 }
1595
1596 fn full_scan(&mut self) -> Self::Output {
1597 AdmissionAccessSummary::non_index(QueryAdmissionAccessKind::FullScan, None)
1598 }
1599
1600 fn union(&mut self, _children: Vec<Self::Output>) -> Self::Output {
1601 AdmissionAccessSummary::composite(QueryAdmissionAccessKind::Union)
1602 }
1603
1604 fn intersection(&mut self, _children: Vec<Self::Output>) -> Self::Output {
1605 AdmissionAccessSummary::composite(QueryAdmissionAccessKind::Intersection)
1606 }
1607}
1608
1609fn summarize_access_plan(plan: &AccessPlannedQuery) -> AdmissionAccessSummary {
1610 project_access_plan(&plan.access, &mut AdmissionAccessProjection)
1611}
1612
1613fn summarize_grouped_plan(plan: &GroupPlan) -> QueryAdmissionGroupedSummary {
1614 QueryAdmissionGroupedSummary::new(
1615 u32::try_from(plan.group.group_fields.len()).unwrap_or(u32::MAX),
1616 u32::try_from(plan.group.aggregates.len()).unwrap_or(u32::MAX),
1617 u32::try_from(
1618 plan.group
1619 .aggregates
1620 .iter()
1621 .filter(|aggregate| aggregate.distinct)
1622 .count(),
1623 )
1624 .unwrap_or(u32::MAX),
1625 plan.group.execution.max_groups(),
1626 plan.group.execution.max_group_bytes(),
1627 plan.having_expr.is_some(),
1628 )
1629}
1630
1631const fn scalar_limit_and_offset(plan: &ScalarPlan) -> (Option<u32>, u32) {
1632 match plan.mode {
1633 QueryMode::Load(load) => match &plan.page {
1634 Some(page) => (page.limit, page.offset),
1635 None => (load.limit(), load.offset()),
1636 },
1637 QueryMode::Delete(delete) => match plan.delete_limit {
1638 Some(delete_limit) => (delete_limit.limit, delete_limit.offset),
1639 None => (delete.limit(), delete.offset()),
1640 },
1641 }
1642}
1643
1644fn returned_row_bound_from_plan(
1645 limit: Option<u32>,
1646 grouped: Option<QueryAdmissionGroupedSummary>,
1647) -> (Option<u32>, QueryBoundKind) {
1648 if let Some(limit) = limit {
1649 return (Some(limit), QueryBoundKind::EnforcedRuntimeCap);
1650 }
1651
1652 let Some(grouped) = grouped else {
1653 return (None, QueryBoundKind::Unavailable);
1654 };
1655 if grouped.max_groups() == u64::MAX {
1656 return (None, QueryBoundKind::Unavailable);
1657 }
1658
1659 (
1660 Some(u32::try_from(grouped.max_groups()).unwrap_or(u32::MAX)),
1661 QueryBoundKind::ConservativeUpperBound,
1662 )
1663}
1664
1665fn returned_row_bound_from_exact_access(
1666 access: &AdmissionAccessSummary,
1667) -> (Option<u32>, QueryBoundKind) {
1668 match (access.kind, access.exact_scan_bound) {
1669 (QueryAdmissionAccessKind::ByKey | QueryAdmissionAccessKind::ByKeys, Some(bound)) => (
1670 Some(clamp_u64_to_u32(bound)),
1671 QueryBoundKind::ConservativeUpperBound,
1672 ),
1673 _ => (None, QueryBoundKind::Unavailable),
1674 }
1675}
1676
1677fn clamp_u64_to_u32(value: u64) -> u32 {
1678 u32::try_from(value).unwrap_or(u32::MAX)
1679}
1680
1681const fn admission_residual_filter(shape: ResidualFilterShape) -> QueryAdmissionResidualFilter {
1682 match shape {
1683 ResidualFilterShape::Absent => QueryAdmissionResidualFilter::Absent,
1684 ResidualFilterShape::Predicate => QueryAdmissionResidualFilter::Predicate,
1685 ResidualFilterShape::Expression => QueryAdmissionResidualFilter::Expression,
1686 ResidualFilterShape::ExpressionAndPredicate => {
1687 QueryAdmissionResidualFilter::ExpressionAndPredicate
1688 }
1689 }
1690}
1691
1692fn admission_ordering(plan: &AccessPlannedQuery) -> QueryAdmissionOrdering {
1693 if plan.scalar_plan().order.is_none() {
1694 return QueryAdmissionOrdering::None;
1695 }
1696
1697 if plan.resolved_order().is_some() {
1698 QueryAdmissionOrdering::Resolved
1699 } else {
1700 QueryAdmissionOrdering::Requested
1701 }
1702}
1703
1704const fn plan_shape(plan: &AccessPlannedQuery) -> QueryAdmissionPlanShape {
1705 if plan.grouped_plan().is_some() {
1706 return QueryAdmissionPlanShape::GroupedAggregate;
1707 }
1708
1709 match plan.scalar_plan().mode {
1710 QueryMode::Load(_) => QueryAdmissionPlanShape::ScalarRead,
1711 QueryMode::Delete(_) => QueryAdmissionPlanShape::Delete,
1712 }
1713}
1714
1715#[cfg(test)]
1716mod tests {
1717 use std::num::{NonZeroU32, NonZeroU64};
1718
1719 use crate::{
1720 db::{
1721 access::{AccessPath, SemanticIndexAccessContract},
1722 predicate::{MissingRowPolicy, Predicate},
1723 query::plan::{
1724 AccessPlannedQuery, AggregateKind, DeleteLimitSpec, DeleteSpec, FieldSlot,
1725 GroupAggregateSpec, GroupSpec, GroupedExecutionConfig, OrderDirection, OrderSpec,
1726 OrderTerm, PageSpec, QueryMode,
1727 expr::{Expr, FieldId},
1728 },
1729 },
1730 model::index::IndexModel,
1731 value::Value,
1732 };
1733
1734 use super::{
1735 GroupedAdmissionPolicy, QueryAdmissionAccessKind, QueryAdmissionDecision,
1736 QueryAdmissionLane, QueryAdmissionOrdering, QueryAdmissionPlanShape, QueryAdmissionPolicy,
1737 QueryAdmissionRejection, QueryAdmissionResidualFilter, QueryAdmissionSummary,
1738 QueryBoundKind, QueryMaterializationSummary,
1739 };
1740
1741 const ADMISSION_INDEX_FIELDS: [&str; 1] = ["tag"];
1742 const ADMISSION_INDEX: IndexModel = IndexModel::generated(
1743 "admission::tag",
1744 "admission::tag_store",
1745 &ADMISSION_INDEX_FIELDS,
1746 false,
1747 );
1748
1749 #[test]
1750 fn public_read_policy_has_safe_finite_defaults() {
1751 let max_rows = NonZeroU32::new(50).expect("test max rows is non-zero");
1752 let max_bytes = NonZeroU32::new(32_768).expect("test max bytes is non-zero");
1753 let policy = QueryAdmissionPolicy::public_read(max_rows, max_bytes);
1754
1755 assert_eq!(policy.lane(), QueryAdmissionLane::PublicRead);
1756 assert!(policy.require_limit());
1757 assert!(policy.require_index());
1758 assert!(policy.reject_non_zero_offset());
1759 assert!(!policy.allow_full_scan());
1760 assert!(!policy.allow_materialized_sort());
1761 assert_eq!(policy.max_returned_rows(), Some(max_rows));
1762 assert_eq!(policy.max_response_bytes(), Some(max_bytes));
1763 assert!(policy.public_caps_are_finite());
1764 assert!(!policy.grouped().has_hard_limits());
1765 }
1766
1767 #[test]
1768 fn admin_policy_is_broader_but_still_budgeted() {
1769 let max_rows = NonZeroU32::new(100).expect("test max rows is non-zero");
1770 let max_scanned = NonZeroU64::new(1_000).expect("test scan cap is non-zero");
1771 let max_bytes = NonZeroU32::new(65_536).expect("test max bytes is non-zero");
1772 let policy = QueryAdmissionPolicy::admin_ad_hoc(max_rows, max_scanned, max_bytes);
1773
1774 assert_eq!(policy.lane(), QueryAdmissionLane::AdminAdHoc);
1775 assert!(!policy.require_limit());
1776 assert!(!policy.require_index());
1777 assert!(policy.allow_full_scan());
1778 assert!(policy.allow_materialized_sort());
1779 assert_eq!(policy.max_scanned_rows(), Some(max_scanned));
1780 assert_eq!(policy.max_materialized_rows(), Some(max_rows));
1781 }
1782
1783 #[test]
1784 fn diagnostic_explain_lane_does_not_execute_rows() {
1785 let policy = QueryAdmissionPolicy::diagnostic_explain();
1786
1787 assert_eq!(policy.lane().as_str(), "diagnostic_explain");
1788 assert!(!policy.lane().executes_rows());
1789 }
1790
1791 #[test]
1792 fn grouped_policy_requires_group_and_memory_budgets() {
1793 let max_groups = NonZeroU32::new(8).expect("test group cap is non-zero");
1794 let max_bytes = NonZeroU32::new(4096).expect("test byte cap is non-zero");
1795 let policy = GroupedAdmissionPolicy::bounded(max_groups, max_bytes, None);
1796
1797 assert!(policy.has_hard_limits());
1798 assert_eq!(policy.max_groups(), Some(max_groups));
1799 assert_eq!(policy.max_group_bytes(), Some(max_bytes));
1800 }
1801
1802 #[test]
1803 fn only_proven_or_enforced_bounds_admit_public_reads() {
1804 assert!(QueryBoundKind::Exact.admits_public_read());
1805 assert!(QueryBoundKind::ConservativeUpperBound.admits_public_read());
1806 assert!(QueryBoundKind::EnforcedRuntimeCap.admits_public_read());
1807 assert!(!QueryBoundKind::EstimateOnly.admits_public_read());
1808 assert!(!QueryBoundKind::Unavailable.admits_public_read());
1809 }
1810
1811 #[test]
1812 fn access_kind_classifies_secondary_indexes_and_full_scans() {
1813 assert!(QueryAdmissionAccessKind::IndexPrefix.is_secondary_index());
1814 assert!(QueryAdmissionAccessKind::FullScan.is_full_scan());
1815 assert!(!QueryAdmissionAccessKind::ByKey.is_secondary_index());
1816 }
1817
1818 #[test]
1819 fn rejection_maps_to_stable_diagnostic() {
1820 let rejection = QueryAdmissionRejection::PublicQueryRequiresLimit;
1821 let diagnostic = rejection.diagnostic();
1822
1823 assert_eq!(
1824 rejection.error_code(),
1825 icydb_diagnostic_code::ErrorCode::QUERY_READ_PUBLIC_REQUIRES_LIMIT
1826 );
1827 assert_eq!(
1828 diagnostic.code(),
1829 icydb_diagnostic_code::DiagnosticCode::QueryReadAdmission
1830 );
1831 }
1832
1833 #[test]
1834 fn summaries_keep_decision_and_rejection_aligned() {
1835 let admitted = QueryAdmissionSummary::admitted(
1836 QueryAdmissionLane::PublicRead,
1837 QueryAdmissionAccessKind::ByKey,
1838 );
1839 let rejected = QueryAdmissionSummary::rejected(
1840 QueryAdmissionLane::PublicRead,
1841 QueryAdmissionAccessKind::FullScan,
1842 QueryAdmissionRejection::UnboundedFullScanRejected,
1843 );
1844
1845 assert_eq!(admitted.decision(), QueryAdmissionDecision::Admitted);
1846 assert_eq!(admitted.rejection(), None);
1847 assert_eq!(rejected.decision(), QueryAdmissionDecision::Rejected);
1848 assert_eq!(
1849 rejected.rejection(),
1850 Some(QueryAdmissionRejection::UnboundedFullScanRejected)
1851 );
1852 }
1853
1854 #[test]
1855 fn admission_summary_renders_stable_verbose_explain_block() {
1856 let summary = QueryAdmissionSummary::rejected(
1857 QueryAdmissionLane::PublicRead,
1858 QueryAdmissionAccessKind::FullScan,
1859 QueryAdmissionRejection::UnboundedFullScanRejected,
1860 );
1861
1862 let rendered = summary.render_text_block();
1863
1864 assert!(
1865 rendered.starts_with("admission:\n lane=public_read\n decision=rejected"),
1866 "admission block should start with stable lane and decision fields: {rendered}",
1867 );
1868 assert!(
1869 rendered.contains("\n reason=unbounded_full_scan_rejected"),
1870 "admission block should include a stable rejection reason: {rendered}",
1871 );
1872 assert!(
1873 rendered.contains("\n selected_access=full_scan"),
1874 "admission block should include the selected access class: {rendered}",
1875 );
1876 assert!(
1877 rendered.contains("\n grouped=false"),
1878 "admission block should include grouped classification: {rendered}",
1879 );
1880 }
1881
1882 #[test]
1883 fn plan_summary_classifies_full_scan_without_overclaiming_bounds() {
1884 let plan = AccessPlannedQuery::new(AccessPath::<Value>::FullScan, MissingRowPolicy::Ignore);
1885
1886 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan);
1887
1888 assert_eq!(summary.plan_shape(), QueryAdmissionPlanShape::ScalarRead);
1889 assert_eq!(
1890 summary.selected_access(),
1891 QueryAdmissionAccessKind::FullScan
1892 );
1893 assert_eq!(summary.selected_index(), None);
1894 assert_eq!(summary.limit(), None);
1895 assert_eq!(summary.offset(), Some(0));
1896 assert_eq!(summary.scan_bound(), None);
1897 assert_eq!(summary.scan_bound_kind(), QueryBoundKind::Unavailable);
1898 assert_eq!(summary.returned_row_bound(), None);
1899 assert_eq!(
1900 summary.returned_row_bound_kind(),
1901 QueryBoundKind::Unavailable
1902 );
1903 assert_eq!(
1904 summary.residual_filter(),
1905 QueryAdmissionResidualFilter::Absent
1906 );
1907 assert_eq!(summary.ordering(), QueryAdmissionOrdering::None);
1908 }
1909
1910 #[test]
1911 fn plan_summary_uses_point_lookup_and_limit_as_proven_bounds() {
1912 let mut plan =
1913 AccessPlannedQuery::new(AccessPath::ByKey(Value::Nat64(7)), MissingRowPolicy::Ignore);
1914 plan.scalar_plan_mut().page = Some(PageSpec {
1915 limit: Some(5),
1916 offset: 2,
1917 });
1918
1919 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan);
1920
1921 assert_eq!(summary.selected_access(), QueryAdmissionAccessKind::ByKey);
1922 assert_eq!(summary.limit(), Some(5));
1923 assert_eq!(summary.offset(), Some(2));
1924 assert_eq!(summary.scan_bound(), Some(1));
1925 assert_eq!(summary.scan_bound_kind(), QueryBoundKind::Exact);
1926 assert_eq!(summary.returned_row_bound(), Some(5));
1927 assert_eq!(
1928 summary.returned_row_bound_kind(),
1929 QueryBoundKind::EnforcedRuntimeCap
1930 );
1931 }
1932
1933 #[test]
1934 fn plan_summary_uses_exact_primary_key_access_as_returned_row_bound_without_limit() {
1935 let plan =
1936 AccessPlannedQuery::new(AccessPath::ByKey(Value::Nat64(7)), MissingRowPolicy::Ignore);
1937
1938 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan);
1939
1940 assert_eq!(summary.selected_access(), QueryAdmissionAccessKind::ByKey);
1941 assert_eq!(summary.limit(), None);
1942 assert_eq!(summary.scan_bound(), Some(1));
1943 assert_eq!(summary.scan_bound_kind(), QueryBoundKind::Exact);
1944 assert_eq!(summary.returned_row_bound(), Some(1));
1945 assert_eq!(
1946 summary.returned_row_bound_kind(),
1947 QueryBoundKind::ConservativeUpperBound
1948 );
1949 }
1950
1951 #[test]
1952 fn plan_summary_uses_exact_primary_key_set_as_returned_row_bound_without_limit() {
1953 let plan = AccessPlannedQuery::new(
1954 AccessPath::ByKeys(vec![Value::Nat64(7), Value::Nat64(8)]),
1955 MissingRowPolicy::Ignore,
1956 );
1957
1958 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan);
1959
1960 assert_eq!(summary.selected_access(), QueryAdmissionAccessKind::ByKeys);
1961 assert_eq!(summary.limit(), None);
1962 assert_eq!(summary.scan_bound(), Some(2));
1963 assert_eq!(summary.scan_bound_kind(), QueryBoundKind::Exact);
1964 assert_eq!(summary.returned_row_bound(), Some(2));
1965 assert_eq!(
1966 summary.returned_row_bound_kind(),
1967 QueryBoundKind::ConservativeUpperBound
1968 );
1969 }
1970
1971 #[test]
1972 fn plan_summary_preserves_selected_index_identity() {
1973 let plan = AccessPlannedQuery::new(
1974 AccessPath::IndexPrefix {
1975 index: SemanticIndexAccessContract::model_only_from_generated_index(
1976 ADMISSION_INDEX,
1977 ),
1978 values: vec![Value::Text("alpha".to_string())],
1979 },
1980 MissingRowPolicy::Ignore,
1981 );
1982
1983 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan);
1984
1985 assert_eq!(
1986 summary.selected_access(),
1987 QueryAdmissionAccessKind::IndexPrefix
1988 );
1989 assert_eq!(summary.selected_index(), Some("admission::tag"));
1990 assert_eq!(summary.scan_bound(), None);
1991 assert_eq!(summary.scan_bound_kind(), QueryBoundKind::Unavailable);
1992 }
1993
1994 #[test]
1995 fn plan_summary_classifies_residual_and_requested_ordering() {
1996 let mut plan =
1997 AccessPlannedQuery::new(AccessPath::<Value>::FullScan, MissingRowPolicy::Ignore);
1998 plan.scalar_plan_mut().predicate = Some(Predicate::eq(
1999 "tag".to_string(),
2000 Value::Text("alpha".to_string()),
2001 ));
2002 plan.scalar_plan_mut().order = Some(OrderSpec {
2003 fields: vec![OrderTerm::field("tag", OrderDirection::Asc)],
2004 });
2005
2006 let summary = QueryAdmissionSummary::from_plan(QueryAdmissionLane::AdminAdHoc, &plan);
2007
2008 assert_eq!(
2009 summary.residual_filter(),
2010 QueryAdmissionResidualFilter::Predicate
2011 );
2012 assert_eq!(summary.ordering(), QueryAdmissionOrdering::Requested);
2013 assert!(!summary.materialization().materialized_sort());
2014 assert_eq!(summary.materialization().materialized_rows(), None);
2015 assert_eq!(
2016 summary.materialization().row_bound_kind(),
2017 QueryBoundKind::Unavailable
2018 );
2019 }
2020
2021 #[test]
2022 fn plan_summary_carries_grouped_execution_budgets() {
2023 let grouped =
2024 AccessPlannedQuery::new(AccessPath::<Value>::FullScan, MissingRowPolicy::Ignore)
2025 .into_grouped_with_having_expr(
2026 GroupSpec {
2027 group_fields: vec![FieldSlot::from_test_slot(0, "tag")],
2028 aggregates: vec![GroupAggregateSpec {
2029 kind: AggregateKind::Count,
2030 input_expr: None,
2031 filter_expr: None,
2032 distinct: false,
2033 }],
2034 execution: GroupedExecutionConfig::with_hard_limits(12, 4096),
2035 },
2036 Some(Expr::Field(FieldId::new("tag"))),
2037 );
2038
2039 let summary =
2040 QueryAdmissionSummary::from_plan(QueryAdmissionLane::DiagnosticExplain, &grouped);
2041 let grouped = summary
2042 .grouped()
2043 .expect("summary should include grouped facts");
2044
2045 assert_eq!(
2046 summary.plan_shape(),
2047 QueryAdmissionPlanShape::GroupedAggregate
2048 );
2049 assert_eq!(grouped.group_field_count(), 1);
2050 assert_eq!(grouped.aggregate_count(), 1);
2051 assert_eq!(grouped.distinct_aggregate_count(), 0);
2052 assert_eq!(grouped.max_groups(), 12);
2053 assert_eq!(grouped.max_group_bytes(), 4096);
2054 assert!(grouped.has_having_filter());
2055 assert_eq!(summary.returned_row_bound(), Some(12));
2056 assert_eq!(
2057 summary.returned_row_bound_kind(),
2058 QueryBoundKind::ConservativeUpperBound
2059 );
2060 }
2061
2062 #[test]
2063 fn plan_summary_reads_delete_window_without_executing_it() {
2064 let mut plan =
2065 AccessPlannedQuery::new(AccessPath::<Value>::FullScan, MissingRowPolicy::Ignore);
2066 plan.scalar_plan_mut().mode = QueryMode::Delete(DeleteSpec::new());
2067 plan.scalar_plan_mut().delete_limit = Some(DeleteLimitSpec {
2068 limit: Some(3),
2069 offset: 1,
2070 });
2071
2072 let summary =
2073 QueryAdmissionSummary::from_plan(QueryAdmissionLane::DiagnosticExplain, &plan);
2074
2075 assert_eq!(summary.plan_shape(), QueryAdmissionPlanShape::Delete);
2076 assert_eq!(summary.limit(), Some(3));
2077 assert_eq!(summary.offset(), Some(1));
2078 assert_eq!(summary.returned_row_bound(), Some(3));
2079 }
2080
2081 #[test]
2082 fn public_read_evaluation_rejects_missing_limit_before_access_shape() {
2083 let policy = public_read_policy();
2084 let summary = summary_for_index_prefix(None, 0);
2085
2086 let evaluated = policy.evaluate(summary);
2087
2088 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2089 assert_eq!(
2090 evaluated.rejection(),
2091 Some(QueryAdmissionRejection::PublicQueryRequiresLimit)
2092 );
2093 }
2094
2095 #[test]
2096 fn public_read_evaluation_rejects_full_scan_even_with_limit() {
2097 let policy = public_read_policy();
2098 let summary = summary_for_path(AccessPath::<Value>::FullScan, Some(5), 0);
2099
2100 let evaluated = policy.evaluate(summary);
2101
2102 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2103 assert_eq!(
2104 evaluated.rejection(),
2105 Some(QueryAdmissionRejection::UnboundedFullScanRejected)
2106 );
2107 }
2108
2109 #[test]
2110 fn public_read_evaluation_admits_indexed_bounded_scalar_read() {
2111 let policy = public_read_policy();
2112 let summary = summary_for_index_prefix(Some(5), 0);
2113
2114 let evaluated = policy.evaluate(summary);
2115
2116 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Admitted);
2117 assert_eq!(evaluated.rejection(), None);
2118 }
2119
2120 #[test]
2121 fn public_read_evaluation_admits_exact_primary_key_read() {
2122 let policy = public_read_policy();
2123 let summary = summary_for_path(
2124 AccessPath::ByKey(Value::Text("primary".to_string())),
2125 None,
2126 0,
2127 );
2128
2129 let evaluated = policy.evaluate(summary);
2130
2131 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Admitted);
2132 assert_eq!(evaluated.limit(), None);
2133 assert_eq!(evaluated.scan_bound(), Some(1));
2134 assert_eq!(evaluated.returned_row_bound(), Some(1));
2135 }
2136
2137 #[test]
2138 fn public_read_evaluation_rejects_primary_key_set_above_returned_row_policy() {
2139 let policy = public_read_policy();
2140 let keys = (0..=50).map(Value::Nat64).collect();
2141 let summary = summary_for_path(AccessPath::ByKeys(keys), None, 0);
2142
2143 let evaluated = policy.evaluate(summary);
2144
2145 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2146 assert_eq!(
2147 evaluated.rejection(),
2148 Some(QueryAdmissionRejection::ReturnedRowBoundExceedsPolicy)
2149 );
2150 }
2151
2152 #[test]
2153 fn public_read_evaluation_rejects_non_zero_offset() {
2154 let policy = public_read_policy();
2155 let summary = summary_for_index_prefix(Some(5), 1);
2156
2157 let evaluated = policy.evaluate(summary);
2158
2159 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2160 assert_eq!(
2161 evaluated.rejection(),
2162 Some(QueryAdmissionRejection::PublicQueryOffsetRejected)
2163 );
2164 }
2165
2166 #[test]
2167 fn public_read_evaluation_rejects_returned_row_cap_overflow() {
2168 let policy = public_read_policy();
2169 let summary = summary_for_index_prefix(Some(51), 0);
2170
2171 let evaluated = policy.evaluate(summary);
2172
2173 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2174 assert_eq!(
2175 evaluated.rejection(),
2176 Some(QueryAdmissionRejection::ReturnedRowBoundExceedsPolicy)
2177 );
2178 }
2179
2180 #[test]
2181 fn public_read_evaluation_rejects_unresolved_order_materialized_sort() {
2182 let policy = public_read_policy();
2183 let summary = summary_for_index_prefix(Some(5), 0);
2184 let returned_row_bound = summary.returned_row_bound();
2185 let returned_row_bound_kind = summary.returned_row_bound_kind();
2186 let summary = summary.with_materialization(QueryMaterializationSummary::sort(
2187 returned_row_bound,
2188 returned_row_bound_kind,
2189 ));
2190
2191 let evaluated = policy.evaluate(summary);
2192
2193 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2194 assert_eq!(
2195 evaluated.rejection(),
2196 Some(QueryAdmissionRejection::SortRequiresMaterialization)
2197 );
2198 }
2199
2200 #[test]
2201 fn public_read_evaluation_admits_exact_key_set_materialized_sort() {
2202 let policy = public_read_policy();
2203 let summary = summary_for_path(
2204 AccessPath::ByKeys(vec![Value::Nat64(1), Value::Nat64(2), Value::Nat64(3)]),
2205 None,
2206 0,
2207 );
2208 let returned_row_bound = summary.returned_row_bound();
2209 let returned_row_bound_kind = summary.returned_row_bound_kind();
2210 let summary = summary.with_materialization(QueryMaterializationSummary::sort(
2211 returned_row_bound,
2212 returned_row_bound_kind,
2213 ));
2214
2215 let evaluated = policy.evaluate(summary);
2216
2217 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Admitted);
2218 assert_eq!(evaluated.rejection(), None);
2219 assert_eq!(
2220 evaluated.selected_access(),
2221 QueryAdmissionAccessKind::ByKeys
2222 );
2223 assert_eq!(evaluated.scan_bound(), Some(3));
2224 assert_eq!(evaluated.returned_row_bound(), Some(3));
2225 assert!(evaluated.materialization().materialized_sort());
2226 }
2227
2228 #[test]
2229 fn public_read_evaluation_rejects_underbounded_key_set_materialized_sort() {
2230 let policy = public_read_policy();
2231 let summary = summary_for_path(
2232 AccessPath::ByKeys(vec![Value::Nat64(1), Value::Nat64(2), Value::Nat64(3)]),
2233 Some(1),
2234 0,
2235 );
2236 let returned_row_bound = summary.returned_row_bound();
2237 let returned_row_bound_kind = summary.returned_row_bound_kind();
2238 let summary = summary.with_materialization(QueryMaterializationSummary::sort(
2239 returned_row_bound,
2240 returned_row_bound_kind,
2241 ));
2242
2243 let evaluated = policy.evaluate(summary);
2244
2245 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2246 assert_eq!(
2247 evaluated.rejection(),
2248 Some(QueryAdmissionRejection::SortRequiresMaterialization)
2249 );
2250 }
2251
2252 #[test]
2253 fn public_read_evaluation_fails_closed_when_bounded_route_falls_back_to_materialized_order() {
2254 let policy = public_read_policy();
2255 let bounded = summary_for_index_prefix(Some(1), 0);
2256 let admitted = policy.evaluate(bounded.clone());
2257
2258 assert_eq!(admitted.decision(), QueryAdmissionDecision::Admitted);
2259 assert_eq!(admitted.returned_row_bound(), Some(1));
2260
2261 let returned_row_bound = bounded.returned_row_bound();
2262 let returned_row_bound_kind = bounded.returned_row_bound_kind();
2263 let fallback = bounded.with_materialization(QueryMaterializationSummary::sort(
2264 returned_row_bound,
2265 returned_row_bound_kind,
2266 ));
2267
2268 let evaluated = policy.evaluate(fallback);
2269
2270 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2271 assert_eq!(
2272 evaluated.rejection(),
2273 Some(QueryAdmissionRejection::SortRequiresMaterialization)
2274 );
2275 }
2276
2277 #[test]
2278 fn public_read_evaluation_rejects_grouped_query_without_group_budgets() {
2279 let policy = public_read_policy();
2280 let summary = grouped_summary_for_index_prefix(12, 4096, false);
2281
2282 let evaluated = policy.evaluate(summary);
2283
2284 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2285 assert_eq!(
2286 evaluated.rejection(),
2287 Some(QueryAdmissionRejection::GroupedQueryRequiresLimits)
2288 );
2289 }
2290
2291 #[test]
2292 fn public_read_evaluation_admits_grouped_query_with_group_budgets_without_limit() {
2293 let policy = public_grouped_read_policy(None);
2294 let summary = grouped_summary_for_index_prefix(12, 4096, false);
2295
2296 let evaluated = policy.evaluate(summary);
2297
2298 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Admitted);
2299 assert_eq!(evaluated.limit(), None);
2300 assert_eq!(evaluated.returned_row_bound(), Some(12));
2301 assert_eq!(evaluated.rejection(), None);
2302 }
2303
2304 #[test]
2305 fn public_read_evaluation_rejects_grouped_query_above_policy_budget() {
2306 let policy = public_grouped_read_policy(None);
2307 let summary = grouped_summary_for_index_prefix(51, 4096, false);
2308
2309 let evaluated = policy.evaluate(summary);
2310
2311 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2312 assert_eq!(
2313 evaluated.rejection(),
2314 Some(QueryAdmissionRejection::GroupedQueryExceedsBudget)
2315 );
2316 }
2317
2318 #[test]
2319 fn public_read_evaluation_rejects_distinct_grouped_query_without_distinct_budget() {
2320 let policy = public_grouped_read_policy(None);
2321 let summary = grouped_summary_for_index_prefix(12, 4096, true);
2322
2323 let evaluated = policy.evaluate(summary);
2324
2325 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2326 assert_eq!(
2327 evaluated.rejection(),
2328 Some(QueryAdmissionRejection::GroupedQueryRequiresLimits)
2329 );
2330 }
2331
2332 #[test]
2333 fn diagnostic_explain_policy_rejects_row_execution() {
2334 let policy = QueryAdmissionPolicy::diagnostic_explain();
2335 let summary = summary_for_index_prefix(Some(5), 0);
2336
2337 let evaluated = policy.evaluate(summary);
2338
2339 assert_eq!(evaluated.decision(), QueryAdmissionDecision::Rejected);
2340 assert_eq!(
2341 evaluated.rejection(),
2342 Some(QueryAdmissionRejection::DiagnosticLaneDoesNotExecute)
2343 );
2344 }
2345
2346 fn public_read_policy() -> QueryAdmissionPolicy {
2347 QueryAdmissionPolicy::public_read(
2348 NonZeroU32::new(50).expect("test public row cap is non-zero"),
2349 NonZeroU32::new(32_768).expect("test public byte cap is non-zero"),
2350 )
2351 }
2352
2353 fn public_grouped_read_policy(distinct_entries: Option<NonZeroU32>) -> QueryAdmissionPolicy {
2354 public_read_policy().with_grouped_policy(GroupedAdmissionPolicy::bounded(
2355 NonZeroU32::new(50).expect("test public group cap is non-zero"),
2356 NonZeroU32::new(8192).expect("test public group byte cap is non-zero"),
2357 distinct_entries,
2358 ))
2359 }
2360
2361 fn summary_for_index_prefix(limit: Option<u32>, offset: u32) -> QueryAdmissionSummary {
2362 summary_for_path(
2363 AccessPath::IndexPrefix {
2364 index: SemanticIndexAccessContract::model_only_from_generated_index(
2365 ADMISSION_INDEX,
2366 ),
2367 values: vec![Value::Text("alpha".to_string())],
2368 },
2369 limit,
2370 offset,
2371 )
2372 }
2373
2374 fn summary_for_path(
2375 path: AccessPath<Value>,
2376 limit: Option<u32>,
2377 offset: u32,
2378 ) -> QueryAdmissionSummary {
2379 let mut plan = AccessPlannedQuery::new(path, MissingRowPolicy::Ignore);
2380 plan.scalar_plan_mut().page = Some(PageSpec { limit, offset });
2381
2382 QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &plan)
2383 }
2384
2385 fn grouped_summary_for_index_prefix(
2386 max_groups: u64,
2387 max_group_bytes: u64,
2388 distinct: bool,
2389 ) -> QueryAdmissionSummary {
2390 let grouped = AccessPlannedQuery::new(index_prefix_path(), MissingRowPolicy::Ignore)
2391 .into_grouped(GroupSpec {
2392 group_fields: vec![FieldSlot::from_test_slot(0, "tag")],
2393 aggregates: vec![GroupAggregateSpec {
2394 kind: AggregateKind::Count,
2395 input_expr: Some(Box::new(Expr::Field(FieldId::new("tag")))),
2396 filter_expr: None,
2397 distinct,
2398 }],
2399 execution: GroupedExecutionConfig::with_hard_limits(max_groups, max_group_bytes),
2400 });
2401
2402 QueryAdmissionSummary::from_plan(QueryAdmissionLane::PublicRead, &grouped)
2403 }
2404
2405 fn index_prefix_path() -> AccessPath<Value> {
2406 AccessPath::IndexPrefix {
2407 index: SemanticIndexAccessContract::model_only_from_generated_index(ADMISSION_INDEX),
2408 values: vec![Value::Text("alpha".to_string())],
2409 }
2410 }
2411}