max / makenotwork

4.4 KB · 129 lines History Blame Raw

1	//! Durable queue for S3 object deletions that must survive server crashes.
2
3	use sqlx::PgPool;
4	use uuid::Uuid;
5	use crate::error::Result;
6
7	/// A pending S3 deletion record.
8	#[derive(Debug, sqlx::FromRow)]
9	pub struct PendingS3Deletion {
10	pub id: Uuid,
11	pub s3_key: String,
12	pub bucket: String,
13	pub source: String,
14	pub attempts: i32,
15	}
16
17	/// Enqueue S3 keys for deletion. Each key is (s3_key, bucket).
18	#[tracing::instrument(skip_all)]
19	pub async fn enqueue_deletions(
20	pool: &PgPool,
21	keys: &[(String, String)],
22	source: &str,
23	) -> Result<()> {
24	if keys.is_empty() {
25	return Ok(());
26	}
27	let s3_keys: Vec<&str> = keys.iter().map(\|(k, _)\| k.as_str()).collect();
28	let buckets: Vec<&str> = keys.iter().map(\|(_, b)\| b.as_str()).collect();
29	sqlx::query(
30	"INSERT INTO pending_s3_deletions (s3_key, bucket, source) SELECT * FROM unnest($1::text[], $2::text[], $3::text[]) ON CONFLICT DO NOTHING",
31	)
32	.bind(&s3_keys)
33	.bind(&buckets)
34	.bind(vec![source; keys.len()])
35	.execute(pool)
36	.await?;
37	Ok(())
38	}
39
40	/// Remove completed deletions by ID.
41	#[tracing::instrument(skip_all)]
42	pub async fn remove_completed(pool: &PgPool, ids: &[Uuid]) -> Result<()> {
43	if ids.is_empty() {
44	return Ok(());
45	}
46	sqlx::query("DELETE FROM pending_s3_deletions WHERE id = ANY($1)")
47	.bind(ids)
48	.execute(pool)
49	.await?;
50	Ok(())
51	}
52
53	/// Returns true if any live row in the given bucket's storage tables still
54	/// references `s3_key`. Used by the deletion worker to detect the
55	/// delete-then-reupload race: a freshly-uploaded object reusing a queued key
56	/// must not be torpedoed by the worker draining the queue.
57	///
58	/// `projects.cover_image_url` stores a full URL (not a bare s3_key) so we
59	/// check via `LIKE %s3_key`. Project image keys live under the
60	/// `projects/{id}/image/` prefix and don't collide with item/media keys, so
61	/// a positive match here is genuine. Without this clause a project cover
62	/// could be deleted out from under a still-live row.
63	#[tracing::instrument(skip_all)]
64	pub async fn is_s3_key_live(pool: &PgPool, bucket: &str, s3_key: &str) -> Result<bool> {
65	let live = if bucket == "synckit" {
66	sqlx::query_scalar::<_, bool>(
67	"SELECT EXISTS(SELECT 1 FROM sync_blobs WHERE s3_key = $1)",
68	)
69	.bind(s3_key)
70	.fetch_one(pool)
71	.await?
72	} else {
73	// Escape the SQL LIKE wildcards (`_` matches one char, `%` matches
74	// any) so an s3_key containing literal underscores from a user-
75	// supplied media folder name (e.g. `cool_stuff/file.png`) doesn't
76	// false-positive against neighbouring rows.
77	let escaped = s3_key.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_");
78	let url_suffix = format!("%{escaped}");
79	sqlx::query_scalar::<_, bool>(
80	r#"
81	SELECT
82	EXISTS(SELECT 1 FROM media_files WHERE s3_key = $1)
83	OR EXISTS(SELECT 1 FROM versions WHERE s3_key = $1)
84	OR EXISTS(SELECT 1 FROM ota_artifacts WHERE s3_key = $1)
85	OR EXISTS(SELECT 1 FROM items
86	WHERE audio_s3_key = $1 OR cover_s3_key = $1 OR video_s3_key = $1)
87	OR EXISTS(SELECT 1 FROM items WHERE cover_image_url LIKE $2 ESCAPE '\')
88	OR EXISTS(SELECT 1 FROM projects WHERE cover_image_url LIKE $2 ESCAPE '\')
89	OR EXISTS(SELECT 1 FROM content_insertions WHERE storage_key = $1)
90	"#,
91	)
92	.bind(s3_key)
93	.bind(&url_suffix)
94	.fetch_one(pool)
95	.await?
96	};
97	Ok(live)
98	}
99
100	/// Fetch stale pending deletions (older than min_age, up to limit).
101	/// Atomically increments attempt count.
102	#[tracing::instrument(skip_all)]
103	pub async fn get_stale_pending(
104	pool: &PgPool,
105	min_age: chrono::Duration,
106	limit: i64,
107	) -> Result<Vec<PendingS3Deletion>> {
108	let cutoff = chrono::Utc::now() - min_age;
109	let rows = sqlx::query_as::<_, PendingS3Deletion>(
110	r#"
111	UPDATE pending_s3_deletions
112	SET attempts = attempts + 1, last_attempted_at = NOW()
113	WHERE id IN (
114	SELECT id FROM pending_s3_deletions
115	WHERE created_at < $1
116	ORDER BY created_at
117	LIMIT $2
118	FOR UPDATE SKIP LOCKED
119	)
120	RETURNING id, s3_key, bucket, source, attempts
121	"#,
122	)
123	.bind(cutoff)
124	.bind(limit)
125	.fetch_all(pool)
126	.await?;
127	Ok(rows)
128	}
129