diff --git a/.github/workflows/basic-tests.yml b/.github/workflows/basic-tests.yml index f5b1d54084b86..b59baa6728fda 100644 --- a/.github/workflows/basic-tests.yml +++ b/.github/workflows/basic-tests.yml @@ -154,8 +154,11 @@ jobs: env: HATCH_ENV: "test" working-directory: ./clients/python + - name: "Prepare FAB provider packages: wheel" + run: > + breeze release-management prepare-provider-packages fab --package-format wheel --skip-tag-check - name: "Install Airflow with fab for webserver tests" - run: pip install ".[fab]" + run: pip install . dist/apache_airflow_providers_fab-*.whl - name: "Install Python client" run: pip install ./dist/apache_airflow_client-*.whl - name: "Initialize Airflow DB and start webserver" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b4b6ec9c3d04b..c89ee2b1a25e1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1295,7 +1295,8 @@ repos: language: python entry: ./scripts/ci/pre_commit/migration_reference.py pass_filenames: false - files: ^airflow/migrations/versions/.*\.py$|^docs/apache-airflow/migrations-ref\.rst$ + files: | + ^airflow/migrations/versions/.*\.py$|^docs/apache-airflow/migrations-ref\.rst$|^airflow/providers/fab/alembic/versions/.*\.py$ additional_dependencies: ['rich>=12.4.4'] - id: update-er-diagram name: Update ER diagram diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 621589d2a5877..70859aebcd594 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -720,6 +720,14 @@ database: type: string example: ~ default: "True" + external_db_managers: + description: | + List of DB managers to use to migrate external tables in airflow database. The managers must inherit + from BaseDBManager + version_added: 3.0.0 + type: string + example: ~ + default: "airflow.providers.fab.auth_manager.models.db.FABDBManager" logging: description: ~ options: diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index c335aa74da830..4be8ac32ca252 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -32,8 +32,8 @@ def include_object(_, name, type_, *args): # Ignore the sqlite_sequence table, which is an internal SQLite construct if name == "sqlite_sequence": return False - # Ignore _anything_ to do with Celery, or FlaskSession's tables - if type_ == "table" and (name.startswith("celery_") or name == "session"): + # Only create migrations for objects that are in the target metadata + if type_ == "table" and name not in target_metadata.tables: return False else: return True @@ -61,6 +61,9 @@ def include_object(_, name, type_, *args): # my_important_option = config.get_main_option("my_important_option") # ... etc. +# version table +version_table = "alembic_version" + def run_migrations_offline(): """ @@ -82,6 +85,8 @@ def run_migrations_offline(): compare_type=compare_type, compare_server_default=compare_server_default, render_as_batch=True, + include_object=include_object, + version_table=version_table, ) with context.begin_transaction(): @@ -119,6 +124,7 @@ def process_revision_directives(context, revision, directives): include_object=include_object, render_as_batch=True, process_revision_directives=process_revision_directives, + version_table=version_table, ) with context.begin_transaction(): diff --git a/airflow/providers/fab/alembic.ini b/airflow/providers/fab/alembic.ini new file mode 100644 index 0000000000000..55c7338f9bd9f --- /dev/null +++ b/airflow/providers/fab/alembic.ini @@ -0,0 +1,133 @@ +# 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. + +# A generic, single database configuration. + +[alembic] +# path to migration scripts +# Use forward slashes (/) also on windows to provide an os agnostic path +script_location = migrations + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python>=3.9 or backports.zoneinfo library. +# Any required deps can installed by adding `alembic[tz]` to the pip requirements +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +sqlalchemy.url = scheme://localhost/airflow + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, execute a binary +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = --fix REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/airflow/providers/fab/auth_manager/models/__init__.py b/airflow/providers/fab/auth_manager/models/__init__.py index bf4e43f275fab..2587d7034d04c 100644 --- a/airflow/providers/fab/auth_manager/models/__init__.py +++ b/airflow/providers/fab/auth_manager/models/__init__.py @@ -23,6 +23,7 @@ # Copyright 2013, Daniel Vaz Gaspar from typing import TYPE_CHECKING +import packaging.version from flask import current_app, g from flask_appbuilder.models.sqla import Model from sqlalchemy import ( @@ -32,6 +33,7 @@ ForeignKey, Index, Integer, + MetaData, String, Table, UniqueConstraint, @@ -39,16 +41,11 @@ func, select, ) -from sqlalchemy.orm import backref, declared_attr, relationship +from sqlalchemy.orm import backref, declared_attr, registry, relationship +from airflow import __version__ as airflow_version from airflow.auth.managers.models.base_user import BaseUser -from airflow.models.base import Base - -""" -Compatibility note: The models in this file are duplicated from Flask AppBuilder. -""" -# Use airflow metadata to create the tables -Model.metadata = Base.metadata +from airflow.models.base import _get_schema, naming_convention if TYPE_CHECKING: try: @@ -56,6 +53,22 @@ except Exception: Identity = None +""" +Compatibility note: The models in this file are duplicated from Flask AppBuilder. +""" + +metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) +mapper_registry = registry(metadata=metadata) + +if packaging.version.parse(packaging.version.parse(airflow_version).base_version) >= packaging.version.parse( + "3.0.0" +): + Model.metadata = metadata +else: + from airflow.models.base import Base + + Model.metadata = Base.metadata + class Action(Model): """Represents permission actions such as `can_read`.""" diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py new file mode 100644 index 0000000000000..a971ea29a3f68 --- /dev/null +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -0,0 +1,35 @@ +# 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 + +import os + +import airflow +from airflow.providers.fab.auth_manager.models import metadata +from airflow.utils.db_manager import BaseDBManager + +PACKAGE_DIR = os.path.dirname(airflow.__file__) + + +class FABDBManager(BaseDBManager): + """Manages FAB database.""" + + metadata = metadata + version_table_name = "fab_alembic_version" + migration_dir = os.path.join(PACKAGE_DIR, "providers/fab/migrations") + alembic_file = os.path.join(PACKAGE_DIR, "providers/fab/alembic.ini") + supports_table_dropping = True diff --git a/airflow/providers/fab/migrations/README b/airflow/providers/fab/migrations/README new file mode 100644 index 0000000000000..2500aa1bcf726 --- /dev/null +++ b/airflow/providers/fab/migrations/README @@ -0,0 +1 @@ +Generic single-database configuration. diff --git a/airflow/providers/fab/migrations/__init__.py b/airflow/providers/fab/migrations/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/fab/migrations/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow/providers/fab/migrations/env.py b/airflow/providers/fab/migrations/env.py new file mode 100644 index 0000000000000..7b31f86024a31 --- /dev/null +++ b/airflow/providers/fab/migrations/env.py @@ -0,0 +1,125 @@ +# 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 + +import contextlib +from logging.config import fileConfig + +from alembic import context + +from airflow import settings +from airflow.providers.fab.auth_manager.models.db import FABDBManager + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +version_table = FABDBManager.version_table_name + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name, disable_existing_loggers=False) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = FABDBManager.metadata + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def include_object(_, name, type_, *args): + if type_ == "table" and name not in target_metadata.tables: + return False + return True + + +def run_migrations_offline(): + """ + Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + context.configure( + url=settings.SQL_ALCHEMY_CONN, + target_metadata=target_metadata, + literal_binds=True, + compare_type=True, + compare_server_default=True, + render_as_batch=True, + version_table=version_table, + include_object=include_object, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online(): + """ + Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + + def process_revision_directives(context, revision, directives): + if getattr(config.cmd_opts, "autogenerate", False): + script = directives[0] + if script.upgrade_ops and script.upgrade_ops.is_empty(): + directives[:] = [] + print("No change detected in ORM schema, skipping revision.") + + with contextlib.ExitStack() as stack: + connection = config.attributes.get("connection", None) + + if not connection: + connection = stack.push(settings.engine.connect()) + + context.configure( + connection=connection, + transaction_per_migration=True, + target_metadata=target_metadata, + compare_type=True, + compare_server_default=True, + include_object=include_object, + render_as_batch=True, + process_revision_directives=process_revision_directives, + version_table=version_table, + ) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/airflow/providers/fab/migrations/script.py.mako b/airflow/providers/fab/migrations/script.py.mako new file mode 100644 index 0000000000000..4d0928fcc09ad --- /dev/null +++ b/airflow/providers/fab/migrations/script.py.mako @@ -0,0 +1,44 @@ +# +# 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. + +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/airflow/providers/fab/migrations/versions/__init__.py b/airflow/providers/fab/migrations/versions/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/fab/migrations/versions/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow/utils/db.py b/airflow/utils/db.py index a86ca5dbf5628..b963eef99a230 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -65,6 +65,7 @@ from airflow.exceptions import AirflowException from airflow.models import import_all_models from airflow.utils import helpers +from airflow.utils.db_manager import RunDBManager # TODO: remove create_session once we decide to break backward compatibility from airflow.utils.session import NEW_SESSION, create_session, provide_session # noqa: F401 @@ -748,7 +749,6 @@ def _create_db_from_orm(session): from alembic import command from airflow.models.base import Base - from airflow.providers.fab.auth_manager.models import Model def _create_flask_session_tbl(sql_database_uri): db = _get_flask_db(sql_database_uri) @@ -757,7 +757,6 @@ def _create_flask_session_tbl(sql_database_uri): with create_global_lock(session=session, lock=DBLocks.MIGRATIONS): engine = session.get_bind().engine Base.metadata.create_all(engine) - Model.metadata.create_all(engine) _create_flask_session_tbl(engine.url) # stamp the migration head config = _get_alembic_config() @@ -767,6 +766,10 @@ def _create_flask_session_tbl(sql_database_uri): @provide_session def initdb(session: Session = NEW_SESSION, load_connections: bool = True): """Initialize Airflow database.""" + # First validate external DB managers before running migration + external_db_manager = RunDBManager() + external_db_manager.validate() + import_all_models() db_exists = _get_current_revision(session) @@ -774,6 +777,8 @@ def initdb(session: Session = NEW_SESSION, load_connections: bool = True): upgradedb(session=session) else: _create_db_from_orm(session=session) + + external_db_manager.initdb(session) if conf.getboolean("database", "LOAD_DEFAULT_CONNECTIONS") and load_connections: create_default_connections(session=session) # Add default pool & sync log_template @@ -1636,7 +1641,6 @@ def upgradedb( # New DB; initialize and exit initdb(session=session, load_connections=False) return - with create_global_lock(session=session, lock=DBLocks.MIGRATIONS): import sqlalchemy.pool @@ -1650,6 +1654,7 @@ def upgradedb( os.environ["AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE"] = "1" settings.reconfigure_orm(pool_class=sqlalchemy.pool.SingletonThreadPool) command.upgrade(config, revision=to_revision or "heads") + finally: if val is None: os.environ.pop("AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE") @@ -1679,6 +1684,8 @@ def resetdb(session: Session = NEW_SESSION, skip_init: bool = False): with create_global_lock(session=session, lock=DBLocks.MIGRATIONS), connection.begin(): drop_airflow_models(connection) drop_airflow_moved_tables(connection) + external_db_manager = RunDBManager() + external_db_manager.drop_tables(connection) if not skip_init: initdb(session=session) @@ -1745,10 +1752,8 @@ def drop_airflow_models(connection): :return: None """ from airflow.models.base import Base - from airflow.providers.fab.auth_manager.models import Model Base.metadata.drop_all(connection) - Model.metadata.drop_all(connection) db = _get_flask_db(connection.engine.url) db.drop_all() # alembic adds significant import time, so we import it lazily diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py new file mode 100644 index 0000000000000..07d1ee9a61be1 --- /dev/null +++ b/airflow/utils/db_manager.py @@ -0,0 +1,167 @@ +# 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 + +import os +from abc import abstractmethod +from typing import TYPE_CHECKING + +from alembic import command + +from airflow.configuration import conf +from airflow.exceptions import AirflowException +from airflow.utils.log.logging_mixin import LoggingMixin +from airflow.utils.module_loading import import_string + +if TYPE_CHECKING: + from sqlalchemy import MetaData + + +class BaseDBManager(LoggingMixin): + """Abstract Base DB manager for external DBs.""" + + metadata: MetaData + migration_dir: str + alembic_file: str + version_table_name: str + # Whether the database supports dropping tables when airflow tables are dropped + supports_table_dropping: bool = False + + def __init__(self, session): + super().__init__() + self.session = session + + def get_alembic_config(self): + from alembic.config import Config + + from airflow import settings + + config = Config(self.alembic_file) + config.set_main_option("script_location", self.migration_dir.replace("%", "%%")) + config.set_main_option("sqlalchemy.url", settings.SQL_ALCHEMY_CONN.replace("%", "%%")) + return config + + def get_current_revision(self): + from alembic.migration import MigrationContext + + conn = self.session.connection() + + migration_ctx = MigrationContext.configure(conn, opts={"version_table": self.version_table_name}) + + return migration_ctx.get_current_revision() + + def _create_db_from_orm(self): + """Create database from ORM.""" + engine = self.session.get_bind().engine + self.metadata.create_all(engine) + config = self.get_alembic_config() + command.stamp(config, "head") + + def initdb(self): + """Initialize the database.""" + db_exists = self.get_current_revision() + if db_exists: + self.upgradedb() + else: + self._create_db_from_orm() + + def upgradedb(self, to_version=None, from_version=None, show_sql_only=False): + """Upgrade the database.""" + self.log.info("Upgrading the %s database", self.__class__.__name__) + + config = self.get_alembic_config() + command.upgrade(config, revision=to_version or "heads", sql=show_sql_only) + + @abstractmethod + def downgradedb(self, to_version, from_version=None, show_sql_only=False): + """Downgrade the database.""" + + +class RunDBManager(LoggingMixin): + """ + Run External DB Managers. + + Validates and runs the external database managers. + """ + + def __init__(self): + super().__init__() + self._managers: list[BaseDBManager] = [] + managers = conf.get("database", "external_db_managers").split(",") + for module in managers: + manager = import_string(module) + self._managers.append(manager) + + def validate(self): + """Validate the external database managers.""" + for manager in self._managers: + RunDBManager._validate(manager) + + @staticmethod + def _validate(manager: BaseDBManager): + """Validate the external database migration.""" + import ast + + from airflow.models.base import metadata as airflow_metadata + + external_metadata = manager.metadata + airflow_m = airflow_metadata + # validate tables are not airflow tables in metadata + for table_ in external_metadata.tables: + if table_ in airflow_m.tables: + raise AirflowException(f"Table '{table_}' already exists in the Airflow metadata") + # validate the version table schema is set appropriately in env.py + migration_dir = manager.migration_dir + env_file = os.path.join(migration_dir, "env.py") + with open(env_file) as f: + tree = ast.parse(f.read(), filename=env_file) + + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "context.configure" + ): + if "version_table" not in node.keywords: + raise AirflowException(f"version_table not set in {env_file}") + # validate the version table is not airflow version table + if manager.version_table_name == "alembic_version": + raise AirflowException(f"{manager}.version_table_name cannot be 'alembic_version'") + + def initdb(self, session): + """Initialize the external database managers.""" + for manager in self._managers: + m = manager(session) + m.initdb() + + def upgradedb(self, session): + """Upgrade the external database managers.""" + for manager in self._managers: + m = manager(session) + m.upgradedb() + + def downgradedb(self, session): + """Downgrade the external database managers.""" + for manager in self._managers: + m = manager(session) + m.downgradedb() + + def drop_tables(self, connection): + """Drop the external database managers.""" + for manager in self._managers: + if manager.supports_table_dropping: + manager.metadata.drop_all(connection) diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256 index e7e8667275258..db64fe3d3b86e 100644 --- a/docs/apache-airflow/img/airflow_erd.sha256 +++ b/docs/apache-airflow/img/airflow_erd.sha256 @@ -1 +1 @@ -37de7143e49532e7650ba09d0172788c23e132b57123a5f6e863e7234cf4f79e \ No newline at end of file +361896a7de8afd8678cf5811309fd242f1426b0d56fee7796444638a019fa6d5 \ No newline at end of file diff --git a/docs/apache-airflow/img/airflow_erd.svg b/docs/apache-airflow/img/airflow_erd.svg index 9f253a53db27f..12eded433404c 100644 --- a/docs/apache-airflow/img/airflow_erd.svg +++ b/docs/apache-airflow/img/airflow_erd.svg @@ -4,11 +4,11 @@ - - + + %3 - + log @@ -239,2033 +239,2033 @@ [BIGINT] - - -ab_user - -ab_user - -id - - [INTEGER] - NOT NULL - -active - - [BOOLEAN] - -changed_by_fk - - [INTEGER] - -changed_on - - [TIMESTAMP] - -created_by_fk - - [INTEGER] - -created_on - - [TIMESTAMP] - -email - - [VARCHAR(512)] - NOT NULL - -fail_login_count - - [INTEGER] - -first_name - - [VARCHAR(256)] - NOT NULL - -last_login - - [TIMESTAMP] - -last_name - - [VARCHAR(256)] - NOT NULL - -login_count - - [INTEGER] - -password - - [VARCHAR(256)] - -username - - [VARCHAR(512)] - NOT NULL - - - -ab_user--ab_user - -0..N -{0,1} - - - -ab_user--ab_user - -0..N -{0,1} - - - -ab_user_role - -ab_user_role - -id - - [INTEGER] - NOT NULL - -role_id - - [INTEGER] - -user_id - - [INTEGER] - - - -ab_user--ab_user_role - -0..N -{0,1} - - - -ab_register_user - -ab_register_user - -id - - [INTEGER] - NOT NULL - -email - - [VARCHAR(512)] - NOT NULL - -first_name - - [VARCHAR(256)] - NOT NULL - -last_name - - [VARCHAR(256)] - NOT NULL - -password - - [VARCHAR(256)] - -registration_date - - [TIMESTAMP] - -registration_hash - - [VARCHAR(256)] - -username - - [VARCHAR(512)] - NOT NULL - - + connection - -connection - -id - - [INTEGER] - NOT NULL - -conn_id - - [VARCHAR(250)] - NOT NULL - -conn_type - - [VARCHAR(500)] - NOT NULL - -description - - [TEXT] - -extra - - [TEXT] - -host - - [VARCHAR(500)] - -is_encrypted - - [BOOLEAN] - -is_extra_encrypted - - [BOOLEAN] - -login - - [TEXT] - -password - - [TEXT] - -port - - [INTEGER] - -schema - - [VARCHAR(500)] + +connection + +id + + [INTEGER] + NOT NULL + +conn_id + + [VARCHAR(250)] + NOT NULL + +conn_type + + [VARCHAR(500)] + NOT NULL + +description + + [TEXT] + +extra + + [TEXT] + +host + + [VARCHAR(500)] + +is_encrypted + + [BOOLEAN] + +is_extra_encrypted + + [BOOLEAN] + +login + + [TEXT] + +password + + [TEXT] + +port + + [INTEGER] + +schema + + [VARCHAR(500)] - + sla_miss - -sla_miss - -dag_id - - [VARCHAR(250)] - NOT NULL - -execution_date - - [TIMESTAMP] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -description - - [TEXT] - -email_sent - - [BOOLEAN] - -notification_sent - - [BOOLEAN] - -timestamp - - [TIMESTAMP] + +sla_miss + +dag_id + + [VARCHAR(250)] + NOT NULL + +execution_date + + [TIMESTAMP] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +description + + [TEXT] + +email_sent + + [BOOLEAN] + +notification_sent + + [BOOLEAN] + +timestamp + + [TIMESTAMP] - + variable - -variable - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -is_encrypted - - [BOOLEAN] - -key - - [VARCHAR(250)] - -val - - [TEXT] + +variable + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +is_encrypted + + [BOOLEAN] + +key + + [VARCHAR(250)] + +val + + [TEXT] - + import_error - -import_error - -id - - [INTEGER] - NOT NULL - -filename - - [VARCHAR(1024)] - -processor_subdir - - [VARCHAR(2000)] - -stacktrace - - [TEXT] - -timestamp - - [TIMESTAMP] + +import_error + +id + + [INTEGER] + NOT NULL + +filename + + [VARCHAR(1024)] + +processor_subdir + + [VARCHAR(2000)] + +stacktrace + + [TEXT] + +timestamp + + [TIMESTAMP] - + serialized_dag - -serialized_dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_hash - - [VARCHAR(32)] - NOT NULL - -data - - [JSON] - -data_compressed - - [BYTEA] - -fileloc - - [VARCHAR(2000)] - NOT NULL - -fileloc_hash - - [BIGINT] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -processor_subdir - - [VARCHAR(2000)] + +serialized_dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_hash + + [VARCHAR(32)] + NOT NULL + +data + + [JSON] + +data_compressed + + [BYTEA] + +fileloc + + [VARCHAR(2000)] + NOT NULL + +fileloc_hash + + [BIGINT] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +processor_subdir + + [VARCHAR(2000)] - + dataset_alias - -dataset_alias - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(3000)] - NOT NULL + +dataset_alias + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(3000)] + NOT NULL - + dataset_alias_dataset - -dataset_alias_dataset - -alias_id - - [INTEGER] - NOT NULL - -dataset_id - - [INTEGER] - NOT NULL + +dataset_alias_dataset + +alias_id + + [INTEGER] + NOT NULL + +dataset_id + + [INTEGER] + NOT NULL - + dataset_alias--dataset_alias_dataset - -0..N -1 + +0..N +1 - + dataset_alias--dataset_alias_dataset - -0..N -1 + +0..N +1 - + dataset_alias_dataset_event - -dataset_alias_dataset_event - -alias_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +dataset_alias_dataset_event + +alias_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL - + dataset_alias--dataset_alias_dataset_event - -0..N -1 + +0..N +1 - + dataset_alias--dataset_alias_dataset_event - -0..N -1 + +0..N +1 - + dag_schedule_dataset_alias_reference - -dag_schedule_dataset_alias_reference - -alias_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_dataset_alias_reference + +alias_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + dataset_alias--dag_schedule_dataset_alias_reference - -0..N -1 + +0..N +1 - + dataset - -dataset - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -extra - - [JSON] - NOT NULL - -is_orphaned - - [BOOLEAN] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -uri - - [VARCHAR(3000)] - NOT NULL + +dataset + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +extra + + [JSON] + NOT NULL + +is_orphaned + + [BOOLEAN] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +uri + + [VARCHAR(3000)] + NOT NULL - + dataset--dataset_alias_dataset - -0..N -1 + +0..N +1 - + dataset--dataset_alias_dataset - -0..N -1 + +0..N +1 - + dag_schedule_dataset_reference - -dag_schedule_dataset_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -dataset_id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_dataset_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +dataset_id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + dataset--dag_schedule_dataset_reference - -0..N -1 + +0..N +1 - + task_outlet_dataset_reference - -task_outlet_dataset_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -dataset_id - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_outlet_dataset_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +dataset_id + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + dataset--task_outlet_dataset_reference - -0..N -1 + +0..N +1 - + dataset_dag_run_queue - -dataset_dag_run_queue - -dataset_id - - [INTEGER] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +dataset_dag_run_queue + +dataset_id + + [INTEGER] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL - + dataset--dataset_dag_run_queue - -0..N -1 + +0..N +1 - + dataset_event - -dataset_event - -id - - [INTEGER] - NOT NULL - -dataset_id - - [INTEGER] - NOT NULL - -extra - - [JSON] - NOT NULL - -source_dag_id - - [VARCHAR(250)] - -source_map_index - - [INTEGER] - -source_run_id - - [VARCHAR(250)] - -source_task_id - - [VARCHAR(250)] - -timestamp - - [TIMESTAMP] - NOT NULL + +dataset_event + +id + + [INTEGER] + NOT NULL + +dataset_id + + [INTEGER] + NOT NULL + +extra + + [JSON] + NOT NULL + +source_dag_id + + [VARCHAR(250)] + +source_map_index + + [INTEGER] + +source_run_id + + [VARCHAR(250)] + +source_task_id + + [VARCHAR(250)] + +timestamp + + [TIMESTAMP] + NOT NULL - + dataset_event--dataset_alias_dataset_event - -0..N -1 + +0..N +1 - + dataset_event--dataset_alias_dataset_event - -0..N -1 + +0..N +1 - + dagrun_dataset_event - -dagrun_dataset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +dagrun_dataset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL - + dataset_event--dagrun_dataset_event - -0..N -1 + +0..N +1 - + dag - -dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_display_name - - [VARCHAR(2000)] - -dataset_expression - - [JSON] - -default_view - - [VARCHAR(25)] - -description - - [TEXT] - -fileloc - - [VARCHAR(2000)] - -has_import_errors - - [BOOLEAN] - -has_task_concurrency_limits - - [BOOLEAN] - NOT NULL - -is_active - - [BOOLEAN] - -is_paused - - [BOOLEAN] - -last_expired - - [TIMESTAMP] - -last_parsed_time - - [TIMESTAMP] - -last_pickled - - [TIMESTAMP] - -max_active_runs - - [INTEGER] - -max_active_tasks - - [INTEGER] - NOT NULL - -max_consecutive_failed_dag_runs - - [INTEGER] - NOT NULL - -next_dagrun - - [TIMESTAMP] - -next_dagrun_create_after - - [TIMESTAMP] - -next_dagrun_data_interval_end - - [TIMESTAMP] - -next_dagrun_data_interval_start - - [TIMESTAMP] - -owners - - [VARCHAR(2000)] - -pickle_id - - [INTEGER] - -processor_subdir - - [VARCHAR(2000)] - -schedule_interval - - [TEXT] - -scheduler_lock - - [BOOLEAN] - -timetable_description - - [VARCHAR(1000)] + +dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_display_name + + [VARCHAR(2000)] + +dataset_expression + + [JSON] + +default_view + + [VARCHAR(25)] + +description + + [TEXT] + +fileloc + + [VARCHAR(2000)] + +has_import_errors + + [BOOLEAN] + +has_task_concurrency_limits + + [BOOLEAN] + NOT NULL + +is_active + + [BOOLEAN] + +is_paused + + [BOOLEAN] + +last_expired + + [TIMESTAMP] + +last_parsed_time + + [TIMESTAMP] + +last_pickled + + [TIMESTAMP] + +max_active_runs + + [INTEGER] + +max_active_tasks + + [INTEGER] + NOT NULL + +max_consecutive_failed_dag_runs + + [INTEGER] + NOT NULL + +next_dagrun + + [TIMESTAMP] + +next_dagrun_create_after + + [TIMESTAMP] + +next_dagrun_data_interval_end + + [TIMESTAMP] + +next_dagrun_data_interval_start + + [TIMESTAMP] + +owners + + [VARCHAR(2000)] + +pickle_id + + [INTEGER] + +processor_subdir + + [VARCHAR(2000)] + +schedule_interval + + [TEXT] + +scheduler_lock + + [BOOLEAN] + +timetable_description + + [VARCHAR(1000)] - + dag--dag_schedule_dataset_alias_reference - -0..N -1 + +0..N +1 - + dag--dag_schedule_dataset_reference - -0..N -1 + +0..N +1 - + dag--task_outlet_dataset_reference - -0..N -1 + +0..N +1 - + dag--dataset_dag_run_queue - -0..N -1 + +0..N +1 - + dag_tag - -dag_tag - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL + +dag_tag + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL - + dag--dag_tag - -0..N -1 + +0..N +1 - + dag_owner_attributes - -dag_owner_attributes - -dag_id - - [VARCHAR(250)] - NOT NULL - -owner - - [VARCHAR(500)] - NOT NULL - -link - - [VARCHAR(500)] - NOT NULL + +dag_owner_attributes + +dag_id + + [VARCHAR(250)] + NOT NULL + +owner + + [VARCHAR(500)] + NOT NULL + +link + + [VARCHAR(500)] + NOT NULL - + dag--dag_owner_attributes - -0..N -1 + +0..N +1 - + dag_warning - -dag_warning - -dag_id - - [VARCHAR(250)] - NOT NULL - -warning_type - - [VARCHAR(50)] - NOT NULL - -message - - [TEXT] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL + +dag_warning + +dag_id + + [VARCHAR(250)] + NOT NULL + +warning_type + + [VARCHAR(50)] + NOT NULL + +message + + [TEXT] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL - + dag--dag_warning - -0..N -1 + +0..N +1 - + log_template - -log_template - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -elasticsearch_id - - [TEXT] - NOT NULL - -filename - - [TEXT] - NOT NULL + +log_template + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +elasticsearch_id + + [TEXT] + NOT NULL + +filename + + [TEXT] + NOT NULL - + dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -clear_number - - [INTEGER] - NOT NULL - -conf - - [BYTEA] - -creating_job_id - - [INTEGER] - -dag_hash - - [VARCHAR(32)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -execution_date - - [TIMESTAMP] - NOT NULL - -external_trigger - - [BOOLEAN] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - -queued_at - - [TIMESTAMP] - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - -updated_at - - [TIMESTAMP] + +dag_run + +id + + [INTEGER] + NOT NULL + +clear_number + + [INTEGER] + NOT NULL + +conf + + [BYTEA] + +creating_job_id + + [INTEGER] + +dag_hash + + [VARCHAR(32)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +execution_date + + [TIMESTAMP] + NOT NULL + +external_trigger + + [BOOLEAN] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + +queued_at + + [TIMESTAMP] + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + +updated_at + + [TIMESTAMP] - + log_template--dag_run - -0..N -{0,1} + +0..N +{0,1} - + dag_run--dagrun_dataset_event - -0..N -1 + +0..N +1 - + task_instance - -task_instance - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -custom_operator_name - - [VARCHAR(1000)] - -duration - - [DOUBLE_PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -job_id - - [INTEGER] - -max_tries - - [INTEGER] - -next_kwargs - - [JSON] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] + +task_instance + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +custom_operator_name + + [VARCHAR(1000)] + +duration + + [DOUBLE_PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +job_id + + [INTEGER] + +max_tries + + [INTEGER] + +next_kwargs + + [JSON] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] - + dag_run--task_instance - -0..N -1 + +0..N +1 - + dag_run--task_instance - -0..N -1 + +0..N +1 - + dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [INTEGER] + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [INTEGER] - + dag_run--dag_run_note - -1 -1 + +1 +1 - + task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -try_number - - [INTEGER] - NOT NULL + +task_reschedule + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +try_number + + [INTEGER] + NOT NULL - + dag_run--task_reschedule - -0..N -1 + +0..N +1 - + dag_run--task_reschedule - -0..N -1 + +0..N +1 - + task_instance--task_reschedule - -0..N -1 + +0..N +1 - + task_instance--task_reschedule - -0..N -1 + +0..N +1 - + task_instance--task_reschedule - -0..N -1 + +0..N +1 - + task_instance--task_reschedule - -0..N -1 + +0..N +1 - + rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL - + task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 - + task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 - + task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 - + task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 - + task_fail - -task_fail - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -duration - - [INTEGER] - -end_date - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -task_id - - [VARCHAR(250)] - NOT NULL + +task_fail + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +duration + + [INTEGER] + +end_date + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +task_id + + [VARCHAR(250)] + NOT NULL - + task_instance--task_fail - -0..N -1 + +0..N +1 - + task_instance--task_fail - -0..N -1 + +0..N +1 - + task_instance--task_fail - -0..N -1 + +0..N +1 - + task_instance--task_fail - -0..N -1 + +0..N +1 - + task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSON] - -length - - [INTEGER] - NOT NULL + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSON] + +length + + [INTEGER] + NOT NULL - + task_instance--task_map - -0..N -1 + +0..N +1 - + task_instance--task_map - -0..N -1 + +0..N +1 - + task_instance--task_map - -0..N -1 + +0..N +1 - + task_instance--task_map - -0..N -1 + +0..N +1 - + xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [BYTEA] + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [BYTEA] - + task_instance--xcom - -0..N -1 + +0..N +1 - + task_instance--xcom - -0..N -1 + +0..N +1 - + task_instance--xcom - -0..N -1 + +0..N +1 - + task_instance--xcom - -0..N -1 + +0..N +1 - + task_instance_note - -task_instance_note - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [INTEGER] + +task_instance_note + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [INTEGER] - + task_instance--task_instance_note - -0..N -1 + +0..N +1 - + task_instance--task_instance_note - -0..N -1 + +0..N +1 - + task_instance--task_instance_note - -0..N -1 + +0..N +1 - + task_instance--task_instance_note - -0..N -1 + +0..N +1 - + task_instance_history - -task_instance_history - -id - - [INTEGER] - NOT NULL - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -duration - - [DOUBLE_PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -job_id - - [INTEGER] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSON] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] + +task_instance_history + +id + + [INTEGER] + NOT NULL + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +duration + + [DOUBLE_PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +job_id + + [INTEGER] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSON] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] - + task_instance--task_instance_history - -0..N -1 + +0..N +1 - + task_instance--task_instance_history - -0..N -1 + +0..N +1 - + task_instance--task_instance_history - -0..N -1 + +0..N +1 - + task_instance--task_instance_history - -0..N -1 + +0..N +1 + + + +trigger + +trigger + +id + + [INTEGER] + NOT NULL + +classpath + + [VARCHAR(1000)] + NOT NULL + +created_date + + [TIMESTAMP] + NOT NULL + +kwargs + + [TEXT] + NOT NULL + +triggerer_id + + [INTEGER] + + + +trigger--task_instance + +0..N +{0,1} + + + +alembic_version + +alembic_version + +version_num + + [VARCHAR(32)] + NOT NULL + + + +session + +session + +id + + [INTEGER] + NOT NULL + +data + + [BYTEA] + +expiry + + [TIMESTAMP] + +session_id + + [VARCHAR(255)] + + + +ab_user + +ab_user + +id + + [INTEGER] + NOT NULL + +active + + [BOOLEAN] + +changed_by_fk + + [INTEGER] + +changed_on + + [TIMESTAMP] + +created_by_fk + + [INTEGER] + +created_on + + [TIMESTAMP] + +email + + [VARCHAR(512)] + NOT NULL + +fail_login_count + + [INTEGER] + +first_name + + [VARCHAR(256)] + NOT NULL + +last_login + + [TIMESTAMP] + +last_name + + [VARCHAR(256)] + NOT NULL + +login_count + + [INTEGER] + +password + + [VARCHAR(256)] + +username + + [VARCHAR(512)] + NOT NULL + + + +ab_user--ab_user + +0..N +{0,1} + + + +ab_user--ab_user + +0..N +{0,1} + + + +ab_user_role + +ab_user_role + +id + + [INTEGER] + NOT NULL + +role_id + + [INTEGER] + +user_id + + [INTEGER] + + + +ab_user--ab_user_role + +0..N +{0,1} + + + +ab_register_user + +ab_register_user + +id + + [INTEGER] + NOT NULL + +email + + [VARCHAR(512)] + NOT NULL + +first_name + + [VARCHAR(256)] + NOT NULL + +last_name + + [VARCHAR(256)] + NOT NULL + +password + + [VARCHAR(256)] + +registration_date + + [TIMESTAMP] + +registration_hash + + [VARCHAR(256)] + +username + + [VARCHAR(512)] + NOT NULL - + ab_permission - -ab_permission - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL + +ab_permission + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL - + ab_permission_view - -ab_permission_view - -id - - [INTEGER] - NOT NULL - -permission_id - - [INTEGER] - -view_menu_id - - [INTEGER] + +ab_permission_view + +id + + [INTEGER] + NOT NULL + +permission_id + + [INTEGER] + +view_menu_id + + [INTEGER] - + ab_permission--ab_permission_view - -0..N -{0,1} + +0..N +{0,1} - + ab_permission_view_role - -ab_permission_view_role - -id - - [INTEGER] - NOT NULL - -permission_view_id - - [INTEGER] - -role_id - - [INTEGER] + +ab_permission_view_role + +id + + [INTEGER] + NOT NULL + +permission_view_id + + [INTEGER] + +role_id + + [INTEGER] - + ab_permission_view--ab_permission_view_role - -0..N -{0,1} + +0..N +{0,1} - + ab_view_menu - -ab_view_menu - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(250)] - NOT NULL + +ab_view_menu + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(250)] + NOT NULL - + ab_view_menu--ab_permission_view - -0..N -{0,1} + +0..N +{0,1} - + ab_role - -ab_role - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(64)] - NOT NULL + +ab_role + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(64)] + NOT NULL - + ab_role--ab_user_role - -0..N -{0,1} + +0..N +{0,1} - + ab_role--ab_permission_view_role - -0..N -{0,1} - - - -trigger - -trigger - -id - - [INTEGER] - NOT NULL - -classpath - - [VARCHAR(1000)] - NOT NULL - -created_date - - [TIMESTAMP] - NOT NULL - -kwargs - - [TEXT] - NOT NULL - -triggerer_id - - [INTEGER] - - - -trigger--task_instance - -0..N -{0,1} - - - -session - -session - -id - - [INTEGER] - NOT NULL - -data - - [BYTEA] - -expiry - - [TIMESTAMP] - -session_id - - [VARCHAR(255)] - - - -alembic_version - -alembic_version - -version_num - - [VARCHAR(32)] - NOT NULL + +0..N +{0,1} diff --git a/scripts/in_container/run_migration_reference.py b/scripts/in_container/run_migration_reference.py index 6dcc08357e4b0..204db7a5ea908 100755 --- a/scripts/in_container/run_migration_reference.py +++ b/scripts/in_container/run_migration_reference.py @@ -34,6 +34,7 @@ from tabulate import tabulate from airflow import __version__ as airflow_version +from airflow.providers.fab import __version__ as fab_version from airflow.utils.db import _get_alembic_config if TYPE_CHECKING: @@ -42,6 +43,7 @@ console = Console(width=400, color_system="standard") airflow_version = re.match(r"(\d+\.\d+\.\d+).*", airflow_version).group(1) # type: ignore +fab_version = re.match(r"(\d+\.\d+\.\d+).*", fab_version).group(1) # type: ignore project_root = Path(__file__).parents[2].resolve() @@ -81,14 +83,25 @@ def update_doc(file, data): ) -def has_version(content): - return re.search(r"^airflow_version\s*=.*", content, flags=re.MULTILINE) is not None +def has_version(content, app): + if app == "airflow": + return re.search(r"^airflow_version\s*=.*", content, flags=re.MULTILINE) is not None + return re.search(r"^fab_version\s*=.*", content, flags=re.MULTILINE) is not None -def insert_version(old_content, file): +def insert_version(old_content, file, app): + if app == "airflow": + new_content = re.sub( + r"(^depends_on.*)", + lambda x: f'{x.group(1)}\nairflow_version = "{airflow_version}"', + old_content, + flags=re.MULTILINE, + ) + return + new_content = re.sub( r"(^depends_on.*)", - lambda x: f'{x.group(1)}\nairflow_version = "{airflow_version}"', + lambda x: f'{x.group(1)}\nfab_version = "{fab_version}"', old_content, flags=re.MULTILINE, ) @@ -107,20 +120,27 @@ def revision_suffix(rev: Script): return "" -def ensure_airflow_version(revisions: Iterable[Script]): +def ensure_version(revisions: Iterable[Script], app): for rev in revisions: if TYPE_CHECKING: # For mypy assert rev.module.__file__ is not None file = Path(rev.module.__file__) content = file.read_text() - if not has_version(content): - insert_version(content, file) + if not has_version(content, app=app): + insert_version(content, file, app=app) -def get_revisions() -> Iterable[Script]: - config = _get_alembic_config() - script = ScriptDirectory.from_config(config) - yield from script.walk_revisions() +def get_revisions(app="airflow") -> Iterable[Script]: + if app == "airflow": + config = _get_alembic_config() + script = ScriptDirectory.from_config(config) + yield from script.walk_revisions() + else: + from airflow.providers.fab.auth_manager.models.db import FABDBManager + + config = FABDBManager(session="").get_alembic_config() + script = ScriptDirectory.from_config(config) + yield from script.walk_revisions() def update_docs(revisions: Iterable[Script]): @@ -155,13 +175,16 @@ def ensure_mod_prefix(mod_name, idx, version): return "_".join(parts) -def ensure_filenames_are_sorted(revisions): +def ensure_filenames_are_sorted(revisions, app): renames = [] is_branched = False unmerged_heads = [] for idx, rev in enumerate(revisions): mod_path = Path(rev.module.__file__) - version = rev.module.airflow_version.split(".")[0:3] # only first 3 tokens + if app == "airflow": + version = rev.module.airflow_version.split(".")[0:3] # only first 3 tokens + else: + version = rev.module.fab_version.split(".")[0:3] # only first 3 tokens correct_mod_basename = ensure_mod_prefix(mod_path.name, idx, version) if mod_path.name != correct_mod_basename: renames.append((mod_path, Path(mod_path.parent, correct_mod_basename))) @@ -209,16 +232,19 @@ def correct_mismatching_revision_nums(revisions: Iterable[Script]): if __name__ == "__main__": - console.print("[bright_blue]Updating migration reference") - revisions = list(reversed(list(get_revisions()))) - console.print("[bright_blue]Making sure airflow version updated") - ensure_airflow_version(revisions=revisions) - console.print("[bright_blue]Making sure there's no mismatching revision numbers") - correct_mismatching_revision_nums(revisions=revisions) - revisions = list(reversed(list(get_revisions()))) - console.print("[bright_blue]Making sure filenames are sorted") - ensure_filenames_are_sorted(revisions=revisions) - revisions = list(get_revisions()) - console.print("[bright_blue]Updating documentation") - update_docs(revisions=revisions) - console.print("[green]Migrations OK") + apps = ["airflow", "fab"] + for app in apps: + console.print("[bright_blue]Updating migration reference") + revisions = list(reversed(list(get_revisions(app)))) + console.print("[bright_blue]Making sure airflow version updated") + ensure_version(revisions=revisions, app=app) + console.print("[bright_blue]Making sure there's no mismatching revision numbers") + correct_mismatching_revision_nums(revisions=revisions) + revisions = list(reversed(list(get_revisions(app=app)))) + console.print("[bright_blue]Making sure filenames are sorted") + ensure_filenames_are_sorted(revisions=revisions, app=app) + revisions = list(get_revisions(app=app)) + if app == "airflow": + console.print("[bright_blue]Updating documentation") + update_docs(revisions=revisions) + console.print("[green]Migrations OK") diff --git a/scripts/in_container/verify_providers.py b/scripts/in_container/verify_providers.py index a7a97d78ca4a2..b3602dead84b4 100755 --- a/scripts/in_container/verify_providers.py +++ b/scripts/in_container/verify_providers.py @@ -212,9 +212,14 @@ def onerror(_): ... except Exception as e: # skip the check as we are temporary vendoring in the google ads client with wrong package - if "No module named 'google.ads.googleads.v12'" not in str(e): + # skip alembic.context which is only available when alembic command is executed from a folder + # containing the alembic.ini file + if "No module named 'google.ads.googleads.v12'" not in str( + e + ) and "module 'alembic.context' has no attribute 'config'" not in str(e): exception_str = traceback.format_exc() tracebacks.append((modinfo.name, exception_str)) + if tracebacks: if IS_AIRFLOW_VERSION_PROVIDED: console.print( diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index 15813ca9ca296..7bb664b068a3a 100644 --- a/tests/always/test_project_structure.py +++ b/tests/always/test_project_structure.py @@ -155,6 +155,7 @@ def test_providers_modules_should_have_tests(self): "tests/providers/snowflake/triggers/test_snowflake_trigger.py", "tests/providers/yandex/hooks/test_yandexcloud_dataproc.py", "tests/providers/yandex/operators/test_yandexcloud_dataproc.py", + "tests/providers/fab/migrations/test_env.py", ] # TODO: Should we extend this test to cover other directories? diff --git a/tests/providers/fab/auth_manager/models/test_db.py b/tests/providers/fab/auth_manager/models/test_db.py new file mode 100644 index 0000000000000..7b0dc345c7188 --- /dev/null +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -0,0 +1,76 @@ +# 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 + +import os + +import pytest +from alembic.autogenerate import compare_metadata +from alembic.migration import MigrationContext +from sqlalchemy import MetaData + +import airflow +from airflow.settings import engine +from airflow.utils.db import ( + compare_server_default, + compare_type, +) + +pytestmark = [pytest.mark.db_test] +try: + from airflow.providers.fab.auth_manager.models.db import FABDBManager + + class TestFABDBManager: + def setup_method(self, session): + self.airflow_dir = os.path.dirname(airflow.__file__) + self.db_manager = FABDBManager(session=session) + + def test_version_table_name_set(self): + assert self.db_manager.version_table_name == "fab_alembic_version" + + def test_migration_dir_set(self): + assert self.db_manager.migration_dir == f"{self.airflow_dir}/providers/fab/migrations" + + def test_alembic_file_set(self): + assert self.db_manager.alembic_file == f"{self.airflow_dir}/providers/fab/alembic.ini" + + def test_supports_table_dropping_set(self): + assert self.db_manager.supports_table_dropping is True + + def test_database_schema_and_sqlalchemy_model_are_in_sync(self): + def include_object(_, name, type_, *args): + if type_ == "table" and name not in self.db_manager.metadata.tables: + return False + return True + + all_meta_data = MetaData() + for table_name, table in self.db_manager.metadata.tables.items(): + all_meta_data._add_table(table_name, table.schema, table) + # create diff between database schema and SQLAlchemy model + mctx = MigrationContext.configure( + engine.connect(), + opts={ + "compare_type": compare_type, + "compare_server_default": compare_server_default, + "include_object": include_object, + }, + ) + diff = compare_metadata(mctx, all_meta_data) + + assert not diff, "Database schema and SQLAlchemy model are not in sync: " + str(diff) +except ModuleNotFoundError: + pass diff --git a/tests/providers/fab/auth_manager/test_security.py b/tests/providers/fab/auth_manager/test_security.py index 8966b5b5f6c67..0cd0cf4af86e2 100644 --- a/tests/providers/fab/auth_manager/test_security.py +++ b/tests/providers/fab/auth_manager/test_security.py @@ -34,13 +34,12 @@ from airflow.configuration import initialize_config from airflow.exceptions import AirflowException from airflow.models import DagModel -from airflow.models.base import Base from airflow.models.dag import DAG from tests.test_utils.compat import ignore_provider_compatibility_error with ignore_provider_compatibility_error("2.9.0+", __file__): from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager - from airflow.providers.fab.auth_manager.models import User, assoc_permission_role + from airflow.providers.fab.auth_manager.models import assoc_permission_role from airflow.providers.fab.auth_manager.models.anonymous_user import AnonymousUser from airflow.security import permissions @@ -1042,12 +1041,6 @@ def test_permissions_work_for_dags_with_dot_in_dagname( session.query(DagModel).delete() -def test_fab_models_use_airflow_base_meta(): - # TODO: move this test to appropriate place when we have more tests for FAB models - user = User() - assert user.metadata is Base.metadata - - @pytest.fixture def mock_security_manager(app_builder): mocked_security_manager = MockSecurityManager(appbuilder=app_builder) diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 3ed2bdcbfae8c..d7fc73aa85783 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -1121,7 +1121,6 @@ def test_extra_serialized_field_and_operator_links( link = simple_task.get_extra_links(ti, GoogleLink.name) assert "https://www.google.com" == link - @pytest.mark.db_test def test_extra_operator_links_logs_error_for_non_registered_extra_links(self, caplog): """ Assert OperatorLinks not registered via Plugins and if it is not an inbuilt Operator Link, diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 287a5a1fdbdb0..4d32d132b2c96 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -50,6 +50,7 @@ resetdb, upgradedb, ) +from airflow.utils.db_manager import RunDBManager from airflow.utils.session import NEW_SESSION pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -61,8 +62,14 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): airflow.models.import_all_models() all_meta_data = MetaData() + # Airflow DB for table_name, table in airflow_base.metadata.tables.items(): all_meta_data._add_table(table_name, table.schema, table) + # External DB Managers + external_db_managers = RunDBManager() + for dbmanager in external_db_managers._managers: + for table_name, table in dbmanager.metadata.tables.items(): + all_meta_data._add_table(table_name, table.schema, table) # create diff between database schema and SQLAlchemy model mctx = MigrationContext.configure( @@ -70,6 +77,7 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): opts={"compare_type": compare_type, "compare_server_default": compare_server_default}, ) diff = compare_metadata(mctx, all_meta_data) + # known diffs to ignore ignores = [ # ignore tables created by celery @@ -86,6 +94,8 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): lambda t: (t[0] == "remove_index" and t[1].name == "session_session_id_uq"), # sqlite sequence is used for autoincrementing columns created with `sqlite_autoincrement` option lambda t: (t[0] == "remove_table" and t[1].name == "sqlite_sequence"), + # fab version table + lambda t: (t[0] == "remove_table" and t[1].name == "fab_alembic_version"), ] for ignore in ignores: diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py new file mode 100644 index 0000000000000..d6fc91f93934a --- /dev/null +++ b/tests/utils/test_db_manager.py @@ -0,0 +1,134 @@ +# 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 unittest import mock + +import pytest +from sqlalchemy import Table + +from airflow.exceptions import AirflowException +from airflow.models import Base +from airflow.utils.db import downgrade, initdb, upgradedb +from airflow.utils.db_manager import BaseDBManager, RunDBManager +from tests.test_utils.config import conf_vars + +pytestmark = [pytest.mark.db_test] + + +class TestRunDBManager: + @conf_vars( + {("database", "external_db_managers"): "airflow.providers.fab.auth_manager.models.db.FABDBManager"} + ) + def test_fab_db_manager_is_default(self): + from airflow.providers.fab.auth_manager.models.db import FABDBManager + + run_db_manager = RunDBManager() + assert run_db_manager._managers == [FABDBManager] + + @conf_vars( + {("database", "external_db_managers"): "airflow.providers.fab.auth_manager.models.db.FABDBManager"} + ) + def test_defining_table_same_name_as_airflow_table_name_raises(self): + from sqlalchemy import Column, Integer, String + + run_db_manager = RunDBManager() + metadata = run_db_manager._managers[0].metadata + # Add dag_run table to metadata + mytable = Table( + "dag_run", metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) + ) + metadata._add_table("dag_run", None, mytable) + with pytest.raises(AirflowException, match="Table 'dag_run' already exists in the Airflow metadata"): + run_db_manager.validate() + metadata._remove_table("dag_run", None) + + @mock.patch.object(RunDBManager, "downgradedb") + @mock.patch.object(RunDBManager, "upgradedb") + @mock.patch.object(RunDBManager, "initdb") + def test_init_db_calls_rundbmanager(self, mock_initdb, mock_upgrade_db, mock_downgrade_db, session): + initdb(session=session) + mock_initdb.assert_called() + mock_initdb.assert_called_once_with(session) + mock_upgrade_db.assert_not_called() + mock_downgrade_db.assert_not_called() + + @mock.patch.object(RunDBManager, "downgradedb") + @mock.patch.object(RunDBManager, "upgradedb") + @mock.patch.object(RunDBManager, "initdb") + @mock.patch("alembic.command") + def test_upgradedb_or_downgrade_dont_call_rundbmanager( + self, mock_alembic_command, mock_initdb, mock_upgrade_db, mock_downgrade_db, session + ): + upgradedb(session=session) + mock_alembic_command.upgrade.assert_called_once_with(mock.ANY, revision="heads") + downgrade(to_revision="base") + mock_alembic_command.downgrade.assert_called_once_with(mock.ANY, revision="base", sql=False) + mock_initdb.assert_not_called() + mock_upgrade_db.assert_not_called() + mock_downgrade_db.assert_not_called() + + @conf_vars( + {("database", "external_db_managers"): "airflow.providers.fab.auth_manager.models.db.FABDBManager"} + ) + @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") + def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session): + mock_fabdb_manager.supports_table_dropping = True + fabdb_manager = mock_fabdb_manager.return_value + ext_db = RunDBManager() + # initdb + ext_db.initdb(session=session) + fabdb_manager.initdb.assert_called_once() + # upgradedb + ext_db.upgradedb(session=session) + fabdb_manager.upgradedb.assert_called_once() + # downgradedb + ext_db.downgradedb(session=session) + mock_fabdb_manager.return_value.downgradedb.assert_called_once() + connection = mock.MagicMock() + ext_db.drop_tables(connection) + mock_fabdb_manager.metadata.drop_all.assert_called_once_with(connection) + + +class MockDBManager(BaseDBManager): + metadata = Base.metadata + version_table_name = "mock_alembic_version" + migration_dir = "mock_migration_dir" + alembic_file = "mock_alembic.ini" + supports_table_dropping = True + + +class TestBaseDBManager: + @mock.patch.object(BaseDBManager, "get_alembic_config") + @mock.patch.object(BaseDBManager, "get_current_revision") + @mock.patch.object(BaseDBManager, "_create_db_from_orm") + def test_create_db_from_orm_called_from_init( + self, mock_create_db_from_orm, mock_current_revision, mock_config, session + ): + mock_current_revision.return_value = None + + manager = MockDBManager(session) + manager.initdb() + mock_create_db_from_orm.assert_called_once() + + @mock.patch.object(BaseDBManager, "get_alembic_config") + @mock.patch("alembic.command.upgrade") + def test_upgradedb(self, mock_alembic_cmd, mock_alembic_config, session, caplog): + manager = MockDBManager(session) + manager.upgradedb() + mock_alembic_cmd.assert_called_once() + assert "Upgrading the MockDBManager database" in caplog.text