reifydb_engine/
subscription.rs1use std::{result::Result as StdResult, sync::Arc};
5
6use reifydb_core::{
7 common::CommitVersion,
8 interface::{catalog::id::SubscriptionId, change::StagedBatch},
9 metrics::execution::ExecutionMetrics,
10};
11use reifydb_evaluate::stack::SymbolTable;
12use reifydb_rql::flow::flow::FlowDag;
13use reifydb_transaction::{multi::lease::VersionLeaseGuard, transaction::Transaction};
14use reifydb_value::{Result, error::Error as TypeError, params::Params, value::identity::IdentityId};
15
16use crate::engine::StandardEngine;
17
18#[derive(Debug, Clone)]
19pub struct SubscriptionContext {
20 pub id: SubscriptionId,
21 pub identity: IdentityId,
22 pub symbols: SymbolTable,
23 pub params: Params,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum HydrationBound {
28 Pushed,
29 Absent,
30 Blocked {
31 operator: String,
32 },
33}
34
35impl HydrationBound {
36 pub fn advice(&self) -> String {
37 match self {
38 Self::Absent => "add `TAKE N` upstream, raise with WITH { hydration: { max_rows: ... } }, or disable with WITH { hydration: { enabled: false } }".to_string(),
39 Self::Blocked {
40 operator,
41 } => format!(
42 "the query's `TAKE` sits below `{}`, which the hydration pushdown cannot see through, so the source was read unbounded; move the `TAKE` above `{}`, raise with WITH {{ hydration: {{ max_rows: ... }} }}, or disable with WITH {{ hydration: {{ enabled: false }} }}",
43 operator, operator
44 ),
45 Self::Pushed => "the query's `TAKE` was already applied at the source and it still returns more rows than the cap, so raise it with WITH { hydration: { max_rows: ... } } or disable with WITH { hydration: { enabled: false } }".to_string(),
46 }
47 }
48}
49
50#[derive(Debug)]
51pub enum HydrateError {
52 SubscriptionNotFound,
53 UnsupportedSourceType,
54 RowCapExceeded {
55 cap: u64,
56 bound: HydrationBound,
57 },
58 Engine(TypeError),
59 Internal(String),
60}
61
62impl From<TypeError> for HydrateError {
63 fn from(e: TypeError) -> Self {
64 HydrateError::Engine(e)
65 }
66}
67
68impl HydrateError {
69 pub fn is_version_evicted(&self) -> bool {
70 matches!(self, HydrateError::Engine(e) if e.0.code == "TXN_012")
71 }
72
73 pub fn wire_code(&self) -> &'static str {
74 match self {
75 Self::SubscriptionNotFound => "HYDRATION_FAILED",
76 Self::UnsupportedSourceType => "HYDRATION_UNSUPPORTED_SOURCE",
77 Self::RowCapExceeded {
78 ..
79 } => "HYDRATION_TOO_LARGE",
80 Self::Engine(_) => {
81 if self.is_version_evicted() {
82 "HYDRATION_VERSION_EVICTED"
83 } else {
84 "HYDRATION_FAILED"
85 }
86 }
87 Self::Internal(_) => "HYDRATION_FAILED",
88 }
89 }
90
91 pub fn wire_message(&self, rql: &str, cap: u64) -> String {
92 match self {
93 Self::SubscriptionNotFound => "Subscription not found at hydration time".to_string(),
94 Self::UnsupportedSourceType => "hydration is not supported for SourceSeries / SourceInlineData; use WITH { hydration: { enabled: false } } to subscribe without it".to_string(),
95 Self::RowCapExceeded {
96 bound,
97 ..
98 } => format!(
99 "Hydration exceeds subscribe.max_hydration_rows={}; {}. Query: {}",
100 cap,
101 bound.advice(),
102 rql
103 ),
104 Self::Engine(e) => {
105 if self.is_version_evicted() {
106 e.0.message.clone()
107 } else {
108 e.to_string()
109 }
110 }
111 Self::Internal(s) => s.clone(),
112 }
113 }
114}
115
116#[derive(Debug)]
117pub struct HydrateOutcome {
118 pub version: CommitVersion,
119 pub batches: Vec<StagedBatch>,
120 pub metrics: ExecutionMetrics,
121}
122
123pub trait SubscriptionService: Send + Sync {
124 fn next_id(&self) -> SubscriptionId;
125
126 fn register_subscription(
127 &self,
128 flow_dag: FlowDag,
129 column_names: Vec<String>,
130 hydration_enabled: bool,
131 ctx: SubscriptionContext,
132 txn: &mut Transaction<'_>,
133 ) -> Result<()>;
134
135 fn unregister_subscription(&self, id: &SubscriptionId) -> Result<()>;
136
137 fn hydrate(
138 &self,
139 sub_id: SubscriptionId,
140 engine: &StandardEngine,
141 identity: IdentityId,
142 lease: VersionLeaseGuard,
143 max_rows: u64,
144 ) -> StdResult<HydrateOutcome, HydrateError>;
145}
146
147pub type SubscriptionServiceRef = Arc<dyn SubscriptionService>;