Add airflow db autogenerate command for generating migration files

This PR customizes the alembic autogenerate command. The DB is reset
using the migration files helping to determine new additions in the ORM
accurately

airflow/cli/cli_config.py

@@ -519,6 +519,13 @@ def string_lower_type(val):
     type=positive_int(allow_zero=False),
     help="Wait time between retries in seconds",
 )
+ARG_DB_MESSAGE = Arg(
+    (
+        "-m",
+        "--message",
+    ),
+    help="Message to apply to the migration",
+)
 
 # pool
 ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
@@ -1644,6 +1651,12 @@ class GroupCommand(NamedTuple):
         func=lazy_load_command("airflow.cli.commands.db_command.drop_archived"),
         args=(ARG_DB_TABLES, ARG_YES),
     ),
+    ActionCommand(
+        name="autogenerate",
+        help="Reset the DB and generate a migration script for the database",
+        func=lazy_load_command("airflow.cli.commands.db_command.autogenerate"),
+        args=(ARG_DB_MESSAGE, ARG_YES),
+    ),
 )
 CONNECTIONS_COMMANDS = (
     ActionCommand(

airflow/cli/commands/db_command.py

@@ -301,3 +301,13 @@ def drop_archived(args):
         table_names=args.tables,
         needs_confirm=not args.yes,
     )
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def autogenerate(args):
+    """Generate migration script."""
+    print(f"DB: {settings.engine.url!r}")
+    if not (args.yes or input("This will drop existing tables if they exist. Proceed? (y/n)").upper() == "Y"):
+        raise SystemExit("Cancelled")
+    db.autogenerate(message=args.message)

airflow/utils/db.py

@@ -1784,6 +1784,20 @@ def drop_airflow_moved_tables(connection):
         Base.metadata.remove(tbl)
 
 
+def autogenerate(message: str | None = None) -> None:
+    """Automatically generate the migration scripts.
+
+    The database is first reset and built using the migration files so that autogenerate
+    would accurately reflect the current state of the database and ORM
+    """
+    from alembic import command
+
+    resetdb(use_migration_files=True)
+
+    config = _get_alembic_config()
+    command.revision(config, message=message, autogenerate=True)
+
+
 @provide_session
 def check(session: Session = NEW_SESSION):
     """

tests/utils/test_db.py

@@ -38,6 +38,7 @@
 from airflow.settings import engine
 from airflow.utils.db import (
     _get_alembic_config,
+    autogenerate,
     check_bad_references,
     check_migrations,
     compare_server_default,
@@ -235,6 +236,16 @@ def test_resetdb(
         else:
             mock_init.assert_called_once_with(session=session_mock, use_migration_files=False)
 
+    @mock.patch("airflow.utils.db._get_alembic_config")
+    @mock.patch("airflow.utils.db.resetdb")
+    @mock.patch("alembic.command")
+    def test_db_autogenerate(self, mock_alembic_command, mock_resetdb, mock_config):
+        autogenerate(message="test migration")
+        mock_resetdb.assert_called_once_with(use_migration_files=True)
+        mock_alembic_command.revision.assert_called_once_with(
+            mock_config.return_value, message="test migration", autogenerate=True
+        )
+
     def test_alembic_configuration(self):
         with mock.patch.dict(
             os.environ, {"AIRFLOW__DATABASE__ALEMBIC_INI_FILE_PATH": "/tmp/alembic.ini"}, clear=True