From b3cfefbd9141c9cecc04de8cc197023847ae7571 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Wed, 21 Oct 2026 21:34:06 +1300 Subject: [PATCH 01/10] fix(queue-storage): insert lane head rows only when the lane is missing The enqueue and claim head triggers fire BEFORE INSERT, ahead of ON CONFLICT, and read the lane's sequence. A lane a 0.6 runtime created has a sequence that runtime owns, so awa storage prepare-queue run as the migrator failed with 42501 on any such lane, and every runtime cache miss re-synced an existing sequence. Head rows are now inserted only when absent, under the lane lock runtimes already take for these inserts. --- awa-model/src/queue_storage.rs | 147 +++++++++++++++++++-------------- 1 file changed, 84 insertions(+), 63 deletions(-) diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 0dba62a00..bd18a8cdd 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -8,7 +8,7 @@ use crate::{InsertParams, JobRow, JobState}; use chrono::TimeDelta; use chrono::{DateTime, Utc}; use sqlx::Acquire; -use sqlx::{PgPool, Postgres, QueryBuilder}; +use sqlx::{PgConnection, PgPool, Postgres, QueryBuilder}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; @@ -673,7 +673,7 @@ async fn set_prune_lock_timeout_tx( Ok(()) } -fn validate_ident(ident: &str) -> Result<(), AwaError> { +pub(crate) fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { Some(first) if first.is_ascii_lowercase() || first == '_' => {} @@ -4021,67 +4021,7 @@ impl QueueStorage { priority: i16, enqueue_shard: i16, ) -> Result<(), AwaError> { - let schema = self.schema(); - let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(lane_lock_key) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_lanes (queue, priority) - VALUES ($1, $2) - ON CONFLICT (queue, priority) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) - VALUES ($1, $2, $3) - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) - VALUES ($1, $2, $3) - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - SELECT {schema}.ensure_lane_sequences($1, $2, $3) - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - + insert_lane_rows(tx.as_mut(), self.schema(), queue, priority, enqueue_shard).await?; self.mark_lane_ensured(queue, priority, enqueue_shard); Ok(()) } @@ -18027,3 +17967,84 @@ impl QueueStorage { format!("{}.{}", self.schema(), self.attempt_state_relname()) } } + +/// Insert a lane's rows and create its sequences, under the lane lock that +/// runtimes take for the same inserts. Head rows are inserted only when +/// missing, so their sequence triggers never fire for an existing lane, whose +/// sequences may belong to the runtime role that created them. +pub(crate) async fn insert_lane_rows( + conn: &mut PgConnection, + schema: &str, + queue: &str, + priority: i16, + enqueue_shard: i16, +) -> Result<(), AwaError> { + let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(lane_lock_key) + .execute(&mut *conn) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_lanes (queue, priority) + VALUES ($1, $2) + ON CONFLICT (queue, priority) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .execute(&mut *conn) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT $1, $2, $3 + WHERE NOT EXISTS ( + SELECT 1 FROM {schema}.queue_enqueue_heads + WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(&mut *conn) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) + SELECT $1, $2, $3 + WHERE NOT EXISTS ( + SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(&mut *conn) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + SELECT {schema}.ensure_lane_sequences($1, $2, $3) + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(&mut *conn) + .await + .map_err(map_sqlx_error)?; + Ok(()) +} From 2a9a3b6a96002e1decb6654e545dc5394eafa7c1 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Wed, 21 Oct 2026 21:34:06 +1300 Subject: [PATCH 02/10] feat(storage): prepare-create-revoke for revoking 0.6 runtimes' schema CREATE without a pause The migration trust check refuses a schema where another role holds CREATE, but 0.6 lane helpers run CREATE SEQUENCE IF NOT EXISTS on every reservation, so revoking the grant before migrating stopped running 0.6 producers and workers with 42501. The documented path was a pause. awa storage prepare-create-revoke, run as the migrator under a running 0.6 fleet, replaces each queue-storage schema's lane helpers with the ones v049 installs (identical to every 0.6 release's helpers apart from creating a sequence only when it is missing), provisions all four priorities of every shard of every queue in use, and grants the roles holding CREATE the privileges they need on the sequences it creates. It applies the migration's own trust rules apart from the CREATE grant, leaves a schema with any other violation unchanged, checks the rules again before committing, and prints the REVOKE statements plus the sequence grants only the sequences' runtime owners can make, which the migration needs to read and set those sequences. The upgrade guide replaces the pause procedure with these steps. The new test installs a v040 schema with the 0.6.9 lane helpers and 0.6 grants, shows revoking CREATE first breaks a 0.6 reservation, and then runs the command, grants and revokes under a 2 s lock_timeout while a 0.6 enqueue transaction stays open, before migrating to the current version. Each part of the fix fails the test when reverted. --- awa-cli/src/main.rs | 22 ++ awa-model/src/migrations.rs | 44 +++ awa-model/src/storage.rs | 309 ++++++++++++++++++ awa-model/tests/create_revoke_upgrade_test.rs | 296 +++++++++++++++++ .../fixtures/v040_lane_sequence_helpers.sql | 136 ++++++++ docs/security/database-roles.md | 4 +- docs/upgrade-0.6-to-0.7.md | 59 ++-- skills/awa-operations/SKILL.md | 3 + 8 files changed, 845 insertions(+), 28 deletions(-) create mode 100644 awa-model/tests/create_revoke_upgrade_test.rs create mode 100644 awa-model/tests/fixtures/v040_lane_sequence_helpers.sql diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index f89003114..b22e386ec 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -328,6 +328,7 @@ fn access_for(command: &Commands) -> context::Access { StorageCommands::Prepare { .. } | StorageCommands::PrepareQueueStorageSchema { .. } | StorageCommands::PrepareQueue { .. } + | StorageCommands::PrepareCreateRevoke | StorageCommands::Abort | StorageCommands::EnterMixedTransition { .. } | StorageCommands::Finalize { .. } @@ -821,6 +822,11 @@ enum StorageCommands { #[arg(long, default_value_t = 1)] queue_stripe_count: usize, }, + /// Before revoking runtime roles' schema CREATE under a running 0.6 fleet: + /// switch the lane helpers to provisioned sequences and provision every + /// lane already in use. Run as the migrator; prints the REVOKE statements + /// to run next + PrepareCreateRevoke, /// Abort a prepared or mixed-transition storage rollout before final activation Abort, /// Enter mixed transition and begin routing new writes to the prepared engine @@ -1903,6 +1909,22 @@ async fn run() -> Result<(), Box> { }))? ); } + StorageCommands::PrepareCreateRevoke => { + let schemas = awa_model::storage::prepare_create_revoke(&pool).await?; + println!("{}", serde_json::to_string_pretty(&schemas)?); + let blocked: Vec<&str> = schemas + .iter() + .filter(|schema| !schema.prepared) + .map(|schema| schema.schema.as_str()) + .collect(); + if !blocked.is_empty() { + return Err(format!( + "not prepared: {}; fix the listed blocking violations and re-run", + blocked.join(", ") + ) + .into()); + } + } StorageCommands::PrepareQueue { schema, queue, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 6d13a6db5..fd0179fc9 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -314,6 +314,32 @@ const V20_UP: &str = include_str!("../migrations/v020_active_queue_storage_schem const V21_UP: &str = include_str!("../migrations/v021_shard_aware_lane_indexes.sql"); const V22_UP: &str = include_str!("../migrations/v022_delete_compat_terminal_counter.sql"); const SCHEMA_DISCOVERY_UP: &str = include_str!("../migrations/queue_storage_schema_discovery.sql"); + +/// The trust and discovery functions every pending bundle installs, without +/// the preflight that refuses an untrusted configured schema. +pub(crate) fn schema_trust_functions_sql() -> &'static str { + SCHEMA_DISCOVERY_UP + .split_once("DO $preflight$") + .map(|(functions, _)| functions) + .expect("the discovery bootstrap ends with its preflight") +} + +/// `awa.refresh_lane_sequence_helpers` as v023's installer defines it. +pub(crate) fn lane_sequence_helper_refresh_sql() -> &'static str { + const START: &str = + "CREATE OR REPLACE FUNCTION awa.refresh_lane_sequence_helpers(p_schema TEXT)"; + const END: &str = + "REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC;"; + let installer = include_str!("../migrations/v023_install_queue_storage_substrate.sql"); + let start = installer + .find(START) + .expect("v023 defines refresh_lane_sequence_helpers"); + let end = installer[start..] + .find(END) + .map(|offset| start + offset + END.len()) + .expect("v023 revokes refresh_lane_sequence_helpers from PUBLIC"); + &installer[start..end] +} const V23_UP: &str = concat!( include_str!("../migrations/queue_storage_schema_discovery.sql"), "\n", @@ -1093,6 +1119,24 @@ pub fn migration_sql_range(from: i32, to: i32) -> Vec<(i32, &'static str, String mod tests { use super::*; + #[test] + fn schema_trust_functions_stop_before_the_preflight() { + let functions = schema_trust_functions_sql(); + assert!(functions.contains("FUNCTION awa.queue_storage_schema_trust_violations")); + assert!(functions.contains("FUNCTION awa.queue_storage_schemas()")); + assert!(!functions.contains("$preflight$")); + assert_eq!(SCHEMA_DISCOVERY_UP.matches("DO $preflight$").count(), 1); + } + + #[test] + fn lane_sequence_helper_refresh_is_one_function_and_its_revoke() { + let refresh = lane_sequence_helper_refresh_sql(); + assert!(refresh.starts_with("CREATE OR REPLACE FUNCTION awa.refresh_lane_sequence_helpers")); + assert!(refresh.ends_with("FROM PUBLIC;")); + assert_eq!(refresh.matches("$refresh$;").count(), 1); + assert!(!refresh.contains("install_queue_storage_substrate")); + } + #[test] fn discovery_bundles_bootstrap_before_discovery_and_default_installation() { for (version, _, sql) in migration_sql() { diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index b7a0ff1ad..d00083bad 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -469,3 +469,312 @@ pub async fn status_report(pool: &PgPool) -> Result, + /// Grants each sequence owner must run so the migrator can read and set + /// the lane sequences that runtime roles created. + pub sequence_grants: Vec, + /// Trust violations other than a `CREATE` grant. They must be fixed + /// before the schema can be prepared or migrated. + pub blocking: Vec, +} + +/// A grant that only a sequence's owner can make. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct SequenceGrant { + pub run_as: String, + pub statement: String, +} + +/// Prepare a database still used by 0.6 runtimes for revoking those roles' +/// `CREATE` on its queue-storage schemas, without stopping them. +/// +/// 0.6 lane helpers issue `CREATE SEQUENCE IF NOT EXISTS` on every +/// reservation, which needs `CREATE` even when the sequence exists. This +/// replaces them with the helpers v049 installs, which only create a +/// sequence that is missing, and provisions every priority and shard of +/// every queue already in use. Run it as the migrator. A schema with any +/// trust violation other than a `CREATE` grant is reported and left +/// unchanged. +pub async fn prepare_create_revoke(pool: &PgPool) -> Result, AwaError> { + let mut tx = pool.begin().await?; + sqlx::raw_sql(crate::migrations::schema_trust_functions_sql()) + .execute(tx.as_mut()) + .await?; + sqlx::raw_sql(crate::migrations::lane_sequence_helper_refresh_sql()) + .execute(tx.as_mut()) + .await?; + + let foreign_helpers: Vec = sqlx::query_scalar( + r#" + SELECT routine.oid::regprocedure::text + FROM pg_proc AS routine + JOIN pg_namespace AS namespace ON namespace.oid = routine.pronamespace + WHERE namespace.nspname = 'awa' + AND routine.proname IN ( + 'queue_storage_canonical_triggers', 'queue_storage_schema_trust_violations', + 'queue_storage_schema_is_trusted', 'queue_storage_schema_trust_summary', + 'queue_storage_schemas', 'refresh_lane_sequence_helpers' + ) + AND routine.proowner NOT IN ( + SELECT nspowner FROM pg_namespace WHERE nspname = 'awa' + UNION SELECT oid FROM pg_roles WHERE rolname = current_user + ) + ORDER BY 1 + "#, + ) + .fetch_all(tx.as_mut()) + .await?; + if !foreign_helpers.is_empty() { + return Err(AwaError::Validation(format!( + "awa helper functions are owned by another role: {}", + foreign_helpers.join(", ") + ))); + } + + let schemas: Vec = sqlx::query_scalar( + r#" + SELECT namespace.nspname::text + FROM pg_namespace AS namespace + WHERE has_schema_privilege(current_user, namespace.oid, 'USAGE') + AND EXISTS ( + SELECT 1 FROM pg_proc AS claim_fn + WHERE claim_fn.pronamespace = namespace.oid + AND claim_fn.proname = 'claim_ready_runtime' + AND oidvectortypes(claim_fn.proargtypes) + = 'text, bigint, double precision, double precision' + ) + AND EXISTS ( + SELECT 1 FROM pg_class AS heads + WHERE heads.relnamespace = namespace.oid AND heads.relname = 'queue_enqueue_heads' + ) + AND EXISTS ( + SELECT 1 FROM pg_class AS heads + WHERE heads.relnamespace = namespace.oid AND heads.relname = 'queue_claim_heads' + ) + ORDER BY 1 + "#, + ) + .fetch_all(tx.as_mut()) + .await?; + + let mut outcomes = Vec::with_capacity(schemas.len()); + for schema in schemas { + let mut blocking = non_create_trust_violations(&mut tx, &schema).await?; + if crate::queue_storage::validate_ident(&schema).is_err() { + blocking.push(format!( + "schema {}: name is not a lowercase identifier", + quote_identifier(&schema) + )); + } + let create_grantees: Vec = sqlx::query_scalar( + r#" + SELECT CASE privilege.grantee WHEN 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(privilege.grantee)) END + FROM pg_namespace AS namespace + CROSS JOIN LATERAL aclexplode( + COALESCE(namespace.nspacl, acldefault('n', namespace.nspowner)) + ) AS privilege + WHERE namespace.nspname = $1 + AND privilege.privilege_type = 'CREATE' + AND privilege.grantee NOT IN ( + SELECT nspowner FROM pg_namespace WHERE nspname = 'awa' + UNION SELECT oid FROM pg_roles WHERE rolname = current_user + ) + ORDER BY 1 + "#, + ) + .bind(&schema) + .fetch_all(tx.as_mut()) + .await?; + let revoke = create_grantees + .iter() + .map(|grantee| { + format!( + "REVOKE CREATE ON SCHEMA {} FROM {grantee};", + quote_identifier(&schema) + ) + }) + .collect(); + + let sequence_grants = sqlx::query_as::<_, (String, String)>( + r#" + SELECT pg_get_userbyid(sequence.relowner)::text, + format( + 'GRANT SELECT, UPDATE ON SEQUENCE %s TO %I;', + string_agg(format('%I.%I', namespace.nspname, sequence.relname), ', ' + ORDER BY sequence.relname), + current_user + ) + FROM pg_class AS sequence + JOIN pg_namespace AS namespace ON namespace.oid = sequence.relnamespace + WHERE namespace.nspname = $1 + AND sequence.relkind = 'S' + AND NOT (has_sequence_privilege(current_user, sequence.oid, 'SELECT') + AND has_sequence_privilege(current_user, sequence.oid, 'UPDATE')) + GROUP BY sequence.relowner + ORDER BY 1 + "#, + ) + .bind(&schema) + .fetch_all(tx.as_mut()) + .await? + .into_iter() + .map(|(run_as, statement)| SequenceGrant { run_as, statement }) + .collect(); + + if !blocking.is_empty() { + outcomes.push(CreateRevokeSchema { + schema, + prepared: false, + provisioned_lanes: 0, + revoke, + sequence_grants, + blocking, + }); + continue; + } + + sqlx::query("SELECT awa.refresh_lane_sequence_helpers($1)") + .bind(&schema) + .execute(tx.as_mut()) + .await?; + let provisioned_lanes = provision_known_lanes(&mut tx, &schema, &create_grantees).await?; + outcomes.push(CreateRevokeSchema { + schema, + prepared: true, + provisioned_lanes, + revoke, + sequence_grants, + blocking, + }); + } + + for outcome in outcomes.iter().filter(|outcome| outcome.prepared) { + let appeared = non_create_trust_violations(&mut tx, &outcome.schema).await?; + if !appeared.is_empty() { + return Err(AwaError::Validation(format!( + "schema {} changed while it was being prepared: {}", + outcome.schema, + appeared.join(", ") + ))); + } + } + tx.commit().await?; + Ok(outcomes) +} + +async fn non_create_trust_violations( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + schema: &str, +) -> Result, AwaError> { + Ok(sqlx::query_scalar( + r#" + SELECT violation.object || ': ' || violation.rule + FROM awa.queue_storage_schema_trust_violations($1) AS violation + WHERE violation.rule NOT LIKE 'CREATE granted to %' + ORDER BY 1 + "#, + ) + .bind(schema) + .fetch_all(tx.as_mut()) + .await?) +} + +/// Provision all four priorities of every shard a runtime may enqueue to for +/// each queue that already has lanes: the configured `enqueue_shards` or the +/// highest shard in use, whichever is larger. The runtimes losing `CREATE` +/// get the sequence privileges they need on the sequences this creates. +/// Returns how many lanes were new. +async fn provision_known_lanes( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + schema: &str, + runtime_roles: &[String], +) -> Result { + let lanes: Vec<(String, i16)> = sqlx::query_as(crate::audited_sql(format!( + r#" + SELECT known.queue, + GREATEST( + COALESCE(MAX(heads.enqueue_shard) + 1, 1), + COALESCE(( + SELECT MAX(meta.enqueue_shards) + FROM awa.queue_meta AS meta + WHERE meta.queue IN (known.queue, split_part(known.queue, '#', 1)) + ), 1) + )::smallint AS shard_count + FROM ( + SELECT queue FROM {schema}.queue_lanes + UNION + SELECT queue FROM {schema}.queue_enqueue_heads + ) AS known + LEFT JOIN {schema}.queue_enqueue_heads AS heads ON heads.queue = known.queue + GROUP BY known.queue + ORDER BY known.queue + "# + ))) + .fetch_all(tx.as_mut()) + .await?; + + let sequences_sql = "SELECT relation.relname::text FROM pg_class AS relation \ + JOIN pg_namespace AS namespace ON namespace.oid = relation.relnamespace \ + WHERE namespace.nspname = $1 AND relation.relkind = 'S'"; + let sequences_before: std::collections::BTreeSet = sqlx::query_scalar(sequences_sql) + .bind(schema) + .fetch_all(tx.as_mut()) + .await? + .into_iter() + .collect(); + let count_heads = format!("SELECT count(*) FROM {schema}.queue_enqueue_heads"); + let before: i64 = sqlx::query_scalar(crate::audited_sql(count_heads.clone())) + .fetch_one(tx.as_mut()) + .await?; + for (queue, shard_count) in lanes { + for priority in 1..=4 { + for shard in 0..shard_count { + crate::queue_storage::insert_lane_rows( + tx.as_mut(), + schema, + &queue, + priority, + shard, + ) + .await?; + } + } + } + let after: i64 = sqlx::query_scalar(crate::audited_sql(count_heads)) + .fetch_one(tx.as_mut()) + .await?; + + let sequences_after: Vec = sqlx::query_scalar(sequences_sql) + .bind(schema) + .fetch_all(tx.as_mut()) + .await?; + for sequence in sequences_after + .iter() + .filter(|sequence| !sequences_before.contains(*sequence)) + { + for role in runtime_roles { + sqlx::query(crate::audited_sql(format!( + "GRANT USAGE, SELECT, UPDATE ON SEQUENCE {schema}.{} TO {role}", + quote_identifier(sequence) + ))) + .execute(tx.as_mut()) + .await?; + } + } + Ok(after - before) +} + +fn quote_identifier(identifier: &str) -> String { + format!("\"{}\"", identifier.replace('"', "\"\"")) +} diff --git a/awa-model/tests/create_revoke_upgrade_test.rs b/awa-model/tests/create_revoke_upgrade_test.rs new file mode 100644 index 000000000..e91cd4b83 --- /dev/null +++ b/awa-model/tests/create_revoke_upgrade_test.rs @@ -0,0 +1,296 @@ +use awa_model::{audited_sql, migrations, storage}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::{Connection, Executor, PgConnection, PgPool}; +use std::str::FromStr; + +type TestResult = Result<(), Box>; +const SCHEMAS: [&str; 2] = ["awa", "custom_lanes"]; +const RELEASED_0_6_VERSION: i32 = 40; + +#[tokio::test] +async fn create_is_revoked_under_running_0_6_runtimes_without_failures() -> TestResult { + let url = std::env::var("DATABASE_URL")?; + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await?; + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("revoke_owner_{suffix}"); + let runtime = format!("revoke_runtime_{suffix}"); + let database = format!("revoke_db_{suffix}"); + admin + .execute(sqlx::raw_sql(audited_sql(format!( + "CREATE ROLE {owner} LOGIN PASSWORD 'revoke_test'; \ + CREATE ROLE {runtime} LOGIN PASSWORD 'revoke_test';" + )))) + .await?; + admin + .execute(sqlx::raw_sql(audited_sql(format!( + "CREATE DATABASE {database} OWNER {owner}" + )))) + .await?; + let options = PgConnectOptions::from_str(&url)? + .database(&database) + .password("revoke_test"); + let result = exercise(options, &owner, &runtime).await; + admin + .execute(sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + )))) + .await?; + admin + .execute(sqlx::raw_sql(audited_sql(format!( + "DROP ROLE {runtime}; DROP ROLE {owner};" + )))) + .await?; + admin.close().await; + result +} + +async fn exercise(options: PgConnectOptions, owner: &str, runtime: &str) -> TestResult { + let owner_options = options.clone().username(owner); + let owner = PgPoolOptions::new() + .max_connections(2) + .connect_with(owner_options.clone()) + .await?; + for (_, _, statement) in migrations::migration_sql_range(0, RELEASED_0_6_VERSION) { + owner.execute(sqlx::raw_sql(audited_sql(statement))).await?; + } + owner + .execute("SELECT awa.install_queue_storage_substrate('custom_lanes')") + .await?; + owner + .execute( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ('existing', 2) \ + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = 2", + ) + .await?; + for schema in SCHEMAS { + let released_helpers = + include_str!("fixtures/v040_lane_sequence_helpers.sql").replace("__schema__", schema); + owner + .execute(sqlx::raw_sql(audited_sql(released_helpers))) + .await?; + owner + .execute(sqlx::raw_sql(audited_sql(format!( + r#" + GRANT USAGE, CREATE ON SCHEMA {schema} TO {runtime}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA {schema} TO {runtime}; + GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA {schema} TO {runtime}; + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA {schema} TO {runtime}; + "# + )))) + .await?; + } + let runtime_options = options.username(runtime); + let mut long_lived = PgConnection::connect_with(&runtime_options).await?; + for schema in SCHEMAS { + create_lane_as_0_6(&mut long_lived, schema, "existing", 2, 0).await?; + reserve(&mut long_lived, schema, "existing", 2, 0).await?; + } + + revoke_create(&owner, runtime).await?; + let error = reserve(&mut long_lived, "awa", "existing", 2, 0) + .await + .expect_err("0.6 helpers need CREATE even for an existing sequence"); + assert_eq!(sqlstate(&error), Some("42501".into())); + grant_create(&owner, runtime).await?; + + owner + .execute(sqlx::raw_sql(audited_sql(format!( + "GRANT TRIGGER ON awa.queue_lanes TO {runtime}" + )))) + .await?; + let blocked = storage::prepare_create_revoke(&owner).await?; + let awa_outcome = blocked + .iter() + .find(|outcome| outcome.schema == "awa") + .unwrap(); + assert!(!awa_outcome.prepared); + assert!(awa_outcome + .blocking + .iter() + .any(|violation| violation + == &format!("relation awa.queue_lanes: TRIGGER granted to {runtime}"))); + assert!(!helper_is_guarded(&owner, "awa").await?); + owner + .execute(sqlx::raw_sql(audited_sql(format!( + "REVOKE TRIGGER ON awa.queue_lanes FROM {runtime}" + )))) + .await?; + + let impatient = PgPoolOptions::new() + .max_connections(1) + .after_connect(|conn, _| { + Box::pin(async move { + conn.execute("SET lock_timeout = '2s'").await?; + Ok(()) + }) + }) + .connect_with(owner_options.clone()) + .await?; + let mut in_flight = PgConnection::connect_with(&runtime_options).await?; + let mut open_enqueue = in_flight.begin().await?; + sqlx::query("SELECT awa.reserve_enqueue_seq('existing', 2::smallint, 0::smallint, 1::bigint)") + .execute(open_enqueue.as_mut()) + .await?; + + let prepared = storage::prepare_create_revoke(&impatient).await?; + assert_eq!( + prepared + .iter() + .map(|outcome| outcome.schema.as_str()) + .collect::>(), + SCHEMAS + ); + for outcome in &prepared { + assert!(outcome.prepared, "{outcome:?}"); + assert!(outcome.blocking.is_empty()); + assert_eq!( + outcome.revoke, + vec![format!( + "REVOKE CREATE ON SCHEMA \"{}\" FROM {runtime};", + outcome.schema + )] + ); + assert!(helper_is_guarded(&owner, &outcome.schema).await?); + } + let awa_outcome = prepared + .iter() + .find(|outcome| outcome.schema == "awa") + .unwrap(); + assert_eq!(awa_outcome.provisioned_lanes, 4 * 2 - 1); + let mut sequence_owner = PgConnection::connect_with(&runtime_options).await?; + for outcome in &prepared { + assert_eq!(outcome.sequence_grants.len(), 1, "{outcome:?}"); + let grant = &outcome.sequence_grants[0]; + assert_eq!(grant.run_as, runtime); + sequence_owner + .execute(sqlx::raw_sql(audited_sql(grant.statement.clone()))) + .await?; + impatient + .execute(sqlx::raw_sql(audited_sql(outcome.revoke.join("\n")))) + .await?; + } + sequence_owner.close().await?; + open_enqueue.commit().await?; + in_flight.close().await?; + impatient.close().await; + + for schema in SCHEMAS { + reserve(&mut long_lived, schema, "existing", 2, 0).await?; + } + let mut restarted = PgConnection::connect_with(&runtime_options).await?; + create_lane_as_0_6(&mut restarted, "awa", "existing", 1, 1).await?; + reserve(&mut restarted, "awa", "existing", 1, 1).await?; + let error = create_lane_as_0_6(&mut restarted, "awa", "unprovisioned", 2, 0) + .await + .expect_err("a queue never used before needs prepare-queue"); + assert_eq!(sqlstate(&error), Some("42501".into())); + + migrations::run(&owner).await?; + assert_eq!( + migrations::current_version(&owner).await?, + migrations::CURRENT_VERSION + ); + for schema in SCHEMAS { + reserve(&mut long_lived, schema, "existing", 2, 0).await?; + } + long_lived.close().await?; + restarted.close().await?; + owner.close().await; + Ok(()) +} + +/// The lane rows a 0.6 runtime inserts on its first enqueue to a lane. Their +/// triggers create the lane's sequences through the lane helpers. +async fn create_lane_as_0_6( + conn: &mut PgConnection, + schema: &str, + queue: &str, + priority: i16, + shard: i16, +) -> Result<(), sqlx::Error> { + let mut tx = conn.begin().await?; + for insert in [ + format!( + "INSERT INTO {schema}.queue_lanes (queue, priority) \ + SELECT $1, $2 WHERE $3::smallint IS NOT NULL ON CONFLICT (queue, priority) DO NOTHING" + ), + format!( + "INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) \ + VALUES ($1, $2, $3) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING" + ), + format!( + "INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) \ + VALUES ($1, $2, $3) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING" + ), + ] { + sqlx::query(audited_sql(insert)) + .bind(queue) + .bind(priority) + .bind(shard) + .execute(tx.as_mut()) + .await?; + } + tx.commit().await +} + +async fn reserve( + conn: &mut PgConnection, + schema: &str, + queue: &str, + priority: i16, + shard: i16, +) -> Result { + let mut tx = conn.begin().await?; + let start = sqlx::query_scalar(audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq($1, $2, $3, 1::bigint)" + ))) + .bind(queue) + .bind(priority) + .bind(shard) + .fetch_one(tx.as_mut()) + .await?; + tx.commit().await?; + Ok(start) +} + +async fn revoke_create(owner: &PgPool, runtime: &str) -> Result<(), sqlx::Error> { + for schema in SCHEMAS { + owner + .execute(sqlx::raw_sql(audited_sql(format!( + "REVOKE CREATE ON SCHEMA {schema} FROM {runtime}" + )))) + .await?; + } + Ok(()) +} + +async fn grant_create(owner: &PgPool, runtime: &str) -> Result<(), sqlx::Error> { + for schema in SCHEMAS { + owner + .execute(sqlx::raw_sql(audited_sql(format!( + "GRANT CREATE ON SCHEMA {schema} TO {runtime}" + )))) + .await?; + } + Ok(()) +} + +async fn helper_is_guarded(owner: &PgPool, schema: &str) -> Result { + let definition: String = sqlx::query_scalar( + "SELECT pg_get_functiondef(to_regprocedure(format('%I.ensure_lane_sequences(text, smallint, smallint)', $1::text)))", + ) + .bind(schema) + .fetch_one(owner) + .await?; + Ok(definition.contains("has not been provisioned")) +} + +fn sqlstate(error: &sqlx::Error) -> Option { + error + .as_database_error() + .and_then(|database| database.code()) + .map(|code| code.into_owned()) +} diff --git a/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql b/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql new file mode 100644 index 000000000..4b4b5beee --- /dev/null +++ b/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql @@ -0,0 +1,136 @@ +-- The lane helpers every 0.6 release installs (v0.6.0 through v0.6.9). + +CREATE OR REPLACE FUNCTION __schema__.ensure_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT +) +RETURNS VOID +LANGUAGE plpgsql +SET search_path = pg_catalog +AS $func$ +DECLARE + v_enqueue_seq TEXT := __schema__.queue_lane_sequence_name( + 'queue_enqueue_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + v_claim_seq TEXT := __schema__.queue_lane_sequence_name( + 'queue_claim_seq', + p_queue, + p_priority, + p_enqueue_shard + ); +BEGIN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + '__schema__', + v_enqueue_seq + ); + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + '__schema__', + v_claim_seq + ); + + UPDATE __schema__.queue_enqueue_heads + SET seq_name = v_enqueue_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_enqueue_seq; + + UPDATE __schema__.queue_claim_heads + SET seq_name = v_claim_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_claim_seq; +END; +$func$; + +CREATE OR REPLACE FUNCTION __schema__.queue_enqueue_head_sequence_sync() +RETURNS TRIGGER +LANGUAGE plpgsql +SET search_path = pg_catalog +AS $func$ +DECLARE + v_seq_name TEXT := __schema__.queue_lane_sequence_name( + 'queue_enqueue_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + v_qualified_seq TEXT := format('%I.%I', '__schema__', v_seq_name); + v_count BIGINT; + v_start BIGINT; +BEGIN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + '__schema__', + v_seq_name + ); + NEW.seq_name := v_seq_name; + + IF TG_OP = 'UPDATE' + AND NEW.next_seq IS DISTINCT FROM OLD.next_seq + AND NEW.next_seq > OLD.next_seq + THEN + v_count := NEW.next_seq - OLD.next_seq; + PERFORM pg_advisory_lock(hashtextextended(v_qualified_seq, 0)); + BEGIN + EXECUTE format( + 'SELECT nextval(%L::regclass)::bigint', + v_qualified_seq + ) + INTO v_start; + + IF v_count > 1 THEN + EXECUTE format( + 'SELECT setval(%L::regclass, %s, true)', + v_qualified_seq, + v_start + v_count - 1 + ); + END IF; + EXCEPTION WHEN OTHERS THEN + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + RAISE; + END; + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + NEW.next_seq := v_start + v_count; + ELSE + PERFORM __schema__.set_sequence_next(v_seq_name, NEW.next_seq); + END IF; + + RETURN NEW; +END; +$func$; + +CREATE OR REPLACE FUNCTION __schema__.queue_claim_head_sequence_sync() +RETURNS TRIGGER +LANGUAGE plpgsql +SET search_path = pg_catalog +AS $func$ +DECLARE + v_seq_name TEXT := __schema__.queue_lane_sequence_name( + 'queue_claim_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); +BEGIN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + '__schema__', + v_seq_name + ); + NEW.seq_name := v_seq_name; + + IF TG_OP = 'INSERT' THEN + PERFORM __schema__.set_sequence_next(v_seq_name, NEW.claim_seq); + END IF; + + RETURN NEW; +END; +$func$; diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index b6c1557e9..28348db24 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -86,7 +86,9 @@ Apply the sequence grants above after provisioning, or set the owner's default privileges first. Lane cursors need `UPDATE` for `setval`, in addition to `USAGE` and `SELECT`. `CREATE SEQUENCE IF NOT EXISTS` requires schema CREATE even when the sequence already exists, so provisioning alone does not repair pre-v049 -helpers: apply the migration first. Schema owners retain lazy creation for +helpers: apply the migration first, or, under a running 0.6 fleet, run +`awa storage prepare-create-revoke` as described in the +[0.6→0.7 upgrade guide](../upgrade-0.6-to-0.7.md#custom-schema-discovery-during-migration). Schema owners retain lazy creation for single-role installations; a restricted runtime encountering an unprepared lane fails with a provisioning hint. diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 3a46052a9..f690cfdd7 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -22,6 +22,9 @@ External migration runners must apply the complete generated `awa migrate --sql --from 48 --to 49` bundle: it refreshes the installer definition before replacing lane helpers in the installed schemas. See [Database roles](security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +To take `CREATE` away from a running 0.6 fleet before migrating, use +`awa storage prepare-create-revoke`; see +[Custom-schema discovery](#custom-schema-discovery-during-migration). ## Custom-schema discovery during migration @@ -71,38 +74,40 @@ schema owner: ``` - **A 0.6 split-role install whose runtime holds `CREATE` on the schema.** 0.6 - runtimes created lane sequences themselves, so they needed `CREATE`. - Sequences they own are accepted; the `CREATE` grant is not, because it lets - the runtime add objects while a migration runs. Until the migration refreshes - them, the 0.6 lane helpers and enqueue-head triggers run - `CREATE SEQUENCE IF NOT EXISTS` on every reservation, and that checks - `CREATE` even for a sequence that exists, so revoking the grant makes running - 0.6 producers and workers fail with `42501` until the migration commits. - Provisioning lanes in advance does not avoid this. Take a short pause: - - 1. Stop 0.6 producers and workers, or hold them at a maintenance gate, so no - enqueue or claim runs. - 2. Revoke the grant and commit: - - ```sql - REVOKE CREATE ON SCHEMA awa FROM awa_runtime; - ``` + runtimes create lane sequences themselves, so they need `CREATE`, and the + sequences they created belong to them. Sequences they own are accepted; the + `CREATE` grant is not, because it lets the runtime add objects while a + migration runs. 0.6 lane helpers run `CREATE SEQUENCE IF NOT EXISTS` on + every reservation, which checks `CREATE` even when the sequence exists, so + revoking the grant first would make running 0.6 producers and workers fail + with `42501`. Keep them running and prepare the database first: - 3. Run `awa migrate` as the migrator. The bundle refreshes the lane helpers - to use provisioned sequences without `CREATE` (v049). - 4. Provision any lanes the 0.7 runtimes will need, as the migrator: + 1. As the migrator, with a 0.7 binary: ```bash - PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue --enqueue-shards + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-create-revoke ``` - 5. Start the 0.7 runtimes. - - An uninterrupted path would need the revoke and the helper refresh to become - visible to every session at once; that has not been verified, so do not - combine them into one transaction as a substitute for the pause. + For each queue-storage schema it replaces the lane helpers with the ones + v049 installs, which create a sequence only when it is missing, and + provisions all four priorities of every shard of every queue already in + use (the larger of `awa.queue_meta.enqueue_shards` and the highest shard + in use). It grants the roles holding `CREATE` `USAGE, SELECT, UPDATE` on + the sequences it creates. It changes no table definitions and takes no + lock that producers or workers wait on. A schema with any other trust + violation is + listed under `blocking` and left unchanged; fix those first and re-run. + 2. Run each `sequence_grants` statement as the role in its `run_as`, usually + the runtime role. The migration reads and sets every lane sequence, and + only a sequence's owner can grant that. + 3. Run the printed `revoke` statements. Running 0.6 binaries keep working on + every provisioned lane. + 4. Run `awa migrate` as the migrator, then roll out 0.7 as usual. + + Between steps 1 and 4, a queue that has never been used, or a raised + `enqueue_shards`, needs `awa storage prepare-queue` before its first + enqueue; until then that enqueue fails with `42501` and a hint naming the + command. - **A custom schema prepared by a runtime or another role.** A worker that auto-prepares a schema, or a DBA who ran the preparation as another user, diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index a06981c7f..4690e0eb9 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -96,6 +96,9 @@ queue with `awa storage prepare-queue` as the migrator before starting producers or workers. Match shard/stripe counts; provision additional lanes before changing routing. Use the sequence grants and owner defaults in [Database roles](../../docs/security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +To revoke `CREATE` from a running 0.6 fleet without stopping it, run +`awa storage prepare-create-revoke` as the migrator first, then the grants and +revokes it prints, then `awa migrate`. ## Rolling Upgrades From 8b65e26bdfb889cf421473fecec8598e353bab2d Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Wed, 21 Oct 2026 21:39:51 +1300 Subject: [PATCH 03/10] test(upgrade): rehearse revoking CREATE under a released 0.6 fleet Runs an awa-pg 0.6.9 producer and worker as the runtime role while CREATE is revoked, with and without prepare-create-revoke, then migrates to the current version and reports failed inserts and slow inserts with lock waits. Against PostgreSQL 17: revoking without preparing fails every enqueue for the rest of the run (471 errors in 10 s); preparing first gives 0 failures and every job completes, with the prepare, grants and revokes taking 0.2 s and no insert slower than 250 ms until migrate, which blocks enqueues for its whole 5.2 s transaction. --- scripts/compat/create_revoke_workload.py | 84 ++++++++++++++++++++++++ scripts/rehearse-create-revoke.sh | 73 ++++++++++++++++++++ 2 files changed, 157 insertions(+) create mode 100644 scripts/compat/create_revoke_workload.py create mode 100755 scripts/rehearse-create-revoke.sh diff --git a/scripts/compat/create_revoke_workload.py b/scripts/compat/create_revoke_workload.py new file mode 100644 index 000000000..6af5c6dd2 --- /dev/null +++ b/scripts/compat/create_revoke_workload.py @@ -0,0 +1,84 @@ +"""A released 0.6 producer and worker that record every failure while the +database is prepared, has CREATE revoked, and is migrated underneath them.""" + +import asyncio +import json +import os +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +import awa + + +@dataclass +class RehearsalJob: + seq: int + + +async def main() -> int: + state_dir = Path(os.environ["REHEARSAL_DIR"]) + client = awa.AsyncClient(os.environ["RUNTIME_DATABASE_URL"]) + completed: set[int] = set() + inserted: list[int] = [] + errors: list[dict] = [] + started = time.monotonic() + last_insert = started + max_insert_gap = 0.0 + slow_inserts: list[dict] = [] + + @client.task(RehearsalJob, queue="rehearsal") + async def handle(job: RehearsalJob): + completed.add(job.args.seq) + + await client.start([("rehearsal", 4)]) + seq = 0 + try: + while not (state_dir / "stop").exists(): + priorities = [2] if not (state_dir / "all-priorities").exists() else [1, 2, 3, 4] + priority = priorities[seq % len(priorities)] + try: + attempt_started = time.monotonic() + await client.insert(RehearsalJob(seq=seq), queue="rehearsal", priority=priority) + now = time.monotonic() + if now - attempt_started > 0.25: + slow_inserts.append( + {"at": round(attempt_started - started, 2), "took": round(now - attempt_started, 2)} + ) + max_insert_gap = max(max_insert_gap, now - last_insert) + last_insert = now + inserted.append(seq) + except Exception as error: # noqa: BLE001 - every failure is the finding + errors.append( + { + "at": round(time.monotonic() - started, 2), + "priority": priority, + "error": f"{type(error).__name__}: {error}"[:300], + } + ) + seq += 1 + await asyncio.sleep(0.02) + + deadline = time.monotonic() + 60 + while time.monotonic() < deadline and not set(inserted) <= completed: + await asyncio.sleep(0.25) + finally: + await client.shutdown() + await client.close() + + summary = { + "attempted": seq, + "inserted": len(inserted), + "completed": len(set(inserted) & completed), + "errors": len(errors), + "first_errors": errors[:3], + "error_window_s": [errors[0]["at"], errors[-1]["at"]] if errors else None, + "max_insert_gap_s": round(max_insert_gap, 3), + "slow_inserts": slow_inserts[:10], + } + print(json.dumps(summary)) + return 0 + + +sys.exit(asyncio.run(main())) diff --git a/scripts/rehearse-create-revoke.sh b/scripts/rehearse-create-revoke.sh new file mode 100755 index 000000000..0bf0fff46 --- /dev/null +++ b/scripts/rehearse-create-revoke.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# A released awa-pg 0.6 producer and worker run through revoking their schema +# CREATE and `awa migrate`. `control` revokes without preparing; `prepared` +# follows the upgrade guide. Prints the workload's failure and stall summary. +# usage: ADMIN_URL=postgres://postgres:pw@host:port/postgres rehearse-create-revoke.sh control|prepared +set -euo pipefail +mode=${1:?control or prepared} +: "${ADMIN_URL:?set ADMIN_URL to a superuser URL on a disposable PostgreSQL server}" +admin=$ADMIN_URL +repo=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) +dir=$(mktemp -d) +SQLX_OFFLINE=true cargo build --quiet --manifest-path "$repo/Cargo.toml" -p awa-cli +awa_new=$(cargo metadata --no-deps --format-version 1 --manifest-path "$repo/Cargo.toml" | python3 -c 'import json,sys; print(json.load(sys.stdin)["target_directory"])')/debug/awa +uv venv --quiet --clear "$dir/venv" +uv pip install --quiet --python "$dir/venv" "awa-pg==${AWA_RELEASED_VERSION:-0.6.9}" +py06=$dir/venv/bin/python +suffix=$(date +%s%N | tail -c 8) +owner=rh_owner_$suffix runtime=rh_runtime_$suffix db=rh_db_$suffix +host=${admin#*@}; host=${host%%/*} +owner_url=postgres://$owner:rh@$host/$db +runtime_url=postgres://$runtime:rh@$host/$db +log() { echo "$(date +%T) [$mode] $*" >&2; } + +psql -q "$admin" -c "CREATE ROLE $owner LOGIN PASSWORD 'rh'" -c "CREATE ROLE $runtime LOGIN PASSWORD 'rh'" \ + -c "CREATE DATABASE $db OWNER $owner" +trap 'psql -q "$admin" -c "DROP DATABASE IF EXISTS $db WITH (FORCE)" -c "DROP ROLE IF EXISTS $runtime" -c "DROP ROLE IF EXISTS $owner" >/dev/null 2>&1 || true' EXIT + +log "migrating with released 0.6: $("$py06" -c 'import importlib.metadata as m; print(m.version("awa-pg"))')" +"$py06" -m awa --database-url "$owner_url" migrate >/dev/null +psql -q "$owner_url" <"$dir/summary.json" 2>"$dir/workload.err" & +workload=$! +t0=$(date +%s.%N) +mark() { log "t=$(echo "$(date +%s.%N) - $t0" | bc | cut -c1-5) $*"; } +sleep 10 + +if [ "$mode" = prepared ]; then + mark "prepare-create-revoke start" + "$awa_new" --database-url "$owner_url" storage prepare-create-revoke >"$dir/prepare.json" + jq -c '.[] | {schema, prepared, provisioned_lanes, revoke, grants: (.sequence_grants | length), blocking}' "$dir/prepare.json" >&2 + jq -r '.[].sequence_grants[] | select(.run_as == "'"$runtime"'") | .statement' "$dir/prepare.json" | psql -q "$runtime_url" + jq -r '.[].revoke[]' "$dir/prepare.json" | psql -q "$owner_url" + mark "grants and revokes done" +else + log "revoking CREATE without preparing" + psql -q "$owner_url" -c "REVOKE CREATE ON SCHEMA awa FROM $runtime" +fi +touch "$dir/all-priorities" +sleep 10 + +if [ "$mode" = prepared ]; then + mark "awa migrate start" + ( while [ ! -e "$dir/migrated" ]; do psql -At "$owner_url" -c "SELECT string_agg(DISTINCT format('%s waits %s on %s', w.pid, w.mode, COALESCE(w.relation::regclass::text, w.locktype)), '; ') FROM pg_locks w WHERE NOT w.granted" 2>/dev/null; sleep 0.1; done | grep -v '^$' | sort | uniq -c | sort -rn | head -8 >"$dir/waits.txt" ) & + "$awa_new" --database-url "$owner_url" migrate >"$dir/migrate.log" 2>&1 || { log "migrate failed"; tail -5 "$dir/migrate.log" >&2; } + touch "$dir/migrated"; sleep 0.3; cat "$dir/waits.txt" >&2 + mark "awa migrate done, schema v$(psql -At "$owner_url" -c 'SELECT max(version) FROM awa.schema_version')" + sleep 10 +fi +touch "$dir/stop" +wait "$workload" || true +tail -3 "$dir/workload.err" >&2 || true +echo "$mode $(cat "$dir/summary.json")" From 9f279a7a300587fa54f9db136cb71296f7841fc7 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Wed, 21 Oct 2026 21:40:28 +1300 Subject: [PATCH 04/10] docs(changelog): describe the zero-pause CREATE revoke --- CHANGELOG.md | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1fe8d79a1..83692b139 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,8 +24,12 @@ for each step. runtime role keeps working; only the migration is refused. Before upgrading, remove `TRIGGER` from runtime-role grants, which `GRANT ALL ON ALL TABLES` includes, remove runtime roles' `CREATE` on the schema, and remove triggers - of your own from awa tables. Revoking a running 0.6 fleet's `CREATE` needs - the order in the upgrade guide. + of your own from awa tables. A 0.6 fleet whose runtime role holds `CREATE` + keeps running throughout: as the migrator, run + `awa storage prepare-create-revoke`, run the sequence grants it prints as + the role it names, then its `REVOKE` statements, then `awa migrate`. A queue + first used after that needs `awa storage prepare-queue` before its first + enqueue. - **External migration runners.** Generate the pending SQL with an upgraded binary, for example `awa migrate --pending --extract-to ./pending`, apply complete bundles, and keep the files already applied unchanged. Apply @@ -43,6 +47,14 @@ for each step. ### Added +- **`awa storage prepare-create-revoke`** prepares a database still used by + 0.6 runtimes for revoking their schema `CREATE` without stopping them. It + swaps 0.6's lane helpers for ones that create a lane sequence only when it + is missing, provisions every priority and shard of each queue in use, and + prints the `REVOKE` statements and the sequence grants to run next. A schema + that breaks any other migration trust rule is reported and left unchanged. + In a rehearsal with released awa-pg 0.6.9, revoking without it failed every + enqueue; with it, nothing failed. - **Owner-scoped periodic schedules (opt-in, v045, [#481](https://github.com/hardbyte/awa/issues/481)).** A client can declare the complete set of periodic schedules it owns. awa then adds, updates and @@ -178,6 +190,9 @@ for each step. `running = 0` while workers held receipt claims, and scanned every ready partition with a catalog lookup per row (about 90 ms to 6 ms for that part on a 20,000-row backlog). Released 0.6.x workers keep working against v044. +- **`awa storage prepare-queue` works on lanes a 0.6 runtime created.** It + failed with `42501` on any such lane, because inserting the lane's existing + head rows fired a trigger that read a sequence the runtime role owns. - **Queue pruning no longer deadlocks with claims.** It takes child-partition locks without waiting and, on contention, retries on the next maintenance tick. From cbf5ebf5788fd97d300b66da79f7f418b8b3ce0f Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 02:49:47 +1300 Subject: [PATCH 05/10] Exclude tombstoned ready rows from point reads Apply the reviewed standalone read correction without changing the CREATE-revoke preparation protocol. Cover priority and queue moves across ready-ring wrap with the same two deterministic regressions. --- CHANGELOG.md | 5 ++ awa-model/src/queue_storage.rs | 13 ++- awa/tests/queue_storage_runtime_test.rs | 111 ++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2291f0328..9bc83fc57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **Point reads show the live queue and priority after a lane move.** Retained + ready rows from before a priority or queue change could win a read's tie + after the ring wrapped, returning the old fields even though the batch + operation succeeded. Point reads now exclude tombstoned source rows. + - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index b868b580a..6f2e74956 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -14197,8 +14197,17 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, ))) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index bf51546fe..a89cff79e 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -21093,3 +21093,114 @@ async fn test_queue_storage_legacy_row_claims_are_locked_before_the_heartbeat_st .expect("join heartbeat") .expect("heartbeat after the rewrite"); } + +/// A lane move leaves immutable source rows behind. Wrap the ring so the +/// destination is in an earlier partition: point reads must use live evidence, +/// rather than depending on which equal-timestamp ready row the plan emits last. +async fn ready_point_read_wrap_fixture() -> (TestDbGuard, sqlx::PgPool, QueueStorage, i64) { + let (guard, pool) = setup_pool(4).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_point_read_wrap".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + ..Default::default() + }, + ) + .await; + // A fresh empty ring can be positioned just before wrap without retaining + // unrelated jobs. This is the state reached after three rotations/prunes. + let schema = store.schema(); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.queue_ring_rotations (generation, slot) VALUES (3, 3)" + ))) + .execute(&pool) + .await + .expect("position empty ring before wrap"); + sqlx::query(audited_sql(format!( + "UPDATE {schema}.queue_ring_slots SET generation = 3 WHERE slot = 3" + ))) + .execute(&pool) + .await + .expect("set current slot generation"); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: "point_read_source".to_string(), + priority: 4, + ..Default::default() + }, + ) + .await; + assert!(matches!( + store.rotate(&pool).await.expect("wrap to slot zero"), + RotateOutcome::Rotated { slot: 0, .. } + )); + (guard, pool, store, job_id) +} + +async fn assert_live_ready_row( + pool: &sqlx::PgPool, + store: &QueueStorage, + job_id: i64, + expected_queue: &str, + expected_priority: i16, +) { + let schema = store.schema(); + let live: Vec<(String, i16, i32)> = sqlx::query_as(audited_sql(format!( + "SELECT ready.queue, ready.priority, ready.ready_slot + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + )" + ))) + .bind(job_id) + .fetch_all(pool) + .await + .expect("read live ready evidence"); + assert_eq!( + live, + vec![(expected_queue.to_string(), expected_priority, 0)] + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_priority_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .set_priority(&pool, job_id, 1) + .await + .expect("set priority")); + assert_live_ready_row(&pool, &store, job_id, "point_read_source", 1).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.priority, 1, + "point read must ignore the tombstoned source" + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .move_queue(&pool, job_id, "point_read_destination", None) + .await + .expect("move queue")); + assert_live_ready_row(&pool, &store, job_id, "point_read_destination", 4).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.queue, "point_read_destination", + "point read must ignore the tombstoned source" + ); +} From be2739ea40d0680a20b8c0ef36177311a7646367 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 05:28:05 +1300 Subject: [PATCH 06/10] test(storage): allow transient blocked prune in carry fixtures --- awa/tests/queue_storage_runtime_test.rs | 51 ++++++++++++++++--------- 1 file changed, 32 insertions(+), 19 deletions(-) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a89cff79e..b55b54fed 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -6,6 +6,8 @@ #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; +mod ci_timing; + use awa::audited_sql; use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, @@ -449,6 +451,32 @@ async fn create_store_with_config(pool: &sqlx::PgPool, config: QueueStorageConfi store } +/// Prune's success/accounting assertions allow a transient competing lock. +/// Retry only the expected slot's Blocked result; errors and other outcomes +/// still reach the caller's original assertions without being retried. +async fn prune_until_not_blocked( + store: &QueueStorage, + pool: &sqlx::PgPool, + retention: Duration, + slot: i32, +) -> PruneOutcome { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { + loop { + let outcome = store + .prune_oldest(pool, retention) + .await + .expect("prune_oldest"); + if matches!(outcome, PruneOutcome::Blocked { slot: blocked } if blocked == slot) { + tokio::time::sleep(Duration::from_millis(10)).await; + } else { + return outcome; + } + } + }) + .await + .expect("prune remained Blocked past the assertion deadline") +} + async fn create_store(pool: &sqlx::PgPool, schema: &str) -> QueueStorage { // Tests that go through this helper exercise the legacy // (non-receipts) lease-materialization path. The receipts mode @@ -13332,11 +13360,7 @@ async fn test_queue_storage_prune_carry_forward_survives_retry_and_rebuild() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest") - { + match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 0).await { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 2, @@ -13614,11 +13638,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest") - { + match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 0).await { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 1, @@ -13662,11 +13682,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { // Second prune: current slot is 0, so the oldest sealed non-current // slot is slot 1 (the carried row). It is carried again into slot 0. - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest re-carry") - { + match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 1).await { PruneOutcome::Pruned { slot: 1, carried_failed_rows: 1, @@ -15622,10 +15638,7 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { .await .expect("Failed to release reader lock"); - let pruned = store - .prune_oldest(&pool, Duration::ZERO) - .await - .expect("Failed to prune after reader lock release"); + let pruned = prune_until_not_blocked(&store, &pool, Duration::ZERO, 0).await; assert!( matches!(pruned, PruneOutcome::Pruned { slot: 0, .. }), "unexpected prune outcome after reader lock release: {pruned:?}" From ba5849f45370b4fecd1d8fc9addb12bbad566f15 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 13 Nov 2026 00:35:07 +0000 Subject: [PATCH 07/10] Remove prepare-create-revoke and the duplicated point-read fix The 0.6 line now ships what prepare-create-revoke did. On a v40 database, the prepared_lane_sequences schema patch (0.6.10) makes the lane helpers skip CREATE SEQUENCE when a lane's sequence exists, and the definer_lane_provisioning patch (0.6.14) creates a missing lane's sequences through the schema owner's SECURITY DEFINER functions. Running `awa migrate` with a 0.6.14 or later CLI therefore lets runtime roles lose schema CREATE while 0.6 binaries keep running, without a 0.7 command. The merged command would also have installed main's lane helpers, whose provision_lane_sequences calls awa.lane_sequence_problems, on 0.6 databases that lack it. The point-read tombstone fix, its two ring-wrap regressions and the carry-forward prune retry reached main through #629 and e2fbfa6b. --- CHANGELOG.md | 24 +- awa-cli/src/main.rs | 22 -- awa-model/src/migrations.rs | 44 --- awa-model/src/queue_storage.rs | 172 +++++----- awa-model/src/storage.rs | 309 ------------------ awa-model/tests/create_revoke_upgrade_test.rs | 296 ----------------- .../fixtures/v040_lane_sequence_helpers.sql | 136 -------- awa/tests/queue_storage_runtime_test.rs | 162 ++------- docs/security/database-roles.md | 4 +- docs/upgrade-0.6-to-0.7.md | 61 ++-- scripts/compat/create_revoke_workload.py | 84 ----- scripts/rehearse-create-revoke.sh | 73 ----- skills/awa-operations/SKILL.md | 3 - 13 files changed, 126 insertions(+), 1264 deletions(-) delete mode 100644 awa-model/tests/create_revoke_upgrade_test.rs delete mode 100644 awa-model/tests/fixtures/v040_lane_sequence_helpers.sql delete mode 100644 scripts/compat/create_revoke_workload.py delete mode 100755 scripts/rehearse-create-revoke.sh diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bc83fc57..065d50db1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,11 +6,6 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed -- **Point reads show the live queue and priority after a lane move.** Retained - ready rows from before a priority or queue change could win a read's tie - after the ring wrapped, returning the old fields even though the batch - operation succeeded. Point reads now exclude tombstoned source rows. - - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane @@ -259,12 +254,8 @@ for each step. runtime role keeps working; only the migration is refused. Before upgrading, remove `TRIGGER` from runtime-role grants, which `GRANT ALL ON ALL TABLES` includes, remove runtime roles' `CREATE` on the schema, and remove triggers - of your own from awa tables. A 0.6 fleet whose runtime role holds `CREATE` - keeps running throughout: as the migrator, run - `awa storage prepare-create-revoke`, run the sequence grants it prints as - the role it names, then its `REVOKE` statements, then `awa migrate`. A queue - first used after that needs `awa storage prepare-queue` before its first - enqueue. + of your own from awa tables. Revoking a running 0.6 fleet's `CREATE` needs + the order in the upgrade guide. - **External migration runners.** Generate the pending SQL with an upgraded binary, for example `awa migrate --pending --extract-to ./pending`, apply complete bundles, and keep the files already applied unchanged. Apply @@ -362,14 +353,6 @@ for each step. ### Added -- **`awa storage prepare-create-revoke`** prepares a database still used by - 0.6 runtimes for revoking their schema `CREATE` without stopping them. It - swaps 0.6's lane helpers for ones that create a lane sequence only when it - is missing, provisions every priority and shard of each queue in use, and - prints the `REVOKE` statements and the sequence grants to run next. A schema - that breaks any other migration trust rule is reported and left unchanged. - In a rehearsal with released awa-pg 0.6.9, revoking without it failed every - enqueue; with it, nothing failed. - **Job output ([ADR-048](docs/adr/048-retained-job-output.md), [guide](docs/job-output.md)).** A handler can record a JSON output of up to 4 KiB with `ctx.set_output` (`ctx.output_slot()` for `register` closures) or `job.set_output` in Python. @@ -780,9 +763,6 @@ for each step. `running = 0` while workers held receipt claims, and scanned every ready partition with a catalog lookup per row (about 90 ms to 6 ms for that part on a 20,000-row backlog). Released 0.6.x workers keep working against v044. -- **`awa storage prepare-queue` works on lanes a 0.6 runtime created.** It - failed with `42501` on any such lane, because inserting the lane's existing - head rows fired a trigger that read a sequence the runtime role owns. - **Queue pruning no longer deadlocks with claims.** It takes child-partition locks without waiting and, on contention, retries on the next maintenance tick. diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index c49ecb3ec..405fc3023 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -333,7 +333,6 @@ fn access_for(command: &Commands) -> context::Access { StorageCommands::Prepare { .. } | StorageCommands::PrepareQueueStorageSchema { .. } | StorageCommands::PrepareQueue { .. } - | StorageCommands::PrepareCreateRevoke | StorageCommands::Abort | StorageCommands::EnterMixedTransition { .. } | StorageCommands::Finalize { .. } @@ -827,11 +826,6 @@ enum StorageCommands { #[arg(long, default_value_t = 1)] queue_stripe_count: usize, }, - /// Before revoking runtime roles' schema CREATE under a running 0.6 fleet: - /// switch the lane helpers to provisioned sequences and provision every - /// lane already in use. Run as the migrator; prints the REVOKE statements - /// to run next - PrepareCreateRevoke, /// Abort a prepared or mixed-transition storage rollout before final activation Abort, /// Enter mixed transition and begin routing new writes to the prepared engine @@ -1919,22 +1913,6 @@ async fn run() -> Result<(), Box> { }))? ); } - StorageCommands::PrepareCreateRevoke => { - let schemas = awa_model::storage::prepare_create_revoke(&pool).await?; - println!("{}", serde_json::to_string_pretty(&schemas)?); - let blocked: Vec<&str> = schemas - .iter() - .filter(|schema| !schema.prepared) - .map(|schema| schema.schema.as_str()) - .collect(); - if !blocked.is_empty() { - return Err(format!( - "not prepared: {}; fix the listed blocking violations and re-run", - blocked.join(", ") - ) - .into()); - } - } StorageCommands::PrepareQueue { schema, queue, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 1dfadcadb..385035b65 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -316,32 +316,6 @@ const V20_UP: &str = include_str!("../migrations/v020_active_queue_storage_schem const V21_UP: &str = include_str!("../migrations/v021_shard_aware_lane_indexes.sql"); const V22_UP: &str = include_str!("../migrations/v022_delete_compat_terminal_counter.sql"); const SCHEMA_DISCOVERY_UP: &str = include_str!("../migrations/queue_storage_schema_discovery.sql"); - -/// The trust and discovery functions every pending bundle installs, without -/// the preflight that refuses an untrusted configured schema. -pub(crate) fn schema_trust_functions_sql() -> &'static str { - SCHEMA_DISCOVERY_UP - .split_once("DO $preflight$") - .map(|(functions, _)| functions) - .expect("the discovery bootstrap ends with its preflight") -} - -/// `awa.refresh_lane_sequence_helpers` as v023's installer defines it. -pub(crate) fn lane_sequence_helper_refresh_sql() -> &'static str { - const START: &str = - "CREATE OR REPLACE FUNCTION awa.refresh_lane_sequence_helpers(p_schema TEXT)"; - const END: &str = - "REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC;"; - let installer = include_str!("../migrations/v023_install_queue_storage_substrate.sql"); - let start = installer - .find(START) - .expect("v023 defines refresh_lane_sequence_helpers"); - let end = installer[start..] - .find(END) - .map(|offset| start + offset + END.len()) - .expect("v023 revokes refresh_lane_sequence_helpers from PUBLIC"); - &installer[start..end] -} const V23_UP: &str = concat!( include_str!("../migrations/queue_storage_schema_discovery.sql"), "\n", @@ -1123,24 +1097,6 @@ pub fn migration_sql_range(from: i32, to: i32) -> Vec<(i32, &'static str, String mod tests { use super::*; - #[test] - fn schema_trust_functions_stop_before_the_preflight() { - let functions = schema_trust_functions_sql(); - assert!(functions.contains("FUNCTION awa.queue_storage_schema_trust_violations")); - assert!(functions.contains("FUNCTION awa.queue_storage_schemas()")); - assert!(!functions.contains("$preflight$")); - assert_eq!(SCHEMA_DISCOVERY_UP.matches("DO $preflight$").count(), 1); - } - - #[test] - fn lane_sequence_helper_refresh_is_one_function_and_its_revoke() { - let refresh = lane_sequence_helper_refresh_sql(); - assert!(refresh.starts_with("CREATE OR REPLACE FUNCTION awa.refresh_lane_sequence_helpers")); - assert!(refresh.ends_with("FROM PUBLIC;")); - assert_eq!(refresh.matches("$refresh$;").count(), 1); - assert!(!refresh.contains("install_queue_storage_substrate")); - } - #[test] fn discovery_bundles_bootstrap_before_discovery_and_default_installation() { for (version, _, sql) in migration_sql() { diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 6f2e74956..339001705 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -9,7 +9,7 @@ use crate::{InsertParams, JobRow, JobState}; use chrono::TimeDelta; use chrono::{DateTime, Utc}; use sqlx::Acquire; -use sqlx::{PgConnection, PgPool, Postgres}; +use sqlx::{PgPool, Postgres}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; @@ -713,7 +713,7 @@ async fn set_prune_lock_timeout_tx( Ok(()) } -pub(crate) fn validate_ident(ident: &str) -> Result<(), AwaError> { +fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { Some(first) if first.is_ascii_lowercase() || first == '_' => {} @@ -4635,7 +4635,79 @@ impl QueueStorage { priority: i16, enqueue_shard: i16, ) -> Result<(), AwaError> { - insert_lane_rows(tx.as_mut(), self.schema(), queue, priority, enqueue_shard).await?; + let schema = self.schema(); + let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(lane_lock_key) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_lanes (queue, priority) + VALUES ($1, $2) + ON CONFLICT (queue, priority) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + // The head tables' BEFORE INSERT triggers run even for a row that + // ON CONFLICT then skips, and they read and set the lane's sequence. + // Skipping rows that exist keeps provisioning a live lane off a + // sequence another role owns, such as one a 0.6 runtime created. + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT $1, $2, $3 + WHERE NOT EXISTS ( + SELECT 1 FROM {schema}.queue_enqueue_heads + WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) + SELECT $1, $2, $3 + WHERE NOT EXISTS ( + SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + sqlx::query(audited_sql(format!( + r#" + SELECT {schema}.ensure_lane_sequences($1, $2, $3) + "# + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.mark_lane_ensured(queue, priority, enqueue_shard); Ok(()) } @@ -14197,17 +14269,8 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries AS ready - WHERE ready.job_id = $1 - AND NOT EXISTS ( - SELECT 1 FROM {schema}.ready_tombstones AS tomb - WHERE tomb.ready_slot = ready.ready_slot - AND tomb.ready_generation = ready.ready_generation - AND tomb.queue = ready.queue - AND tomb.priority = ready.priority - AND tomb.enqueue_shard = ready.enqueue_shard - AND tomb.lane_seq = ready.lane_seq - ) + FROM {schema}.ready_entries + WHERE job_id = $1 ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, ))) @@ -19700,87 +19763,6 @@ impl QueueStorage { } } -/// Insert a lane's rows and create its sequences, under the lane lock that -/// runtimes take for the same inserts. Head rows are inserted only when -/// missing, so their sequence triggers never fire for an existing lane, whose -/// sequences may belong to the runtime role that created them. -pub(crate) async fn insert_lane_rows( - conn: &mut PgConnection, - schema: &str, - queue: &str, - priority: i16, - enqueue_shard: i16, -) -> Result<(), AwaError> { - let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(lane_lock_key) - .execute(&mut *conn) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_lanes (queue, priority) - VALUES ($1, $2) - ON CONFLICT (queue, priority) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .execute(&mut *conn) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) - SELECT $1, $2, $3 - WHERE NOT EXISTS ( - SELECT 1 FROM {schema}.queue_enqueue_heads - WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 - ) - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(&mut *conn) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) - SELECT $1, $2, $3 - WHERE NOT EXISTS ( - SELECT 1 FROM {schema}.queue_claim_heads - WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 - ) - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(&mut *conn) - .await - .map_err(map_sqlx_error)?; - - sqlx::query(audited_sql(format!( - r#" - SELECT {schema}.ensure_lane_sequences($1, $2, $3) - "# - ))) - .bind(queue) - .bind(priority) - .bind(enqueue_shard) - .execute(&mut *conn) - .await - .map_err(map_sqlx_error)?; - Ok(()) -} - #[cfg(test)] mod writer_bind_tests { /// Multi-row `VALUES` writers bind one parameter per column per row, so a diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index d00083bad..b7a0ff1ad 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -469,312 +469,3 @@ pub async fn status_report(pool: &PgPool) -> Result, - /// Grants each sequence owner must run so the migrator can read and set - /// the lane sequences that runtime roles created. - pub sequence_grants: Vec, - /// Trust violations other than a `CREATE` grant. They must be fixed - /// before the schema can be prepared or migrated. - pub blocking: Vec, -} - -/// A grant that only a sequence's owner can make. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct SequenceGrant { - pub run_as: String, - pub statement: String, -} - -/// Prepare a database still used by 0.6 runtimes for revoking those roles' -/// `CREATE` on its queue-storage schemas, without stopping them. -/// -/// 0.6 lane helpers issue `CREATE SEQUENCE IF NOT EXISTS` on every -/// reservation, which needs `CREATE` even when the sequence exists. This -/// replaces them with the helpers v049 installs, which only create a -/// sequence that is missing, and provisions every priority and shard of -/// every queue already in use. Run it as the migrator. A schema with any -/// trust violation other than a `CREATE` grant is reported and left -/// unchanged. -pub async fn prepare_create_revoke(pool: &PgPool) -> Result, AwaError> { - let mut tx = pool.begin().await?; - sqlx::raw_sql(crate::migrations::schema_trust_functions_sql()) - .execute(tx.as_mut()) - .await?; - sqlx::raw_sql(crate::migrations::lane_sequence_helper_refresh_sql()) - .execute(tx.as_mut()) - .await?; - - let foreign_helpers: Vec = sqlx::query_scalar( - r#" - SELECT routine.oid::regprocedure::text - FROM pg_proc AS routine - JOIN pg_namespace AS namespace ON namespace.oid = routine.pronamespace - WHERE namespace.nspname = 'awa' - AND routine.proname IN ( - 'queue_storage_canonical_triggers', 'queue_storage_schema_trust_violations', - 'queue_storage_schema_is_trusted', 'queue_storage_schema_trust_summary', - 'queue_storage_schemas', 'refresh_lane_sequence_helpers' - ) - AND routine.proowner NOT IN ( - SELECT nspowner FROM pg_namespace WHERE nspname = 'awa' - UNION SELECT oid FROM pg_roles WHERE rolname = current_user - ) - ORDER BY 1 - "#, - ) - .fetch_all(tx.as_mut()) - .await?; - if !foreign_helpers.is_empty() { - return Err(AwaError::Validation(format!( - "awa helper functions are owned by another role: {}", - foreign_helpers.join(", ") - ))); - } - - let schemas: Vec = sqlx::query_scalar( - r#" - SELECT namespace.nspname::text - FROM pg_namespace AS namespace - WHERE has_schema_privilege(current_user, namespace.oid, 'USAGE') - AND EXISTS ( - SELECT 1 FROM pg_proc AS claim_fn - WHERE claim_fn.pronamespace = namespace.oid - AND claim_fn.proname = 'claim_ready_runtime' - AND oidvectortypes(claim_fn.proargtypes) - = 'text, bigint, double precision, double precision' - ) - AND EXISTS ( - SELECT 1 FROM pg_class AS heads - WHERE heads.relnamespace = namespace.oid AND heads.relname = 'queue_enqueue_heads' - ) - AND EXISTS ( - SELECT 1 FROM pg_class AS heads - WHERE heads.relnamespace = namespace.oid AND heads.relname = 'queue_claim_heads' - ) - ORDER BY 1 - "#, - ) - .fetch_all(tx.as_mut()) - .await?; - - let mut outcomes = Vec::with_capacity(schemas.len()); - for schema in schemas { - let mut blocking = non_create_trust_violations(&mut tx, &schema).await?; - if crate::queue_storage::validate_ident(&schema).is_err() { - blocking.push(format!( - "schema {}: name is not a lowercase identifier", - quote_identifier(&schema) - )); - } - let create_grantees: Vec = sqlx::query_scalar( - r#" - SELECT CASE privilege.grantee WHEN 0 THEN 'PUBLIC' - ELSE quote_ident(pg_get_userbyid(privilege.grantee)) END - FROM pg_namespace AS namespace - CROSS JOIN LATERAL aclexplode( - COALESCE(namespace.nspacl, acldefault('n', namespace.nspowner)) - ) AS privilege - WHERE namespace.nspname = $1 - AND privilege.privilege_type = 'CREATE' - AND privilege.grantee NOT IN ( - SELECT nspowner FROM pg_namespace WHERE nspname = 'awa' - UNION SELECT oid FROM pg_roles WHERE rolname = current_user - ) - ORDER BY 1 - "#, - ) - .bind(&schema) - .fetch_all(tx.as_mut()) - .await?; - let revoke = create_grantees - .iter() - .map(|grantee| { - format!( - "REVOKE CREATE ON SCHEMA {} FROM {grantee};", - quote_identifier(&schema) - ) - }) - .collect(); - - let sequence_grants = sqlx::query_as::<_, (String, String)>( - r#" - SELECT pg_get_userbyid(sequence.relowner)::text, - format( - 'GRANT SELECT, UPDATE ON SEQUENCE %s TO %I;', - string_agg(format('%I.%I', namespace.nspname, sequence.relname), ', ' - ORDER BY sequence.relname), - current_user - ) - FROM pg_class AS sequence - JOIN pg_namespace AS namespace ON namespace.oid = sequence.relnamespace - WHERE namespace.nspname = $1 - AND sequence.relkind = 'S' - AND NOT (has_sequence_privilege(current_user, sequence.oid, 'SELECT') - AND has_sequence_privilege(current_user, sequence.oid, 'UPDATE')) - GROUP BY sequence.relowner - ORDER BY 1 - "#, - ) - .bind(&schema) - .fetch_all(tx.as_mut()) - .await? - .into_iter() - .map(|(run_as, statement)| SequenceGrant { run_as, statement }) - .collect(); - - if !blocking.is_empty() { - outcomes.push(CreateRevokeSchema { - schema, - prepared: false, - provisioned_lanes: 0, - revoke, - sequence_grants, - blocking, - }); - continue; - } - - sqlx::query("SELECT awa.refresh_lane_sequence_helpers($1)") - .bind(&schema) - .execute(tx.as_mut()) - .await?; - let provisioned_lanes = provision_known_lanes(&mut tx, &schema, &create_grantees).await?; - outcomes.push(CreateRevokeSchema { - schema, - prepared: true, - provisioned_lanes, - revoke, - sequence_grants, - blocking, - }); - } - - for outcome in outcomes.iter().filter(|outcome| outcome.prepared) { - let appeared = non_create_trust_violations(&mut tx, &outcome.schema).await?; - if !appeared.is_empty() { - return Err(AwaError::Validation(format!( - "schema {} changed while it was being prepared: {}", - outcome.schema, - appeared.join(", ") - ))); - } - } - tx.commit().await?; - Ok(outcomes) -} - -async fn non_create_trust_violations( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - schema: &str, -) -> Result, AwaError> { - Ok(sqlx::query_scalar( - r#" - SELECT violation.object || ': ' || violation.rule - FROM awa.queue_storage_schema_trust_violations($1) AS violation - WHERE violation.rule NOT LIKE 'CREATE granted to %' - ORDER BY 1 - "#, - ) - .bind(schema) - .fetch_all(tx.as_mut()) - .await?) -} - -/// Provision all four priorities of every shard a runtime may enqueue to for -/// each queue that already has lanes: the configured `enqueue_shards` or the -/// highest shard in use, whichever is larger. The runtimes losing `CREATE` -/// get the sequence privileges they need on the sequences this creates. -/// Returns how many lanes were new. -async fn provision_known_lanes( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - schema: &str, - runtime_roles: &[String], -) -> Result { - let lanes: Vec<(String, i16)> = sqlx::query_as(crate::audited_sql(format!( - r#" - SELECT known.queue, - GREATEST( - COALESCE(MAX(heads.enqueue_shard) + 1, 1), - COALESCE(( - SELECT MAX(meta.enqueue_shards) - FROM awa.queue_meta AS meta - WHERE meta.queue IN (known.queue, split_part(known.queue, '#', 1)) - ), 1) - )::smallint AS shard_count - FROM ( - SELECT queue FROM {schema}.queue_lanes - UNION - SELECT queue FROM {schema}.queue_enqueue_heads - ) AS known - LEFT JOIN {schema}.queue_enqueue_heads AS heads ON heads.queue = known.queue - GROUP BY known.queue - ORDER BY known.queue - "# - ))) - .fetch_all(tx.as_mut()) - .await?; - - let sequences_sql = "SELECT relation.relname::text FROM pg_class AS relation \ - JOIN pg_namespace AS namespace ON namespace.oid = relation.relnamespace \ - WHERE namespace.nspname = $1 AND relation.relkind = 'S'"; - let sequences_before: std::collections::BTreeSet = sqlx::query_scalar(sequences_sql) - .bind(schema) - .fetch_all(tx.as_mut()) - .await? - .into_iter() - .collect(); - let count_heads = format!("SELECT count(*) FROM {schema}.queue_enqueue_heads"); - let before: i64 = sqlx::query_scalar(crate::audited_sql(count_heads.clone())) - .fetch_one(tx.as_mut()) - .await?; - for (queue, shard_count) in lanes { - for priority in 1..=4 { - for shard in 0..shard_count { - crate::queue_storage::insert_lane_rows( - tx.as_mut(), - schema, - &queue, - priority, - shard, - ) - .await?; - } - } - } - let after: i64 = sqlx::query_scalar(crate::audited_sql(count_heads)) - .fetch_one(tx.as_mut()) - .await?; - - let sequences_after: Vec = sqlx::query_scalar(sequences_sql) - .bind(schema) - .fetch_all(tx.as_mut()) - .await?; - for sequence in sequences_after - .iter() - .filter(|sequence| !sequences_before.contains(*sequence)) - { - for role in runtime_roles { - sqlx::query(crate::audited_sql(format!( - "GRANT USAGE, SELECT, UPDATE ON SEQUENCE {schema}.{} TO {role}", - quote_identifier(sequence) - ))) - .execute(tx.as_mut()) - .await?; - } - } - Ok(after - before) -} - -fn quote_identifier(identifier: &str) -> String { - format!("\"{}\"", identifier.replace('"', "\"\"")) -} diff --git a/awa-model/tests/create_revoke_upgrade_test.rs b/awa-model/tests/create_revoke_upgrade_test.rs deleted file mode 100644 index e91cd4b83..000000000 --- a/awa-model/tests/create_revoke_upgrade_test.rs +++ /dev/null @@ -1,296 +0,0 @@ -use awa_model::{audited_sql, migrations, storage}; -use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; -use sqlx::{Connection, Executor, PgConnection, PgPool}; -use std::str::FromStr; - -type TestResult = Result<(), Box>; -const SCHEMAS: [&str; 2] = ["awa", "custom_lanes"]; -const RELEASED_0_6_VERSION: i32 = 40; - -#[tokio::test] -async fn create_is_revoked_under_running_0_6_runtimes_without_failures() -> TestResult { - let url = std::env::var("DATABASE_URL")?; - let admin = PgPoolOptions::new() - .max_connections(2) - .connect(&url) - .await?; - let suffix = uuid::Uuid::new_v4().simple().to_string(); - let owner = format!("revoke_owner_{suffix}"); - let runtime = format!("revoke_runtime_{suffix}"); - let database = format!("revoke_db_{suffix}"); - admin - .execute(sqlx::raw_sql(audited_sql(format!( - "CREATE ROLE {owner} LOGIN PASSWORD 'revoke_test'; \ - CREATE ROLE {runtime} LOGIN PASSWORD 'revoke_test';" - )))) - .await?; - admin - .execute(sqlx::raw_sql(audited_sql(format!( - "CREATE DATABASE {database} OWNER {owner}" - )))) - .await?; - let options = PgConnectOptions::from_str(&url)? - .database(&database) - .password("revoke_test"); - let result = exercise(options, &owner, &runtime).await; - admin - .execute(sqlx::raw_sql(audited_sql(format!( - "DROP DATABASE {database} WITH (FORCE)" - )))) - .await?; - admin - .execute(sqlx::raw_sql(audited_sql(format!( - "DROP ROLE {runtime}; DROP ROLE {owner};" - )))) - .await?; - admin.close().await; - result -} - -async fn exercise(options: PgConnectOptions, owner: &str, runtime: &str) -> TestResult { - let owner_options = options.clone().username(owner); - let owner = PgPoolOptions::new() - .max_connections(2) - .connect_with(owner_options.clone()) - .await?; - for (_, _, statement) in migrations::migration_sql_range(0, RELEASED_0_6_VERSION) { - owner.execute(sqlx::raw_sql(audited_sql(statement))).await?; - } - owner - .execute("SELECT awa.install_queue_storage_substrate('custom_lanes')") - .await?; - owner - .execute( - "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ('existing', 2) \ - ON CONFLICT (queue) DO UPDATE SET enqueue_shards = 2", - ) - .await?; - for schema in SCHEMAS { - let released_helpers = - include_str!("fixtures/v040_lane_sequence_helpers.sql").replace("__schema__", schema); - owner - .execute(sqlx::raw_sql(audited_sql(released_helpers))) - .await?; - owner - .execute(sqlx::raw_sql(audited_sql(format!( - r#" - GRANT USAGE, CREATE ON SCHEMA {schema} TO {runtime}; - GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA {schema} TO {runtime}; - GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA {schema} TO {runtime}; - GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA {schema} TO {runtime}; - "# - )))) - .await?; - } - let runtime_options = options.username(runtime); - let mut long_lived = PgConnection::connect_with(&runtime_options).await?; - for schema in SCHEMAS { - create_lane_as_0_6(&mut long_lived, schema, "existing", 2, 0).await?; - reserve(&mut long_lived, schema, "existing", 2, 0).await?; - } - - revoke_create(&owner, runtime).await?; - let error = reserve(&mut long_lived, "awa", "existing", 2, 0) - .await - .expect_err("0.6 helpers need CREATE even for an existing sequence"); - assert_eq!(sqlstate(&error), Some("42501".into())); - grant_create(&owner, runtime).await?; - - owner - .execute(sqlx::raw_sql(audited_sql(format!( - "GRANT TRIGGER ON awa.queue_lanes TO {runtime}" - )))) - .await?; - let blocked = storage::prepare_create_revoke(&owner).await?; - let awa_outcome = blocked - .iter() - .find(|outcome| outcome.schema == "awa") - .unwrap(); - assert!(!awa_outcome.prepared); - assert!(awa_outcome - .blocking - .iter() - .any(|violation| violation - == &format!("relation awa.queue_lanes: TRIGGER granted to {runtime}"))); - assert!(!helper_is_guarded(&owner, "awa").await?); - owner - .execute(sqlx::raw_sql(audited_sql(format!( - "REVOKE TRIGGER ON awa.queue_lanes FROM {runtime}" - )))) - .await?; - - let impatient = PgPoolOptions::new() - .max_connections(1) - .after_connect(|conn, _| { - Box::pin(async move { - conn.execute("SET lock_timeout = '2s'").await?; - Ok(()) - }) - }) - .connect_with(owner_options.clone()) - .await?; - let mut in_flight = PgConnection::connect_with(&runtime_options).await?; - let mut open_enqueue = in_flight.begin().await?; - sqlx::query("SELECT awa.reserve_enqueue_seq('existing', 2::smallint, 0::smallint, 1::bigint)") - .execute(open_enqueue.as_mut()) - .await?; - - let prepared = storage::prepare_create_revoke(&impatient).await?; - assert_eq!( - prepared - .iter() - .map(|outcome| outcome.schema.as_str()) - .collect::>(), - SCHEMAS - ); - for outcome in &prepared { - assert!(outcome.prepared, "{outcome:?}"); - assert!(outcome.blocking.is_empty()); - assert_eq!( - outcome.revoke, - vec![format!( - "REVOKE CREATE ON SCHEMA \"{}\" FROM {runtime};", - outcome.schema - )] - ); - assert!(helper_is_guarded(&owner, &outcome.schema).await?); - } - let awa_outcome = prepared - .iter() - .find(|outcome| outcome.schema == "awa") - .unwrap(); - assert_eq!(awa_outcome.provisioned_lanes, 4 * 2 - 1); - let mut sequence_owner = PgConnection::connect_with(&runtime_options).await?; - for outcome in &prepared { - assert_eq!(outcome.sequence_grants.len(), 1, "{outcome:?}"); - let grant = &outcome.sequence_grants[0]; - assert_eq!(grant.run_as, runtime); - sequence_owner - .execute(sqlx::raw_sql(audited_sql(grant.statement.clone()))) - .await?; - impatient - .execute(sqlx::raw_sql(audited_sql(outcome.revoke.join("\n")))) - .await?; - } - sequence_owner.close().await?; - open_enqueue.commit().await?; - in_flight.close().await?; - impatient.close().await; - - for schema in SCHEMAS { - reserve(&mut long_lived, schema, "existing", 2, 0).await?; - } - let mut restarted = PgConnection::connect_with(&runtime_options).await?; - create_lane_as_0_6(&mut restarted, "awa", "existing", 1, 1).await?; - reserve(&mut restarted, "awa", "existing", 1, 1).await?; - let error = create_lane_as_0_6(&mut restarted, "awa", "unprovisioned", 2, 0) - .await - .expect_err("a queue never used before needs prepare-queue"); - assert_eq!(sqlstate(&error), Some("42501".into())); - - migrations::run(&owner).await?; - assert_eq!( - migrations::current_version(&owner).await?, - migrations::CURRENT_VERSION - ); - for schema in SCHEMAS { - reserve(&mut long_lived, schema, "existing", 2, 0).await?; - } - long_lived.close().await?; - restarted.close().await?; - owner.close().await; - Ok(()) -} - -/// The lane rows a 0.6 runtime inserts on its first enqueue to a lane. Their -/// triggers create the lane's sequences through the lane helpers. -async fn create_lane_as_0_6( - conn: &mut PgConnection, - schema: &str, - queue: &str, - priority: i16, - shard: i16, -) -> Result<(), sqlx::Error> { - let mut tx = conn.begin().await?; - for insert in [ - format!( - "INSERT INTO {schema}.queue_lanes (queue, priority) \ - SELECT $1, $2 WHERE $3::smallint IS NOT NULL ON CONFLICT (queue, priority) DO NOTHING" - ), - format!( - "INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) \ - VALUES ($1, $2, $3) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING" - ), - format!( - "INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) \ - VALUES ($1, $2, $3) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING" - ), - ] { - sqlx::query(audited_sql(insert)) - .bind(queue) - .bind(priority) - .bind(shard) - .execute(tx.as_mut()) - .await?; - } - tx.commit().await -} - -async fn reserve( - conn: &mut PgConnection, - schema: &str, - queue: &str, - priority: i16, - shard: i16, -) -> Result { - let mut tx = conn.begin().await?; - let start = sqlx::query_scalar(audited_sql(format!( - "SELECT {schema}.reserve_enqueue_seq($1, $2, $3, 1::bigint)" - ))) - .bind(queue) - .bind(priority) - .bind(shard) - .fetch_one(tx.as_mut()) - .await?; - tx.commit().await?; - Ok(start) -} - -async fn revoke_create(owner: &PgPool, runtime: &str) -> Result<(), sqlx::Error> { - for schema in SCHEMAS { - owner - .execute(sqlx::raw_sql(audited_sql(format!( - "REVOKE CREATE ON SCHEMA {schema} FROM {runtime}" - )))) - .await?; - } - Ok(()) -} - -async fn grant_create(owner: &PgPool, runtime: &str) -> Result<(), sqlx::Error> { - for schema in SCHEMAS { - owner - .execute(sqlx::raw_sql(audited_sql(format!( - "GRANT CREATE ON SCHEMA {schema} TO {runtime}" - )))) - .await?; - } - Ok(()) -} - -async fn helper_is_guarded(owner: &PgPool, schema: &str) -> Result { - let definition: String = sqlx::query_scalar( - "SELECT pg_get_functiondef(to_regprocedure(format('%I.ensure_lane_sequences(text, smallint, smallint)', $1::text)))", - ) - .bind(schema) - .fetch_one(owner) - .await?; - Ok(definition.contains("has not been provisioned")) -} - -fn sqlstate(error: &sqlx::Error) -> Option { - error - .as_database_error() - .and_then(|database| database.code()) - .map(|code| code.into_owned()) -} diff --git a/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql b/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql deleted file mode 100644 index 4b4b5beee..000000000 --- a/awa-model/tests/fixtures/v040_lane_sequence_helpers.sql +++ /dev/null @@ -1,136 +0,0 @@ --- The lane helpers every 0.6 release installs (v0.6.0 through v0.6.9). - -CREATE OR REPLACE FUNCTION __schema__.ensure_lane_sequences( - p_queue TEXT, - p_priority SMALLINT, - p_enqueue_shard SMALLINT -) -RETURNS VOID -LANGUAGE plpgsql -SET search_path = pg_catalog -AS $func$ -DECLARE - v_enqueue_seq TEXT := __schema__.queue_lane_sequence_name( - 'queue_enqueue_seq', - p_queue, - p_priority, - p_enqueue_shard - ); - v_claim_seq TEXT := __schema__.queue_lane_sequence_name( - 'queue_claim_seq', - p_queue, - p_priority, - p_enqueue_shard - ); -BEGIN - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - '__schema__', - v_enqueue_seq - ); - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - '__schema__', - v_claim_seq - ); - - UPDATE __schema__.queue_enqueue_heads - SET seq_name = v_enqueue_seq - WHERE queue = p_queue - AND priority = p_priority - AND enqueue_shard = p_enqueue_shard - AND seq_name IS DISTINCT FROM v_enqueue_seq; - - UPDATE __schema__.queue_claim_heads - SET seq_name = v_claim_seq - WHERE queue = p_queue - AND priority = p_priority - AND enqueue_shard = p_enqueue_shard - AND seq_name IS DISTINCT FROM v_claim_seq; -END; -$func$; - -CREATE OR REPLACE FUNCTION __schema__.queue_enqueue_head_sequence_sync() -RETURNS TRIGGER -LANGUAGE plpgsql -SET search_path = pg_catalog -AS $func$ -DECLARE - v_seq_name TEXT := __schema__.queue_lane_sequence_name( - 'queue_enqueue_seq', - NEW.queue, - NEW.priority, - NEW.enqueue_shard - ); - v_qualified_seq TEXT := format('%I.%I', '__schema__', v_seq_name); - v_count BIGINT; - v_start BIGINT; -BEGIN - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - '__schema__', - v_seq_name - ); - NEW.seq_name := v_seq_name; - - IF TG_OP = 'UPDATE' - AND NEW.next_seq IS DISTINCT FROM OLD.next_seq - AND NEW.next_seq > OLD.next_seq - THEN - v_count := NEW.next_seq - OLD.next_seq; - PERFORM pg_advisory_lock(hashtextextended(v_qualified_seq, 0)); - BEGIN - EXECUTE format( - 'SELECT nextval(%L::regclass)::bigint', - v_qualified_seq - ) - INTO v_start; - - IF v_count > 1 THEN - EXECUTE format( - 'SELECT setval(%L::regclass, %s, true)', - v_qualified_seq, - v_start + v_count - 1 - ); - END IF; - EXCEPTION WHEN OTHERS THEN - PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); - RAISE; - END; - PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); - NEW.next_seq := v_start + v_count; - ELSE - PERFORM __schema__.set_sequence_next(v_seq_name, NEW.next_seq); - END IF; - - RETURN NEW; -END; -$func$; - -CREATE OR REPLACE FUNCTION __schema__.queue_claim_head_sequence_sync() -RETURNS TRIGGER -LANGUAGE plpgsql -SET search_path = pg_catalog -AS $func$ -DECLARE - v_seq_name TEXT := __schema__.queue_lane_sequence_name( - 'queue_claim_seq', - NEW.queue, - NEW.priority, - NEW.enqueue_shard - ); -BEGIN - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %I.%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - '__schema__', - v_seq_name - ); - NEW.seq_name := v_seq_name; - - IF TG_OP = 'INSERT' THEN - PERFORM __schema__.set_sequence_next(v_seq_name, NEW.claim_seq); - END IF; - - RETURN NEW; -END; -$func$; diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index b55b54fed..bf51546fe 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -6,8 +6,6 @@ #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; -mod ci_timing; - use awa::audited_sql; use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, @@ -451,32 +449,6 @@ async fn create_store_with_config(pool: &sqlx::PgPool, config: QueueStorageConfi store } -/// Prune's success/accounting assertions allow a transient competing lock. -/// Retry only the expected slot's Blocked result; errors and other outcomes -/// still reach the caller's original assertions without being retried. -async fn prune_until_not_blocked( - store: &QueueStorage, - pool: &sqlx::PgPool, - retention: Duration, - slot: i32, -) -> PruneOutcome { - tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { - loop { - let outcome = store - .prune_oldest(pool, retention) - .await - .expect("prune_oldest"); - if matches!(outcome, PruneOutcome::Blocked { slot: blocked } if blocked == slot) { - tokio::time::sleep(Duration::from_millis(10)).await; - } else { - return outcome; - } - } - }) - .await - .expect("prune remained Blocked past the assertion deadline") -} - async fn create_store(pool: &sqlx::PgPool, schema: &str) -> QueueStorage { // Tests that go through this helper exercise the legacy // (non-receipts) lease-materialization path. The receipts mode @@ -13360,7 +13332,11 @@ async fn test_queue_storage_prune_carry_forward_survives_retry_and_rebuild() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 0).await { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("prune_oldest") + { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 2, @@ -13638,7 +13614,11 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 0).await { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("prune_oldest") + { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 1, @@ -13682,7 +13662,11 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { // Second prune: current slot is 0, so the oldest sealed non-current // slot is slot 1 (the carried row). It is carried again into slot 0. - match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 1).await { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("prune_oldest re-carry") + { PruneOutcome::Pruned { slot: 1, carried_failed_rows: 1, @@ -15638,7 +15622,10 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { .await .expect("Failed to release reader lock"); - let pruned = prune_until_not_blocked(&store, &pool, Duration::ZERO, 0).await; + let pruned = store + .prune_oldest(&pool, Duration::ZERO) + .await + .expect("Failed to prune after reader lock release"); assert!( matches!(pruned, PruneOutcome::Pruned { slot: 0, .. }), "unexpected prune outcome after reader lock release: {pruned:?}" @@ -21106,114 +21093,3 @@ async fn test_queue_storage_legacy_row_claims_are_locked_before_the_heartbeat_st .expect("join heartbeat") .expect("heartbeat after the rewrite"); } - -/// A lane move leaves immutable source rows behind. Wrap the ring so the -/// destination is in an earlier partition: point reads must use live evidence, -/// rather than depending on which equal-timestamp ready row the plan emits last. -async fn ready_point_read_wrap_fixture() -> (TestDbGuard, sqlx::PgPool, QueueStorage, i64) { - let (guard, pool) = setup_pool(4).await; - let store = create_store_with_config( - &pool, - QueueStorageConfig { - schema: "awa_qs_point_read_wrap".to_string(), - queue_slot_count: 4, - lease_slot_count: 2, - claim_slot_count: 2, - ..Default::default() - }, - ) - .await; - // A fresh empty ring can be positioned just before wrap without retaining - // unrelated jobs. This is the state reached after three rotations/prunes. - let schema = store.schema(); - sqlx::query(audited_sql(format!( - "INSERT INTO {schema}.queue_ring_rotations (generation, slot) VALUES (3, 3)" - ))) - .execute(&pool) - .await - .expect("position empty ring before wrap"); - sqlx::query(audited_sql(format!( - "UPDATE {schema}.queue_ring_slots SET generation = 3 WHERE slot = 3" - ))) - .execute(&pool) - .await - .expect("set current slot generation"); - let job_id = enqueue_job( - &pool, - &store, - &CompleteJob { id: 1 }, - InsertOpts { - queue: "point_read_source".to_string(), - priority: 4, - ..Default::default() - }, - ) - .await; - assert!(matches!( - store.rotate(&pool).await.expect("wrap to slot zero"), - RotateOutcome::Rotated { slot: 0, .. } - )); - (guard, pool, store, job_id) -} - -async fn assert_live_ready_row( - pool: &sqlx::PgPool, - store: &QueueStorage, - job_id: i64, - expected_queue: &str, - expected_priority: i16, -) { - let schema = store.schema(); - let live: Vec<(String, i16, i32)> = sqlx::query_as(audited_sql(format!( - "SELECT ready.queue, ready.priority, ready.ready_slot - FROM {schema}.ready_entries AS ready - WHERE ready.job_id = $1 - AND NOT EXISTS ( - SELECT 1 FROM {schema}.ready_tombstones AS tomb - WHERE tomb.ready_slot = ready.ready_slot - AND tomb.ready_generation = ready.ready_generation - AND tomb.queue = ready.queue - AND tomb.priority = ready.priority - AND tomb.enqueue_shard = ready.enqueue_shard - AND tomb.lane_seq = ready.lane_seq - )" - ))) - .bind(job_id) - .fetch_all(pool) - .await - .expect("read live ready evidence"); - assert_eq!( - live, - vec![(expected_queue.to_string(), expected_priority, 0)] - ); -} - -#[tokio::test] -async fn test_point_read_uses_live_ready_priority_after_ring_wrap() { - let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; - assert!(store - .set_priority(&pool, job_id, 1) - .await - .expect("set priority")); - assert_live_ready_row(&pool, &store, job_id, "point_read_source", 1).await; - let job = admin::get_job(&pool, job_id).await.expect("read moved job"); - assert_eq!( - job.priority, 1, - "point read must ignore the tombstoned source" - ); -} - -#[tokio::test] -async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { - let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; - assert!(store - .move_queue(&pool, job_id, "point_read_destination", None) - .await - .expect("move queue")); - assert_live_ready_row(&pool, &store, job_id, "point_read_destination", 4).await; - let job = admin::get_job(&pool, job_id).await.expect("read moved job"); - assert_eq!( - job.queue, "point_read_destination", - "point read must ignore the tombstoned source" - ); -} diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 28348db24..b6c1557e9 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -86,9 +86,7 @@ Apply the sequence grants above after provisioning, or set the owner's default privileges first. Lane cursors need `UPDATE` for `setval`, in addition to `USAGE` and `SELECT`. `CREATE SEQUENCE IF NOT EXISTS` requires schema CREATE even when the sequence already exists, so provisioning alone does not repair pre-v049 -helpers: apply the migration first, or, under a running 0.6 fleet, run -`awa storage prepare-create-revoke` as described in the -[0.6→0.7 upgrade guide](../upgrade-0.6-to-0.7.md#custom-schema-discovery-during-migration). Schema owners retain lazy creation for +helpers: apply the migration first. Schema owners retain lazy creation for single-role installations; a restricted runtime encountering an unprepared lane fails with a provisioning hint. diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index cec6b451c..0afc48afb 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -22,9 +22,6 @@ External migration runners must apply the complete generated `awa migrate --sql --from 48 --to 49` bundle: it refreshes the installer definition before replacing lane helpers in the installed schemas. See [Database roles](security/database-roles.md#provision-queues-before-starting-restricted-runtimes). -To take `CREATE` away from a running 0.6 fleet before migrating, use -`awa storage prepare-create-revoke`; see -[Custom-schema discovery](#custom-schema-discovery-during-migration). ## Custom-schema discovery during migration @@ -74,42 +71,38 @@ schema owner: ``` - **A 0.6 split-role install whose runtime holds `CREATE` on the schema.** 0.6 - runtimes create lane sequences themselves, so they need `CREATE`, and the - sequences they created belong to them. Sequences they own are accepted; the - `CREATE` grant is not, because it lets the runtime add objects while a - migration runs. 0.6 lane helpers run `CREATE SEQUENCE IF NOT EXISTS` on - every reservation, which checks `CREATE` even when the sequence exists, so - revoking the grant first would make running 0.6 producers and workers fail - with `42501`. Keep them running and prepare the database first: + runtimes created lane sequences themselves, so they needed `CREATE`. + Sequences they own are accepted; the `CREATE` grant is not, because it lets + the runtime add objects while a migration runs. Until the migration refreshes + them, the 0.6 lane helpers and enqueue-head triggers run + `CREATE SEQUENCE IF NOT EXISTS` on every reservation, and that checks + `CREATE` even for a sequence that exists, so revoking the grant makes running + 0.6 producers and workers fail with `42501` until the migration commits. + Provisioning lanes in advance does not avoid this. Take a short pause: + + 1. Stop 0.6 producers and workers, or hold them at a maintenance gate, so no + enqueue or claim runs. + 2. Revoke the grant and commit: + + ```sql + REVOKE CREATE ON SCHEMA awa FROM awa_runtime; + ``` - 1. As the migrator, with a 0.7 binary: + 3. Run `awa migrate` as the migrator. The bundle refreshes the lane helpers + to use provisioned sequences without `CREATE` (v049). + 4. Provision any lanes the 0.7 runtimes will need, as the migrator: ```bash - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-create-revoke + PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ + --queue --enqueue-shards ``` - For each queue-storage schema it replaces the lane helpers with the ones - v049 installs, which create a sequence only when it is missing, and - provisions all four priorities of every shard of every queue already in - use (the larger of `awa.queue_meta.enqueue_shards` and the highest shard - in use). It grants the roles holding `CREATE` `USAGE, SELECT, UPDATE` on - the sequences it creates. It changes no table definitions. Producers and - workers using already-provisioned lanes can continue reserving sequences; - a runtime provisioning a lane may wait on the preparation transaction's - lane advisory lock until preparation commits. A schema with any other - trust violation is listed under `blocking` and left unchanged; fix those - first and re-run. - 2. Run each `sequence_grants` statement as the role in its `run_as`, usually - the runtime role. The migration reads and sets every lane sequence, and - only a sequence's owner can grant that. - 3. Run the printed `revoke` statements. Running 0.6 binaries keep working on - every provisioned lane. - 4. Run `awa migrate` as the migrator, then roll out 0.7 as usual. - - Between steps 1 and 4, a queue that has never been used, or a raised - `enqueue_shards`, needs `awa storage prepare-queue` before its first - enqueue; until then that enqueue fails with `42501` and a hint naming the - command. + 5. Start the 0.7 runtimes. + + An uninterrupted path would need the revoke and the helper refresh to become + visible to every session at once; that has not been verified, so do not + combine them into one transaction as a substitute for the pause. - **A custom schema prepared by a runtime or another role.** A worker that auto-prepares a schema, or a DBA who ran the preparation as another user, diff --git a/scripts/compat/create_revoke_workload.py b/scripts/compat/create_revoke_workload.py deleted file mode 100644 index 6af5c6dd2..000000000 --- a/scripts/compat/create_revoke_workload.py +++ /dev/null @@ -1,84 +0,0 @@ -"""A released 0.6 producer and worker that record every failure while the -database is prepared, has CREATE revoked, and is migrated underneath them.""" - -import asyncio -import json -import os -import sys -import time -from dataclasses import dataclass -from pathlib import Path - -import awa - - -@dataclass -class RehearsalJob: - seq: int - - -async def main() -> int: - state_dir = Path(os.environ["REHEARSAL_DIR"]) - client = awa.AsyncClient(os.environ["RUNTIME_DATABASE_URL"]) - completed: set[int] = set() - inserted: list[int] = [] - errors: list[dict] = [] - started = time.monotonic() - last_insert = started - max_insert_gap = 0.0 - slow_inserts: list[dict] = [] - - @client.task(RehearsalJob, queue="rehearsal") - async def handle(job: RehearsalJob): - completed.add(job.args.seq) - - await client.start([("rehearsal", 4)]) - seq = 0 - try: - while not (state_dir / "stop").exists(): - priorities = [2] if not (state_dir / "all-priorities").exists() else [1, 2, 3, 4] - priority = priorities[seq % len(priorities)] - try: - attempt_started = time.monotonic() - await client.insert(RehearsalJob(seq=seq), queue="rehearsal", priority=priority) - now = time.monotonic() - if now - attempt_started > 0.25: - slow_inserts.append( - {"at": round(attempt_started - started, 2), "took": round(now - attempt_started, 2)} - ) - max_insert_gap = max(max_insert_gap, now - last_insert) - last_insert = now - inserted.append(seq) - except Exception as error: # noqa: BLE001 - every failure is the finding - errors.append( - { - "at": round(time.monotonic() - started, 2), - "priority": priority, - "error": f"{type(error).__name__}: {error}"[:300], - } - ) - seq += 1 - await asyncio.sleep(0.02) - - deadline = time.monotonic() + 60 - while time.monotonic() < deadline and not set(inserted) <= completed: - await asyncio.sleep(0.25) - finally: - await client.shutdown() - await client.close() - - summary = { - "attempted": seq, - "inserted": len(inserted), - "completed": len(set(inserted) & completed), - "errors": len(errors), - "first_errors": errors[:3], - "error_window_s": [errors[0]["at"], errors[-1]["at"]] if errors else None, - "max_insert_gap_s": round(max_insert_gap, 3), - "slow_inserts": slow_inserts[:10], - } - print(json.dumps(summary)) - return 0 - - -sys.exit(asyncio.run(main())) diff --git a/scripts/rehearse-create-revoke.sh b/scripts/rehearse-create-revoke.sh deleted file mode 100755 index 0bf0fff46..000000000 --- a/scripts/rehearse-create-revoke.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env bash -# A released awa-pg 0.6 producer and worker run through revoking their schema -# CREATE and `awa migrate`. `control` revokes without preparing; `prepared` -# follows the upgrade guide. Prints the workload's failure and stall summary. -# usage: ADMIN_URL=postgres://postgres:pw@host:port/postgres rehearse-create-revoke.sh control|prepared -set -euo pipefail -mode=${1:?control or prepared} -: "${ADMIN_URL:?set ADMIN_URL to a superuser URL on a disposable PostgreSQL server}" -admin=$ADMIN_URL -repo=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) -dir=$(mktemp -d) -SQLX_OFFLINE=true cargo build --quiet --manifest-path "$repo/Cargo.toml" -p awa-cli -awa_new=$(cargo metadata --no-deps --format-version 1 --manifest-path "$repo/Cargo.toml" | python3 -c 'import json,sys; print(json.load(sys.stdin)["target_directory"])')/debug/awa -uv venv --quiet --clear "$dir/venv" -uv pip install --quiet --python "$dir/venv" "awa-pg==${AWA_RELEASED_VERSION:-0.6.9}" -py06=$dir/venv/bin/python -suffix=$(date +%s%N | tail -c 8) -owner=rh_owner_$suffix runtime=rh_runtime_$suffix db=rh_db_$suffix -host=${admin#*@}; host=${host%%/*} -owner_url=postgres://$owner:rh@$host/$db -runtime_url=postgres://$runtime:rh@$host/$db -log() { echo "$(date +%T) [$mode] $*" >&2; } - -psql -q "$admin" -c "CREATE ROLE $owner LOGIN PASSWORD 'rh'" -c "CREATE ROLE $runtime LOGIN PASSWORD 'rh'" \ - -c "CREATE DATABASE $db OWNER $owner" -trap 'psql -q "$admin" -c "DROP DATABASE IF EXISTS $db WITH (FORCE)" -c "DROP ROLE IF EXISTS $runtime" -c "DROP ROLE IF EXISTS $owner" >/dev/null 2>&1 || true' EXIT - -log "migrating with released 0.6: $("$py06" -c 'import importlib.metadata as m; print(m.version("awa-pg"))')" -"$py06" -m awa --database-url "$owner_url" migrate >/dev/null -psql -q "$owner_url" <"$dir/summary.json" 2>"$dir/workload.err" & -workload=$! -t0=$(date +%s.%N) -mark() { log "t=$(echo "$(date +%s.%N) - $t0" | bc | cut -c1-5) $*"; } -sleep 10 - -if [ "$mode" = prepared ]; then - mark "prepare-create-revoke start" - "$awa_new" --database-url "$owner_url" storage prepare-create-revoke >"$dir/prepare.json" - jq -c '.[] | {schema, prepared, provisioned_lanes, revoke, grants: (.sequence_grants | length), blocking}' "$dir/prepare.json" >&2 - jq -r '.[].sequence_grants[] | select(.run_as == "'"$runtime"'") | .statement' "$dir/prepare.json" | psql -q "$runtime_url" - jq -r '.[].revoke[]' "$dir/prepare.json" | psql -q "$owner_url" - mark "grants and revokes done" -else - log "revoking CREATE without preparing" - psql -q "$owner_url" -c "REVOKE CREATE ON SCHEMA awa FROM $runtime" -fi -touch "$dir/all-priorities" -sleep 10 - -if [ "$mode" = prepared ]; then - mark "awa migrate start" - ( while [ ! -e "$dir/migrated" ]; do psql -At "$owner_url" -c "SELECT string_agg(DISTINCT format('%s waits %s on %s', w.pid, w.mode, COALESCE(w.relation::regclass::text, w.locktype)), '; ') FROM pg_locks w WHERE NOT w.granted" 2>/dev/null; sleep 0.1; done | grep -v '^$' | sort | uniq -c | sort -rn | head -8 >"$dir/waits.txt" ) & - "$awa_new" --database-url "$owner_url" migrate >"$dir/migrate.log" 2>&1 || { log "migrate failed"; tail -5 "$dir/migrate.log" >&2; } - touch "$dir/migrated"; sleep 0.3; cat "$dir/waits.txt" >&2 - mark "awa migrate done, schema v$(psql -At "$owner_url" -c 'SELECT max(version) FROM awa.schema_version')" - sleep 10 -fi -touch "$dir/stop" -wait "$workload" || true -tail -3 "$dir/workload.err" >&2 || true -echo "$mode $(cat "$dir/summary.json")" diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index d0017a9f7..4cfb039b0 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -96,9 +96,6 @@ queue with `awa storage prepare-queue` as the migrator before starting producers or workers. Match shard/stripe counts; provision additional lanes before changing routing. Use the sequence grants and owner defaults in [Database roles](../../docs/security/database-roles.md#provision-queues-before-starting-restricted-runtimes). -To revoke `CREATE` from a running 0.6 fleet without stopping it, run -`awa storage prepare-create-revoke` as the migrator first, then the grants and -revokes it prints, then `awa migrate`. ## Rolling Upgrades From e2de04f868fb33954aa1ae60e150c70a29408b31 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 13 Nov 2026 00:56:34 +0000 Subject: [PATCH 08/10] test(prune): retry a Blocked prune in the re-carry and reader-lock tests The carry-forward test already retries prune while a background session holds the slot. The re-carry test's two prunes and the reader-lock test's prune after the lock is released had the same exposure; route all three through one helper that retries only Blocked for the expected slot. --- awa/tests/queue_storage_runtime_test.rs | 45 +++++++++++++++++-------- 1 file changed, 31 insertions(+), 14 deletions(-) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 66fb9c5b6..7a9e19b5e 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -455,6 +455,34 @@ async fn create_store_with_config(pool: &sqlx::PgPool, config: QueueStorageConfi store } +/// Prune `slot`, retrying while it reports `Blocked` for that slot. Prune takes +/// its locks NOWAIT, so a background session on the slot (autovacuum, or a +/// hosted Postgres's own workers) blocks it; the maintenance loop retries the +/// same way. Any other outcome, or `Blocked` past the deadline, is returned for +/// the caller to assert on. +async fn prune_retrying_blocked( + store: &QueueStorage, + pool: &sqlx::PgPool, + retention: Duration, + slot: i32, +) -> PruneOutcome { + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(5)); + loop { + let outcome = store + .prune_oldest(pool, retention) + .await + .expect("prune_oldest"); + match outcome { + PruneOutcome::Blocked { slot: blocked } + if blocked == slot && Instant::now() < deadline => + { + tokio::time::sleep(Duration::from_millis(50)).await; + } + other => return other, + } + } +} + async fn create_store(pool: &sqlx::PgPool, schema: &str) -> QueueStorage { // Tests that go through this helper exercise the legacy // (non-receipts) lease-materialization path. The receipts mode @@ -15711,11 +15739,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest") - { + match prune_retrying_blocked(&store, &pool, Duration::from_secs(3600), 0).await { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 1, @@ -15759,11 +15783,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { // Second prune: current slot is 0, so the oldest sealed non-current // slot is slot 1 (the carried row). It is carried again into slot 0. - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest re-carry") - { + match prune_retrying_blocked(&store, &pool, Duration::from_secs(3600), 1).await { PruneOutcome::Pruned { slot: 1, carried_failed_rows: 1, @@ -18048,10 +18068,7 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { .await .expect("Failed to release reader lock"); - let pruned = store - .prune_oldest(&pool, Duration::ZERO) - .await - .expect("Failed to prune after reader lock release"); + let pruned = prune_retrying_blocked(&store, &pool, Duration::ZERO, 0).await; assert!( matches!(pruned, PruneOutcome::Pruned { slot: 0, .. }), "unexpected prune outcome after reader lock release: {pruned:?}" From 21fb119a2f2f191266bfe3a6cac3270792d15b72 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 13 Nov 2026 00:56:35 +0000 Subject: [PATCH 09/10] docs(upgrade): revoke a running 0.6 fleet's schema CREATE after the 0.6 lane patches The guide asked for a pause because unpatched 0.6 lane helpers check schema CREATE on every reservation. awa migrate from 0.6.14 or later applies prepared_lane_sequences and definer_lane_provisioning to a v40 database, after which no runtime role needs CREATE, so the revoke can happen under traffic. Point at the 0.6 rehearsal that covers it. --- CHANGELOG.md | 5 +++-- docs/upgrade-0.6-to-0.7.md | 33 ++++++--------------------------- 2 files changed, 9 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dd033804d..d2fb5716c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -430,8 +430,9 @@ for each step. runtime role keeps working; only the migration is refused. Before upgrading, remove `TRIGGER` from runtime-role grants, which `GRANT ALL ON ALL TABLES` includes, remove runtime roles' `CREATE` on the schema, and remove triggers - of your own from awa tables. Revoking a running 0.6 fleet's `CREATE` needs - the order in the upgrade guide. + of your own from awa tables. A running 0.6 fleet keeps running if + `awa migrate` from 0.6.14 or later has applied its lane schema patches + before the `CREATE` revoke; see the upgrade guide. - **External migration runners.** Generate the pending SQL with an upgraded binary, for example `awa migrate --pending --extract-to ./pending`, apply complete bundles, and keep the files already applied unchanged. Apply diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index f7a5efa9e..ae106ea0b 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -104,39 +104,18 @@ schema owner: REVOKE TRIGGER, REFERENCES ON TABLES FROM awa_runtime; ``` -- **A 0.6 split-role install whose runtime holds `CREATE` on the schema.** 0.6 - runtimes created lane sequences themselves, so they needed `CREATE`. - Sequences they own are accepted; the `CREATE` grant is not, because it lets - the runtime add objects while a migration runs. Until the migration refreshes - them, the 0.6 lane helpers and enqueue-head triggers run - `CREATE SEQUENCE IF NOT EXISTS` on every reservation, and that checks - `CREATE` even for a sequence that exists, so revoking the grant makes running - 0.6 producers and workers fail with `42501` until the migration commits. - Provisioning lanes in advance does not avoid this. Take a short pause: - - 1. Stop 0.6 producers and workers, or hold them at a maintenance gate, so no - enqueue or claim runs. - 2. Revoke the grant and commit: +- **A 0.6 split-role install whose runtime holds `CREATE` on the schema.** Before the 0.6.14 schema patch, 0.6 runtimes created lane sequences themselves, so they needed `CREATE`. Sequences they own are accepted; the `CREATE` grant is not, because it lets the runtime add objects while a migration runs. Remove it while 0.6 producers and workers keep running: + + 1. Run `awa migrate` with a 0.6.14 or later CLI, with the schema owner as the effective role. The schema stays at v40. Its `prepared_lane_sequences` patch makes the lane helpers and head-table triggers skip `CREATE SEQUENCE` when the lane's sequence exists, and its `definer_lane_provisioning` patch creates a missing lane's sequences through the schema owner's `SECURITY DEFINER` functions, so no runtime role needs `CREATE`. The patches only replace functions, in `awa` and every other queue-storage schema, and earlier 0.6 binaries call the patched helpers with the same arguments. External migration runners apply the `R__` patches as the [0.6.13 to 0.6.14 upgrade guide](https://github.com/hardbyte/awa/blob/v0.6.14/docs/upgrade-0.6.13-to-0.6.14.md#external-migration-runners) describes. + 2. Revoke the grant in each queue-storage schema: ```sql REVOKE CREATE ON SCHEMA awa FROM awa_runtime; ``` - 3. Run `awa migrate` as the migrator. The bundle refreshes the lane helpers - to use provisioned sequences without `CREATE` (v049). - 4. Provision any lanes the 0.7 runtimes will need, as the migrator: - - ```bash - PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue --enqueue-shards - ``` - - 5. Start the 0.7 runtimes. + 3. Continue with the 0.7 migration. - An uninterrupted path would need the revoke and the helper refresh to become - visible to every session at once; that has not been verified, so do not - combine them into one transaction as a substitute for the pause. + Revoking before the patches makes running 0.6 producers and workers fail with `42501`: the unpatched helpers run `CREATE SEQUENCE IF NOT EXISTS` on every reservation, which checks `CREATE` even for a sequence that exists. The 0.6 release branch's `awa/tests/patch_upgrade_rehearsal_test.rs` applies the patches under a released 0.6.9 or 0.6.15 fleet, provisions its queues, grants sequence privileges and revokes `CREATE` from its runtime role, then checks that restarted released workers complete every job and a released producer creates a new queue's lanes through the schema owner's functions. - **A custom schema prepared by a runtime or another role.** A worker that auto-prepares a schema, or a DBA who ran the preparation as another user, From bb9865ebb1b865d8d8fbf8a730ef7a91373b1e79 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 13 Nov 2026 01:11:13 +0000 Subject: [PATCH 10/10] docs(upgrade): name a 0.6.x CLI for the lane patches before the CREATE revoke --- CHANGELOG.md | 2 +- docs/upgrade-0.6-to-0.7.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d2fb5716c..cca63fe19 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -431,7 +431,7 @@ for each step. remove `TRIGGER` from runtime-role grants, which `GRANT ALL ON ALL TABLES` includes, remove runtime roles' `CREATE` on the schema, and remove triggers of your own from awa tables. A running 0.6 fleet keeps running if - `awa migrate` from 0.6.14 or later has applied its lane schema patches + `awa migrate` from a 0.6.x CLI, 0.6.14 or later, has applied its lane schema patches before the `CREATE` revoke; see the upgrade guide. - **External migration runners.** Generate the pending SQL with an upgraded binary, for example `awa migrate --pending --extract-to ./pending`, apply diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index ae106ea0b..ade299b59 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -106,7 +106,7 @@ schema owner: - **A 0.6 split-role install whose runtime holds `CREATE` on the schema.** Before the 0.6.14 schema patch, 0.6 runtimes created lane sequences themselves, so they needed `CREATE`. Sequences they own are accepted; the `CREATE` grant is not, because it lets the runtime add objects while a migration runs. Remove it while 0.6 producers and workers keep running: - 1. Run `awa migrate` with a 0.6.14 or later CLI, with the schema owner as the effective role. The schema stays at v40. Its `prepared_lane_sequences` patch makes the lane helpers and head-table triggers skip `CREATE SEQUENCE` when the lane's sequence exists, and its `definer_lane_provisioning` patch creates a missing lane's sequences through the schema owner's `SECURITY DEFINER` functions, so no runtime role needs `CREATE`. The patches only replace functions, in `awa` and every other queue-storage schema, and earlier 0.6 binaries call the patched helpers with the same arguments. External migration runners apply the `R__` patches as the [0.6.13 to 0.6.14 upgrade guide](https://github.com/hardbyte/awa/blob/v0.6.14/docs/upgrade-0.6.13-to-0.6.14.md#external-migration-runners) describes. + 1. Run `awa migrate` with a 0.6.x CLI, 0.6.14 or later, with the schema owner as the effective role. The schema stays at v40. A 0.7 CLI cannot do this step: it applies the versioned 0.7 migrations, whose schema check refuses the runtime role's `CREATE` grant. Its `prepared_lane_sequences` patch makes the lane helpers and head-table triggers skip `CREATE SEQUENCE` when the lane's sequence exists, and its `definer_lane_provisioning` patch creates a missing lane's sequences through the schema owner's `SECURITY DEFINER` functions, so no runtime role needs `CREATE`. The patches only replace functions, in `awa` and every other queue-storage schema, and earlier 0.6 binaries call the patched helpers with the same arguments. External migration runners apply the `R__` patches as the [0.6.13 to 0.6.14 upgrade guide](https://github.com/hardbyte/awa/blob/v0.6.14/docs/upgrade-0.6.13-to-0.6.14.md#external-migration-runners) describes. 2. Revoke the grant in each queue-storage schema: ```sql