rust_tg_bot_ext/
ext_bot.rs

1//! Extended bot with convenience features.
2//!
3//! Ported from `python-telegram-bot/src/telegram/ext/_extbot.py`.
4//!
5//! [`ExtBot`] wraps the low-level [`rust_tg_bot_raw::bot::Bot`] and adds:
6//!
7//! * [`Defaults`](crate::defaults::Defaults) injection into API calls
8//! * [`CallbackDataCache`](crate::callback_data_cache::CallbackDataCache) for arbitrary
9//!   callback data
10//! * A [`DynRateLimiter`](crate::rate_limiter::DynRateLimiter) that intercepts all HTTP
11//!   requests through a [`RateLimitedRequest`](crate::rate_limiter::RateLimitedRequest)
12//!   wrapper.
13//!
14//! # Construction
15//!
16//! Use [`ExtBot::builder`] for the full option set, or [`ExtBot::from_bot`] when you
17//! only have a raw `Bot` and need no extras:
18//!
19//! ```rust,ignore
20//! // Minimal:
21//! let ext = ExtBot::from_bot(bot);
22//!
23//! // Full control:
24//! let ext = ExtBot::builder("token", request)
25//!     .defaults(defaults)
26//!     .arbitrary_callback_data(512)
27//!     .rate_limiter(Arc::new(AioRateLimiter::default_limits()))
28//!     .build();
29//! ```
30//!
31//! # Rate limiting
32//!
33//! When a rate limiter is provided, the builder wraps the HTTP request backend
34//! in a [`RateLimitedRequest`](crate::rate_limiter::RateLimitedRequest), so
35//! **all** API calls flow through the limiter transparently.  No changes to
36//! handler code are needed.
37//!
38//! # `Deref` to `Bot`
39//!
40//! `ExtBot` implements `Deref<Target = Bot>`, so all `Bot` methods are accessible
41//! directly without calling `.inner()`:
42//!
43//! ```rust,ignore
44//! // Instead of: ext_bot.inner().send_message(chat_id, text)
45//! ext_bot.send_message(chat_id, text)
46//! ```
47
48use std::sync::Arc;
49
50use tokio::sync::RwLock;
51
52use rust_tg_bot_raw::bot::Bot;
53use rust_tg_bot_raw::request::base::BaseRequest;
54
55use crate::callback_data_cache::CallbackDataCache;
56use crate::defaults::Defaults;
57
58#[cfg(feature = "rate-limiter")]
59use crate::rate_limiter::{DynRateLimiter, RateLimitedRequest};
60
61/// Extended bot that adds defaults, arbitrary callback data, and a rate-limiter slot on top
62/// of the raw [`Bot`].
63///
64/// # Construction
65///
66/// Use [`ExtBot::builder`] for the full option set, or [`ExtBot::from_bot`] for the
67/// simplest case (no defaults, no cache, no rate limiter).
68///
69/// # `Deref` to `Bot`
70///
71/// `ExtBot` implements [`Deref<Target = Bot>`](std::ops::Deref), making all `Bot` methods
72/// accessible directly. This is a zero-cost abstraction -- no allocation or indirection
73/// beyond what `Bot` already provides.
74///
75/// # Rate limiter
76///
77/// When a rate limiter is set, the inner `Bot`'s HTTP request backend is wrapped in a
78/// [`RateLimitedRequest`](crate::rate_limiter::RateLimitedRequest) so that all API calls
79/// are throttled transparently at the transport layer.
80pub struct ExtBot {
81    /// The underlying raw bot.
82    bot: Bot,
83
84    /// User-defined defaults for API calls.
85    defaults: Option<Defaults>,
86
87    /// Cache for arbitrary inline keyboard callback data.
88    callback_data_cache: Option<Arc<RwLock<CallbackDataCache>>>,
89
90    /// The rate limiter, if any.  Stored here for introspection; the actual
91    /// throttling is done by the `RateLimitedRequest` wrapper inside `bot`.
92    #[cfg(feature = "rate-limiter")]
93    rate_limiter: Option<Arc<dyn DynRateLimiter>>,
94
95    /// Placeholder when the rate-limiter feature is disabled.
96    #[cfg(not(feature = "rate-limiter"))]
97    rate_limiter: Option<()>,
98}
99
100// ---------------------------------------------------------------------------
101// Deref<Target = Bot> -- zero-cost access to all Bot methods
102// ---------------------------------------------------------------------------
103
104impl std::ops::Deref for ExtBot {
105    type Target = Bot;
106
107    fn deref(&self) -> &Bot {
108        &self.bot
109    }
110}
111
112impl std::fmt::Debug for ExtBot {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct("ExtBot")
115            .field("token", &self.bot.token())
116            .field("defaults", &self.defaults)
117            .field(
118                "has_callback_data_cache",
119                &self.callback_data_cache.is_some(),
120            )
121            .field("has_rate_limiter", &self.rate_limiter.is_some())
122            .finish()
123    }
124}
125
126impl ExtBot {
127    /// Creates a new `ExtBot`.
128    ///
129    /// # Arguments
130    ///
131    /// * `bot` - The underlying raw bot.
132    /// * `defaults` - Optional user-defined defaults for API calls.
133    /// * `arbitrary_callback_data` - Pass `Some(maxsize)` to enable the callback data cache,
134    ///   or `None` to disable it.  `Some(0)` uses the default maxsize of 1024.
135    /// * `rate_limiter` - An optional rate limiter.  When `Some`, the bot's request backend
136    ///   is wrapped in a `RateLimitedRequest` so all API calls are throttled.
137    ///
138    /// Prefer [`ExtBot::builder`] or [`ExtBot::from_bot`] for public construction.
139    #[cfg(feature = "rate-limiter")]
140    #[must_use]
141    pub(crate) fn new(
142        bot: Bot,
143        defaults: Option<Defaults>,
144        arbitrary_callback_data: Option<usize>,
145        rate_limiter: Option<Arc<dyn DynRateLimiter>>,
146    ) -> Self {
147        let callback_data_cache = arbitrary_callback_data.map(|maxsize| {
148            let effective = if maxsize == 0 { 1024 } else { maxsize };
149            Arc::new(RwLock::new(CallbackDataCache::new(effective)))
150        });
151
152        Self {
153            bot,
154            defaults,
155            callback_data_cache,
156            rate_limiter,
157        }
158    }
159
160    /// Creates a new `ExtBot` (no rate-limiter feature).
161    #[cfg(not(feature = "rate-limiter"))]
162    #[must_use]
163    pub(crate) fn new(
164        bot: Bot,
165        defaults: Option<Defaults>,
166        arbitrary_callback_data: Option<usize>,
167        rate_limiter: Option<()>,
168    ) -> Self {
169        let callback_data_cache = arbitrary_callback_data.map(|maxsize| {
170            let effective = if maxsize == 0 { 1024 } else { maxsize };
171            Arc::new(RwLock::new(CallbackDataCache::new(effective)))
172        });
173
174        Self {
175            bot,
176            defaults,
177            callback_data_cache,
178            rate_limiter,
179        }
180    }
181
182    /// Creates an `ExtBot` from a raw `Bot` with no extras.
183    ///
184    /// This is the simplest construction path -- no defaults, no callback data
185    /// cache, and no rate limiter.
186    ///
187    /// # Example
188    ///
189    /// ```rust,ignore
190    /// let bot = Bot::new("token", request);
191    /// let ext = ExtBot::from_bot(bot);
192    /// ```
193    #[must_use]
194    pub fn from_bot(bot: Bot) -> Self {
195        Self::new(bot, None, None, None)
196    }
197
198    /// Returns a reference to the underlying raw bot.
199    ///
200    /// Note: With the `Deref<Target = Bot>` implementation, you can call `Bot`
201    /// methods directly on `ExtBot` without using `.inner()`. This method is
202    /// retained for backward compatibility.
203    #[must_use]
204    pub fn inner(&self) -> &Bot {
205        &self.bot
206    }
207
208    /// Returns the bot token (delegates to the inner bot).
209    #[must_use]
210    pub fn token(&self) -> &str {
211        self.bot.token()
212    }
213
214    /// Returns the user-defined defaults, if any.
215    #[must_use]
216    pub fn defaults(&self) -> Option<&Defaults> {
217        self.defaults.as_ref()
218    }
219
220    /// Returns a reference to the callback data cache, if enabled.
221    #[must_use]
222    pub fn callback_data_cache(&self) -> Option<&Arc<RwLock<CallbackDataCache>>> {
223        self.callback_data_cache.as_ref()
224    }
225
226    /// Returns `true` if arbitrary callback data is enabled.
227    #[must_use]
228    pub fn has_callback_data_cache(&self) -> bool {
229        self.callback_data_cache.is_some()
230    }
231
232    /// Returns `true` if a rate limiter is configured.
233    #[must_use]
234    pub fn has_rate_limiter(&self) -> bool {
235        self.rate_limiter.is_some()
236    }
237
238    /// Returns a reference to the rate limiter, if set.
239    #[cfg(feature = "rate-limiter")]
240    #[must_use]
241    pub fn rate_limiter(&self) -> Option<&Arc<dyn DynRateLimiter>> {
242        self.rate_limiter.as_ref()
243    }
244
245    /// Returns the rate-limiter placeholder (always `None` when feature is disabled).
246    #[cfg(not(feature = "rate-limiter"))]
247    #[must_use]
248    pub fn rate_limiter(&self) -> Option<()> {
249        self.rate_limiter
250    }
251
252    /// Convenience builder entry point.
253    #[must_use]
254    pub fn builder(token: impl Into<String>, request: Arc<dyn BaseRequest>) -> ExtBotBuilder {
255        ExtBotBuilder::new(token, request)
256    }
257
258    /// Initializes the bot.
259    ///
260    /// If a rate limiter is present, it is initialized here.
261    pub async fn initialize(&self) -> rust_tg_bot_raw::error::Result<()> {
262        #[cfg(feature = "rate-limiter")]
263        if let Some(ref rl) = self.rate_limiter {
264            rl.initialize().await;
265        }
266        Ok(())
267    }
268
269    /// Shuts down the bot.
270    ///
271    /// If a rate limiter is present, it is shut down here.
272    pub async fn shutdown(&self) -> rust_tg_bot_raw::error::Result<()> {
273        #[cfg(feature = "rate-limiter")]
274        if let Some(ref rl) = self.rate_limiter {
275            rl.shutdown().await;
276        }
277        Ok(())
278    }
279}
280
281// ---------------------------------------------------------------------------
282// ExtBotBuilder
283// ---------------------------------------------------------------------------
284
285/// Builder for [`ExtBot`].
286///
287/// # Example
288///
289/// ```rust,ignore
290/// let ext = ExtBot::builder("my_token", request)
291///     .defaults(defaults)
292///     .arbitrary_callback_data(256)
293///     .rate_limiter(Arc::new(AioRateLimiter::default_limits()))
294///     .build();
295/// ```
296pub struct ExtBotBuilder {
297    token: String,
298    request: Arc<dyn BaseRequest>,
299    base_url: Option<String>,
300    base_file_url: Option<String>,
301    defaults: Option<Defaults>,
302    arbitrary_callback_data: Option<usize>,
303    #[cfg(feature = "rate-limiter")]
304    rate_limiter: Option<Arc<dyn DynRateLimiter>>,
305    #[cfg(not(feature = "rate-limiter"))]
306    rate_limiter: Option<()>,
307}
308
309impl ExtBotBuilder {
310    /// Creates a new builder with the required token and HTTP request backend.
311    #[must_use]
312    pub fn new(token: impl Into<String>, request: Arc<dyn BaseRequest>) -> Self {
313        Self {
314            token: token.into(),
315            request,
316            base_url: None,
317            base_file_url: None,
318            defaults: None,
319            arbitrary_callback_data: None,
320            rate_limiter: None,
321        }
322    }
323
324    /// Sets a custom base URL (e.g. for a local Bot API server).
325    #[must_use]
326    pub fn base_url(mut self, url: impl Into<String>) -> Self {
327        self.base_url = Some(url.into());
328        self
329    }
330
331    /// Sets a custom base file URL.
332    #[must_use]
333    pub fn base_file_url(mut self, url: impl Into<String>) -> Self {
334        self.base_file_url = Some(url.into());
335        self
336    }
337
338    /// Sets the user-defined defaults.
339    #[must_use]
340    pub fn defaults(mut self, defaults: Defaults) -> Self {
341        self.defaults = Some(defaults);
342        self
343    }
344
345    /// Enables arbitrary callback data with the given cache size.
346    ///
347    /// Pass `0` to use the default maxsize of 1024.
348    #[must_use]
349    pub fn arbitrary_callback_data(mut self, maxsize: usize) -> Self {
350        self.arbitrary_callback_data = Some(maxsize);
351        self
352    }
353
354    /// Sets the rate limiter.
355    ///
356    /// When set, the builder wraps the request backend in a
357    /// [`RateLimitedRequest`](crate::rate_limiter::RateLimitedRequest) so all
358    /// API calls are throttled transparently.
359    #[cfg(feature = "rate-limiter")]
360    #[must_use]
361    pub fn rate_limiter(mut self, rl: Arc<dyn DynRateLimiter>) -> Self {
362        self.rate_limiter = Some(rl);
363        self
364    }
365
366    /// Sets the rate-limiter placeholder (feature disabled).
367    #[cfg(not(feature = "rate-limiter"))]
368    #[must_use]
369    pub fn rate_limiter(mut self, _rl: ()) -> Self {
370        self.rate_limiter = Some(());
371        self
372    }
373
374    /// Builds the [`ExtBot`].
375    ///
376    /// If a rate limiter was provided, the HTTP request backend is wrapped in a
377    /// `RateLimitedRequest` before being passed to the inner `Bot`.
378    #[must_use]
379    pub fn build(self) -> ExtBot {
380        #[cfg(feature = "rate-limiter")]
381        let (request, rate_limiter) = if let Some(ref rl) = self.rate_limiter {
382            let wrapped: Arc<dyn BaseRequest> =
383                Arc::new(RateLimitedRequest::new(self.request.clone(), rl.clone()));
384            (wrapped, self.rate_limiter)
385        } else {
386            (self.request, None)
387        };
388
389        #[cfg(not(feature = "rate-limiter"))]
390        let (request, rate_limiter) = (self.request, self.rate_limiter);
391
392        let bot = Bot::new(&self.token, request);
393
394        ExtBot::new(
395            bot,
396            self.defaults,
397            self.arbitrary_callback_data,
398            rate_limiter,
399        )
400    }
401}
402
403impl std::fmt::Debug for ExtBotBuilder {
404    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
405        f.debug_struct("ExtBotBuilder")
406            .field("token", &"[REDACTED]")
407            .field("has_rate_limiter", &self.rate_limiter.is_some())
408            .finish()
409    }
410}
411
412// ---------------------------------------------------------------------------
413// Test-only mock request
414// ---------------------------------------------------------------------------
415
416/// A minimal mock [`BaseRequest`] for use in tests throughout the ext crate.
417///
418/// Returns `{"ok": true, "result": []}` for every request.
419#[cfg(test)]
420pub(crate) mod test_support {
421    use std::time::Duration;
422
423    use rust_tg_bot_raw::request::base::{HttpMethod, TimeoutOverride};
424    use rust_tg_bot_raw::request::request_data::RequestData;
425
426    use super::*;
427
428    #[derive(Debug)]
429    pub struct MockRequest;
430
431    #[async_trait::async_trait]
432    impl BaseRequest for MockRequest {
433        async fn initialize(&self) -> rust_tg_bot_raw::error::Result<()> {
434            Ok(())
435        }
436
437        async fn shutdown(&self) -> rust_tg_bot_raw::error::Result<()> {
438            Ok(())
439        }
440
441        fn default_read_timeout(&self) -> Option<Duration> {
442            Some(Duration::from_secs(5))
443        }
444
445        async fn do_request(
446            &self,
447            _url: &str,
448            _method: HttpMethod,
449            _request_data: Option<&RequestData>,
450            _timeouts: TimeoutOverride,
451        ) -> rust_tg_bot_raw::error::Result<(u16, bytes::Bytes)> {
452            let body = br#"{"ok":true,"result":[]}"#;
453            Ok((200, bytes::Bytes::from_static(body)))
454        }
455
456        async fn do_request_json_bytes(
457            &self,
458            _url: &str,
459            _body: &[u8],
460            _timeouts: TimeoutOverride,
461        ) -> rust_tg_bot_raw::error::Result<(u16, bytes::Bytes)> {
462            let body = br#"{"ok":true,"result":[]}"#;
463            Ok((200, bytes::Bytes::from_static(body)))
464        }
465    }
466
467    pub fn mock_request() -> Arc<dyn BaseRequest> {
468        Arc::new(MockRequest)
469    }
470}
471
472#[cfg(test)]
473mod tests {
474    use super::*;
475    use test_support::*;
476
477    #[test]
478    fn ext_bot_creation() {
479        let bot = Bot::new("test_token", mock_request());
480        let ext = ExtBot::from_bot(bot);
481
482        assert_eq!(ext.token(), "test_token");
483        assert!(ext.defaults().is_none());
484        assert!(!ext.has_callback_data_cache());
485        assert!(!ext.has_rate_limiter());
486    }
487
488    #[test]
489    fn ext_bot_with_callback_cache() {
490        let bot = Bot::new("token", mock_request());
491        let ext = ExtBot::new(bot, None, Some(512), None);
492
493        assert!(ext.has_callback_data_cache());
494    }
495
496    #[test]
497    fn ext_bot_with_defaults() {
498        let defaults = Defaults::builder().parse_mode("HTML").build();
499        let bot = Bot::new("token", mock_request());
500        let ext = ExtBot::new(bot, Some(defaults), None, None);
501
502        assert_eq!(ext.defaults().unwrap().parse_mode(), Some("HTML"));
503    }
504
505    #[test]
506    fn ext_bot_builder() {
507        let ext = ExtBot::builder("my_token", mock_request())
508            .arbitrary_callback_data(256)
509            .build();
510
511        assert_eq!(ext.token(), "my_token");
512        assert!(ext.has_callback_data_cache());
513    }
514
515    #[tokio::test]
516    async fn ext_bot_lifecycle() {
517        let bot = Bot::new("token", mock_request());
518        let ext = ExtBot::from_bot(bot);
519        assert!(ext.initialize().await.is_ok());
520        assert!(ext.shutdown().await.is_ok());
521    }
522
523    #[test]
524    fn ext_bot_debug() {
525        let bot = Bot::new("token", mock_request());
526        let ext = ExtBot::from_bot(bot);
527        let s = format!("{ext:?}");
528        assert!(s.contains("ExtBot"));
529        assert!(s.contains("token"));
530    }
531
532    #[test]
533    fn ext_bot_from_bot_convenience() {
534        let bot = Bot::new("tk", mock_request());
535        let ext = ExtBot::from_bot(bot);
536        assert_eq!(ext.token(), "tk");
537        assert!(ext.defaults().is_none());
538        assert!(!ext.has_callback_data_cache());
539        assert!(!ext.has_rate_limiter());
540    }
541
542    #[test]
543    fn ext_bot_deref_provides_bot_methods() {
544        let bot = Bot::new("deref_token", mock_request());
545        let ext = ExtBot::from_bot(bot);
546
547        // token() is available on Bot via Deref (same as ext.inner().token())
548        let deref_token: &str = (*ext).token();
549        assert_eq!(deref_token, "deref_token");
550        assert_eq!(ext.token(), deref_token);
551    }
552
553    #[cfg(feature = "rate-limiter")]
554    #[test]
555    fn ext_bot_builder_with_rate_limiter() {
556        use crate::rate_limiter::NoRateLimiter;
557
558        let limiter: Arc<dyn DynRateLimiter> = Arc::new(NoRateLimiter);
559        let ext = ExtBot::builder("rl_token", mock_request())
560            .rate_limiter(limiter)
561            .build();
562
563        assert_eq!(ext.token(), "rl_token");
564        assert!(ext.has_rate_limiter());
565        assert!(ext.rate_limiter().is_some());
566    }
567
568    #[cfg(feature = "rate-limiter")]
569    #[test]
570    fn ext_bot_builder_without_rate_limiter() {
571        let ext = ExtBot::builder("no_rl", mock_request()).build();
572
573        assert!(!ext.has_rate_limiter());
574        assert!(ext.rate_limiter().is_none());
575    }
576
577    #[cfg(feature = "rate-limiter")]
578    #[tokio::test]
579    async fn ext_bot_lifecycle_with_rate_limiter() {
580        use crate::rate_limiter::NoRateLimiter;
581
582        let limiter: Arc<dyn DynRateLimiter> = Arc::new(NoRateLimiter);
583        let ext = ExtBot::builder("rl_lc", mock_request())
584            .rate_limiter(limiter)
585            .build();
586
587        assert!(ext.initialize().await.is_ok());
588        assert!(ext.shutdown().await.is_ok());
589    }
590}
rust_tg_bot_ext/ext_bot.rs

rust_tg_bot_ext/
ext_bot.rs
