shardline_server/postgres_backend/
backend.rs1use std::{num::NonZeroUsize, path::PathBuf};
2
3use shardline_index::{PostgresIndexStore, PostgresRecordStore};
4
5use super::connect_postgres_metadata_pool;
6use crate::{
7 ServerError, ServerFrontend, config::default_upload_max_in_flight_chunks,
8 object_store::ServerObjectStore,
9};
10
11#[derive(Debug, Clone)]
13pub struct PostgresBackend {
14 pub(super) public_base_url: String,
15 pub(super) chunk_size: NonZeroUsize,
16 pub(super) upload_max_in_flight_chunks: NonZeroUsize,
17 pub(super) server_frontends: Vec<ServerFrontend>,
18 pub(super) index_store: PostgresIndexStore,
19 pub(super) record_store: PostgresRecordStore,
20 pub(super) object_store: ServerObjectStore,
21}
22
23impl PostgresBackend {
24 pub async fn new(
31 root: PathBuf,
32 public_base_url: String,
33 chunk_size: NonZeroUsize,
34 index_postgres_url: &str,
35 ) -> Result<Self, ServerError> {
36 let object_store = ServerObjectStore::local(root.join("chunks"))?;
37 Self::new_with_object_store(
38 root,
39 public_base_url,
40 chunk_size,
41 index_postgres_url,
42 object_store,
43 )
44 .await
45 }
46
47 pub(crate) async fn new_with_object_store(
48 root: PathBuf,
49 public_base_url: String,
50 chunk_size: NonZeroUsize,
51 index_postgres_url: &str,
52 object_store: ServerObjectStore,
53 ) -> Result<Self, ServerError> {
54 Self::new_with_object_store_and_upload_parallelism_with_frontends(
55 root,
56 public_base_url,
57 chunk_size,
58 default_upload_max_in_flight_chunks(),
59 index_postgres_url,
60 object_store,
61 &[ServerFrontend::Xet],
62 )
63 .await
64 }
65
66 #[cfg(test)]
67 pub(crate) async fn new_with_object_store_and_upload_parallelism(
68 _root: PathBuf,
69 public_base_url: String,
70 chunk_size: NonZeroUsize,
71 upload_max_in_flight_chunks: NonZeroUsize,
72 index_postgres_url: &str,
73 object_store: ServerObjectStore,
74 ) -> Result<Self, ServerError> {
75 Self::new_with_object_store_and_upload_parallelism_with_frontends(
76 _root,
77 public_base_url,
78 chunk_size,
79 upload_max_in_flight_chunks,
80 index_postgres_url,
81 object_store,
82 &[ServerFrontend::Xet],
83 )
84 .await
85 }
86
87 pub(crate) async fn new_with_object_store_and_upload_parallelism_with_frontends(
88 _root: PathBuf,
89 public_base_url: String,
90 chunk_size: NonZeroUsize,
91 upload_max_in_flight_chunks: NonZeroUsize,
92 index_postgres_url: &str,
93 object_store: ServerObjectStore,
94 server_frontends: &[ServerFrontend],
95 ) -> Result<Self, ServerError> {
96 let pool = connect_postgres_metadata_pool(index_postgres_url, 10)?;
97
98 Ok(Self {
99 public_base_url,
100 chunk_size,
101 upload_max_in_flight_chunks,
102 server_frontends: server_frontends.to_vec(),
103 index_store: PostgresIndexStore::new(pool.clone()),
104 record_store: PostgresRecordStore::new(pool),
105 object_store,
106 })
107 }
108
109 #[must_use]
111 pub fn public_base_url(&self) -> &str {
112 &self.public_base_url
113 }
114
115 pub(crate) const fn object_backend_name(&self) -> &'static str {
116 self.object_store.backend_name()
117 }
118
119 pub(crate) fn object_store(&self) -> ServerObjectStore {
120 self.object_store.clone()
121 }
122
123 pub(crate) async fn probe_metadata(&self) -> Result<(), String> {
125 self.index_store.probe().await
126 }
127}
128
129#[cfg(test)]
130mod tests {
131 use super::*;
132 use crate::object_store::ServerObjectStore;
133 use std::num::NonZeroUsize;
134 use tempfile::tempdir;
135
136 const TEST_POSTGRES_URL: &str = "postgres://localhost:5432/test";
137
138 async fn make_backend() -> PostgresBackend {
139 let root = tempdir().unwrap();
140 let object_store = ServerObjectStore::local(root.path().join("chunks")).unwrap();
141 PostgresBackend::new_with_object_store_and_upload_parallelism(
142 root.path().to_path_buf(),
143 "http://127.0.0.1:8080".to_owned(),
144 NonZeroUsize::new(65536).unwrap(),
145 NonZeroUsize::new(64).unwrap(),
146 TEST_POSTGRES_URL,
147 object_store,
148 )
149 .await
150 .expect("constructor should succeed with lazy pool")
151 }
152
153 #[tokio::test]
154 async fn constructor_succeeds() {
155 let root = tempdir().unwrap();
156 let object_store = ServerObjectStore::local(root.path().join("chunks")).unwrap();
157 let backend = PostgresBackend::new_with_object_store_and_upload_parallelism(
158 root.path().to_path_buf(),
159 "http://127.0.0.1:8080".to_owned(),
160 NonZeroUsize::new(65536).unwrap(),
161 NonZeroUsize::new(64).unwrap(),
162 TEST_POSTGRES_URL,
163 object_store,
164 )
165 .await;
166 assert!(backend.is_ok(), "constructor should succeed with lazy pool");
167 }
168
169 #[tokio::test]
170 async fn constructor_fails_empty_url() {
171 let root = tempdir().unwrap();
172 let object_store = ServerObjectStore::local(root.path().join("chunks")).unwrap();
173 let backend = PostgresBackend::new_with_object_store_and_upload_parallelism(
174 root.path().to_path_buf(),
175 "http://127.0.0.1:8080".to_owned(),
176 NonZeroUsize::new(65536).unwrap(),
177 NonZeroUsize::new(64).unwrap(),
178 "",
179 object_store,
180 )
181 .await;
182 assert!(
183 backend.is_err(),
184 "constructor should fail with empty postgres URL"
185 );
186 }
187
188 #[tokio::test]
189 async fn public_base_url() {
190 let backend = make_backend().await;
191 assert_eq!(backend.public_base_url(), "http://127.0.0.1:8080");
192 }
193
194 #[tokio::test]
195 async fn object_backend_name() {
196 let backend = make_backend().await;
197 assert_eq!(backend.object_backend_name(), "local");
198 }
199
200 #[tokio::test]
201 async fn object_store() {
202 let root = tempdir().unwrap();
203 let object_store = ServerObjectStore::local(root.path().join("chunks")).unwrap();
204 let backend = PostgresBackend::new_with_object_store_and_upload_parallelism(
205 root.path().to_path_buf(),
206 "http://127.0.0.1:8080".to_owned(),
207 NonZeroUsize::new(65536).unwrap(),
208 NonZeroUsize::new(64).unwrap(),
209 TEST_POSTGRES_URL,
210 object_store.clone(),
211 )
212 .await
213 .unwrap();
214 let retrieved = backend.object_store();
215 assert_eq!(retrieved.backend_name(), object_store.backend_name());
217 }
218}