Skip to main content

sz_orm_stream/
config.rs

1//! 流式结果集配置与策略枚举
2
3use serde::{Deserialize, Serialize};
4
5use sz_orm_core::DbType;
6
7/// 分页策略
8#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
9pub enum PaginationStrategy {
10    /// keyset 分页(WHERE key > last_key ORDER BY key LIMIT batch)
11    Keyset,
12    /// OFFSET 分页(LIMIT batch OFFSET n)
13    LimitOffset,
14    /// 服务端游标(DECLARE CURSOR + FETCH)
15    ServerCursor,
16}
17
18/// 排序方向
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
20pub enum OrderDirection {
21    /// 升序(默认)
22    #[default]
23    Asc,
24    /// 降序
25    Desc,
26}
27
28/// 流式结果集配置
29#[derive(Debug, Clone)]
30pub struct StreamResultSetConfig {
31    /// 批次大小(默认 1000)
32    pub batch_size: usize,
33    /// 背压阈值(默认 10000)
34    pub backpressure_threshold: usize,
35    /// 分页策略
36    pub pagination_strategy: PaginationStrategy,
37    /// keyset 列名(Keyset 策略时必须设置)
38    pub keyset_column: Option<String>,
39    /// 排序方向
40    pub order_direction: OrderDirection,
41    /// 数据库类型
42    pub db_type: DbType,
43}
44
45impl Default for StreamResultSetConfig {
46    fn default() -> Self {
47        Self {
48            batch_size: 1000,
49            backpressure_threshold: 10000,
50            pagination_strategy: PaginationStrategy::LimitOffset,
51            keyset_column: None,
52            order_direction: OrderDirection::Asc,
53            db_type: DbType::PostgreSQL,
54        }
55    }
56}
57
58impl StreamResultSetConfig {
59    pub fn new(db_type: DbType) -> Self {
60        Self {
61            db_type,
62            ..Self::default()
63        }
64    }
65
66    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
67        self.batch_size = batch_size.max(1);
68        self
69    }
70
71    pub fn with_backpressure_threshold(mut self, threshold: usize) -> Self {
72        self.backpressure_threshold = threshold;
73        self
74    }
75
76    pub fn with_pagination_strategy(mut self, strategy: PaginationStrategy) -> Self {
77        self.pagination_strategy = strategy;
78        self
79    }
80
81    pub fn with_keyset_column(mut self, column: impl Into<String>) -> Self {
82        self.keyset_column = Some(column.into());
83        self
84    }
85
86    pub fn with_order_direction(mut self, direction: OrderDirection) -> Self {
87        self.order_direction = direction;
88        self
89    }
90
91    /// 校验配置合法性
92    pub fn validate(&self) -> Result<(), String> {
93        if self.batch_size == 0 {
94            return Err("batch_size must be > 0".into());
95        }
96        if self.pagination_strategy == PaginationStrategy::Keyset && self.keyset_column.is_none() {
97            return Err("keyset pagination requires keyset_column".into());
98        }
99        Ok(())
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn config_default() {
109        let config = StreamResultSetConfig::new(DbType::PostgreSQL);
110        assert_eq!(config.batch_size, 1000);
111        assert_eq!(config.backpressure_threshold, 10000);
112        assert_eq!(config.pagination_strategy, PaginationStrategy::LimitOffset);
113        assert_eq!(config.order_direction, OrderDirection::Asc);
114    }
115
116    #[test]
117    fn config_builder_chain() {
118        let config = StreamResultSetConfig::new(DbType::MySQL)
119            .with_batch_size(500)
120            .with_backpressure_threshold(5000)
121            .with_pagination_strategy(PaginationStrategy::Keyset)
122            .with_keyset_column("id")
123            .with_order_direction(OrderDirection::Desc);
124        assert_eq!(config.batch_size, 500);
125        assert_eq!(config.backpressure_threshold, 5000);
126        assert_eq!(config.pagination_strategy, PaginationStrategy::Keyset);
127        assert_eq!(config.keyset_column.as_deref(), Some("id"));
128        assert_eq!(config.order_direction, OrderDirection::Desc);
129    }
130
131    #[test]
132    fn config_validate_keyset_requires_column() {
133        let config = StreamResultSetConfig::new(DbType::PostgreSQL)
134            .with_pagination_strategy(PaginationStrategy::Keyset);
135        assert!(config.validate().is_err());
136    }
137
138    #[test]
139    fn config_validate_keyset_with_column() {
140        let config = StreamResultSetConfig::new(DbType::PostgreSQL)
141            .with_pagination_strategy(PaginationStrategy::Keyset)
142            .with_keyset_column("id");
143        assert!(config.validate().is_ok());
144    }
145
146    #[test]
147    fn config_validate_limit_offset_ok() {
148        let config = StreamResultSetConfig::new(DbType::PostgreSQL);
149        assert!(config.validate().is_ok());
150    }
151
152    #[test]
153    fn pagination_strategy_serde() {
154        let s = serde_json::to_string(&PaginationStrategy::Keyset).unwrap();
155        let d: PaginationStrategy = serde_json::from_str(&s).unwrap();
156        assert_eq!(d, PaginationStrategy::Keyset);
157    }
158
159    #[test]
160    fn order_direction_serde() {
161        let s = serde_json::to_string(&OrderDirection::Desc).unwrap();
162        let d: OrderDirection = serde_json::from_str(&s).unwrap();
163        assert_eq!(d, OrderDirection::Desc);
164    }
165
166    #[test]
167    fn batch_size_zero_clamped() {
168        let config = StreamResultSetConfig::new(DbType::PostgreSQL).with_batch_size(0);
169        assert_eq!(config.batch_size, 1);
170    }
171}