1use sift_query_core::{QueryPlan, SearchEntity, SearchOptions};
2use sqlx::FromRow;
3use sqlx::Postgres;
4use sqlx::pool::PoolConnection;
5use sqlx::postgres::PgRow;
6use tokio::time::{Instant, timeout_at};
7
8use crate::execution::io::{PreparedPgQuery, close_hard, fetch_cluster, fetch_maps, fetch_typed};
9use crate::{Page, PgSearchError, PgSiftQuery, RowMap, TimeoutStage};
10
11#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
12pub enum ExecutionStrategy {
13 #[default]
14 SequentialBestEffort,
15 ParallelBestEffort,
16 Snapshot(SnapshotIsolation),
17}
18
19#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
20pub enum SnapshotIsolation {
21 #[default]
22 RepeatableRead,
23 Serializable,
24}
25
26#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
27pub enum CountStrategy {
28 #[default]
29 Separate,
30 WindowWhenPossible,
31}
32
33impl PgSiftQuery {
34 pub(crate) async fn execute_typed_pair<E, R>(
35 &self,
36 list: Option<PreparedPgQuery>,
37 cluster: PreparedPgQuery,
38 options: Option<&SearchOptions<E>>,
39 deadline: Instant,
40 ) -> Result<Page<R>, PgSearchError>
41 where
42 E: SearchEntity,
43 R: Send + Unpin + for<'row> FromRow<'row, PgRow>,
44 {
45 match self.execution.strategy {
46 ExecutionStrategy::SequentialBestEffort => {
47 let connection = self.acquire(deadline).await?;
48 let (connection, data) = match list {
49 Some(rendered) => {
50 fetch_typed::<R>(connection, rendered, deadline, TimeoutStage::ListQuery)
51 .await?
52 }
53 None => (connection, Vec::new()),
54 };
55 let (_, (total_count, summaries)) =
56 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
57 Ok(Page {
58 data,
59 total_count,
60 summaries,
61 next_cursor: None,
62 })
63 }
64 ExecutionStrategy::ParallelBestEffort if list.is_some() => {
65 let permit = timeout_at(deadline, self.parallel_gate.acquire())
66 .await
67 .map_err(|_| PgSearchError::Timeout(TimeoutStage::PoolAcquire))?
68 .map_err(|_| PgSearchError::ParallelGateClosed)?;
69 let list = list.expect("checked above");
70 let list_future = async {
71 let connection = self.acquire(deadline).await?;
72 fetch_typed::<R>(connection, list, deadline, TimeoutStage::ListQuery).await
73 };
74 let cluster_future = async {
75 let connection = self.acquire(deadline).await?;
76 fetch_cluster(connection, cluster, options, E::META, deadline).await
77 };
78 let (list_result, cluster_result) =
79 futures_util::future::join(list_future, cluster_future).await;
80 drop(permit);
81 let (_, data) = list_result?;
82 let (_, (total_count, summaries)) = cluster_result?;
83 Ok(Page {
84 data,
85 total_count,
86 summaries,
87 next_cursor: None,
88 })
89 }
90 ExecutionStrategy::ParallelBestEffort => {
91 let connection = self.acquire(deadline).await?;
92 let (_, (total_count, summaries)) =
93 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
94 Ok(Page {
95 data: Vec::new(),
96 total_count,
97 summaries,
98 next_cursor: None,
99 })
100 }
101 ExecutionStrategy::Snapshot(isolation) => {
102 let mut connection = self.acquire(deadline).await?;
103 connection.close_on_drop();
104 let connection = begin_snapshot(connection, isolation, deadline).await?;
105 let (connection, data) = match list {
106 Some(rendered) => {
107 fetch_typed::<R>(connection, rendered, deadline, TimeoutStage::ListQuery)
108 .await?
109 }
110 None => (connection, Vec::new()),
111 };
112 let (connection, (total_count, summaries)) =
113 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
114 let mut connection = commit(connection, deadline).await?;
115 connection.return_to_pool().await;
116 Ok(Page {
117 data,
118 total_count,
119 summaries,
120 next_cursor: None,
121 })
122 }
123 }
124 }
125
126 pub(crate) async fn execute_map_pair<E: SearchEntity>(
127 &self,
128 plan: &QueryPlan<E>,
129 list: Option<PreparedPgQuery>,
130 cluster: PreparedPgQuery,
131 options: Option<&SearchOptions<E>>,
132 deadline: Instant,
133 ) -> Result<Page<RowMap>, PgSearchError> {
134 match self.execution.strategy {
135 ExecutionStrategy::SequentialBestEffort => {
136 let connection = self.acquire(deadline).await?;
137 let (connection, data) = match list {
138 Some(rendered) => fetch_maps(connection, rendered, plan, deadline).await?,
139 None => (connection, Vec::new()),
140 };
141 let (_, (total_count, summaries)) =
142 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
143 Ok(Page {
144 data,
145 total_count,
146 summaries,
147 next_cursor: None,
148 })
149 }
150 ExecutionStrategy::ParallelBestEffort if list.is_some() => {
151 let permit = timeout_at(deadline, self.parallel_gate.acquire())
152 .await
153 .map_err(|_| PgSearchError::Timeout(TimeoutStage::PoolAcquire))?
154 .map_err(|_| PgSearchError::ParallelGateClosed)?;
155 let list = list.expect("checked above");
156 let list_future = async {
157 let connection = self.acquire(deadline).await?;
158 fetch_maps(connection, list, plan, deadline).await
159 };
160 let cluster_future = async {
161 let connection = self.acquire(deadline).await?;
162 fetch_cluster(connection, cluster, options, E::META, deadline).await
163 };
164 let (list_result, cluster_result) =
165 futures_util::future::join(list_future, cluster_future).await;
166 drop(permit);
167 let (_, data) = list_result?;
168 let (_, (total_count, summaries)) = cluster_result?;
169 Ok(Page {
170 data,
171 total_count,
172 summaries,
173 next_cursor: None,
174 })
175 }
176 ExecutionStrategy::ParallelBestEffort => {
177 let connection = self.acquire(deadline).await?;
178 let (_, (total_count, summaries)) =
179 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
180 Ok(Page {
181 data: Vec::new(),
182 total_count,
183 summaries,
184 next_cursor: None,
185 })
186 }
187 ExecutionStrategy::Snapshot(isolation) => {
188 let mut connection = self.acquire(deadline).await?;
189 connection.close_on_drop();
190 let connection = begin_snapshot(connection, isolation, deadline).await?;
191 let (connection, data) = match list {
192 Some(rendered) => fetch_maps(connection, rendered, plan, deadline).await?,
193 None => (connection, Vec::new()),
194 };
195 let (connection, (total_count, summaries)) =
196 fetch_cluster(connection, cluster, options, E::META, deadline).await?;
197 let mut connection = commit(connection, deadline).await?;
198 connection.return_to_pool().await;
199 Ok(Page {
200 data,
201 total_count,
202 summaries,
203 next_cursor: None,
204 })
205 }
206 }
207 }
208}
209
210async fn begin_snapshot(
211 connection: PoolConnection<Postgres>,
212 isolation: SnapshotIsolation,
213 deadline: Instant,
214) -> Result<PoolConnection<Postgres>, PgSearchError> {
215 let sql = match isolation {
216 SnapshotIsolation::RepeatableRead => "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY",
217 SnapshotIsolation::Serializable => "BEGIN ISOLATION LEVEL SERIALIZABLE READ ONLY",
218 };
219 execute_control(connection, sql, deadline).await
220}
221
222async fn commit(
223 connection: PoolConnection<Postgres>,
224 deadline: Instant,
225) -> Result<PoolConnection<Postgres>, PgSearchError> {
226 execute_control(connection, "COMMIT", deadline).await
227}
228
229async fn execute_control(
230 mut connection: PoolConnection<Postgres>,
231 sql: &'static str,
232 deadline: Instant,
233) -> Result<PoolConnection<Postgres>, PgSearchError> {
234 let result = timeout_at(deadline, sqlx::query(sql).execute(&mut *connection)).await;
235 let Ok(result) = result else {
236 close_hard(connection).await;
237 return Err(PgSearchError::Timeout(TimeoutStage::ClusterQuery));
238 };
239 match result {
240 Ok(_) => Ok(connection),
241 Err(error) => {
242 close_hard(connection).await;
243 Err(PgSearchError::Sql(error))
244 }
245 }
246}
247
248#[cfg(test)]
249mod tests {
250 use super::*;
251 use std::time::Duration;
252
253 #[tokio::test]
254 #[ignore = "requires a disposable PostgreSQL instance in DATABASE_URL"]
255 async fn snapshot_control_faults_discard_connection_and_release_pool_slot() {
256 let url = std::env::var("DATABASE_URL").unwrap();
257 let pool = sqlx::postgres::PgPoolOptions::new()
258 .max_connections(1)
259 .connect(&url)
260 .await
261 .unwrap();
262 let admin = sqlx::postgres::PgPoolOptions::new()
263 .max_connections(1)
264 .connect(&url)
265 .await
266 .unwrap();
267 for control in ["COMMIT", "ROLLBACK"] {
268 let mut connection = pool.acquire().await.unwrap();
269 connection.close_on_drop();
270 let mut connection = begin_snapshot(
271 connection,
272 SnapshotIsolation::RepeatableRead,
273 Instant::now() + Duration::from_secs(5),
274 )
275 .await
276 .unwrap();
277 let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
278 .fetch_one(&mut *connection)
279 .await
280 .unwrap();
281 assert!(
283 sqlx::query_scalar::<_, bool>("SELECT pg_terminate_backend($1)")
284 .bind(pid)
285 .fetch_one(&admin)
286 .await
287 .unwrap()
288 );
289 assert!(
290 execute_control(connection, control, Instant::now() + Duration::from_secs(5))
291 .await
292 .is_err()
293 );
294 assert_clean(&pool).await;
295 }
296 let mut connection = pool.acquire().await.unwrap();
297 connection.close_on_drop();
298 let connection = begin_snapshot(
299 connection,
300 SnapshotIsolation::Serializable,
301 Instant::now() + Duration::from_secs(5),
302 )
303 .await
304 .unwrap();
305 assert!(matches!(
306 commit(connection, Instant::now() - Duration::from_secs(1)).await,
307 Err(PgSearchError::Timeout(_))
308 ));
309 assert_clean(&pool).await;
310 pool.close().await;
311 admin.close().await;
312 }
313
314 async fn assert_clean(pool: &sqlx::PgPool) {
315 let read_only: String = tokio::time::timeout(
316 Duration::from_secs(5),
317 sqlx::query_scalar("SHOW transaction_read_only").fetch_one(pool),
318 )
319 .await
320 .unwrap()
321 .unwrap();
322 assert_eq!(read_only, "off");
323 }
324}