Skip to main content

sift_query_sqlx/
router.rs

1use std::any::TypeId;
2use std::collections::HashMap;
3use std::sync::Arc;
4
5use sift_query_core::{EntityMeta, QueryInput, SearchEntity, ServerScope};
6use sqlx::postgres::PgPool;
7
8use crate::{Page, PgSearchError, PgSiftQuery, RowMap};
9
10#[derive(Debug, thiserror::Error)]
11pub enum PoolRouteError {
12    #[error("no PostgreSQL pool is registered for entity {0}")]
13    MissingEntity(&'static str),
14    #[error("invalid PostgreSQL pool name")]
15    InvalidName,
16    #[error("duplicate PostgreSQL pool name: {0}")]
17    DuplicateName(String),
18    #[error("PostgreSQL pool is not registered: {0}")]
19    MissingPool(String),
20    #[error("server pool selection rejected: {0}")]
21    Rejected(&'static str),
22    #[error("query failed on routed PostgreSQL pool: {0}")]
23    Query(#[from] PgSearchError),
24}
25
26/// Application-owned, nonblocking resolver for trusted request context.
27/// Client query parameters are never supplied to this hook.
28pub trait PgPoolSelector<C>: Send + Sync {
29    fn select_pool(
30        &self,
31        entity: &'static EntityMeta,
32        context: &C,
33    ) -> Result<String, PoolRouteError>;
34}
35
36/// Immutable server-owned entity-to-pool routing table.
37#[derive(Clone, Default)]
38pub struct PgPoolRouter {
39    default: Option<Arc<PgSiftQuery>>,
40    routes: HashMap<TypeId, Arc<PgSiftQuery>>,
41    named: HashMap<String, Arc<PgSiftQuery>>,
42    labels: HashMap<TypeId, String>,
43    default_label: String,
44}
45
46impl std::fmt::Debug for PgPoolRouter {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        formatter
49            .debug_struct("PgPoolRouter")
50            .field("has_default", &self.default.is_some())
51            .field("route_count", &self.routes.len())
52            .field("named_pool_count", &self.named.len())
53            .finish_non_exhaustive()
54    }
55}
56
57impl PgPoolRouter {
58    #[must_use]
59    pub fn new(default: PgPool) -> Self {
60        Self {
61            default: Some(Arc::new(PgSiftQuery::new(default))),
62            routes: HashMap::new(),
63            default_label: "default".to_owned(),
64            ..Self::default()
65        }
66    }
67
68    #[must_use]
69    pub fn with_searcher(default: PgSiftQuery) -> Self {
70        Self {
71            default: Some(Arc::new(default)),
72            routes: HashMap::new(),
73            default_label: "default".to_owned(),
74            ..Self::default()
75        }
76    }
77
78    #[must_use]
79    pub fn register_entity<E: SearchEntity>(mut self, searcher: PgSiftQuery) -> Self {
80        self.routes.insert(TypeId::of::<E>(), Arc::new(searcher));
81        self.labels
82            .insert(TypeId::of::<E>(), E::META.entity_name.to_owned());
83        self
84    }
85
86    #[must_use]
87    pub fn register_pool<E: SearchEntity>(mut self, pool: PgPool) -> Self {
88        self.routes
89            .insert(TypeId::of::<E>(), Arc::new(PgSiftQuery::new(pool)));
90        self.labels
91            .insert(TypeId::of::<E>(), E::META.entity_name.to_owned());
92        self
93    }
94
95    /// Names are stable configuration labels, never tenant/user identifiers.
96    pub fn register_named(
97        mut self,
98        name: &str,
99        searcher: PgSiftQuery,
100    ) -> Result<Self, PoolRouteError> {
101        if name.is_empty()
102            || name.len() > 64
103            || !name
104                .bytes()
105                .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
106        {
107            return Err(PoolRouteError::InvalidName);
108        }
109        if self.named.contains_key(name) {
110            return Err(PoolRouteError::DuplicateName(name.to_owned()));
111        }
112        self.named.insert(name.to_owned(), Arc::new(searcher));
113        Ok(self)
114    }
115
116    pub fn route_entity<E: SearchEntity>(mut self, name: &str) -> Result<Self, PoolRouteError> {
117        let searcher = self
118            .named
119            .get(name)
120            .cloned()
121            .ok_or_else(|| PoolRouteError::MissingPool(name.to_owned()))?;
122        self.routes.insert(TypeId::of::<E>(), searcher);
123        self.labels.insert(TypeId::of::<E>(), name.to_owned());
124        Ok(self)
125    }
126
127    pub fn with_default_named(mut self, name: &str) -> Result<Self, PoolRouteError> {
128        self.default = Some(
129            self.named
130                .get(name)
131                .cloned()
132                .ok_or_else(|| PoolRouteError::MissingPool(name.to_owned()))?,
133        );
134        name.clone_into(&mut self.default_label);
135        Ok(self)
136    }
137
138    /// Selected searcher exposes normal query, scan and export entry points.
139    pub fn searcher<E: SearchEntity>(&self) -> Result<Arc<PgSiftQuery>, PoolRouteError> {
140        let result = self
141            .routes
142            .get(&TypeId::of::<E>())
143            .cloned()
144            .or_else(|| self.default.clone())
145            .ok_or(PoolRouteError::MissingEntity(E::META.entity_name))?;
146        let label = self
147            .labels
148            .get(&TypeId::of::<E>())
149            .unwrap_or(&self.default_label);
150        tracing::debug!(entity = E::META.entity_name, backend = "postgres", pool_route = %label, "selected PostgreSQL pool");
151        Ok(result)
152    }
153
154    pub fn searcher_with_context<E: SearchEntity, C>(
155        &self,
156        context: &C,
157        selector: &impl PgPoolSelector<C>,
158    ) -> Result<Arc<PgSiftQuery>, PoolRouteError> {
159        let name = selector.select_pool(E::META, context)?;
160        if name.is_empty() {
161            return self.searcher::<E>();
162        }
163        // Unknown selection fails closed, with no cross-pool fallback.
164        let searcher = self
165            .named
166            .get(&name)
167            .ok_or_else(|| PoolRouteError::MissingPool(name.clone()))?;
168        tracing::debug!(entity = E::META.entity_name, backend = "postgres", pool_route = %name, "selected PostgreSQL pool from trusted context");
169        Ok(Arc::clone(searcher))
170    }
171
172    pub async fn search_map<E: SearchEntity>(
173        &self,
174        query: impl QueryInput<E>,
175        scope: &ServerScope<E>,
176    ) -> Result<Page<RowMap>, PoolRouteError> {
177        Ok(self.searcher::<E>()?.search_map(query, scope).await?)
178    }
179
180    pub async fn search_with_context<E, C>(
181        &self,
182        query: impl QueryInput<E>,
183        scope: &ServerScope<E>,
184        context: &C,
185        selector: &impl PgPoolSelector<C>,
186    ) -> Result<Page<E>, PoolRouteError>
187    where
188        E: SearchEntity + Send + Unpin + for<'row> sqlx::FromRow<'row, sqlx::postgres::PgRow>,
189    {
190        Ok(self
191            .searcher_with_context::<E, C>(context, selector)?
192            .search(query, scope)
193            .await?)
194    }
195
196    pub async fn search_map_with_context<E: SearchEntity, C>(
197        &self,
198        query: impl QueryInput<E>,
199        scope: &ServerScope<E>,
200        context: &C,
201        selector: &impl PgPoolSelector<C>,
202    ) -> Result<Page<RowMap>, PoolRouteError> {
203        Ok(self
204            .searcher_with_context::<E, C>(context, selector)?
205            .search_map(query, scope)
206            .await?)
207    }
208
209    pub async fn search<E>(
210        &self,
211        query: impl QueryInput<E>,
212        scope: &ServerScope<E>,
213    ) -> Result<Page<E>, PoolRouteError>
214    where
215        E: SearchEntity + Send + Unpin + for<'row> sqlx::FromRow<'row, sqlx::postgres::PgRow>,
216    {
217        Ok(self.searcher::<E>()?.search(query, scope).await?)
218    }
219}