Skip to main content

max / makenotwork

5.7 KB · 203 lines History Blame Raw
1 //! Import job CRUD: create, update progress, complete, list.
2
3 use sqlx::PgPool;
4
5 use super::enums::ImportJobStatus;
6 use super::id_types::{ImportJobId, ProjectId, UserId};
7 use super::models::DbImportJob;
8 use crate::error::Result;
9 use crate::import::ImportSource;
10
11 /// Create a new import job in `pending` status.
12 #[tracing::instrument(skip_all)]
13 pub async fn create_import_job(
14 pool: &PgPool,
15 user_id: UserId,
16 project_id: ProjectId,
17 source: ImportSource,
18 total_rows: i32,
19 ) -> Result<DbImportJob> {
20 let job = sqlx::query_as::<_, DbImportJob>(
21 r"
22 INSERT INTO import_jobs (user_id, project_id, source, total_rows)
23 VALUES ($1, $2, $3, $4)
24 RETURNING *
25 ",
26 )
27 .bind(user_id)
28 .bind(project_id)
29 .bind(source)
30 .bind(total_rows)
31 .fetch_one(pool)
32 .await?;
33
34 Ok(job)
35 }
36
37 /// Update the processing progress of an import job.
38 #[tracing::instrument(skip_all)]
39 pub async fn update_import_progress(
40 pool: &PgPool,
41 job_id: ImportJobId,
42 processed_rows: i32,
43 created_rows: i32,
44 skipped_rows: i32,
45 ) -> Result<()> {
46 sqlx::query(
47 r"
48 UPDATE import_jobs
49 SET processed_rows = $2, created_rows = $3, skipped_rows = $4
50 WHERE id = $1
51 ",
52 )
53 .bind(job_id)
54 .bind(processed_rows)
55 .bind(created_rows)
56 .bind(skipped_rows)
57 .execute(pool)
58 .await?;
59
60 Ok(())
61 }
62
63 /// Update the status of an import job.
64 ///
65 /// Typed on [`ImportJobStatus`] rather than a raw `&str` so a caller can't stamp
66 /// a status the reader ([`DbImportJob::status`]) can't decode. Entering
67 /// `Processing` also stamps `heartbeat_at = NOW()` so the reaper has a liveness
68 /// baseline from the moment the job starts running (the first chunk may take a
69 /// while); other transitions leave the heartbeat untouched.
70 #[tracing::instrument(skip_all)]
71 pub async fn update_import_status(
72 pool: &PgPool,
73 job_id: ImportJobId,
74 status: ImportJobStatus,
75 ) -> Result<()> {
76 sqlx::query(
77 "UPDATE import_jobs \
78 SET status = $2, \
79 heartbeat_at = CASE WHEN $2 = 'processing' THEN NOW() ELSE heartbeat_at END \
80 WHERE id = $1",
81 )
82 .bind(job_id)
83 .bind(status)
84 .execute(pool)
85 .await?;
86
87 Ok(())
88 }
89
90 /// Refresh a running import's liveness heartbeat. Called after every processed
91 /// chunk so [`reap_stuck_import_jobs`] can tell a slow-but-progressing import
92 /// (fresh beat) from one whose owning process died (stale beat).
93 #[tracing::instrument(skip_all)]
94 pub async fn bump_import_heartbeat(pool: &PgPool, job_id: ImportJobId) -> Result<()> {
95 sqlx::query(
96 "UPDATE import_jobs SET heartbeat_at = NOW() WHERE id = $1 AND status = 'processing'",
97 )
98 .bind(job_id)
99 .execute(pool)
100 .await?;
101
102 Ok(())
103 }
104
105 /// Fail every import job stuck in `processing` whose heartbeat has gone stale
106 /// past `max_age_secs`, i.e. the process running it died without reaching
107 /// `complete_import_job` / `fail_import_job`. Without this a crash mid-import
108 /// leaves the job `processing` forever. Returns the number of jobs reaped.
109 ///
110 /// `COALESCE(heartbeat_at, created_at)` covers a job that crashed before its
111 /// first heartbeat (older rows and the window between insert and the first
112 /// `processing` stamp). Mirrors `scan_jobs::reap_stuck`'s liveness definition.
113 #[tracing::instrument(skip_all)]
114 pub async fn reap_stuck_import_jobs(pool: &PgPool, max_age_secs: i64) -> Result<u64> {
115 let affected = sqlx::query(
116 "UPDATE import_jobs \
117 SET status = 'failed', \
118 completed_at = NOW(), \
119 error_log = COALESCE(error_log, '') || \
120 CASE WHEN error_log IS NULL OR error_log = '' THEN '' ELSE E'\\n' END || \
121 'import reaped: no heartbeat for over the stuck-job threshold (process likely crashed)' \
122 WHERE status = 'processing' \
123 AND COALESCE(heartbeat_at, created_at) < NOW() - ($1 || ' seconds')::interval",
124 )
125 .bind(max_age_secs.to_string())
126 .execute(pool)
127 .await?
128 .rows_affected();
129
130 Ok(affected)
131 }
132
133 /// Mark an import job as completed with optional error log.
134 #[tracing::instrument(skip_all)]
135 pub async fn complete_import_job(
136 pool: &PgPool,
137 job_id: ImportJobId,
138 error_log: Option<String>,
139 ) -> Result<()> {
140 sqlx::query(
141 r"
142 UPDATE import_jobs
143 SET status = 'completed', completed_at = NOW(), error_log = $2
144 WHERE id = $1
145 ",
146 )
147 .bind(job_id)
148 .bind(error_log)
149 .execute(pool)
150 .await?;
151
152 Ok(())
153 }
154
155 /// Mark an import job as failed with an error message.
156 #[tracing::instrument(skip_all)]
157 pub async fn fail_import_job(pool: &PgPool, job_id: ImportJobId, error: &str) -> Result<()> {
158 sqlx::query(
159 r"
160 UPDATE import_jobs
161 SET status = 'failed', completed_at = NOW(), error_log = $2
162 WHERE id = $1
163 ",
164 )
165 .bind(job_id)
166 .bind(error)
167 .execute(pool)
168 .await?;
169
170 Ok(())
171 }
172
173 /// Get a single import job by ID, scoped to a user.
174 #[tracing::instrument(skip_all)]
175 pub async fn get_import_job(
176 pool: &PgPool,
177 job_id: ImportJobId,
178 user_id: UserId,
179 ) -> Result<Option<DbImportJob>> {
180 let job = sqlx::query_as::<_, DbImportJob>(
181 "SELECT * FROM import_jobs WHERE id = $1 AND user_id = $2",
182 )
183 .bind(job_id)
184 .bind(user_id)
185 .fetch_optional(pool)
186 .await?;
187
188 Ok(job)
189 }
190
191 /// List import jobs for a user, most recent first.
192 #[tracing::instrument(skip_all)]
193 pub async fn list_import_jobs(pool: &PgPool, user_id: UserId) -> Result<Vec<DbImportJob>> {
194 let jobs = sqlx::query_as::<_, DbImportJob>(
195 "SELECT * FROM import_jobs WHERE user_id = $1 ORDER BY created_at DESC LIMIT 20",
196 )
197 .bind(user_id)
198 .fetch_all(pool)
199 .await?;
200
201 Ok(jobs)
202 }
203