diff --git a/docker/data-restore/data-restore-entry.sh b/docker/data-restore/data-restore-entry.sh
index f7f108d364..3d715b407a 100755
--- a/docker/data-restore/data-restore-entry.sh
+++ b/docker/data-restore/data-restore-entry.sh
@@ -2,16 +2,21 @@
 
 set -e
 
-function migrate() {
+function reset_db() {
+      cd core/lib/storage
       psql "$DATABASE_URL" -c 'DROP OWNED BY CURRENT_USER CASCADE' || /bin/true
       psql "$DATABASE_URL" -c 'DROP SCHEMA IF EXISTS public CASCADE' || /bin/true
       psql "$DATABASE_URL" -c 'CREATE SCHEMA public'
+      diesel database setup
+      cd $ZKSYNC_HOME
+}
 
+function migrate() {
       cd core/lib/storage
-      diesel database setup
       diesel migration run
       cd $ZKSYNC_HOME
 }
+
 cd $ZKSYNC_HOME
 
 # Wait for the database to be ready.
@@ -28,7 +33,7 @@ fi
 case $COMMAND in
   genesis)
       echo "Resetting the database"
-      migrate
+      reset_db
       COMMAND="--genesis"
     ;;
   continue)
@@ -68,7 +73,7 @@ then
   # Do not drop db if the file doesn't exist.
   [ -f $PG_DUMP_PATH/$PG_DUMP ] || { echo "$PG_DUMP_PATH/$PG_DUMP  not found" ; exit 1 ; }
 
-  migrate
+  reset_db
 
   echo "Applying $PG_DUMP"
   pg_restore -j 8 -d $DATABASE_URL --clean --if-exists $PG_DUMP_PATH/$PG_DUMP
@@ -81,4 +86,5 @@ else
   CONFIG_FILE="${CONFIG_PATH}/${NETWORK}.json"
 fi
 
+migrate
 $ZKSYNC_HOME/target/release/zksync_data_restore $COMMAND $MODE --config $CONFIG_FILE --web3 $WEB3_URL || exit 1