Skip to main content

shardline_server/postgres_backend/
backend.rs

1use 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/// Server backend that keeps file metadata in Postgres and object bytes in the selected store.
12#[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    /// Creates a Postgres-backed metadata backend.
25    ///
26    /// # Errors
27    ///
28    /// Returns [`ServerError`] when the local chunk store cannot initialize or the
29    /// Postgres pool configuration is invalid.
30    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    /// Returns the public base URL used in generated download links.
110    #[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    /// Probes the Postgres metadata store for connectivity.
124    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        // ServerObjectStore implements Clone but not PartialEq, so compare backend names
216        assert_eq!(retrieved.backend_name(), object_store.backend_name());
217    }
218}