fix checking for migrations in granian (#5348)

Author: adhami3310

reflex/reflex.py

@@ -201,6 +201,9 @@ def _run(
     # Get the app module.
     app_task = prerequisites.compile_or_validate_app
     args = (frontend,)
+    kwargs = {
+        "check_if_schema_up_to_date": True,
+    }
 
     # Granian fails if the app is already imported.
     if should_use_granian():
@@ -209,16 +212,15 @@ def _run(
         compile_future = concurrent.futures.ProcessPoolExecutor(max_workers=1).submit(
             app_task,
             *args,
+            **kwargs,
         )
         validation_result = compile_future.result()
     else:
-        validation_result = app_task(*args)
+        validation_result = app_task(*args, **kwargs)
+
     if not validation_result:
         raise click.exceptions.Exit(1)
 
-    # Warn if schema is not up to date.
-    prerequisites.check_schema_up_to_date()
-
     # Get the frontend and backend commands, based on the environment.
     setup_frontend = frontend_cmd = backend_cmd = None
     if env == constants.Env.DEV:
@@ -529,9 +531,7 @@ def migrate():
     from reflex import model
     from reflex.utils import prerequisites
 
-    # TODO see if we can use `get_app()` instead (no compile).  Would _skip_compile still be needed then?
-    _skip_compile()
-    prerequisites.get_compiled_app()
+    prerequisites.get_app()
     if not prerequisites.check_db_initialized():
         return
     model.Model.migrate()

reflex/utils/prerequisites.py

@@ -393,11 +393,14 @@ def get_app(reload: bool = False) -> ModuleType:
         return app
 
 
-def get_and_validate_app(reload: bool = False) -> AppInfo:
+def get_and_validate_app(
+    reload: bool = False, check_if_schema_up_to_date: bool = False
+) -> AppInfo:
     """Get the app instance based on the default config and validate it.
 
     Args:
         reload: Re-import the app module from disk
+        check_if_schema_up_to_date: If True, check if the schema is up to date.
 
     Returns:
         The app instance and the app module.
@@ -412,32 +415,47 @@ def get_and_validate_app(reload: bool = False) -> AppInfo:
     if not isinstance(app, App):
         msg = "The app instance in the specified app_module_import in rxconfig must be an instance of rx.App."
         raise RuntimeError(msg)
+
+    if check_if_schema_up_to_date:
+        check_schema_up_to_date()
+
     return AppInfo(app=app, module=app_module)
 
 
-def validate_app(reload: bool = False) -> None:
+def validate_app(
+    reload: bool = False, check_if_schema_up_to_date: bool = False
+) -> None:
     """Validate the app instance based on the default config.
 
     Args:
         reload: Re-import the app module from disk
+        check_if_schema_up_to_date: If True, check if the schema is up to date.
     """
-    get_and_validate_app(reload=reload)
+    get_and_validate_app(
+        reload=reload, check_if_schema_up_to_date=check_if_schema_up_to_date
+    )
 
 
 def get_compiled_app(
-    reload: bool = False, export: bool = False, dry_run: bool = False
+    reload: bool = False,
+    export: bool = False,
+    dry_run: bool = False,
+    check_if_schema_up_to_date: bool = False,
 ) -> ModuleType:
     """Get the app module based on the default config after first compiling it.
 
     Args:
         reload: Re-import the app module from disk
         export: Compile the app for export
         dry_run: If True, do not write the compiled app to disk.
+        check_if_schema_up_to_date: If True, check if the schema is up to date.
 
     Returns:
         The compiled app based on the default config.
     """
-    app, app_module = get_and_validate_app(reload=reload)
+    app, app_module = get_and_validate_app(
+        reload=reload, check_if_schema_up_to_date=check_if_schema_up_to_date
+    )
     # For py3.9 compatibility when redis is used, we MUST add any decorator pages
     # before compiling the app in a thread to avoid event loop error (REF-2172).
     app._apply_decorated_pages()
@@ -446,16 +464,25 @@ def get_compiled_app(
 
 
 def compile_app(
-    reload: bool = False, export: bool = False, dry_run: bool = False
+    reload: bool = False,
+    export: bool = False,
+    dry_run: bool = False,
+    check_if_schema_up_to_date: bool = False,
 ) -> None:
     """Compile the app module based on the default config.
 
     Args:
         reload: Re-import the app module from disk
         export: Compile the app for export
         dry_run: If True, do not write the compiled app to disk.
+        check_if_schema_up_to_date: If True, check if the schema is up to date.
     """
-    get_compiled_app(reload=reload, export=export, dry_run=dry_run)
+    get_compiled_app(
+        reload=reload,
+        export=export,
+        dry_run=dry_run,
+        check_if_schema_up_to_date=check_if_schema_up_to_date,
+    )
 
 
 def _can_colorize() -> bool:
@@ -500,20 +527,23 @@ def _can_colorize() -> bool:
         return file.isatty()
 
 
-def compile_or_validate_app(compile: bool = False) -> bool:
+def compile_or_validate_app(
+    compile: bool = False, check_if_schema_up_to_date: bool = False
+) -> bool:
     """Compile or validate the app module based on the default config.
 
     Args:
         compile: Whether to compile the app.
+        check_if_schema_up_to_date: If True, check if the schema is up to date.
 
     Returns:
         If the app is compiled successfully.
     """
     try:
         if compile:
-            compile_app()
+            compile_app(check_if_schema_up_to_date=check_if_schema_up_to_date)
         else:
-            validate_app()
+            validate_app(check_if_schema_up_to_date=check_if_schema_up_to_date)
     except Exception as e:
         if isinstance(e, click.exceptions.Exit):
             return False