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
26pub 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#[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 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 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 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}