From 353a100bdcc238975321c27fa4cf5486b1193c38 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 13 Aug 2024 16:18:07 +0100 Subject: [PATCH 01/30] Separate FAB migration from Core Airflow migration This PR separates FAB migration from Airflow Core migration and provides a way for apps to integrate into Airflow and run their migrations. --- airflow/config_templates/config.yml | 8 ++ airflow/migrations/env.py | 30 ++-- airflow/models/base.py | 75 +++++++++- airflow/providers/fab/alembic.ini | 133 ++++++++++++++++++ airflow/providers/fab/alembic/README | 1 + airflow/providers/fab/alembic/__init__.py | 16 +++ airflow/providers/fab/alembic/env.py | 124 ++++++++++++++++ airflow/providers/fab/alembic/script.py.mako | 26 ++++ .../fab/alembic/versions/__init__.py | 16 +++ .../fab/auth_manager/models/__init__.py | 12 +- .../providers/fab/auth_manager/models/db.py | 34 +++++ airflow/utils/db.py | 87 +++++++++++- 12 files changed, 539 insertions(+), 23 deletions(-) create mode 100644 airflow/providers/fab/alembic.ini create mode 100644 airflow/providers/fab/alembic/README create mode 100644 airflow/providers/fab/alembic/__init__.py create mode 100644 airflow/providers/fab/alembic/env.py create mode 100644 airflow/providers/fab/alembic/script.py.mako create mode 100644 airflow/providers/fab/alembic/versions/__init__.py create mode 100644 airflow/providers/fab/auth_manager/models/db.py 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..27541b1591be0 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -26,19 +26,6 @@ from airflow import models, settings from airflow.utils.db import compare_server_default, compare_type - -def include_object(_, name, type_, *args): - """Filter objects for autogenerating revisions.""" - # 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"): - return False - else: - return True - - # Make sure everything is imported so that alembic can find it all models.import_all_models() @@ -61,6 +48,21 @@ def include_object(_, name, type_, *args): # my_important_option = config.get_main_option("my_important_option") # ... etc. +# version table +version_table = "alembic_version" + + +def include_object(_, name, type_, *args): + """Filter objects for autogenerating revisions.""" + # 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"): + return False + else: + return True + def run_migrations_offline(): """ @@ -82,6 +84,7 @@ def run_migrations_offline(): compare_type=compare_type, compare_server_default=compare_server_default, render_as_batch=True, + version_table=version_table, ) with context.begin_transaction(): @@ -119,6 +122,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/models/base.py b/airflow/models/base.py index e9f86f8d7e672..8c59c469af985 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -22,7 +22,9 @@ from sqlalchemy import Column, Integer, MetaData, String, text from sqlalchemy.orm import registry +from airflow import settings from airflow.configuration import conf +from airflow.utils.log.logging_mixin import LoggingMixin SQL_ALCHEMY_SCHEMA = conf.get("database", "SQL_ALCHEMY_SCHEMA") @@ -38,13 +40,13 @@ } -def _get_schema(): +def get_schema(): if not SQL_ALCHEMY_SCHEMA or SQL_ALCHEMY_SCHEMA.isspace(): return None return SQL_ALCHEMY_SCHEMA -metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) +metadata = MetaData(schema=get_schema(), naming_convention=naming_convention) mapper_registry = registry(metadata=metadata) _sentinel = object() @@ -94,3 +96,72 @@ class TaskInstanceDependencies(Base): dag_id = Column(StringID(), nullable=False) run_id = Column(StringID(), nullable=False) map_index = Column(Integer, nullable=False, server_default=text("-1")) + + +class AttributeCheckerMeta(type): + """Metaclass to check attributes of subclasses.""" + + def __new__(cls, name, bases, dct): + """Check that subclasses are setting the required attributes.""" + required_attrs = ["metadata", "migration_dir", "alembic_file", "version_table_name"] + for attr in required_attrs: + if attr not in dct: + raise AttributeError(f"{name} is missing required attribute: {attr}") + return super().__new__(cls, name, bases, dct) + + +class BaseDBManager(LoggingMixin, metaclass=AttributeCheckerMeta): + """Base DB manager for external DBs.""" + + metadata: MetaData = None + migration_dir: str = "" + alembic_file: str = "" + version_table_name: str = "" + + def __init__(self, session): + super().__init__() + self.session = session + + def get_alembic_config(self): + from alembic.config import Config + + 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.""" + from alembic import command + + 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): + """Upgrade the database.""" + from alembic import command + + config = self.get_alembic_config() + command.upgrade(config, "heads") + + def downgradedb(self): + """Downgrade the database.""" diff --git a/airflow/providers/fab/alembic.ini b/airflow/providers/fab/alembic.ini new file mode 100644 index 0000000000000..877cab3a1e82d --- /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 = alembic + +# 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/alembic/README b/airflow/providers/fab/alembic/README new file mode 100644 index 0000000000000..2500aa1bcf726 --- /dev/null +++ b/airflow/providers/fab/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. diff --git a/airflow/providers/fab/alembic/__init__.py b/airflow/providers/fab/alembic/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/fab/alembic/__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/alembic/env.py b/airflow/providers/fab/alembic/env.py new file mode 100644 index 0000000000000..9e52dd44f8f42 --- /dev/null +++ b/airflow/providers/fab/alembic/env.py @@ -0,0 +1,124 @@ +# 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 import metadata +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) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = 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, + ) + + 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/alembic/script.py.mako b/airflow/providers/fab/alembic/script.py.mako new file mode 100644 index 0000000000000..fbc4b07dcef98 --- /dev/null +++ b/airflow/providers/fab/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${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/alembic/versions/__init__.py b/airflow/providers/fab/alembic/versions/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/fab/alembic/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/providers/fab/auth_manager/models/__init__.py b/airflow/providers/fab/auth_manager/models/__init__.py index bf4e43f275fab..f9f33f81dc31c 100644 --- a/airflow/providers/fab/auth_manager/models/__init__.py +++ b/airflow/providers/fab/auth_manager/models/__init__.py @@ -32,6 +32,7 @@ ForeignKey, Index, Integer, + MetaData, String, Table, UniqueConstraint, @@ -39,16 +40,19 @@ func, select, ) -from sqlalchemy.orm import backref, declared_attr, relationship +from sqlalchemy.orm import backref, declared_attr, registry, relationship from airflow.auth.managers.models.base_user import BaseUser -from airflow.models.base import Base +from airflow.models.base import get_schema, naming_convention """ Compatibility note: The models in this file are duplicated from Flask AppBuilder. """ -# Use airflow metadata to create the tables -Model.metadata = Base.metadata + +metadata = MetaData(schema=get_schema(), naming_convention=naming_convention) +mapper_registry = registry(metadata=metadata) + +Model.metadata = metadata if TYPE_CHECKING: try: 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..7daf71c0da582 --- /dev/null +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -0,0 +1,34 @@ +# 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.models.base import BaseDBManager +from airflow.providers.fab.auth_manager.models import metadata + +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/alembic") + alembic_file = os.path.join(PACKAGE_DIR, "providers/fab/alembic.ini") diff --git a/airflow/utils/db.py b/airflow/utils/db.py index a86ca5dbf5628..610c3d1510a57 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.module_loading import import_string # 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,8 @@ def _create_flask_session_tbl(sql_database_uri): @provide_session def initdb(session: Session = NEW_SESSION, load_connections: bool = True): """Initialize Airflow database.""" + external_db_manager = ExternalDBManager() + external_db_manager.validate() import_all_models() db_exists = _get_current_revision(session) @@ -774,6 +775,7 @@ 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 @@ -1650,6 +1652,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") @@ -1745,10 +1748,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 @@ -2112,3 +2113,81 @@ def _coerce_slice(key: slice) -> tuple[int, int | None, bool]: else: raise ValueError("non-trivial slice step not supported") return _coerce_index(key.start) or 0, _coerce_index(key.stop), reverse + + +class ExternalDBManager: + """ + External DB Managers. + + This class is a container for external database managers. + """ + + def __init__(self): + self._managers = [] + managers = conf.get("database", "external_db_managers", fallback="").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: + self._validate(manager) + + def _validate(self, manager): + """ + Validate the external database migration. + + :param cls: External database class + """ + 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 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 _create_db_from_orm(self, session): + """Create the external database from ORM.""" + for manager in self._managers: + m = manager(session) + m.create_db_from_orm() From 6e55b5af1c0db31a41b8b532d9b5af20a5d45a28 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 14 Aug 2024 01:04:56 +0100 Subject: [PATCH 02/30] minor fixes and use pre-commit on fab migrations --- .pre-commit-config.yaml | 3 +- airflow/utils/db.py | 22 ++---- .../in_container/run_migration_reference.py | 78 ++++++++++++------- 3 files changed, 61 insertions(+), 42 deletions(-) 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/utils/db.py b/airflow/utils/db.py index 610c3d1510a57..28e42d50f751f 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -79,6 +79,7 @@ from sqlalchemy.sql.elements import ClauseElement, TextClause from sqlalchemy.sql.selectable import Select + from airflow.models.base import BaseDBManager from airflow.models.connection import Connection from airflow.typing_compat import Self @@ -2123,8 +2124,8 @@ class ExternalDBManager: """ def __init__(self): - self._managers = [] - managers = conf.get("database", "external_db_managers", fallback="").split(",") + self._managers: list[type[BaseDBManager]] = [] + managers = conf.get("database", "external_db_managers").split(",") for module in managers: manager = import_string(module) self._managers.append(manager) @@ -2132,14 +2133,11 @@ def __init__(self): def validate(self): """Validate the external database managers.""" for manager in self._managers: - self._validate(manager) + ExternalDBManager._validate(manager) - def _validate(self, manager): - """ - Validate the external database migration. - - :param cls: External database class - """ + @staticmethod + def _validate(manager: type[BaseDBManager]): + """Validate the external database migration.""" import ast from airflow.models.base import metadata as airflow_metadata @@ -2185,9 +2183,3 @@ def downgradedb(self, session): for manager in self._managers: m = manager(session) m.downgradedb() - - def _create_db_from_orm(self, session): - """Create the external database from ORM.""" - for manager in self._managers: - m = manager(session) - m.create_db_from_orm() 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") From 39bb2a15085a01dafc17f0a11eef61eb56dae486 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 14 Aug 2024 01:30:02 +0100 Subject: [PATCH 03/30] support external db managers to drop tables when airflow db reset is run --- airflow/models/base.py | 2 ++ airflow/providers/fab/auth_manager/models/db.py | 1 + airflow/utils/db.py | 15 +++++++++++---- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/airflow/models/base.py b/airflow/models/base.py index 8c59c469af985..a2fa31bccf769 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -117,6 +117,8 @@ class BaseDBManager(LoggingMixin, metaclass=AttributeCheckerMeta): 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__() diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index 7daf71c0da582..9b5b17ef81e8c 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -32,3 +32,4 @@ class FABDBManager(BaseDBManager): version_table_name = "fab_alembic_version" migration_dir = os.path.join(PACKAGE_DIR, "providers/fab/alembic") alembic_file = os.path.join(PACKAGE_DIR, "providers/fab/alembic.ini") + supports_table_dropping = True diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 28e42d50f751f..cf23b1be98c22 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -767,7 +767,7 @@ def _create_flask_session_tbl(sql_database_uri): @provide_session def initdb(session: Session = NEW_SESSION, load_connections: bool = True): """Initialize Airflow database.""" - external_db_manager = ExternalDBManager() + external_db_manager = RunDBManager() external_db_manager.validate() import_all_models() @@ -1683,6 +1683,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) @@ -2116,9 +2118,9 @@ def _coerce_slice(key: slice) -> tuple[int, int | None, bool]: return _coerce_index(key.start) or 0, _coerce_index(key.stop), reverse -class ExternalDBManager: +class RunDBManager: """ - External DB Managers. + Run External DB Managers. This class is a container for external database managers. """ @@ -2133,7 +2135,7 @@ def __init__(self): def validate(self): """Validate the external database managers.""" for manager in self._managers: - ExternalDBManager._validate(manager) + RunDBManager._validate(manager) @staticmethod def _validate(manager: type[BaseDBManager]): @@ -2183,3 +2185,8 @@ def downgradedb(self, session): 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: + manager.metadata.drop_all(connection) From feac59691a6b6a66c7252aeed88f9459c4736c76 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 14 Aug 2024 16:31:04 +0100 Subject: [PATCH 04/30] Add tests --- airflow/providers/fab/alembic/env.py | 1 + airflow/utils/db.py | 2 +- tests/always/test_project_structure.py | 1 + tests/models/test_base.py | 46 +++++++++++- .../fab/auth_manager/models/test_db.py | 71 +++++++++++++++++++ tests/utils/test_db.py | 32 +++++++++ 6 files changed, 151 insertions(+), 2 deletions(-) create mode 100644 tests/providers/fab/auth_manager/models/test_db.py diff --git a/airflow/providers/fab/alembic/env.py b/airflow/providers/fab/alembic/env.py index 9e52dd44f8f42..0b91e54ade399 100644 --- a/airflow/providers/fab/alembic/env.py +++ b/airflow/providers/fab/alembic/env.py @@ -74,6 +74,7 @@ def run_migrations_offline(): compare_server_default=True, render_as_batch=True, version_table=version_table, + include_object=include_object, ) with context.begin_transaction(): diff --git a/airflow/utils/db.py b/airflow/utils/db.py index cf23b1be98c22..754bb48ee4cc8 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -2149,7 +2149,7 @@ def _validate(manager: type[BaseDBManager]): # 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 Airflow metadata") + 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") diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index 15813ca9ca296..f9487c13e224b 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/alembic/test_env.py", ] # TODO: Should we extend this test to cover other directories? diff --git a/tests/models/test_base.py b/tests/models/test_base.py index 3224616404d6a..5da936f9f6650 100644 --- a/tests/models/test_base.py +++ b/tests/models/test_base.py @@ -18,7 +18,7 @@ import pytest -from airflow.models.base import get_id_collation_args +from airflow.models.base import Base, BaseDBManager, get_id_collation_args from tests.test_utils.config import conf_vars pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -47,3 +47,47 @@ def test_collation(dsn, expected, extra): with conf_vars({("database", "sql_alchemy_conn"): dsn, **extra}): assert expected == get_id_collation_args() + + +def test_subclassing_db_manager_with_missing_attrs(): + """Test subclassing BaseDBManager.""" + + with pytest.raises(AttributeError, match="SubclassDBManager is missing required attribute: metadata"): + + class SubclassDBManager(BaseDBManager): ... + + +def test_subclassing_db_manager_with_set_metadata(): + with pytest.raises( + AttributeError, match="SubclassDbManager is missing required attribute: migration_dir" + ): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + + +def test_subclassing_db_manager_with_set_metadata_and_migration_dir(): + with pytest.raises(AttributeError, match="SubclassDbManager is missing required attribute: alembic_file"): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + + +def test_subclassing_db_manager_with_attrs_set_except_version_table_name(): + with pytest.raises( + AttributeError, match="SubclassDbManager is missing required attribute: version_table_name" + ): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + alembic_file = "some_file" + + +def test_subclassing_db_manager_with_attrs_set_dont_raise(session): + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + alembic_file = "some_file" + version_table_name = "some_table" 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..067d9b7f98801 --- /dev/null +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -0,0 +1,71 @@ +# 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 alembic.autogenerate import compare_metadata +from alembic.migration import MigrationContext +from sqlalchemy import MetaData + +import airflow +from airflow.providers.fab.auth_manager.models.db import FABDBManager +from airflow.settings import engine +from airflow.utils.db import ( + compare_server_default, + compare_type, +) + + +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/alembic" + + 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) diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 287a5a1fdbdb0..7e289614120ed 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -34,9 +34,11 @@ from sqlalchemy import MetaData, Table from sqlalchemy.sql import Select +from airflow.exceptions import AirflowException from airflow.models import Base as airflow_base from airflow.settings import engine from airflow.utils.db import ( + RunDBManager, _get_alembic_config, check_bad_references, check_migrations, @@ -57,12 +59,20 @@ class TestDb: def test_database_schema_and_sqlalchemy_model_are_in_sync(self): + external_db_managers = RunDBManager() + import airflow.models + from airflow.utils.db import _get_flask_db 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 + 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 +80,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 @@ -307,3 +318,24 @@ def test_check_bad_references( mock_session, task_fail_table, mock_select, dangling_task_fail_table_name ) mock_session.rollback.assert_called_once() + + +class TestRunDBManager: + 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] + + def test_defining_table_same_name_as_airflow_table_name_raises(self): + from sqlalchemy import Column, Integer, String + + run_db_manager = RunDBManager() + manager = run_db_manager._managers[0] + # Add dag_run table to metadata + mytable = Table( + "dag_run", manager.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) + ) + manager.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() From 37fe5816b6a9908a36d56e4f5481a335c77c6c22 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 09:05:02 +0100 Subject: [PATCH 05/30] fix autogenerate and update function signatures --- airflow/migrations/env.py | 2 +- airflow/models/base.py | 6 +++--- airflow/providers/fab/alembic/env.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index 27541b1591be0..e1531e6722e64 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -58,7 +58,7 @@ def include_object(_, name, type_, *args): 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"): + if type_ == "table" and name not in target_metadata.tables: return False else: return True diff --git a/airflow/models/base.py b/airflow/models/base.py index a2fa31bccf769..64b65428bfcdc 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -158,12 +158,12 @@ def initdb(self): else: self._create_db_from_orm() - def upgradedb(self): + def upgradedb(self, to_version=None, from_version=None, show_sql_only=False): """Upgrade the database.""" from alembic import command config = self.get_alembic_config() - command.upgrade(config, "heads") + command.upgrade(config, revision=to_version or "heads", sql=show_sql_only) - def downgradedb(self): + def downgradedb(self, to_version, from_version=None, show_sql_only=False): """Downgrade the database.""" diff --git a/airflow/providers/fab/alembic/env.py b/airflow/providers/fab/alembic/env.py index 0b91e54ade399..27be8fd17f49b 100644 --- a/airflow/providers/fab/alembic/env.py +++ b/airflow/providers/fab/alembic/env.py @@ -22,12 +22,12 @@ from alembic import context from airflow import settings -from airflow.providers.fab.auth_manager.models import metadata 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. @@ -39,7 +39,7 @@ # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata -target_metadata = metadata +target_metadata = FABDBManager.metadata # other values from the config, defined by the needs of env.py, # can be acquired: From 58e86a2e83fe64815d48bdecb1326feb8eadc543 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 10:00:38 +0100 Subject: [PATCH 06/30] Rename alembic directory to migrations --- airflow/migrations/env.py | 25 ++++++++++--------- airflow/providers/fab/alembic.ini | 2 +- .../providers/fab/auth_manager/models/db.py | 2 +- .../fab/{alembic => migrations}/README | 0 .../fab/{alembic => migrations}/__init__.py | 0 .../fab/{alembic => migrations}/env.py | 0 .../{alembic => migrations}/script.py.mako | 0 .../versions/__init__.py | 0 .../fab/auth_manager/models/test_db.py | 2 +- 9 files changed, 16 insertions(+), 15 deletions(-) rename airflow/providers/fab/{alembic => migrations}/README (100%) rename airflow/providers/fab/{alembic => migrations}/__init__.py (100%) rename airflow/providers/fab/{alembic => migrations}/env.py (100%) rename airflow/providers/fab/{alembic => migrations}/script.py.mako (100%) rename airflow/providers/fab/{alembic => migrations}/versions/__init__.py (100%) diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index e1531e6722e64..9b1c88d99a907 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -26,6 +26,19 @@ from airflow import models, settings from airflow.utils.db import compare_server_default, compare_type + +def include_object(_, name, type_, *args): + """Filter objects for autogenerating revisions.""" + # 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 not in target_metadata.tables: + return False + else: + return True + + # Make sure everything is imported so that alembic can find it all models.import_all_models() @@ -52,18 +65,6 @@ version_table = "alembic_version" -def include_object(_, name, type_, *args): - """Filter objects for autogenerating revisions.""" - # 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 not in target_metadata.tables: - return False - else: - return True - - def run_migrations_offline(): """ Run migrations in 'offline' mode. diff --git a/airflow/providers/fab/alembic.ini b/airflow/providers/fab/alembic.ini index 877cab3a1e82d..55c7338f9bd9f 100644 --- a/airflow/providers/fab/alembic.ini +++ b/airflow/providers/fab/alembic.ini @@ -20,7 +20,7 @@ [alembic] # path to migration scripts # Use forward slashes (/) also on windows to provide an os agnostic path -script_location = alembic +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 diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index 9b5b17ef81e8c..b1c8067e61f7c 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -30,6 +30,6 @@ class FABDBManager(BaseDBManager): metadata = metadata version_table_name = "fab_alembic_version" - migration_dir = os.path.join(PACKAGE_DIR, "providers/fab/alembic") + 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/alembic/README b/airflow/providers/fab/migrations/README similarity index 100% rename from airflow/providers/fab/alembic/README rename to airflow/providers/fab/migrations/README diff --git a/airflow/providers/fab/alembic/__init__.py b/airflow/providers/fab/migrations/__init__.py similarity index 100% rename from airflow/providers/fab/alembic/__init__.py rename to airflow/providers/fab/migrations/__init__.py diff --git a/airflow/providers/fab/alembic/env.py b/airflow/providers/fab/migrations/env.py similarity index 100% rename from airflow/providers/fab/alembic/env.py rename to airflow/providers/fab/migrations/env.py diff --git a/airflow/providers/fab/alembic/script.py.mako b/airflow/providers/fab/migrations/script.py.mako similarity index 100% rename from airflow/providers/fab/alembic/script.py.mako rename to airflow/providers/fab/migrations/script.py.mako diff --git a/airflow/providers/fab/alembic/versions/__init__.py b/airflow/providers/fab/migrations/versions/__init__.py similarity index 100% rename from airflow/providers/fab/alembic/versions/__init__.py rename to airflow/providers/fab/migrations/versions/__init__.py diff --git a/tests/providers/fab/auth_manager/models/test_db.py b/tests/providers/fab/auth_manager/models/test_db.py index 067d9b7f98801..e703a94e6c382 100644 --- a/tests/providers/fab/auth_manager/models/test_db.py +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -40,7 +40,7 @@ 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/alembic" + 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" From d9172128b4aa71f5e9a4ee9bafcb8956675613bb Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 11:50:47 +0100 Subject: [PATCH 07/30] add more tests --- airflow/models/base.py | 5 +++++ airflow/utils/db.py | 1 - tests/utils/test_db.py | 35 ++++++++++++++++++++++++++++++++++- 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/airflow/models/base.py b/airflow/models/base.py index 64b65428bfcdc..9541e4c82ced6 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -46,6 +46,11 @@ def get_schema(): return SQL_ALCHEMY_SCHEMA +def _get_schema(): + # Backcompat + return get_schema() + + metadata = MetaData(schema=get_schema(), naming_convention=naming_convention) mapper_registry = registry(metadata=metadata) _sentinel = object() diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 754bb48ee4cc8..a35a5f4d46498 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -1639,7 +1639,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 diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 7e289614120ed..1d78b9ea57523 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -49,6 +49,7 @@ # guard against removing it from utils.db accidentally create_session, # noqa: F401 downgrade, + initdb, resetdb, upgradedb, ) @@ -62,7 +63,6 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): external_db_managers = RunDBManager() import airflow.models - from airflow.utils.db import _get_flask_db airflow.models.import_all_models() all_meta_data = MetaData() @@ -339,3 +339,36 @@ def test_defining_table_same_name_as_airflow_table_name_raises(self): manager.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() + + @mock.patch("airflow.utils.db.RunDBManager") + def test_init_db_calls_rundbmanager(self, mock_rundbmanager, session): + initdb(session=session) + mock_rundbmanager.return_value.initdb.assert_called_once_with(session) + mock_rundbmanager.return_value.upgradedb.assert_not_called() + mock_rundbmanager.return_value.downgrade.assert_not_called() + + @mock.patch("airflow.utils.db.RunDBManager") + @mock.patch("alembic.command") + def test_upgradedb_or_downgrade_dont_call_rundbmanager( + self, mock_alembic_command, mock_rundbmanager, 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_rundbmanager.return_value.initdb.assert_not_called() + mock_rundbmanager.return_value.upgradedb.assert_not_called() + mock_rundbmanager.return_value.downgrade.assert_not_called() + + @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") + def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session): + ext_db = RunDBManager() + # initdb + ext_db.initdb(session=session) + mock_fabdb_manager.return_value.initdb.assert_called_once() + # upgradedb + ext_db.upgradedb(session=session) + mock_fabdb_manager.return_value.upgradedb.assert_called_once() + # downgradedb + ext_db.downgradedb(session=session) + mock_fabdb_manager.return_value.downgradedb.assert_called_once() From 85a38c2cd09755baf03fc347d61a9b650beb47ae Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 12:25:05 +0100 Subject: [PATCH 08/30] Fix import --- airflow/models/base.py | 9 ++------- airflow/providers/fab/auth_manager/models/__init__.py | 4 ++-- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/airflow/models/base.py b/airflow/models/base.py index 9541e4c82ced6..be4fcfa959345 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -40,18 +40,13 @@ } -def get_schema(): +def _get_schema(): if not SQL_ALCHEMY_SCHEMA or SQL_ALCHEMY_SCHEMA.isspace(): return None return SQL_ALCHEMY_SCHEMA -def _get_schema(): - # Backcompat - return get_schema() - - -metadata = MetaData(schema=get_schema(), naming_convention=naming_convention) +metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) mapper_registry = registry(metadata=metadata) _sentinel = object() diff --git a/airflow/providers/fab/auth_manager/models/__init__.py b/airflow/providers/fab/auth_manager/models/__init__.py index f9f33f81dc31c..9dc5c75747849 100644 --- a/airflow/providers/fab/auth_manager/models/__init__.py +++ b/airflow/providers/fab/auth_manager/models/__init__.py @@ -43,13 +43,13 @@ from sqlalchemy.orm import backref, declared_attr, registry, relationship from airflow.auth.managers.models.base_user import BaseUser -from airflow.models.base import get_schema, naming_convention +from airflow.models.base import _get_schema, naming_convention """ Compatibility note: The models in this file are duplicated from Flask AppBuilder. """ -metadata = MetaData(schema=get_schema(), naming_convention=naming_convention) +metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) mapper_registry = registry(metadata=metadata) Model.metadata = metadata From 481413f7ea1dcaec4bcf54905de917725cf98f14 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 12:44:16 +0100 Subject: [PATCH 09/30] Fix static check --- .../providers/fab/migrations/script.py.mako | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/airflow/providers/fab/migrations/script.py.mako b/airflow/providers/fab/migrations/script.py.mako index fbc4b07dcef98..4d0928fcc09ad 100644 --- a/airflow/providers/fab/migrations/script.py.mako +++ b/airflow/providers/fab/migrations/script.py.mako @@ -1,3 +1,21 @@ +# +# 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} From 3cd8577c16cba1435f5b2a789014506929a5935d Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 13:49:02 +0100 Subject: [PATCH 10/30] Fix tests --- tests/always/test_project_structure.py | 2 +- tests/providers/fab/auth_manager/test_security.py | 9 +-------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index f9487c13e224b..7bb664b068a3a 100644 --- a/tests/always/test_project_structure.py +++ b/tests/always/test_project_structure.py @@ -155,7 +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/alembic/test_env.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/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) From 554b11fec10024042dc3ad7c8bd99adc6286a430 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 14:35:29 +0100 Subject: [PATCH 11/30] fixup! Fix tests --- airflow/migrations/env.py | 3 ++- airflow/providers/fab/auth_manager/models/db.py | 1 + tests/utils/test_db.py | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index 9b1c88d99a907..4be8ac32ca252 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -32,7 +32,7 @@ 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 + # Only create migrations for objects that are in the target metadata if type_ == "table" and name not in target_metadata.tables: return False else: @@ -85,6 +85,7 @@ 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, ) diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index b1c8067e61f7c..c9606ceacfb9d 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -33,3 +33,4 @@ class FABDBManager(BaseDBManager): 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/tests/utils/test_db.py b/tests/utils/test_db.py index 1d78b9ea57523..7f6ac33f61554 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -97,6 +97,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: From 018db91c0eed6cd54651e3794dd8eeb503ea5e29 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Aug 2024 16:40:18 +0100 Subject: [PATCH 12/30] fixup! fixup! Fix tests --- airflow/providers/fab/auth_manager/models/db.py | 1 - airflow/utils/db.py | 4 +++- tests/providers/fab/auth_manager/models/test_db.py | 3 +++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index c9606ceacfb9d..b1c8067e61f7c 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -33,4 +33,3 @@ class FABDBManager(BaseDBManager): 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/utils/db.py b/airflow/utils/db.py index a35a5f4d46498..092857cfdc909 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.log.logging_mixin import LoggingMixin from airflow.utils.module_loading import import_string # TODO: remove create_session once we decide to break backward compatibility @@ -2117,7 +2118,7 @@ def _coerce_slice(key: slice) -> tuple[int, int | None, bool]: return _coerce_index(key.start) or 0, _coerce_index(key.stop), reverse -class RunDBManager: +class RunDBManager(LoggingMixin): """ Run External DB Managers. @@ -2125,6 +2126,7 @@ class RunDBManager: """ def __init__(self): + super().__init__() self._managers: list[type[BaseDBManager]] = [] managers = conf.get("database", "external_db_managers").split(",") for module in managers: diff --git a/tests/providers/fab/auth_manager/models/test_db.py b/tests/providers/fab/auth_manager/models/test_db.py index e703a94e6c382..4b4b621d0cc60 100644 --- a/tests/providers/fab/auth_manager/models/test_db.py +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -18,6 +18,7 @@ import os +import pytest from alembic.autogenerate import compare_metadata from alembic.migration import MigrationContext from sqlalchemy import MetaData @@ -30,6 +31,8 @@ compare_type, ) +pytestmark = [pytest.mark.db_test] + class TestFABDBManager: def setup_method(self, session): From 17457ed14cafbf7f5c569879749c1b3e74c030be Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 19 Aug 2024 12:56:26 +0100 Subject: [PATCH 13/30] Mark extra operator links test as non-db test --- tests/serialization/test_dag_serialization.py | 1 - 1 file changed, 1 deletion(-) 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, From b7fe8601a85f7fa0607bd56ca43425e6ee9e72f9 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 19 Aug 2024 13:01:18 +0100 Subject: [PATCH 14/30] update erd diagram --- docs/apache-airflow/img/airflow_erd.sha256 | 2 +- docs/apache-airflow/img/airflow_erd.svg | 3494 ++++++++++---------- 2 files changed, 1748 insertions(+), 1748 deletions(-) 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} From ea816cdf0bd9b058e0610c31fd8ab12984d685b6 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 19 Aug 2024 15:06:34 +0100 Subject: [PATCH 15/30] import airflow settings at point of use --- airflow/models/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/models/base.py b/airflow/models/base.py index be4fcfa959345..a716ed3d0e29c 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -22,7 +22,6 @@ from sqlalchemy import Column, Integer, MetaData, String, text from sqlalchemy.orm import registry -from airflow import settings from airflow.configuration import conf from airflow.utils.log.logging_mixin import LoggingMixin @@ -127,6 +126,8 @@ def __init__(self, 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("%", "%%")) From df10d288caeaed79439a3630247103f44f3b73f2 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 19 Aug 2024 15:39:39 +0100 Subject: [PATCH 16/30] skip import error for alembic env.py --- scripts/in_container/verify_providers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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( From 0eef772853bfa93861272d5e9b536844cff06380 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 20 Aug 2024 09:15:59 +0100 Subject: [PATCH 17/30] Remove type on typing --- airflow/utils/db.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 092857cfdc909..b59ae34801c5b 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -2127,7 +2127,7 @@ class RunDBManager(LoggingMixin): def __init__(self): super().__init__() - self._managers: list[type[BaseDBManager]] = [] + self._managers: list[BaseDBManager] = [] managers = conf.get("database", "external_db_managers").split(",") for module in managers: manager = import_string(module) @@ -2139,7 +2139,7 @@ def validate(self): RunDBManager._validate(manager) @staticmethod - def _validate(manager: type[BaseDBManager]): + def _validate(manager: BaseDBManager): """Validate the external database migration.""" import ast From 376c24d048e3872abc9e53be77e2257acd32d2c5 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 20 Aug 2024 11:47:18 +0100 Subject: [PATCH 18/30] Move code around --- airflow/models/base.py | 74 -------- .../providers/fab/auth_manager/models/db.py | 5 +- airflow/utils/db.py | 83 +-------- airflow/utils/db_manager.py | 176 ++++++++++++++++++ tests/models/test_base.py | 46 +---- tests/utils/test_db.py | 56 ------ tests/utils/test_db_manager.py | 125 +++++++++++++ 7 files changed, 309 insertions(+), 256 deletions(-) create mode 100644 airflow/utils/db_manager.py create mode 100644 tests/utils/test_db_manager.py diff --git a/airflow/models/base.py b/airflow/models/base.py index a716ed3d0e29c..e9f86f8d7e672 100644 --- a/airflow/models/base.py +++ b/airflow/models/base.py @@ -23,7 +23,6 @@ from sqlalchemy.orm import registry from airflow.configuration import conf -from airflow.utils.log.logging_mixin import LoggingMixin SQL_ALCHEMY_SCHEMA = conf.get("database", "SQL_ALCHEMY_SCHEMA") @@ -95,76 +94,3 @@ class TaskInstanceDependencies(Base): dag_id = Column(StringID(), nullable=False) run_id = Column(StringID(), nullable=False) map_index = Column(Integer, nullable=False, server_default=text("-1")) - - -class AttributeCheckerMeta(type): - """Metaclass to check attributes of subclasses.""" - - def __new__(cls, name, bases, dct): - """Check that subclasses are setting the required attributes.""" - required_attrs = ["metadata", "migration_dir", "alembic_file", "version_table_name"] - for attr in required_attrs: - if attr not in dct: - raise AttributeError(f"{name} is missing required attribute: {attr}") - return super().__new__(cls, name, bases, dct) - - -class BaseDBManager(LoggingMixin, metaclass=AttributeCheckerMeta): - """Base DB manager for external DBs.""" - - metadata: MetaData = None - 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.""" - from alembic import command - - 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.""" - from alembic import command - - config = self.get_alembic_config() - command.upgrade(config, revision=to_version or "heads", sql=show_sql_only) - - def downgradedb(self, to_version, from_version=None, show_sql_only=False): - """Downgrade the database.""" diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index b1c8067e61f7c..e90efd557b0ad 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -19,8 +19,8 @@ import os import airflow -from airflow.models.base import BaseDBManager from airflow.providers.fab.auth_manager.models import metadata +from airflow.utils.db_manager import BaseDBManager PACKAGE_DIR = os.path.dirname(airflow.__file__) @@ -33,3 +33,6 @@ class FABDBManager(BaseDBManager): 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 + + def __init__(self, session): + self.session = session diff --git a/airflow/utils/db.py b/airflow/utils/db.py index b59ae34801c5b..aae973288b7bd 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -65,8 +65,7 @@ from airflow.exceptions import AirflowException from airflow.models import import_all_models from airflow.utils import helpers -from airflow.utils.log.logging_mixin import LoggingMixin -from airflow.utils.module_loading import import_string +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 @@ -80,7 +79,6 @@ from sqlalchemy.sql.elements import ClauseElement, TextClause from sqlalchemy.sql.selectable import Select - from airflow.models.base import BaseDBManager from airflow.models.connection import Connection from airflow.typing_compat import Self @@ -768,8 +766,6 @@ def _create_flask_session_tbl(sql_database_uri): @provide_session def initdb(session: Session = NEW_SESSION, load_connections: bool = True): """Initialize Airflow database.""" - external_db_manager = RunDBManager() - external_db_manager.validate() import_all_models() db_exists = _get_current_revision(session) @@ -777,6 +773,8 @@ def initdb(session: Session = NEW_SESSION, load_connections: bool = True): upgradedb(session=session) else: _create_db_from_orm(session=session) + external_db_manager = RunDBManager() + external_db_manager.validate() external_db_manager.initdb(session) if conf.getboolean("database", "LOAD_DEFAULT_CONNECTIONS") and load_connections: create_default_connections(session=session) @@ -2116,78 +2114,3 @@ def _coerce_slice(key: slice) -> tuple[int, int | None, bool]: else: raise ValueError("non-trivial slice step not supported") return _coerce_index(key.start) or 0, _coerce_index(key.stop), reverse - - -class RunDBManager(LoggingMixin): - """ - Run External DB Managers. - - This class is a container for 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: - manager.metadata.drop_all(connection) diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py new file mode 100644 index 0000000000000..44b440e55f0d0 --- /dev/null +++ b/airflow/utils/db_manager.py @@ -0,0 +1,176 @@ +# 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 typing import TYPE_CHECKING + +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 AttributeCheckerMeta(type): + """Metaclass to check attributes of subclasses.""" + + def __new__(cls, name, bases, dct): + """Check that subclasses are setting the required attributes.""" + required_attrs = ["metadata", "migration_dir", "alembic_file", "version_table_name"] + for attr in required_attrs: + if attr not in dct: + raise AttributeError(f"{name} is missing required attribute: {attr}") + return super().__new__(cls, name, bases, dct) + + +class BaseDBManager(LoggingMixin, metaclass=AttributeCheckerMeta): + """Base DB manager for external DBs.""" + + metadata: MetaData = None + 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.""" + from alembic import command + + 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.""" + from alembic import command + + config = self.get_alembic_config() + command.upgrade(config, revision=to_version or "heads", sql=show_sql_only) + + def downgradedb(self, to_version, from_version=None, show_sql_only=False): + """Downgrade the database.""" + + +class RunDBManager(LoggingMixin): + """ + Run External DB Managers. + + This class is a container for external database managers. + """ + + def __init__(self): + super().__init__() + self._managers = [] + 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): + """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: + manager.metadata.drop_all(connection) diff --git a/tests/models/test_base.py b/tests/models/test_base.py index 5da936f9f6650..3224616404d6a 100644 --- a/tests/models/test_base.py +++ b/tests/models/test_base.py @@ -18,7 +18,7 @@ import pytest -from airflow.models.base import Base, BaseDBManager, get_id_collation_args +from airflow.models.base import get_id_collation_args from tests.test_utils.config import conf_vars pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -47,47 +47,3 @@ def test_collation(dsn, expected, extra): with conf_vars({("database", "sql_alchemy_conn"): dsn, **extra}): assert expected == get_id_collation_args() - - -def test_subclassing_db_manager_with_missing_attrs(): - """Test subclassing BaseDBManager.""" - - with pytest.raises(AttributeError, match="SubclassDBManager is missing required attribute: metadata"): - - class SubclassDBManager(BaseDBManager): ... - - -def test_subclassing_db_manager_with_set_metadata(): - with pytest.raises( - AttributeError, match="SubclassDbManager is missing required attribute: migration_dir" - ): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - - -def test_subclassing_db_manager_with_set_metadata_and_migration_dir(): - with pytest.raises(AttributeError, match="SubclassDbManager is missing required attribute: alembic_file"): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" - - -def test_subclassing_db_manager_with_attrs_set_except_version_table_name(): - with pytest.raises( - AttributeError, match="SubclassDbManager is missing required attribute: version_table_name" - ): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" - alembic_file = "some_file" - - -def test_subclassing_db_manager_with_attrs_set_dont_raise(session): - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" - alembic_file = "some_file" - version_table_name = "some_table" diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 7f6ac33f61554..63f698235fb45 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -34,7 +34,6 @@ from sqlalchemy import MetaData, Table from sqlalchemy.sql import Select -from airflow.exceptions import AirflowException from airflow.models import Base as airflow_base from airflow.settings import engine from airflow.utils.db import ( @@ -49,7 +48,6 @@ # guard against removing it from utils.db accidentally create_session, # noqa: F401 downgrade, - initdb, resetdb, upgradedb, ) @@ -320,57 +318,3 @@ def test_check_bad_references( mock_session, task_fail_table, mock_select, dangling_task_fail_table_name ) mock_session.rollback.assert_called_once() - - -class TestRunDBManager: - 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] - - def test_defining_table_same_name_as_airflow_table_name_raises(self): - from sqlalchemy import Column, Integer, String - - run_db_manager = RunDBManager() - manager = run_db_manager._managers[0] - # Add dag_run table to metadata - mytable = Table( - "dag_run", manager.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) - ) - manager.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() - - @mock.patch("airflow.utils.db.RunDBManager") - def test_init_db_calls_rundbmanager(self, mock_rundbmanager, session): - initdb(session=session) - mock_rundbmanager.return_value.initdb.assert_called_once_with(session) - mock_rundbmanager.return_value.upgradedb.assert_not_called() - mock_rundbmanager.return_value.downgrade.assert_not_called() - - @mock.patch("airflow.utils.db.RunDBManager") - @mock.patch("alembic.command") - def test_upgradedb_or_downgrade_dont_call_rundbmanager( - self, mock_alembic_command, mock_rundbmanager, 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_rundbmanager.return_value.initdb.assert_not_called() - mock_rundbmanager.return_value.upgradedb.assert_not_called() - mock_rundbmanager.return_value.downgrade.assert_not_called() - - @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") - def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session): - ext_db = RunDBManager() - # initdb - ext_db.initdb(session=session) - mock_fabdb_manager.return_value.initdb.assert_called_once() - # upgradedb - ext_db.upgradedb(session=session) - mock_fabdb_manager.return_value.upgradedb.assert_called_once() - # downgradedb - ext_db.downgradedb(session=session) - mock_fabdb_manager.return_value.downgradedb.assert_called_once() diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py new file mode 100644 index 0000000000000..b9e7599af1a9d --- /dev/null +++ b/tests/utils/test_db_manager.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 + +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 + + +class TestRunDBManager: + 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] + + def test_defining_table_same_name_as_airflow_table_name_raises(self): + from sqlalchemy import Column, Integer, String + + run_db_manager = RunDBManager() + manager = run_db_manager._managers[0] + # Add dag_run table to metadata + mytable = Table( + "dag_run", manager.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) + ) + manager.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() + + @mock.patch("airflow.utils.db.RunDBManager") + def test_init_db_calls_rundbmanager(self, mock_rundbmanager, session): + initdb(session=session) + mock_rundbmanager.return_value.initdb.assert_called_once_with(session) + mock_rundbmanager.return_value.upgradedb.assert_not_called() + mock_rundbmanager.return_value.downgrade.assert_not_called() + + @mock.patch("airflow.utils.db.RunDBManager") + @mock.patch("alembic.command") + def test_upgradedb_or_downgrade_dont_call_rundbmanager( + self, mock_alembic_command, mock_rundbmanager, 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_rundbmanager.return_value.initdb.assert_not_called() + mock_rundbmanager.return_value.upgradedb.assert_not_called() + mock_rundbmanager.return_value.downgrade.assert_not_called() + + @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") + def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session): + ext_db = RunDBManager() + # initdb + ext_db.initdb(session=session) + mock_fabdb_manager.return_value.initdb.assert_called_once() + # upgradedb + ext_db.upgradedb(session=session) + mock_fabdb_manager.return_value.upgradedb.assert_called_once() + # downgradedb + ext_db.downgradedb(session=session) + mock_fabdb_manager.return_value.downgradedb.assert_called_once() + + +def test_subclassing_db_manager_with_missing_attrs(): + """Test subclassing BaseDBManager.""" + + with pytest.raises(AttributeError, match="SubclassDBManager is missing required attribute: metadata"): + + class SubclassDBManager(BaseDBManager): ... + + +def test_subclassing_db_manager_with_set_metadata(): + with pytest.raises( + AttributeError, match="SubclassDbManager is missing required attribute: migration_dir" + ): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + + +def test_subclassing_db_manager_with_set_metadata_and_migration_dir(): + with pytest.raises(AttributeError, match="SubclassDbManager is missing required attribute: alembic_file"): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + + +def test_subclassing_db_manager_with_attrs_set_except_version_table_name(): + with pytest.raises( + AttributeError, match="SubclassDbManager is missing required attribute: version_table_name" + ): + + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + alembic_file = "some_file" + + +def test_subclassing_db_manager_with_attrs_set_dont_raise(session): + class SubclassDbManager(BaseDBManager): + metadata = Base.metadata + migration_dir = "some_dir" + alembic_file = "some_file" + version_table_name = "some_table" From 79cf6db560451c2625febe91c78d6bb9af391d5c Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 20 Aug 2024 12:31:49 +0100 Subject: [PATCH 19/30] fixup! Move code around --- tests/utils/test_db.py | 3 +-- tests/utils/test_db_manager.py | 2 ++ 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 63f698235fb45..8f7abf5b325bc 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -58,8 +58,6 @@ class TestDb: def test_database_schema_and_sqlalchemy_model_are_in_sync(self): - external_db_managers = RunDBManager() - import airflow.models airflow.models.import_all_models() @@ -68,6 +66,7 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): 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) diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py index b9e7599af1a9d..73fa47715db70 100644 --- a/tests/utils/test_db_manager.py +++ b/tests/utils/test_db_manager.py @@ -26,6 +26,8 @@ from airflow.utils.db import downgrade, initdb, upgradedb from airflow.utils.db_manager import BaseDBManager, RunDBManager +pytestmark = [pytest.mark.db_test] + class TestRunDBManager: def test_fab_db_manager_is_default(self): From c898bad06b0135854cba9d3f69b2b6bad8a9e55d Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 20 Aug 2024 18:11:52 +0100 Subject: [PATCH 20/30] Move import --- airflow/utils/db.py | 5 ++++- tests/utils/test_db.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index aae973288b7bd..f35edf792aa1f 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -65,7 +65,6 @@ 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 @@ -773,6 +772,8 @@ def initdb(session: Session = NEW_SESSION, load_connections: bool = True): upgradedb(session=session) else: _create_db_from_orm(session=session) + from airflow.utils.db_manager import RunDBManager + external_db_manager = RunDBManager() external_db_manager.validate() external_db_manager.initdb(session) @@ -1681,6 +1682,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) + from airflow.utils.db_manager import RunDBManager + external_db_manager = RunDBManager() external_db_manager.drop_tables(connection) diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 8f7abf5b325bc..4d32d132b2c96 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -37,7 +37,6 @@ from airflow.models import Base as airflow_base from airflow.settings import engine from airflow.utils.db import ( - RunDBManager, _get_alembic_config, check_bad_references, check_migrations, @@ -51,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] From 422234488c401252015879fac19053f4685678a5 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 20 Aug 2024 19:10:12 +0100 Subject: [PATCH 21/30] fixup! Move import --- tests/utils/test_db_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py index 73fa47715db70..c919d904a48b5 100644 --- a/tests/utils/test_db_manager.py +++ b/tests/utils/test_db_manager.py @@ -49,14 +49,14 @@ def test_defining_table_same_name_as_airflow_table_name_raises(self): with pytest.raises(AirflowException, match="Table 'dag_run' already exists in the Airflow metadata"): run_db_manager.validate() - @mock.patch("airflow.utils.db.RunDBManager") + @mock.patch("airflow.utils.db_manager.RunDBManager") def test_init_db_calls_rundbmanager(self, mock_rundbmanager, session): initdb(session=session) mock_rundbmanager.return_value.initdb.assert_called_once_with(session) mock_rundbmanager.return_value.upgradedb.assert_not_called() mock_rundbmanager.return_value.downgrade.assert_not_called() - @mock.patch("airflow.utils.db.RunDBManager") + @mock.patch("airflow.utils.db_manager.RunDBManager") @mock.patch("alembic.command") def test_upgradedb_or_downgrade_dont_call_rundbmanager( self, mock_alembic_command, mock_rundbmanager, session From 2dcf266b601557604747ea7d651a66b1b784a17c Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 08:02:44 +0100 Subject: [PATCH 22/30] Don't disable existing loggers --- airflow/providers/fab/migrations/env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/fab/migrations/env.py b/airflow/providers/fab/migrations/env.py index 27be8fd17f49b..7b31f86024a31 100644 --- a/airflow/providers/fab/migrations/env.py +++ b/airflow/providers/fab/migrations/env.py @@ -33,7 +33,7 @@ # 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) + fileConfig(config.config_file_name, disable_existing_loggers=False) # add your model's MetaData object here # for 'autogenerate' support From d70eccbf8fe3f6417924a447757e16dd7a9a7897 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 13:20:30 +0100 Subject: [PATCH 23/30] Use abstract base and update tests --- airflow/migrations/env.py | 8 +- .../providers/fab/auth_manager/models/db.py | 3 - airflow/utils/db.py | 10 +- airflow/utils/db_manager.py | 37 +++--- tests/utils/test_db_manager.py | 108 +++++++++--------- 5 files changed, 76 insertions(+), 90 deletions(-) diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index 4be8ac32ca252..1233fe215acde 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -33,7 +33,13 @@ def include_object(_, name, type_, *args): if name == "sqlite_sequence": return False # Only create migrations for objects that are in the target metadata - if type_ == "table" and name not in target_metadata.tables: + from airflow import __version__ as airflow_version + + if airflow_version.startswith("2"): # For backwards compatibility with FAB provider + # Ignore _anything_ to do with Celery, or FlaskSession's tables + if type_ == "table" and (name.startswith("celery_") or name == "session"): + return False + elif type_ == "table" and name not in target_metadata.tables: return False else: return True diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index e90efd557b0ad..a971ea29a3f68 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -33,6 +33,3 @@ class FABDBManager(BaseDBManager): 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 - - def __init__(self, session): - self.session = session diff --git a/airflow/utils/db.py b/airflow/utils/db.py index f35edf792aa1f..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 @@ -765,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) @@ -772,10 +777,7 @@ def initdb(session: Session = NEW_SESSION, load_connections: bool = True): upgradedb(session=session) else: _create_db_from_orm(session=session) - from airflow.utils.db_manager import RunDBManager - external_db_manager = RunDBManager() - external_db_manager.validate() external_db_manager.initdb(session) if conf.getboolean("database", "LOAD_DEFAULT_CONNECTIONS") and load_connections: create_default_connections(session=session) @@ -1682,8 +1684,6 @@ 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) - from airflow.utils.db_manager import RunDBManager - external_db_manager = RunDBManager() external_db_manager.drop_tables(connection) diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py index 44b440e55f0d0..3acf69c172666 100644 --- a/airflow/utils/db_manager.py +++ b/airflow/utils/db_manager.py @@ -19,6 +19,8 @@ import os 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 @@ -28,25 +30,13 @@ from sqlalchemy import MetaData -class AttributeCheckerMeta(type): - """Metaclass to check attributes of subclasses.""" - - def __new__(cls, name, bases, dct): - """Check that subclasses are setting the required attributes.""" - required_attrs = ["metadata", "migration_dir", "alembic_file", "version_table_name"] - for attr in required_attrs: - if attr not in dct: - raise AttributeError(f"{name} is missing required attribute: {attr}") - return super().__new__(cls, name, bases, dct) - +class BaseDBManager(LoggingMixin): + """Abstract Base DB manager for external DBs.""" -class BaseDBManager(LoggingMixin, metaclass=AttributeCheckerMeta): - """Base DB manager for external DBs.""" - - metadata: MetaData = None - migration_dir: str = "" - alembic_file: str = "" - version_table_name: str = "" + 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 @@ -75,8 +65,6 @@ def get_current_revision(self): def _create_db_from_orm(self): """Create database from ORM.""" - from alembic import command - engine = self.session.get_bind().engine self.metadata.create_all(engine) config = self.get_alembic_config() @@ -92,25 +80,26 @@ def initdb(self): def upgradedb(self, to_version=None, from_version=None, show_sql_only=False): """Upgrade the database.""" - from alembic import command + 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) def downgradedb(self, to_version, from_version=None, show_sql_only=False): """Downgrade the database.""" + self.log.info("Downgrading the %s database", self.__class__.__name__) class RunDBManager(LoggingMixin): """ Run External DB Managers. - This class is a container for external database managers. + Validates and runs the external database managers. """ def __init__(self): super().__init__() - self._managers = [] + self._managers: list[BaseDBManager] = [] managers = conf.get("database", "external_db_managers").split(",") for module in managers: manager = import_string(module) @@ -122,7 +111,7 @@ def validate(self): RunDBManager._validate(manager) @staticmethod - def _validate(manager): + def _validate(manager: BaseDBManager): """Validate the external database migration.""" import ast diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py index c919d904a48b5..4906177e981de 100644 --- a/tests/utils/test_db_manager.py +++ b/tests/utils/test_db_manager.py @@ -30,44 +30,52 @@ class TestRunDBManager: + def setup_method(self): + self.run_db_manager = RunDBManager() + self.manager = self.run_db_manager._managers[0] + self.metadata = self.manager.metadata + 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] + assert self.run_db_manager._managers == [FABDBManager] def test_defining_table_same_name_as_airflow_table_name_raises(self): from sqlalchemy import Column, Integer, String - run_db_manager = RunDBManager() - manager = run_db_manager._managers[0] # Add dag_run table to metadata mytable = Table( - "dag_run", manager.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) + "dag_run", self.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) ) - manager.metadata._add_table("dag_run", None, mytable) + self.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() + self.run_db_manager.validate() + self.metadata._remove_table("dag_run", schema=None) - @mock.patch("airflow.utils.db_manager.RunDBManager") - def test_init_db_calls_rundbmanager(self, mock_rundbmanager, session): + @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_rundbmanager.return_value.initdb.assert_called_once_with(session) - mock_rundbmanager.return_value.upgradedb.assert_not_called() - mock_rundbmanager.return_value.downgrade.assert_not_called() - - @mock.patch("airflow.utils.db_manager.RunDBManager") + 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_rundbmanager, session + 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_rundbmanager.return_value.initdb.assert_not_called() - mock_rundbmanager.return_value.upgradedb.assert_not_called() - mock_rundbmanager.return_value.downgrade.assert_not_called() + mock_initdb.assert_not_called() + mock_upgrade_db.assert_not_called() + mock_downgrade_db.assert_not_called() @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session): @@ -83,45 +91,31 @@ def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session) mock_fabdb_manager.return_value.downgradedb.assert_called_once() -def test_subclassing_db_manager_with_missing_attrs(): - """Test subclassing BaseDBManager.""" - - with pytest.raises(AttributeError, match="SubclassDBManager is missing required attribute: metadata"): - - class SubclassDBManager(BaseDBManager): ... - - -def test_subclassing_db_manager_with_set_metadata(): - with pytest.raises( - AttributeError, match="SubclassDbManager is missing required attribute: migration_dir" - ): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - - -def test_subclassing_db_manager_with_set_metadata_and_migration_dir(): - with pytest.raises(AttributeError, match="SubclassDbManager is missing required attribute: alembic_file"): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" +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 -def test_subclassing_db_manager_with_attrs_set_except_version_table_name(): - with pytest.raises( - AttributeError, match="SubclassDbManager is missing required attribute: version_table_name" +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 ): - - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" - alembic_file = "some_file" - - -def test_subclassing_db_manager_with_attrs_set_dont_raise(session): - class SubclassDbManager(BaseDBManager): - metadata = Base.metadata - migration_dir = "some_dir" - alembic_file = "some_file" - version_table_name = "some_table" + 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 From b19cfc45250b97dc7b0a2fdfbc48d05d2a90fa9f Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 15:45:35 +0100 Subject: [PATCH 24/30] Revert compat change --- airflow/migrations/env.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index 1233fe215acde..4be8ac32ca252 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -33,13 +33,7 @@ def include_object(_, name, type_, *args): if name == "sqlite_sequence": return False # Only create migrations for objects that are in the target metadata - from airflow import __version__ as airflow_version - - if airflow_version.startswith("2"): # For backwards compatibility with FAB provider - # Ignore _anything_ to do with Celery, or FlaskSession's tables - if type_ == "table" and (name.startswith("celery_") or name == "session"): - return False - elif type_ == "table" and name not in target_metadata.tables: + if type_ == "table" and name not in target_metadata.tables: return False else: return True From 250c59ed119aa49e4e2fee4a3450eba5c3ad6c04 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 16:37:35 +0100 Subject: [PATCH 25/30] Fix backcompat for fab provider --- .../fab/auth_manager/models/test_db.py | 68 ++++++++++--------- 1 file changed, 35 insertions(+), 33 deletions(-) diff --git a/tests/providers/fab/auth_manager/models/test_db.py b/tests/providers/fab/auth_manager/models/test_db.py index 4b4b621d0cc60..7b0dc345c7188 100644 --- a/tests/providers/fab/auth_manager/models/test_db.py +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -24,7 +24,6 @@ from sqlalchemy import MetaData import airflow -from airflow.providers.fab.auth_manager.models.db import FABDBManager from airflow.settings import engine from airflow.utils.db import ( compare_server_default, @@ -32,43 +31,46 @@ ) 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) -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_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_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_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_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 - 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) - 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) + assert not diff, "Database schema and SQLAlchemy model are not in sync: " + str(diff) +except ModuleNotFoundError: + pass From 70b1e2d633ea04fa3f46142b705373b16f25f087 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 18:40:08 +0100 Subject: [PATCH 26/30] Build fab provider for www test instead of installing from pypi --- .github/workflows/basic-tests.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/basic-tests.yml b/.github/workflows/basic-tests.yml index f5b1d54084b86..bb1c9c82c4b23 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 - 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" From 48d8e10de90e54184c0a0a046b11ee6443721201 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 19:16:06 +0100 Subject: [PATCH 27/30] fixup! Fix backcompat for fab provider --- .../fab/auth_manager/models/__init__.py | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/airflow/providers/fab/auth_manager/models/__init__.py b/airflow/providers/fab/auth_manager/models/__init__.py index 9dc5c75747849..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 ( @@ -42,9 +43,16 @@ ) 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 _get_schema, naming_convention +if TYPE_CHECKING: + try: + from sqlalchemy import Identity + except Exception: + Identity = None + """ Compatibility note: The models in this file are duplicated from Flask AppBuilder. """ @@ -52,13 +60,14 @@ metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) mapper_registry = registry(metadata=metadata) -Model.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 -if TYPE_CHECKING: - try: - from sqlalchemy import Identity - except Exception: - Identity = None + Model.metadata = Base.metadata class Action(Model): From 08411de06a48f742f12e0047c3fd08212d00b540 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 21 Aug 2024 19:36:29 +0100 Subject: [PATCH 28/30] Mark downgradedb method as abstract --- airflow/utils/db_manager.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py index 3acf69c172666..4730ac1a9d896 100644 --- a/airflow/utils/db_manager.py +++ b/airflow/utils/db_manager.py @@ -17,6 +17,7 @@ from __future__ import annotations import os +from abc import abstractmethod from typing import TYPE_CHECKING from alembic import command @@ -85,9 +86,9 @@ def upgradedb(self, to_version=None, from_version=None, show_sql_only=False): 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.""" - self.log.info("Downgrading the %s database", self.__class__.__name__) class RunDBManager(LoggingMixin): From 497cd3c14ca81f83011c45039bec8d6c3e964480 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 23 Aug 2024 00:53:58 +0100 Subject: [PATCH 29/30] Ensure supports_table_dropping works and improve tests --- airflow/utils/db_manager.py | 3 ++- tests/utils/test_db_manager.py | 37 +++++++++++++++++++++++----------- 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py index 4730ac1a9d896..07d1ee9a61be1 100644 --- a/airflow/utils/db_manager.py +++ b/airflow/utils/db_manager.py @@ -163,4 +163,5 @@ def downgradedb(self, session): def drop_tables(self, connection): """Drop the external database managers.""" for manager in self._managers: - manager.metadata.drop_all(connection) + if manager.supports_table_dropping: + manager.metadata.drop_all(connection) diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py index 4906177e981de..d6fc91f93934a 100644 --- a/tests/utils/test_db_manager.py +++ b/tests/utils/test_db_manager.py @@ -25,32 +25,37 @@ 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: - def setup_method(self): - self.run_db_manager = RunDBManager() - self.manager = self.run_db_manager._managers[0] - self.metadata = self.manager.metadata - + @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 - assert self.run_db_manager._managers == [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", self.metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) + "dag_run", metadata, Column("id", Integer, primary_key=True), Column("name", String(50)) ) - self.metadata._add_table("dag_run", None, mytable) + metadata._add_table("dag_run", None, mytable) with pytest.raises(AirflowException, match="Table 'dag_run' already exists in the Airflow metadata"): - self.run_db_manager.validate() - self.metadata._remove_table("dag_run", schema=None) + run_db_manager.validate() + metadata._remove_table("dag_run", None) @mock.patch.object(RunDBManager, "downgradedb") @mock.patch.object(RunDBManager, "upgradedb") @@ -77,18 +82,26 @@ def test_upgradedb_or_downgrade_dont_call_rundbmanager( 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) - mock_fabdb_manager.return_value.initdb.assert_called_once() + fabdb_manager.initdb.assert_called_once() # upgradedb ext_db.upgradedb(session=session) - mock_fabdb_manager.return_value.upgradedb.assert_called_once() + 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): From 167ecb6de900dc1503bb2de4197d58c96129b41c Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 23 Aug 2024 08:16:04 +0100 Subject: [PATCH 30/30] Skip tag check when preparing fab provider in CI --- .github/workflows/basic-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/basic-tests.yml b/.github/workflows/basic-tests.yml index bb1c9c82c4b23..b59baa6728fda 100644 --- a/.github/workflows/basic-tests.yml +++ b/.github/workflows/basic-tests.yml @@ -156,7 +156,7 @@ jobs: working-directory: ./clients/python - name: "Prepare FAB provider packages: wheel" run: > - breeze release-management prepare-provider-packages fab --package-format wheel + breeze release-management prepare-provider-packages fab --package-format wheel --skip-tag-check - name: "Install Airflow with fab for webserver tests" run: pip install . dist/apache_airflow_providers_fab-*.whl - name: "Install Python client"