Skip to main content

reifydb_engine/
subscription.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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>;