shahar1 commented on code in PR #41804:
URL: https://github.com/apache/airflow/pull/41804#discussion_r1738132491
##########
airflow/cli/commands/db_command.py:
##########
@@ -72,15 +72,17 @@ def upgradedb(args):
migratedb(args)
-def get_version_revision(version: str, recursion_limit=10) -> str | None:
+def get_version_revision(
+ version: str, recursion_limit=10, revision_heads_map=_REVISION_HEADS_MAP
+) -> str | None:
Review Comment:
1. Add type hints in function signature
2. I'd also suggest making the function private, or adding docstring if it
is expected to be utilized by other modules.
##########
airflow/providers/fab/auth_manager/cli_commands/db_command.py:
##########
@@ -0,0 +1,138 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from packaging.version import InvalidVersion, parse as parse_version
+
+from airflow import settings
+from airflow.cli.commands.db_command import get_version_revision
+from airflow.providers.fab.auth_manager.models.db import _REVISION_HEADS_MAP,
FABDBManager
+from airflow.utils import cli as cli_utils
+from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+
+
+@providers_configuration_loaded
+def resetdb(args):
+ """Reset the metadata database."""
+ 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")
+ FABDBManager(settings.Session()).resetdb(skip_init=args.skip_init)
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def migratedb(args):
+ """Migrates the metadata database."""
+ print(f"DB: {settings.engine.url!r}")
+ session = settings.Session()
+ if args.to_revision and args.to_version:
+ raise SystemExit("Cannot supply both `--to-revision` and
`--to-version`.")
+ if args.from_version and args.from_revision:
+ raise SystemExit("Cannot supply both `--from-revision` and
`--from-version`")
+ if (args.from_revision or args.from_version) and not args.show_sql_only:
+ raise SystemExit(
+ "Args `--from-revision` and `--from-version` may only be used with
`--show-sql-only`"
+ )
+ to_revision = None
+ from_revision = None
+ if args.from_revision:
+ from_revision = args.from_revision
+ elif args.from_version:
+ try:
+ parse_version(args.from_version)
+ except InvalidVersion:
+ raise SystemExit(f"Invalid version {args.from_version!r} supplied
as `--from-version`.")
+ from_revision = get_version_revision(args.from_version,
revision_heads_map=_REVISION_HEADS_MAP)
+ if not from_revision:
+ raise SystemExit(f"Unknown version {args.from_version!r} supplied
as `--from-version`.")
+
+ if args.to_version:
+ try:
+ parse_version(args.to_version)
+ except InvalidVersion:
+ raise SystemExit(f"Invalid version {args.to_version!r} supplied as
`--to-version`.")
+ to_revision = get_version_revision(args.to_version,
revision_heads_map=_REVISION_HEADS_MAP)
+ if not to_revision:
+ raise SystemExit(f"Unknown version {args.to_version!r} supplied as
`--to-version`.")
+ elif args.to_revision:
+ to_revision = args.to_revision
+
+ if not args.show_sql_only:
+ print(f"Performing upgradedb to the FAB metadata database
{settings.engine.url!r}")
+ else:
+ print("Generating sql for upgradedb -- FAB upgradedb commands will
*not* be submitted.")
+
+ FABDBManager(session).upgradedb(
+ to_revision=to_revision,
+ from_revision=from_revision,
+ show_sql_only=args.show_sql_only,
+ )
+ if not args.show_sql_only:
+ print("Database migrating done!")
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def downgrade(args):
+ """Downgrades the metadata database."""
+ session = settings.Session()
+ if args.to_revision and args.to_version:
+ raise SystemExit("Cannot supply both `--to-revision` and
`--to-version`.")
+ if args.from_version and args.from_revision:
+ raise SystemExit("`--from-revision` may not be combined with
`--from-version`")
+ if (args.from_revision or args.from_version) and not args.show_sql_only:
+ raise SystemExit(
+ "Args `--from-revision` and `--from-version` may only be used with
`--show-sql-only`"
+ )
+ if not (args.to_version or args.to_revision):
+ raise SystemExit("Must provide either --to-revision or --to-version.")
Review Comment:
There's some repeating logic from `migratedb` - I think that extracting it
to a seperate function would be a good idea.
##########
airflow/providers/fab/auth_manager/fab_auth_manager.py:
##########
@@ -22,11 +22,13 @@
from pathlib import Path
Review Comment:
General concern - this PR includes both changes to the FAB provider and to
Airflow core. While it seems that the changes to the FAB provider address
future compatbility, I think that it would be safer to make them in a separate
PR to mark the `3.0.0` milestone properly and to avoid confusion.
##########
airflow/cli/commands/db_command.py:
##########
@@ -72,15 +72,17 @@ def upgradedb(args):
migratedb(args)
-def get_version_revision(version: str, recursion_limit=10) -> str | None:
+def get_version_revision(
+ version: str, recursion_limit=10, revision_heads_map=_REVISION_HEADS_MAP
+) -> str | None:
"""
Recursively search for the revision of the given version.
This searches REVISION_HEADS_MAP for the revision of the given version,
recursively
searching for the previous version if the given version is not found.
Review Comment:
Now that `revision_heads_map` is parametrized, it's worth clarifying it in
the description.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
