core/test/rust/runtime/job/postgres/mod.rs

use crate::{
    adapter::Document,
    database::{Binding, Database, Error, Executor, Options, Parameter, Registry},
    job::{self, Claim, Execution, Policy, Submission},
    task::{Runtime, Scope},
    wire::json::Json,
};
use futures_util::FutureExt;
use std::{
    panic::AssertUnwindSafe,
    path::PathBuf,
    sync::{Arc, LazyLock},
    time::Duration,
};

mod heartbeat;
mod lease;
mod queue;
mod schedule;
mod workers;

const NOW: i64 = 1_700_000_000_000;

#[derive(Clone)]
struct Case {
    database: Database,
    schema: String,
    directory: PathBuf,
}

impl Case {
    async fn registry(&self, scope: &Scope<'_>) -> Arc<Registry> {
        Registry::load(
            &self.directory.join("setting.json"),
            &[Binding {
                storage: crate::database::Storage::Global,
                component: "demo",
                database: Some("default"),
                table: job::schema::TABLE,
                schema: job::schema::SQL,
                plan: None,
                references: &[],
            }],
            scope,
        )
        .await
        .unwrap()
    }

    async fn executor(&self, scope: &Scope<'_>) -> Executor {
        let mut executor = self.database.executor().await.unwrap();
        executor
            .query(
                scope,
                "SELECT set_config('search_path',$1,false)",
                vec![Parameter::Text(self.schema.clone())],
                1,
                |row| row.get::<_, String>(0),
            )
            .await
            .unwrap();
        executor
    }

    async fn count(&self, scope: &Scope<'_>, sql: &'static str) -> i64 {
        self.executor(scope)
            .await
            .query(scope, sql, vec![], 1, |row| row.get::<_, i64>(0))
            .await
            .unwrap()[0]
    }
}

fn fixture(check: impl AsyncFnOnce(&Case, &Scope<'_>)) {
    let path =
        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../config/setting.json");
    let document = Document::load(&path).expect("cannot read PostgreSQL test settings");
    let root: serde_json::Value = document.decode().expect("invalid PostgreSQL test settings");
    let value = root
        .get("database")
        .and_then(|database| database.get("postgres_test"))
        .expect("database.postgres_test must select an isolated test database");
    let mut options: Options = serde_json::from_value(value.clone())
        .unwrap_or_else(|_| panic!("invalid PostgreSQL test options"));
    assert_eq!(options.driver, "postgres");
    options.max_connections = 8;
    options.wait_timeout_ms = 1000;
    options.io_timeout_ms = 5000;
    Runtime::new(Default::default())
        .unwrap()
        .run(async |scope| {
            let database = Database::open(std::path::Path::new("."), options, scope)
                .await
                .unwrap();
            let schema = format!("dever_job_test_{}", job::identifier().unwrap());
            let directory = std::env::temp_dir().join(&schema);
            std::fs::create_dir(&directory).unwrap();
            let case = Case {
                database,
                schema,
                directory,
            };
            let mut owner = case.database.executor().await.unwrap();
            owner
                .postgres_operation(async |client| {
                    client
                        .batch_execute(&format!("CREATE SCHEMA {}", case.schema))
                        .await
                        .map_err(Error::from)
                })
                .await
                .unwrap();
            let outcome = AssertUnwindSafe(async {
                use std::{
                    io::Write,
                    os::unix::fs::{OpenOptionsExt, PermissionsExt},
                };
                std::fs::set_permissions(&case.directory, std::fs::Permissions::from_mode(0o700))
                    .unwrap();
                let mut connection = value.clone();
                connection["max_connections"] = 8.into();
                let url = connection["url"].as_str().unwrap();
                assert!(
                    url.parse::<tokio_postgres::Config>()
                        .unwrap()
                        .get_options()
                        .is_none()
                );
                let separator = if url.contains('?') { '&' } else { '?' };
                connection["url"] =
                    format!("{url}{separator}options=-csearch_path%3D{}", case.schema).into();
                let mut settings = std::fs::OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .mode(0o600)
                    .open(case.directory.join("setting.json"))
                    .unwrap();
                settings
                    .write_all(
                        &serde_json::to_vec(
                            &serde_json::json!({"database":{"default":connection}}),
                        )
                        .unwrap(),
                    )
                    .unwrap();
                let mut executor = case.executor(scope).await;
                executor
                    .schema(scope, job::schema::TABLE, job::schema::SQL)
                    .await
                    .unwrap();
                executor
                    .schema(scope, job::schema::TABLE, job::schema::SQL)
                    .await
                    .unwrap();
                drop(executor);
                tokio::time::timeout(Duration::from_secs(30), check(&case, scope))
                    .await
                    .expect("PostgreSQL Job test deadline exceeded");
            })
            .catch_unwind()
            .await;
            let cleanup = owner
                .postgres_operation(async |client| {
                    client
                        .batch_execute(&format!("DROP SCHEMA {} CASCADE", case.schema))
                        .await
                        .map_err(Error::from)
                })
                .await;
            case.database.close();
            std::fs::remove_dir_all(&case.directory).unwrap();
            cleanup.expect("PostgreSQL Job test schema cleanup failed");
            if let Err(panic) = outcome {
                std::panic::resume_unwind(panic);
            }
        })
        .unwrap();
}

fn submission(key: &str) -> Submission<'_> {
    static PAYLOAD: LazyLock<Json> = LazyLock::new(|| Json::parse(r#"{"number":1}"#).unwrap());
    Submission {
        target: "demo.account.publish",
        payload: Some(&PAYLOAD),
        schema: "schema-1",
        execution: &Execution::System { tenant: None },
        key,
        policy: Policy {
            retries: 2,
            timeout_ms: 50,
        },
        at: NOW,
    }
}