diff --git a/snap/hooks/post-refresh b/snap/hooks/post-refresh
index 7de41a2..7429816 100644
--- a/snap/hooks/post-refresh
+++ b/snap/hooks/post-refresh
@@ -77,11 +77,7 @@ if [ "$SNAP_VERSION" != "$CURRENT_SNAP_VERSION" ]; then
         link_mastodon_config
         upgrade_mastodon_config
 
-        echo Update mastodon database
-        mastodon_migrate_db
-
-        echo Remove logs and tempfiles
-        mastodon_rails log:clear tmp:clear
+        migration_required
     fi
 fi
 
diff --git a/snap/hooks/pre-refresh b/snap/hooks/pre-refresh
index 38a6e14..fb54d38 100644
--- a/snap/hooks/pre-refresh
+++ b/snap/hooks/pre-refresh
@@ -13,3 +13,5 @@ if [ "$PG_VERSION" != "$CURRENT_PG_VERSION" ]; then
     echo Create postgres backup
     pg_dumpall.wrapper -f "$SNAP_COMMON/database.sql"
 fi
+
+pre_migration_required
diff --git a/snap/snapcraft.yaml b/snap/snapcraft.yaml
index 77a64fb..0b36b71 100644
--- a/snap/snapcraft.yaml
+++ b/snap/snapcraft.yaml
@@ -49,6 +49,8 @@ apps:
 
   postgres:
     command: bin/run_as_daemon_user postgres.wrapper
+    stop-command: bin/migrate
+    stop-timeout: 30m
     daemon: simple
     plugs:
       - network
diff --git a/src/mastodon/bin/migrate b/src/mastodon/bin/migrate
new file mode 100755
index 0000000..3cf5184
--- /dev/null
+++ b/src/mastodon/bin/migrate
@@ -0,0 +1,18 @@
+#!/bin/bash -e
+
+. "$SNAP/mastodon.env"
+
+if [ "$(snapctl get migration-required)" == true ]; then
+    postgres_waitready
+    mastodon_rails db:migrate
+    mastodon_rails log:clear tmp:clear
+    reset_migration_required
+    reset_pre_migration_required
+elif [ "$(snapctl get pre-migration-required)" == true ]; then
+    postgres_waitready
+    SKIP_POST_DEPLOYMENT_MIGRATIONS=true \
+        mastodon_rails db:migrate
+    mastodon_rails log:clear tmp:clear
+    reset_pre_migration_required
+fi
+
diff --git a/src/mastodon/bin/restore b/src/mastodon/bin/restore
index 1df5f66..38e92d9 100755
--- a/src/mastodon/bin/restore
+++ b/src/mastodon/bin/restore
@@ -53,7 +53,7 @@ echo Restore database
 $SNAP/bin/psql.wrapper -f "$backup_dir/database.sql" template1 > /dev/null
 if [ "$BACKUP_VERSION" != "$SNAP_VERSION" ]; then
     echo Migrate database
-    mastodon_migrate_db > /dev/null
+    $SNAP/bin/migrate > /dev/null
 fi
 
 if [ -f "$backup_dir/redis.rdb" ]; then
diff --git a/src/mastodon/mastodon.env b/src/mastodon/mastodon.env
index e81c322..f30d3bf 100644
--- a/src/mastodon/mastodon.env
+++ b/src/mastodon/mastodon.env
@@ -116,15 +116,6 @@ setup_mastodon_database() {
     psql.wrapper -c "alter user $MASTODON_DBUSER createdb;"
 }
 
-mastodon_migrate_db() {
-    postgres_waitready
-    mastodon_rails db:migrate
-}
-
-mastodon_migrate_db_pre_deployment() {
-    SKIP_POST_DEPLOYMENT_MIGRATIONS=true mastodon_migrate_db
-}
-
 generate_mastodon_vapid_secrets() {
     eval $(mastodon_rails mastodon:webpush:generate_vapid_key)
     export VAPID_PRIVATE_KEY="$VAPID_PRIVATE_KEY"
@@ -233,6 +224,22 @@ recompile_if_required() {
     fi
 }
 
+migration_required() {
+    snapctl set migration-required=true
+}
+
+reset_migration_required() {
+    snapctl set migration-required=false
+}
+
+pre_migration_required() {
+    snapctl set pre-migration-required=true
+}
+
+reset_pre_migration_required() {
+    snapctl set pre-migration-required=false
+}
+
 get_status_max_chars_file() {
     run_as_daemon_user sed -rn "s/^.*MAX_CHARS = ([0-9]+).*$/\1/p" "$status_length_validator"
 }
diff --git a/src/snap/bin/save_file b/src/snap/bin/save_file
old mode 100644
new mode 100755