pipedash_core/infrastructure/storage/
postgres.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use chrono::{
5    DateTime,
6    Utc,
7};
8use sqlx::postgres::PgPool;
9use sqlx::Row;
10
11use super::{
12    ObjectMetadata,
13    StorageBackend,
14};
15use crate::domain::{
16    DomainError,
17    DomainResult,
18};
19
20#[derive(Debug, Clone)]
21pub struct PostgresCacheConfig {
22    pub connection_string: String,
23}
24
25pub struct PostgresStorage {
26    pool: PgPool,
27}
28
29impl PostgresStorage {
30    pub fn new(pool: PgPool) -> Self {
31        Self { pool }
32    }
33
34    pub async fn from_connection_string(connection_string: &str) -> DomainResult<Self> {
35        use sqlx::postgres::PgPoolOptions;
36
37        let pool = PgPoolOptions::new()
38            .max_connections(20)
39            .min_connections(5)
40            .acquire_timeout(std::time::Duration::from_secs(30))
41            .connect(connection_string)
42            .await
43            .map_err(|e| {
44                DomainError::DatabaseError(format!("Failed to connect to PostgreSQL: {}", e))
45            })?;
46
47        Self::init_cache_table(&pool).await?;
48
49        Ok(Self { pool })
50    }
51
52    async fn init_cache_table(pool: &PgPool) -> DomainResult<()> {
53        sqlx::query(
54            r#"
55            CREATE TABLE IF NOT EXISTS cache_objects (
56                key TEXT PRIMARY KEY,
57                data BYTEA NOT NULL,
58                size BIGINT NOT NULL,
59                content_type TEXT,
60                etag TEXT,
61                metadata JSONB DEFAULT '{}',
62                created_at TIMESTAMPTZ DEFAULT NOW(),
63                updated_at TIMESTAMPTZ DEFAULT NOW()
64            )
65            "#,
66        )
67        .execute(pool)
68        .await
69        .map_err(|e| DomainError::DatabaseError(format!("Failed to create cache table: {}", e)))?;
70
71        sqlx::query(
72            r#"
73            CREATE INDEX IF NOT EXISTS idx_cache_objects_key ON cache_objects(key)
74            "#,
75        )
76        .execute(pool)
77        .await
78        .map_err(|e| DomainError::DatabaseError(format!("Failed to create cache index: {}", e)))?;
79
80        Ok(())
81    }
82
83    fn generate_etag(data: &[u8]) -> String {
84        use std::collections::hash_map::DefaultHasher;
85        use std::hash::{
86            Hash,
87            Hasher,
88        };
89
90        let mut hasher = DefaultHasher::new();
91        data.hash(&mut hasher);
92        format!("{:x}", hasher.finish())
93    }
94}
95
96#[async_trait]
97impl StorageBackend for PostgresStorage {
98    fn backend_type(&self) -> &str {
99        "postgres"
100    }
101
102    async fn is_available(&self) -> bool {
103        sqlx::query("SELECT 1").fetch_one(&self.pool).await.is_ok()
104    }
105
106    async fn list(&self, prefix: Option<&str>) -> DomainResult<Vec<ObjectMetadata>> {
107        let rows = if let Some(prefix) = prefix {
108            let pattern = format!("{}%", prefix);
109            sqlx::query(
110                r#"
111                SELECT key, size, content_type, etag, metadata, updated_at
112                FROM cache_objects
113                WHERE key LIKE $1
114                ORDER BY key
115                "#,
116            )
117            .bind(&pattern)
118            .fetch_all(&self.pool)
119            .await
120        } else {
121            sqlx::query(
122                r#"
123                SELECT key, size, content_type, etag, metadata, updated_at
124                FROM cache_objects
125                ORDER BY key
126                "#,
127            )
128            .fetch_all(&self.pool)
129            .await
130        }
131        .map_err(|e| DomainError::DatabaseError(format!("Failed to list cache objects: {}", e)))?;
132
133        let mut objects = Vec::new();
134        for row in rows {
135            let key: String = row.get("key");
136            let size: i64 = row.get("size");
137            let content_type: Option<String> = row.get("content_type");
138            let etag: Option<String> = row.get("etag");
139            let metadata_json: serde_json::Value = row.get("metadata");
140            let updated_at: DateTime<Utc> = row.get("updated_at");
141
142            let metadata: HashMap<String, String> =
143                serde_json::from_value(metadata_json).unwrap_or_default();
144
145            objects.push(ObjectMetadata {
146                key,
147                size: size as u64,
148                last_modified: updated_at,
149                etag,
150                content_type,
151                metadata,
152            });
153        }
154
155        Ok(objects)
156    }
157
158    async fn get(&self, key: &str) -> DomainResult<Vec<u8>> {
159        let row = sqlx::query(
160            r#"
161            SELECT data FROM cache_objects WHERE key = $1
162            "#,
163        )
164        .bind(key)
165        .fetch_optional(&self.pool)
166        .await
167        .map_err(|e| DomainError::DatabaseError(format!("Failed to get cache object: {}", e)))?;
168
169        match row {
170            Some(row) => {
171                let data: Vec<u8> = row.get("data");
172                Ok(data)
173            }
174            None => Err(DomainError::NotFound(format!(
175                "Cache object not found: {}",
176                key
177            ))),
178        }
179    }
180
181    async fn put(
182        &self, key: &str, data: &[u8], content_type: Option<&str>,
183    ) -> DomainResult<ObjectMetadata> {
184        let size = data.len() as i64;
185        let etag = Self::generate_etag(data);
186        let now = Utc::now();
187
188        sqlx::query(
189            r#"
190            INSERT INTO cache_objects (key, data, size, content_type, etag, updated_at)
191            VALUES ($1, $2, $3, $4, $5, $6)
192            ON CONFLICT (key) DO UPDATE SET
193                data = EXCLUDED.data,
194                size = EXCLUDED.size,
195                content_type = EXCLUDED.content_type,
196                etag = EXCLUDED.etag,
197                updated_at = EXCLUDED.updated_at
198            "#,
199        )
200        .bind(key)
201        .bind(data)
202        .bind(size)
203        .bind(content_type)
204        .bind(&etag)
205        .bind(now)
206        .execute(&self.pool)
207        .await
208        .map_err(|e| DomainError::DatabaseError(format!("Failed to put cache object: {}", e)))?;
209
210        Ok(ObjectMetadata {
211            key: key.to_string(),
212            size: size as u64,
213            last_modified: now,
214            etag: Some(etag),
215            content_type: content_type.map(String::from),
216            metadata: HashMap::new(),
217        })
218    }
219
220    async fn delete(&self, key: &str) -> DomainResult<()> {
221        sqlx::query(
222            r#"
223            DELETE FROM cache_objects WHERE key = $1
224            "#,
225        )
226        .bind(key)
227        .execute(&self.pool)
228        .await
229        .map_err(|e| DomainError::DatabaseError(format!("Failed to delete cache object: {}", e)))?;
230
231        Ok(())
232    }
233
234    async fn exists(&self, key: &str) -> DomainResult<bool> {
235        let row = sqlx::query(
236            r#"
237            SELECT 1 FROM cache_objects WHERE key = $1
238            "#,
239        )
240        .bind(key)
241        .fetch_optional(&self.pool)
242        .await
243        .map_err(|e| DomainError::DatabaseError(format!("Failed to check cache object: {}", e)))?;
244
245        Ok(row.is_some())
246    }
247
248    async fn head(&self, key: &str) -> DomainResult<Option<ObjectMetadata>> {
249        let row = sqlx::query(
250            r#"
251            SELECT key, size, content_type, etag, metadata, updated_at
252            FROM cache_objects
253            WHERE key = $1
254            "#,
255        )
256        .bind(key)
257        .fetch_optional(&self.pool)
258        .await
259        .map_err(|e| {
260            DomainError::DatabaseError(format!("Failed to get cache object metadata: {}", e))
261        })?;
262
263        match row {
264            Some(row) => {
265                let key: String = row.get("key");
266                let size: i64 = row.get("size");
267                let content_type: Option<String> = row.get("content_type");
268                let etag: Option<String> = row.get("etag");
269                let metadata_json: serde_json::Value = row.get("metadata");
270                let updated_at: DateTime<Utc> = row.get("updated_at");
271
272                let metadata: HashMap<String, String> =
273                    serde_json::from_value(metadata_json).unwrap_or_default();
274
275                Ok(Some(ObjectMetadata {
276                    key,
277                    size: size as u64,
278                    last_modified: updated_at,
279                    etag,
280                    content_type,
281                    metadata,
282                }))
283            }
284            None => Ok(None),
285        }
286    }
287}
288
289#[cfg(test)]
290mod tests {}
pipedash_core/infrastructure/storage/postgres.rs

pipedash_core/infrastructure/storage/
postgres.rs
