Something went wrong. Try again.
Our Personal Data Server from scratch!
Something went wrong. Try again.
8.9 kB · 298 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299use async_trait::async_trait;use sqlx::PgPool;use tranquil_db_traits::{ BlobForExport, BlobMetadata, BlobRepository, BlobWithTakedown, DbError, MissingBlobInfo,};use tranquil_types::{AtUri, CidLink, Did, Tid};use uuid::Uuid;
use super::col;use super::user::map_sqlx_error;use super::{column, column_vec, opt_column};
pub struct PostgresBlobRepository { pool: PgPool,}
impl PostgresBlobRepository { pub fn new(pool: PgPool) -> Self { Self { pool } }}
#[async_trait]impl BlobRepository for PostgresBlobRepository { async fn insert_blob( &self, cid: &CidLink, mime_type: &str, size_bytes: i64, created_by_user: Uuid, storage_key: &str, ) -> Result<Option<CidLink>, DbError> { let result = sqlx::query_scalar!( r#"INSERT INTO blobs (cid, mime_type, size_bytes, created_by_user, storage_key) VALUES ($1, $2, $3, $4, $5) ON CONFLICT (cid, created_by_user) DO NOTHING RETURNING cid"#, cid.as_str(), mime_type, size_bytes, created_by_user, storage_key ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
opt_column(result, col::BLOBS_CID) }
async fn get_blob_metadata(&self, cid: &CidLink) -> Result<Option<BlobMetadata>, DbError> { let result = sqlx::query!( "SELECT storage_key, mime_type, size_bytes FROM blobs WHERE cid = $1 LIMIT 1", cid.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.map(|r| BlobMetadata { storage_key: r.storage_key, mime_type: r.mime_type, size_bytes: r.size_bytes, })) }
async fn get_blob_with_takedown( &self, cid: &CidLink, ) -> Result<Option<BlobWithTakedown>, DbError> { let result = sqlx::query!( "SELECT cid, takedown_ref FROM blobs WHERE cid = $1 ORDER BY takedown_ref NULLS LAST LIMIT 1", cid.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
result .map(|r| { Ok(BlobWithTakedown { cid: column(r.cid, col::BLOBS_CID)?, takedown_ref: r.takedown_ref, }) }) .transpose() }
async fn get_blob_storage_key(&self, cid: &CidLink) -> Result<Option<String>, DbError> { let result = sqlx::query_scalar!( "SELECT storage_key FROM blobs WHERE cid = $1 LIMIT 1", cid.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result) }
async fn list_blobs_by_user( &self, user_id: Uuid, cursor: Option<&str>, limit: i64, ) -> Result<Vec<CidLink>, DbError> { let cursor_val = cursor.unwrap_or(""); let results = sqlx::query_scalar!( r#"SELECT cid FROM blobs WHERE created_by_user = $1 AND cid > $2 ORDER BY cid ASC LIMIT $3"#, user_id, cursor_val, limit ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
column_vec(results, col::BLOBS_CID) }
async fn list_blobs_since_rev(&self, did: &Did, since: &Tid) -> Result<Vec<CidLink>, DbError> { let results = sqlx::query_scalar!( r#"SELECT DISTINCT unnest(blobs) as "cid!" FROM repo_seq WHERE did = $1 AND rev > $2 AND blobs IS NOT NULL"#, did.as_str(), since.as_str() ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
column_vec(results, col::REPO_SEQ_BLOBS) }
async fn count_blobs_by_user(&self, user_id: Uuid) -> Result<i64, DbError> { let result = sqlx::query_scalar!( r#"SELECT COUNT(*) as "count!" FROM blobs WHERE created_by_user = $1"#, user_id ) .fetch_one(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result) }
async fn sum_blob_storage(&self) -> Result<i64, DbError> { let result = sqlx::query_scalar!( r#"SELECT COALESCE(SUM(size_bytes), 0)::BIGINT as "total!" FROM (SELECT DISTINCT cid, size_bytes FROM blobs) t"# ) .fetch_one(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result) }
async fn update_blob_takedown( &self, cid: &CidLink, takedown_ref: Option<&str>, ) -> Result<bool, DbError> { let result = sqlx::query!( "UPDATE blobs SET takedown_ref = $1 WHERE cid = $2", takedown_ref, cid.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected() > 0) }
async fn delete_blobs_by_user(&self, user_id: Uuid) -> Result<u64, DbError> { let result = sqlx::query!("DELETE FROM blobs WHERE created_by_user = $1", user_id) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn get_blob_storage_keys_by_user(&self, user_id: Uuid) -> Result<Vec<String>, DbError> { let results = sqlx::query_scalar!( r#"SELECT storage_key as "storage_key!" FROM blobs b WHERE created_by_user = $1 AND NOT EXISTS ( SELECT 1 FROM blobs o WHERE o.cid = b.cid AND o.created_by_user <> $1 )"#, user_id ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(results) }
async fn insert_record_blobs( &self, repo_id: Uuid, record_uris: &[AtUri], blob_cids: &[CidLink], ) -> Result<(), DbError> { let uris_str: Vec<&str> = record_uris.iter().map(|u| u.as_str()).collect(); let cids_str: Vec<&str> = blob_cids.iter().map(|c| c.as_str()).collect();
sqlx::query!( r#"INSERT INTO record_blobs (repo_id, record_uri, blob_cid) SELECT $1, record_uri, blob_cid FROM UNNEST($2::text[], $3::text[]) AS t(record_uri, blob_cid) ON CONFLICT (repo_id, record_uri, blob_cid) DO NOTHING"#, repo_id, &uris_str as &[&str], &cids_str as &[&str] ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(()) }
async fn list_missing_blobs( &self, repo_id: Uuid, cursor: Option<&str>, limit: i64, ) -> Result<Vec<MissingBlobInfo>, DbError> { let cursor_val = cursor.unwrap_or(""); let results = sqlx::query!( r#"SELECT rb.blob_cid, rb.record_uri FROM record_blobs rb LEFT JOIN blobs b ON rb.blob_cid = b.cid AND b.created_by_user = $1 WHERE rb.repo_id = $1 AND b.cid IS NULL AND rb.blob_cid > $2 ORDER BY rb.blob_cid LIMIT $3"#, repo_id, cursor_val, limit ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
results .into_iter() .map(|r| { Ok(MissingBlobInfo { blob_cid: column(r.blob_cid, col::RECORD_BLOBS_BLOB_CID)?, record_uri: column(r.record_uri, col::RECORD_BLOBS_RECORD_URI)?, }) }) .collect() }
async fn count_distinct_record_blobs(&self, repo_id: Uuid) -> Result<i64, DbError> { let result = sqlx::query_scalar!( r#"SELECT COUNT(DISTINCT blob_cid) as "count!" FROM record_blobs WHERE repo_id = $1"#, repo_id ) .fetch_one(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result) }
async fn get_blobs_for_export(&self, repo_id: Uuid) -> Result<Vec<BlobForExport>, DbError> { let results = sqlx::query!( r#"SELECT DISTINCT b.cid, b.storage_key, b.mime_type FROM blobs b JOIN record_blobs rb ON rb.blob_cid = b.cid WHERE rb.repo_id = $1"#, repo_id ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
results .into_iter() .map(|r| { Ok(BlobForExport { cid: column(r.cid, col::BLOBS_CID)?, storage_key: r.storage_key, mime_type: r.mime_type, }) }) .collect() }}