Files
crank/crates/crank-registry/src/migrations/schema_guard_v11.rs
T

332 lines
11 KiB
Rust

use sqlx::{PgConnection, Row, query};
use super::authority::MigrationError;
use super::schema_guard::{
column_exists, constraint_expression, normalize_definition, relation_exists, schema_error,
};
pub(super) async fn validate_v11_absent(
connection: &mut PgConnection,
) -> Result<bool, MigrationError> {
Ok(!relation_exists(connection, "product_events").await?
&& !relation_exists(connection, "product_event_daily_rollups").await?
&& !relation_exists(connection, "onboarding_selections").await?
&& !column_exists(connection, "invocation_logs", "platform_api_key_id").await?)
}
pub(super) async fn validate_v11_onboarding_product_events(
connection: &mut PgConnection,
) -> Result<(), MigrationError> {
for relation in [
"product_events",
"product_event_daily_rollups",
"onboarding_selections",
] {
if !relation_exists(connection, relation).await? {
return Err(schema_error(11));
}
}
if !column_exists(connection, "invocation_logs", "platform_api_key_id").await? {
return Err(schema_error(11));
}
for relation in [
"product_events_workspace_id_idempotency_key_key",
"product_events_workspace_occurred_idx",
"invocation_logs_workspace_agent_key_success_idx",
] {
if !relation_exists(connection, relation).await? {
return Err(schema_error(11));
}
}
for (table, constraint, required) in [
(
"product_events",
"product_events_id_check",
&["id~", "^pe_[a-za-z0-9_-]{1,128}$"][..],
),
(
"product_events",
"product_events_name_check",
&[
"event_name=any",
"onboarding_eligible",
"onboarding_started",
"onboarding_resumed",
"onboarding_dismissed",
"onboarding_abandoned",
"onboarding_completed",
][..],
),
(
"product_events",
"product_events_schema_version_check",
&["schema_version=1"][..],
),
(
"product_events",
"product_events_idempotency_key_check",
&[
"octet_lengthidempotency_key>=1",
"octet_lengthidempotency_key<=256",
][..],
),
(
"product_events",
"product_events_properties_check",
&[
"jsonb_typeofproperties_json='object'",
"pg_column_sizeproperties_json<=4096",
"event_name<>'onboarding_eligible'",
"properties_json@>'{\"eligible\":true}'",
"jsonb_typeofproperties_json->'eligible_since'",
"='string'",
][..],
),
(
"product_event_daily_rollups",
"product_event_daily_rollups_name_check",
&[
"event_name=any",
"onboarding_eligible",
"onboarding_completed",
][..],
),
(
"product_event_daily_rollups",
"product_event_daily_rollups_counts_check",
&[
"events_total>=0",
"eligible_total>=0",
"eligible_total<=events_total",
][..],
),
] {
let definition =
normalize_definition(&constraint_expression(connection, table, constraint).await?);
if required.iter().any(|snippet| !definition.contains(snippet)) {
return Err(schema_error(11));
}
}
validate_index(
connection,
"product_events_workspace_occurred_idx",
"product_events",
false,
&["workspace_id", "occurred_at", "id"],
&[],
)
.await?;
validate_index(
connection,
"invocation_logs_workspace_agent_key_success_idx",
"invocation_logs",
false,
&[
"workspace_id",
"agent_id",
"platform_api_key_id",
"created_atdesc",
],
&[
"platform_api_key_idisnotnull",
"source='agent_tool_call'",
"status='ok'",
],
)
.await?;
let key_scope_fk = constraint_definition(
connection,
"invocation_logs",
"invocation_logs_platform_key_scope_fk",
)
.await?;
for required in [
"foreignkeyworkspace_id,agent_id,platform_api_key_id",
"referencesplatform_api_keysworkspace_id,agent_id,id",
"ondeletesetnullplatform_api_key_id",
] {
if !key_scope_fk.contains(required) {
return Err(schema_error(11));
}
}
for (kind, name) in [
("constraint", "invocation_logs_platform_key_scope_fk"),
("trigger", "product_events_append_only_guard"),
] {
let present: bool = match kind {
"constraint" => query(
"select exists (select 1 from pg_constraint c
join pg_namespace n on n.oid = c.connamespace
where n.nspname = current_schema() and c.conname = $1) as present",
),
_ => query(
"select exists (select 1 from pg_trigger t
join pg_class c on c.oid = t.tgrelid
join pg_namespace n on n.oid = c.relnamespace
where n.nspname = current_schema() and t.tgname = $1 and not t.tgisinternal) as present",
),
}
.bind(name)
.fetch_one(&mut *connection)
.await
.map_err(|_| MigrationError::storage("preflight.schema"))?
.try_get("present")
.map_err(|_| MigrationError::storage("preflight.schema"))?;
if !present {
return Err(schema_error(11));
}
}
let nullable = query(
"select is_nullable from information_schema.columns
where table_schema = current_schema()
and table_name = 'invocation_logs'
and column_name = 'platform_api_key_id'",
)
.fetch_optional(&mut *connection)
.await
.map_err(|_| MigrationError::storage("preflight.schema"))?
.and_then(|row| row.try_get::<String, _>("is_nullable").ok());
if nullable.as_deref() != Some("YES") {
return Err(schema_error(11));
}
let shape = constraint_expression(
connection,
"onboarding_selections",
"onboarding_selections_shape_check",
)
.await?;
let normalized_shape = normalize_definition(&shape);
for required in [
"operation_idisnull",
"operation_version>0",
"catalog_revision>0",
"invocation_log_idisnotnull",
"selected_atisnull",
"selected_atisnotnull",
] {
if !normalized_shape.contains(required) {
return Err(schema_error(11));
}
}
let function_definition = query(
"select p.prosrc as definition
from pg_proc p
join pg_namespace n on n.oid = p.pronamespace
where n.nspname = current_schema()
and p.proname = 'crank_reject_product_event_mutation'",
)
.fetch_optional(connection)
.await
.map_err(|_| MigrationError::storage("preflight.schema"))?
.and_then(|row| row.try_get::<String, _>("definition").ok());
if !function_definition.is_some_and(|definition| {
let normalized = normalize_definition(&definition);
[
"notexists",
"fromworkspaces",
"old.workspace_id",
"returnold",
"producteventisappend-only",
]
.into_iter()
.all(|required| normalized.contains(required))
}) {
return Err(schema_error(11));
}
Ok(())
}
async fn constraint_definition(
connection: &mut PgConnection,
table: &str,
constraint: &str,
) -> Result<String, MigrationError> {
query(
"select pg_get_constraintdef(c.oid, true) as definition
from pg_catalog.pg_constraint c
join pg_catalog.pg_class t on t.oid = c.conrelid
join pg_catalog.pg_namespace n on n.oid = t.relnamespace
where n.nspname = current_schema() and t.relname = $1 and c.conname = $2",
)
.bind(table)
.bind(constraint)
.fetch_optional(&mut *connection)
.await
.map_err(|_| MigrationError::storage("preflight.schema"))?
.ok_or_else(|| schema_error(11))?
.try_get::<String, _>("definition")
.map(|definition| normalize_definition(&definition))
.map_err(|_| MigrationError::storage("preflight.schema"))
}
async fn validate_index(
connection: &mut PgConnection,
index: &str,
expected_table: &str,
expected_unique: bool,
expected_columns: &[&str],
predicate_snippets: &[&str],
) -> Result<(), MigrationError> {
let row = query(
"select t.relname as table_name, am.amname as access_method,
i.indisvalid, i.indisready, i.indisunique,
pg_get_indexdef(i.indexrelid, 1, true) as first_column,
pg_get_indexdef(i.indexrelid, 2, true) as second_column,
pg_get_indexdef(i.indexrelid, 3, true) as third_column,
pg_get_indexdef(i.indexrelid, 4, true) as fourth_column,
pg_get_indexdef(i.indexrelid) as definition,
coalesce(pg_get_expr(i.indpred, i.indrelid), '') as predicate
from pg_catalog.pg_index i
join pg_catalog.pg_class idx on idx.oid = i.indexrelid
join pg_catalog.pg_class t on t.oid = i.indrelid
join pg_catalog.pg_namespace n on n.oid = t.relnamespace
join pg_catalog.pg_am am on am.oid = idx.relam
where n.nspname = current_schema() and idx.relname = $1",
)
.bind(index)
.fetch_optional(&mut *connection)
.await
.map_err(|_| MigrationError::storage("preflight.schema"))?;
let valid = row.is_some_and(|row| {
let actual_columns = [
"first_column",
"second_column",
"third_column",
"fourth_column",
]
.into_iter()
.filter_map(|field| row.try_get::<Option<String>, _>(field).ok().flatten())
.map(|value| normalize_definition(&value))
.filter(|value| !value.is_empty())
.collect::<Vec<_>>();
let expected_base_columns = expected_columns
.iter()
.map(|column| column.strip_suffix("desc").unwrap_or(column).to_owned())
.collect::<Vec<_>>();
let definition = row
.try_get::<String, _>("definition")
.ok()
.map(|value| normalize_definition(&value))
.unwrap_or_default();
let predicate = row
.try_get::<String, _>("predicate")
.ok()
.map(|value| normalize_definition(&value))
.unwrap_or_default();
row.try_get::<String, _>("table_name").ok().as_deref() == Some(expected_table)
&& row.try_get::<String, _>("access_method").ok().as_deref() == Some("btree")
&& row.try_get::<bool, _>("indisvalid").ok() == Some(true)
&& row.try_get::<bool, _>("indisready").ok() == Some(true)
&& row.try_get::<bool, _>("indisunique").ok() == Some(expected_unique)
&& actual_columns == expected_base_columns
&& expected_columns
.iter()
.filter(|column| column.ends_with("desc"))
.all(|column| definition.contains(column))
&& predicate_snippets
.iter()
.all(|snippet| predicate.contains(snippet))
});
if valid { Ok(()) } else { Err(schema_error(11)) }
}