Repository: apache/airflow
File: providers/edge3/src/airflow/providers/edge3/models/db.py (Python)

Code before:
# 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 warnings
from pathlib import Path

from sqlalchemy import MetaData, inspect

from airflow.models.base import Base
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
from airflow.utils.db_manager import BaseDBManager

EDGE_DB_MANAGER_FQDN = "airflow.providers.edge3.models.db.EdgeDBManager"

PACKAGE_DIR = Path(__file__).parents[1]

_REVISION_HEADS_MAP: dict[str, str] = {
    "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__)


def check_db_manager_config() -> None:
    """
    Warn if EdgeDBManager is not registered in the external_db_managers config.


Review comment (suggestion):
Not needed to have a constant. You can make this also via `f"{EdgeDBManager.__module__}.{EdgeDBManager.__name__}"`

Code after:
# 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 warnings
from pathlib import Path

from sqlalchemy import inspect

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
from airflow.utils.db_manager import BaseDBManager

PACKAGE_DIR = Path(__file__).parents[1]

_REVISION_HEADS_MAP: dict[str, str] = {
    "3.0.0": "9d34dfc2de06",
}


class EdgeDBManager(BaseDBManager):
    """Manages Edge3 provider database tables."""

    metadata = edge_metadata

    version_table_name = "alembic_version_edge3"
    migration_dir = (PACKAGE_DIR / "migrations").as_posix()
    alembic_file = (PACKAGE_DIR / "alembic.ini").as_posix()
    supports_table_dropping = True
    revision_heads_map = _REVISION_HEADS_MAP

    def drop_tables(self, connection):
        """Drop only edge3 tables in reverse dependency order."""
        if not self.supports_table_dropping:
            return

        inspector = inspect(connection)

        edge_tables = [