reifydb_engine/
subscription.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{collections::BTreeMap, result::Result as StdResult, sync::Arc};
5
6use reifydb_core::{
7	common::CommitVersion,
8	flow::dag::FlowDag,
9	interface::{
10		catalog::{id::SubscriptionId, object::ObjectId},
11		change::StagedBatch,
12	},
13};
14use reifydb_evaluate::stack::SymbolTable;
15use reifydb_transaction::{multi::lease::VersionLeaseGuard, transaction::Transaction};
16use reifydb_value::{
17	Result,
18	error::Error as TypeError,
19	params::Params,
20	value::{duration::Duration, identity::IdentityId, partition::Partition, system_columns::SystemColumn},
21};
22
23use crate::engine::StandardEngine;
24
25#[derive(Debug, Clone)]
26pub struct SubscriptionContext {
27	pub id: SubscriptionId,
28	pub identity: IdentityId,
29	pub symbols: SymbolTable,
30	pub params: Params,
31	pub named_system_columns: Vec<SystemColumn>,
32	pub partitions: BTreeMap<ObjectId, Option<Partition>>,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum HydrationBound {
37	Present,
38	Absent,
39}
40
41impl HydrationBound {
42	pub fn advice(&self) -> String {
43		match self {
44			Self::Absent => "add `TAKE N` upstream, raise the hydration.max_rows subscribe option, or set the hydration.enabled subscribe option to false".to_string(),
45			Self::Present => "the query's `TAKE` still returns more rows than the cap, so raise the hydration.max_rows subscribe option or set the hydration.enabled subscribe option to 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; set the hydration.enabled subscribe option to 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 total: Duration,
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		hydration_enabled: bool,
130		ctx: SubscriptionContext,
131		txn: &mut Transaction<'_>,
132	) -> Result<()>;
133
134	fn unregister_subscription(&self, id: &SubscriptionId) -> Result<bool>;
135
136	fn hydrate(
137		&self,
138		sub_id: SubscriptionId,
139		identity: IdentityId,
140		lease: VersionLeaseGuard,
141		max_rows: u64,
142	) -> StdResult<HydrateOutcome, HydrateError>;
143}
144
145pub type SubscriptionServiceRef = Arc<dyn SubscriptionService>;
146
147#[cfg(feature = "testing")]
148pub trait HandOffHooks: Send + Sync {
149	fn during_hand_off(&self, _subscription: SubscriptionId) {}
150}
151
152#[cfg(feature = "testing")]
153#[derive(Clone)]
154pub struct InstalledHandOffHooks(pub Arc<dyn HandOffHooks>);
155
156pub fn acquire_hand_off_lease(
157	engine: &StandardEngine,
158	_subscriptions: &[SubscriptionId],
159) -> Result<(CommitVersion, VersionLeaseGuard)> {
160	#[cfg(feature = "testing")]
161	if let Some(InstalledHandOffHooks(hooks)) = engine.ioc().try_resolve::<InstalledHandOffHooks>() {
162		for subscription in _subscriptions {
163			hooks.during_hand_off(*subscription);
164		}
165	}
166	engine.acquire_current_snapshot_lease()
167}
reifydb_engine/subscription.rs

reifydb_engine/
subscription.rs
