From 1c19c7039507d8b8ba4b2521103c824ba3e5bcca Mon Sep 17 00:00:00 2001 From: Dheeraj Turaga Date: Thu, 19 Feb 2026 22:40:14 -0600 Subject: [PATCH 1/3] Replace Base.metadata.remove() hack with dedicated Edge3 declarative base MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce models/edge_base.py with its own MetaData, registry, and Base for the Edge3 provider. All three Edge models now inherit from this isolated base, so their tables are never registered in Airflow core's Base.metadata in the first place — eliminating the need for the to_metadata() copy + Base.metadata.remove() workaround in db.py. --- .../src/airflow/providers/edge3/models/db.py | 28 +++-------------- .../providers/edge3/models/edge_base.py | 31 +++++++++++++++++++ .../providers/edge3/models/edge_job.py | 3 +- .../providers/edge3/models/edge_logs.py | 3 +- .../providers/edge3/models/edge_worker.py | 2 +- 5 files changed, 40 insertions(+), 27 deletions(-) create mode 100644 providers/edge3/src/airflow/providers/edge3/models/edge_base.py diff --git a/providers/edge3/src/airflow/providers/edge3/models/db.py b/providers/edge3/src/airflow/providers/edge3/models/db.py index f434da64d4013..ac61ed987c394 100644 --- a/providers/edge3/src/airflow/providers/edge3/models/db.py +++ b/providers/edge3/src/airflow/providers/edge3/models/db.py @@ -18,9 +18,9 @@ from pathlib import Path -from sqlalchemy import MetaData, inspect +from sqlalchemy import inspect -from airflow.models.base import Base +from airflow.providers.edge3.models.edge_base import edge_metadata from airflow.providers.edge3.models.edge_job import EdgeJobModel from airflow.providers.edge3.models.edge_logs import EdgeLogsModel from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel @@ -32,25 +32,11 @@ "3.0.0": "9d34dfc2de06", } -# Create filtered metadata containing only edge3 tables -# This avoids validation issues with shared Base.metadata -_edge_metadata = MetaData() -EdgeWorkerModel.__table__.to_metadata(_edge_metadata) -EdgeJobModel.__table__.to_metadata(_edge_metadata) -EdgeLogsModel.__table__.to_metadata(_edge_metadata) - -# Remove edge tables from Airflow's core metadata to prevent validation conflicts -# The tables are now managed exclusively through _edge_metadata -Base.metadata.remove(EdgeWorkerModel.__table__) -Base.metadata.remove(EdgeJobModel.__table__) -Base.metadata.remove(EdgeLogsModel.__table__) - class EdgeDBManager(BaseDBManager): """Manages Edge3 provider database tables.""" - # Use filtered metadata instead of shared Base.metadata - metadata = _edge_metadata + metadata = edge_metadata version_table_name = "alembic_version_edge3" migration_dir = (PACKAGE_DIR / "migrations").as_posix() @@ -59,17 +45,12 @@ class EdgeDBManager(BaseDBManager): revision_heads_map = _REVISION_HEADS_MAP def drop_tables(self, connection): - """ - Drop only edge3 tables. - - Override base implementation to avoid dropping all tables in shared metadata. - """ + """Drop only edge3 tables in reverse dependency order.""" if not self.supports_table_dropping: return inspector = inspect(connection) - # Drop edge3 tables in reverse dependency order edge_tables = [ EdgeLogsModel.__table__, EdgeJobModel.__table__, @@ -81,7 +62,6 @@ def drop_tables(self, connection): self.log.info("Dropping table %s", table.name) table.drop(connection) - # Drop version table version = self._get_migration_ctx()._version if inspector.has_table(version.name): self.log.info("Dropping version table %s", version.name) diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_base.py b/providers/edge3/src/airflow/providers/edge3/models/edge_base.py new file mode 100644 index 0000000000000..2d97c99759927 --- /dev/null +++ b/providers/edge3/src/airflow/providers/edge3/models/edge_base.py @@ -0,0 +1,31 @@ +# 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 sqlalchemy import MetaData +from sqlalchemy.orm import registry + +from airflow.models.base import _get_schema, naming_convention + +# Isolated metadata for Edge3 provider tables. +# By using a dedicated MetaData + registry + Base, Edge3 tables are never +# registered in Airflow core's Base.metadata, avoiding validation conflicts +# without needing the post-hoc Base.metadata.remove() hack. +edge_metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) +_edge_mapper_registry = registry(metadata=edge_metadata) +Base = _edge_mapper_registry.generate_base() +Base.__allow_unmapped__ = True # match core Base workaround for unmapped v1.4 models diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_job.py b/providers/edge3/src/airflow/providers/edge3/models/edge_job.py index d322790346522..d98f3076a324a 100644 --- a/providers/edge3/src/airflow/providers/edge3/models/edge_job.py +++ b/providers/edge3/src/airflow/providers/edge3/models/edge_job.py @@ -26,9 +26,10 @@ ) from sqlalchemy.orm import Mapped -from airflow.models.base import Base, StringID +from airflow.models.base import StringID from airflow.providers.common.compat.sdk import TaskInstanceKey, timezone from airflow.providers.common.compat.sqlalchemy.orm import mapped_column +from airflow.providers.edge3.models.edge_base import Base from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.sqlalchemy import UtcDateTime diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_logs.py b/providers/edge3/src/airflow/providers/edge3/models/edge_logs.py index f4366756841b9..df9e7cd8db335 100644 --- a/providers/edge3/src/airflow/providers/edge3/models/edge_logs.py +++ b/providers/edge3/src/airflow/providers/edge3/models/edge_logs.py @@ -26,8 +26,9 @@ from sqlalchemy.dialects.mysql import MEDIUMTEXT from sqlalchemy.orm import Mapped -from airflow.models.base import Base, StringID +from airflow.models.base import StringID from airflow.providers.common.compat.sqlalchemy.orm import mapped_column +from airflow.providers.edge3.models.edge_base import Base from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.sqlalchemy import UtcDateTime diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py index 2e739826793d7..e4b0698e7e051 100644 --- a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py +++ b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py @@ -26,8 +26,8 @@ from sqlalchemy import Integer, String, delete, select from sqlalchemy.orm import Mapped -from airflow.models.base import Base from airflow.providers.common.compat.sdk import AirflowException, Stats, timezone +from airflow.providers.edge3.models.edge_base import Base try: from airflow.sdk.observability.stats import DualStatsManager From 67268b238ae89455fd55ac4137b0752194417f39 Mon Sep 17 00:00:00 2001 From: Dheeraj Turaga Date: Thu, 19 Feb 2026 23:25:18 -0600 Subject: [PATCH 2/3] Fix tests --- providers/edge3/tests/conftest.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/edge3/tests/conftest.py b/providers/edge3/tests/conftest.py index 27f6db6427ebd..b5de80cd45a9c 100644 --- a/providers/edge3/tests/conftest.py +++ b/providers/edge3/tests/conftest.py @@ -30,8 +30,8 @@ def _create_edge_tables(): yield return - from airflow.providers.edge3.models.db import _edge_metadata + from airflow.providers.edge3.models.edge_base import edge_metadata - _edge_metadata.create_all(settings.engine) + edge_metadata.create_all(settings.engine) yield - _edge_metadata.drop_all(settings.engine) + edge_metadata.drop_all(settings.engine) From 2fc0d4cd9f742bdd672c5efe3c680357a21bd94f Mon Sep 17 00:00:00 2001 From: Dheeraj Turaga Date: Thu, 19 Feb 2026 23:45:35 -0600 Subject: [PATCH 3/3] Add unit test for edge base --- .../tests/unit/edge3/models/test_edge_base.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 providers/edge3/tests/unit/edge3/models/test_edge_base.py diff --git a/providers/edge3/tests/unit/edge3/models/test_edge_base.py b/providers/edge3/tests/unit/edge3/models/test_edge_base.py new file mode 100644 index 0000000000000..148bb71e92eb8 --- /dev/null +++ b/providers/edge3/tests/unit/edge3/models/test_edge_base.py @@ -0,0 +1,55 @@ +# 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 airflow.models.base import naming_convention +from airflow.providers.edge3.models.edge_base import Base, edge_metadata + + +class TestEdgeBase: + def test_edge_metadata_is_isolated_from_core(self): + """edge_metadata must not be the same object as Airflow core's Base.metadata.""" + from airflow.models.base import Base as CoreBase + + assert edge_metadata is not CoreBase.metadata + + def test_edge_tables_not_in_core_metadata(self): + """Edge3 tables must never appear in Airflow core's Base.metadata.""" + from airflow.models.base import Base as CoreBase + + edge_table_names = {"edge_worker", "edge_job", "edge_logs"} + core_table_names = set(CoreBase.metadata.tables.keys()) + assert not edge_table_names & core_table_names + + def test_edge_metadata_contains_edge_tables(self): + """edge_metadata must contain all three Edge3 tables.""" + # Import models to ensure they are registered + import airflow.providers.edge3.models.edge_job + import airflow.providers.edge3.models.edge_logs + import airflow.providers.edge3.models.edge_worker # noqa: F401 + + assert "edge_worker" in edge_metadata.tables + assert "edge_job" in edge_metadata.tables + assert "edge_logs" in edge_metadata.tables + + def test_edge_metadata_uses_same_naming_convention_as_core(self): + """edge_metadata should use the same naming convention as Airflow core.""" + assert edge_metadata.naming_convention == naming_convention + + def test_base_allow_unmapped(self): + """Base must have __allow_unmapped__ set to match core Base workaround.""" + assert Base.__allow_unmapped__ is True