Commit source sync state and jobs atomically

crates/utopia-server/src/api/google_drive_routes.rs

@@ -469,13 +469,7 @@ pub async fn select(
     )
     .await?;
     if !body.folders.is_empty() {
-        let _ = utopia_store::sources::mark_queued(&state.pool, source_id).await;
-        utopia_store::jobs::enqueue(
-            &state.pool,
-            "sync_source",
-            json!({ "source_id": source_id }),
-        )
-        .await?;
+        utopia_store::sources::enqueue_sync_after_config_change(&state.pool, source_id).await?;
     }
     state.emit_source(kb_id);
     let _ = utopia_store::audit::record(

crates/utopia-server/src/api/google_drive_routes_tests.rs

@@ -1,7 +1,10 @@
 //! 连接流程的几道关卡与「谁连上的谁来选」（0074 决定 3、5），连真库跑。
 //! 测试环境没配 OAuth 客户端：过了所有关卡的回调停在「没配置」，不会去敲 Google
 use super::*;
+use axum::body::{to_bytes, Body};
+use axum::http::{Request, StatusCode};
 use std::sync::Arc;
+use tower::ServiceExt;
 use utopia_core::models::User;
 
 struct Fixture {
@@ -114,6 +117,268 @@ impl Fixture {
             .await?
             .config)
     }
+
+    async fn sync_http(&self) -> anyhow::Result<(StatusCode, Value)> {
+        self.post_http(
+            &format!("/api/v1/kbs/{}/sources/{}/sync", self.kb, self.source),
+            json!({}),
+        )
+        .await
+    }
+
+    async fn post_http(&self, path: &str, body: Value) -> anyhow::Result<(StatusCode, Value)> {
+        let token = crate::auth::issue_token(&self.state, self.connector.id)?;
+        let response = super::super::router(self.state.clone(), &Default::default())
+            .oneshot(
+                Request::post(path)
+                    .header("Authorization", format!("Bearer {token}"))
+                    .header("Content-Type", "application/json")
+                    .body(Body::from(body.to_string()))?,
+            )
+            .await?;
+        let status = response.status();
+        let bytes = to_bytes(response.into_body(), 1 << 20).await?;
+        Ok((status, serde_json::from_slice(&bytes)?))
+    }
+
+    async fn sync_state(&self) -> anyhow::Result<(String, i64)> {
+        Ok(sqlx::query_as(
+            "SELECT last_sync_status,
+                    (SELECT count(*) FROM jobs WHERE kind = 'sync_source'
+                       AND payload->>'source_id' = sources.id::text)
+               FROM sources WHERE id = $1",
+        )
+        .bind(self.source)
+        .fetch_one(&self.pool)
+        .await?)
+    }
+
+    async fn choose_folder(&self, id: &str) -> ApiResult<Json<Value>> {
+        select(
+            State(self.state.clone()),
+            AuthUser(self.connector.clone()),
+            Path((self.kb, self.source)),
+            Json(SelectBody {
+                folders: vec![FolderEntry {
+                    id: id.into(),
+                    name: id.into(),
+                }],
+            }),
+        )
+        .await
+    }
+
+    async fn reject_sync_jobs(&self, reject: bool) -> anyhow::Result<()> {
+        // Install once: dropping a table trigger between tests can deadlock parallel fixtures.
+        static INSTALLED: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::const_new();
+        INSTALLED
+            .get_or_try_init(|| async {
+                sqlx::raw_sql(
+                    "SELECT pg_advisory_xact_lock(hashtext('test_reject_source_sync'));
+                     CREATE OR REPLACE FUNCTION test_reject_source_sync() RETURNS trigger
+                     LANGUAGE plpgsql AS $$ BEGIN
+                       IF NEW.kind = 'sync_source' AND EXISTS(
+                         SELECT 1 FROM sources WHERE id::text = NEW.payload->>'source_id'
+                           AND name = 'reject-source-sync-test')
+                       THEN RAISE EXCEPTION 'test source sync insert failure'; END IF;
+                       RETURN NEW;
+                     END $$;
+                     DO $$ BEGIN
+                       IF NOT EXISTS(SELECT 1 FROM pg_trigger WHERE tgname = 'test_reject_source_sync'
+                                        AND tgrelid = 'jobs'::regclass) THEN
+                         CREATE TRIGGER test_reject_source_sync BEFORE INSERT ON jobs
+                           FOR EACH ROW EXECUTE FUNCTION test_reject_source_sync();
+                       END IF;
+                     END $$;",
+                )
+                .execute(&self.pool)
+                .await
+                .map(drop)
+            })
+            .await?;
+        sqlx::query("UPDATE sources SET name = $2 WHERE id = $1")
+            .bind(self.source)
+            .bind(if reject {
+                "reject-source-sync-test"
+            } else {
+                "Drive"
+            })
+            .execute(&self.pool)
+            .await?;
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn source_sync_enqueue_creation_failure_keeps_source_schedulable() -> anyhow::Result<()> {
+    let Some(mut fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    fx.reject_sync_jobs(false).await?;
+    let config = json!({"url": "https://example.com/source-sync-test"});
+    let (status, body) = fx
+        .post_http(
+            &format!("/api/v1/kbs/{}/sources", fx.kb),
+            json!({"kind": "url", "name": "reject-source-sync-test", "config": config}),
+        )
+        .await?;
+    assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR, "{body}");
+    fx.source = sqlx::query_scalar(
+        "SELECT id FROM sources WHERE kb_id = $1 AND kind = 'url'
+           AND name = 'reject-source-sync-test'",
+    )
+    .bind(fx.kb)
+    .fetch_one(&fx.pool)
+    .await?;
+    assert_eq!(fx.config().await?, config);
+    assert_eq!(fx.sync_state().await?, ("never".into(), 0));
+    fx.reject_sync_jobs(false).await?;
+    let (status, body) = fx.sync_http().await?;
+    assert_eq!(status, StatusCode::OK, "{body}");
+    assert_eq!(body["queued"], true);
+    assert_eq!(fx.sync_state().await?, ("queued".into(), 1));
+    Ok(())
+}
+
+#[tokio::test]
+async fn source_sync_enqueue_manual_failure_rolls_back_and_can_retry() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    fx.reject_sync_jobs(true).await?;
+    let (status, body) = fx.sync_http().await?;
+    assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR, "{body}");
+    assert_eq!(fx.sync_state().await?, ("never".into(), 0));
+
+    fx.reject_sync_jobs(false).await?;
+    let (status, body) = fx.sync_http().await?;
+    assert_eq!(status, StatusCode::OK, "{body}");
+    assert_eq!(body["queued"], true);
+    assert_eq!(fx.sync_state().await?, ("queued".into(), 1));
+    let (payload, max_attempts): (Value, i32) = sqlx::query_as(
+        "SELECT payload, max_attempts FROM jobs WHERE kind = 'sync_source'
+           AND payload->>'source_id' = $1",
+    )
+    .bind(fx.source.to_string())
+    .fetch_one(&fx.pool)
+    .await?;
+    assert_eq!(payload, json!({"source_id": fx.source}));
+    assert_eq!(max_attempts, 3);
+
+    for pending in ["queued", "running"] {
+        sqlx::query("UPDATE sources SET last_sync_status = $2 WHERE id = $1")
+            .bind(fx.source)
+            .bind(pending)
+            .execute(&fx.pool)
+            .await?;
+        let (status, body) = fx.sync_http().await?;
+        assert_eq!(status, StatusCode::OK, "{body}");
+        assert_eq!(body["queued"], false);
+        assert_eq!(fx.sync_state().await?, (pending.into(), 1));
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn source_sync_enqueue_concurrent_manual_requests_enqueue_once() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    // AppState and the router do not start a worker, so the first task stays queued.
+    let (first, second) = tokio::join!(fx.sync_http(), fx.sync_http());
+    let mut queued = Vec::new();
+    for result in [first, second] {
+        let (status, body) = result?;
+        assert_eq!(status, StatusCode::OK, "{body}");
+        queued.push(body["queued"].as_bool().expect("queued is a boolean"));
+    }
+    queued.sort_unstable();
+    assert_eq!(queued, [false, true]);
+    assert_eq!(fx.sync_state().await?, ("queued".into(), 1));
+    Ok(())
+}
+
+#[tokio::test]
+async fn source_sync_enqueue_drive_selection_keeps_follow_up_when_pending() -> anyhow::Result<()> {
+    for pending in ["queued", "running"] {
+        let Some(fx) = Fixture::new().await? else {
+            return Ok(());
+        };
+        let (status, body) = fx.sync_http().await?;
+        assert_eq!(status, StatusCode::OK, "{body}");
+        sqlx::query("UPDATE sources SET last_sync_status = $2 WHERE id = $1")
+            .bind(fx.source)
+            .bind(pending)
+            .execute(&fx.pool)
+            .await?;
+        sqlx::query("UPDATE jobs SET status = $2 WHERE payload->>'source_id' = $1")
+            .bind(fx.source.to_string())
+            .bind(pending)
+            .execute(&fx.pool)
+            .await?;
+        let _ = fx
+            .choose_folder("next-folder")
+            .await
+            .map_err(|ApiErr(error)| error)?;
+        assert_eq!(fx.sync_state().await?, (pending.into(), 2));
+        let latest_status: String = sqlx::query_scalar(
+            "SELECT status FROM jobs WHERE kind = 'sync_source'
+               AND payload->>'source_id' = $1 ORDER BY id DESC LIMIT 1",
+        )
+        .bind(fx.source.to_string())
+        .fetch_one(&fx.pool)
+        .await?;
+        assert_eq!(latest_status, "queued");
+        assert_eq!(fx.config().await?["folders"][0]["id"], "next-folder");
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn source_sync_enqueue_drive_failure_keeps_config_but_rolls_back_queue() -> anyhow::Result<()>
+{
+    for initial in ["never", "running"] {
+        let Some(fx) = Fixture::new().await? else {
+            return Ok(());
+        };
+        let existing_jobs = if initial == "running" {
+            let (status, body) = fx.sync_http().await?;
+            assert_eq!(status, StatusCode::OK, "{body}");
+            sqlx::query("UPDATE sources SET last_sync_status = 'running' WHERE id = $1")
+                .bind(fx.source)
+                .execute(&fx.pool)
+                .await?;
+            sqlx::query("UPDATE jobs SET status = 'running' WHERE payload->>'source_id' = $1")
+                .bind(fx.source.to_string())
+                .execute(&fx.pool)
+                .await?;
+            1
+        } else {
+            0
+        };
+        fx.reject_sync_jobs(true).await?;
+        assert!(matches!(
+            fx.choose_folder("saved-folder").await,
+            Err(ApiErr(AppError::Db(_)))
+        ));
+        assert_eq!(fx.config().await?["folders"][0]["id"], "saved-folder");
+        assert_eq!(fx.sync_state().await?, (initial.into(), existing_jobs));
+        fx.reject_sync_jobs(false).await?;
+        let _ = fx
+            .choose_folder("saved-folder")
+            .await
+            .map_err(|ApiErr(error)| error)?;
+        let expected_status = if initial == "running" {
+            "running"
+        } else {
+            "queued"
+        };
+        assert_eq!(
+            fx.sync_state().await?,
+            (expected_status.into(), existing_jobs + 1)
+        );
+    }
+    Ok(())
 }
 
 fn denied_code<T>(result: ApiResult<T>) -> Option<&'static str> {

crates/utopia-server/src/api/sources_routes.rs

@@ -132,13 +132,7 @@ pub async fn create(
     .await?;
     // 拉取型来源建好立即同步一次（有配置即产出，无需等下一个调度周期）
     if matches!(source.kind.as_str(), "url" | "rss" | "custom") {
-        let _ = utopia_store::sources::mark_queued(&state.pool, source.id).await;
-        utopia_store::jobs::enqueue(
-            &state.pool,
-            "sync_source",
-            json!({ "source_id": source.id }),
-        )
-        .await?;
+        utopia_store::sources::enqueue_sync(&state.pool, source.id).await?;
     }
     // api 来源：生成专属推送密钥（此后可随时经 get_token 查看）
     let mut ingest_token: Option<String> = None;
@@ -396,14 +390,8 @@ pub async fn sync_now(
 ) -> ApiResult<Json<serde_json::Value>> {
     require_kb(&state, &user, kb_id, Role::Editor).await?;
     source_in_kb(&state, kb_id, source_id).await?;
-    let queued = utopia_store::sources::mark_queued(&state.pool, source_id).await?;
+    let queued = utopia_store::sources::enqueue_sync(&state.pool, source_id).await?;
     if queued {
-        utopia_store::jobs::enqueue(
-            &state.pool,
-            "sync_source",
-            json!({ "source_id": source_id }),
-        )
-        .await?;
         state.emit_source(kb_id);
     }
     Ok(Json(json!({ "queued": queued })))

crates/utopia-server/src/main.rs

@@ -232,20 +232,10 @@ async fn main() -> anyhow::Result<()> {
             match utopia_store::sources::due_sources(&sched_state.pool).await {
                 Ok(due) => {
                     for s in due {
-                        match utopia_store::sources::mark_queued(&sched_state.pool, s.id).await {
-                            Ok(true) => {
-                                if let Err(e) = utopia_store::jobs::enqueue(
-                                    &sched_state.pool,
-                                    "sync_source",
-                                    serde_json::json!({ "source_id": s.id }),
-                                )
-                                .await
-                                {
-                                    tracing::warn!(source_id = %s.id, error = %e, "同步任务入队失败");
-                                }
-                            }
-                            Ok(false) => {}
-                            Err(e) => tracing::warn!(source_id = %s.id, error = %e, "标记入队失败"),
+                        if let Err(e) =
+                            utopia_store::sources::enqueue_sync(&sched_state.pool, s.id).await
+                        {
+                            tracing::warn!(source_id = %s.id, error = %e, "同步任务入队失败");
                         }
                     }
                 }