Skip to content
Clone
eeff292
121 lines 4.4 KB Rust
use std::hash::{DefaultHasher, Hash, Hasher};
use sea_orm::{
ConnectionTrait, DatabaseBackend, DatabaseConnection, DbErr, Statement, TransactionTrait,
};
const SCHEMA_PRE_SYNC: &str = include_str!("../../schema_pre_sync.sql");
const SCHEMA_EXTRAS: &str = include_str!("../../schema_extras.sql");
/// Bump whenever `sync_schema` gains a Rust-side step that is not captured by
/// the entity registry or by `schema_pre_sync.sql` / `schema_extras.sql`.
const SYNC_FORMAT_VERSION: u32 = 3;
pub const MIGRATION_LOCK_KEY: i64 = 5_241_915_330_551;
const SYNC_STATE_DDL: &str = "CREATE TABLE IF NOT EXISTS schema_sync_state (
id smallint PRIMARY KEY CHECK (id = 1),
fingerprint text NOT NULL,
live_fingerprint text,
synced_at timestamptz NOT NULL DEFAULT now()
)";
const SYNC_STATE_LIVE_COLUMN: &str =
"ALTER TABLE schema_sync_state ADD COLUMN IF NOT EXISTS live_fingerprint text";
const LIVE_FINGERPRINT_SQL: &str = r#"
SELECT COALESCE(md5(string_agg(entry, E'\n' ORDER BY entry)), '') AS fingerprint
FROM (
SELECT format('table:%s', table_name) AS entry
FROM information_schema.tables
WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
UNION ALL
SELECT format('column:%s.%s:%s:%s:%s',
table_name, column_name, data_type, is_nullable,
COALESCE(column_default, ''))
FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name IN (
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
)
UNION ALL
SELECT format('index:%s:%s', tablename, indexdef)
FROM pg_indexes
WHERE schemaname = 'public'
UNION ALL
SELECT format('constraint:%s:%s', conrelid::regclass, pg_get_constraintdef(oid))
FROM pg_constraint
WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
UNION ALL
SELECT format('enum:%s:%s', t.typname, e.enumlabel)
FROM pg_type t
JOIN pg_enum e ON e.enumtypid = t.oid
JOIN pg_namespace n ON n.oid = t.typnamespace
WHERE n.nspname = 'public'
) AS entries
"#;
pub async fn sync_schema(db: &DatabaseConnection) -> Result<(), DbErr> {
let builder = db.get_schema_registry("models::");
if db.get_database_backend() != DatabaseBackend::Postgres {
return builder.sync(db).await;
}
let mut hasher = DefaultHasher::new();
SYNC_FORMAT_VERSION.hash(&mut hasher);
SCHEMA_PRE_SYNC.hash(&mut hasher);
SCHEMA_EXTRAS.hash(&mut hasher);
format!("{builder:?}").hash(&mut hasher);
let fingerprint = format!("{:016x}", hasher.finish());
let txn = db.begin().await?;
txn.execute_unprepared(&format!(
"SELECT pg_advisory_xact_lock({MIGRATION_LOCK_KEY})"
))
.await?;
txn.execute_unprepared(SYNC_STATE_DDL).await?;
txn.execute_unprepared(SYNC_STATE_LIVE_COLUMN).await?;
let live = live_fingerprint(&txn).await?;
let stored = txn
.query_one_raw(Statement::from_string(
DatabaseBackend::Postgres,
"SELECT fingerprint, live_fingerprint FROM schema_sync_state WHERE id = 1",
))
.await?;
if let Some(row) = stored {
let stored_fingerprint = row.try_get::<String>("", "fingerprint")?;
let stored_live = row.try_get::<Option<String>>("", "live_fingerprint")?;
if stored_fingerprint == fingerprint && stored_live.as_deref() == Some(live.as_str()) {
txn.commit().await?;
return Ok(());
}
}
txn.execute_unprepared(SCHEMA_PRE_SYNC).await?;
builder.sync(&txn).await?;
txn.execute_unprepared(SCHEMA_EXTRAS).await?;
let live = live_fingerprint(&txn).await?;
let update = format!(
"INSERT INTO schema_sync_state (id, fingerprint, live_fingerprint) \
VALUES (1, '{fingerprint}', '{live}') \
ON CONFLICT (id) DO UPDATE SET fingerprint = EXCLUDED.fingerprint, \
live_fingerprint = EXCLUDED.live_fingerprint, synced_at = now()"
);
txn.execute_unprepared(&update).await?;
txn.commit().await
}
async fn live_fingerprint<C>(db: &C) -> Result<String, DbErr>
where
C: ConnectionTrait,
{
let row = db
.query_one_raw(Statement::from_string(
DatabaseBackend::Postgres,
LIVE_FINGERPRINT_SQL,
))
.await?;
match row {
Some(row) => row.try_get::<String>("", "fingerprint"),
None => Ok(String::new()),
}
}