apache-airflow-providers-fab 2.4.4rc1__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

airflow/providers/fab/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "2.4.4"
+__version__ = "3.0.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "3.0.2"

airflow/providers/fab/auth_manager/api_endpoints/role_and_permission_endpoint.py

@@ -72,7 +72,7 @@ def get_role(*, role_name: str) -> APIResponse:
 def get_roles(*, order_by: str = "name", limit: int, offset: int | None = None) -> APIResponse:
     """Get roles."""
     security_manager = cast("FabAuthManager", get_auth_manager()).security_manager
-    session = security_manager.get_session
+    session = security_manager.session
     total_entries = session.scalars(select(func.count(Role.id))).one()
     direction = desc if order_by.startswith("-") else asc
     to_replace = {"role_id": "id"}
@@ -99,7 +99,7 @@ def get_roles(*, order_by: str = "name", limit: int, offset: int | None = None)
 def get_permissions(*, limit: int, offset: int | None = None) -> APIResponse:
     """Get permissions."""
     security_manager = cast("FabAuthManager", get_auth_manager()).security_manager
-    session = security_manager.get_session
+    session = security_manager.session
     total_entries = session.scalars(select(func.count(Action.id))).one()
     query = select(Action)
     actions = session.scalars(query.offset(offset).limit(limit)).all()

airflow/providers/fab/auth_manager/api_endpoints/user_endpoint.py

@@ -56,10 +56,10 @@ def get_user(*, username: str) -> APIResponse:
 
 @requires_access_custom_view("GET", permissions.RESOURCE_USER)
 @format_parameters({"limit": check_limit})
-def get_users(*, limit: int, order_by: str = "id", offset: str | None = None) -> APIResponse:
+def get_users(*, limit: int, order_by: str = "id", offset: int | None = None) -> APIResponse:
     """Get users."""
     security_manager = cast("FabAuthManager", get_auth_manager()).security_manager
-    session = security_manager.get_session
+    session = security_manager.session
     total_entries = session.execute(select(func.count(User.id))).scalar()
     direction = desc if order_by.startswith("-") else asc
     to_replace = {"user_id": "id"}
@@ -212,7 +212,7 @@ def delete_user(*, username: str) -> APIResponse:
         raise NotFound(title="User not found", detail=detail)
 
     user.roles = []  # Clear foreign keys on this user first.
-    security_manager.get_session.delete(user)
-    security_manager.get_session.commit()
+    security_manager.session.delete(user)
+    security_manager.session.commit()
 
     return NoContent, HTTPStatus.NO_CONTENT

airflow/providers/fab/auth_manager/api_fastapi/routes/login.py

@@ -23,6 +23,7 @@ from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_
 from airflow.configuration import conf
 from airflow.providers.fab.auth_manager.api_fastapi.datamodels.login import LoginBody, LoginResponse
 from airflow.providers.fab.auth_manager.api_fastapi.services.login import FABAuthManagerLogin
+from airflow.providers.fab.auth_manager.cli_commands.utils import get_application_builder
 
 login_router = AirflowRouter(tags=["FabAuthManager"])
 
@@ -35,7 +36,8 @@ login_router = AirflowRouter(tags=["FabAuthManager"])
 )
 def create_token(body: LoginBody) -> LoginResponse:
     """Generate a new API token."""
-    return FABAuthManagerLogin.create_token(body=body)
+    with get_application_builder():
+        return FABAuthManagerLogin.create_token(body=body)
 
 
 @login_router.post(
@@ -46,6 +48,7 @@ def create_token(body: LoginBody) -> LoginResponse:
 )
 def create_token_cli(body: LoginBody) -> LoginResponse:
     """Generate a new CLI API token."""
-    return FABAuthManagerLogin.create_token(
-        body=body, expiration_time_in_seconds=conf.getint("api_auth", "jwt_cli_expiration_time")
-    )
+    with get_application_builder():
+        return FABAuthManagerLogin.create_token(
+            body=body, expiration_time_in_seconds=conf.getint("api_auth", "jwt_cli_expiration_time")
+        )

airflow/providers/fab/auth_manager/cli_commands/user_command.py

@@ -144,9 +144,8 @@ def users_delete(args):
 @providers_configuration_loaded
 def users_manage_role(args, remove=False):
     """Delete or appends user roles."""
-    user = _find_user(args)
-
     with get_application_builder() as appbuilder:
+        user = _find_user(args)
         role = appbuilder.sm.find_role(args.role)
         if not role:
             valid_roles = appbuilder.sm.get_all_roles()

airflow/providers/fab/auth_manager/cli_commands/utils.py

@@ -25,6 +25,7 @@ from os.path import isabs
 from typing import TYPE_CHECKING
 
 from flask import Flask
+from flask_sqlalchemy import SQLAlchemy
 from sqlalchemy.engine import make_url
 
 import airflow
@@ -39,11 +40,11 @@ if TYPE_CHECKING:
 
 
 @cache
-def _return_appbuilder(app: Flask) -> AirflowAppBuilder:
+def _return_appbuilder(app: Flask, db) -> AirflowAppBuilder:
     """Return an appbuilder instance for the given app."""
     init_appbuilder(app, enable_plugins=False)
     init_plugins(app)
-    init_airflow_session_interface(app)
+    init_airflow_session_interface(app, db)
     return app.appbuilder  # type: ignore[attr-defined]
 
 
@@ -63,4 +64,7 @@ def get_application_builder() -> Generator[AirflowAppBuilder, None, None]:
                 "Please use absolute path such as `sqlite:////tmp/airflow.db`."
             )
         flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
-        yield _return_appbuilder(flask_app)
+
+        db = SQLAlchemy(flask_app)
+        yield _return_appbuilder(flask_app, db)
+        db.engine.dispose(close=True)

airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -26,7 +26,7 @@ from urllib.parse import urljoin
 import packaging.version
 from connexion import FlaskApi
 from fastapi import FastAPI
-from flask import Blueprint, g
+from flask import Blueprint, current_app, g
 from sqlalchemy import select
 from sqlalchemy.orm import Session, joinedload
 from starlette.middleware.wsgi import WSGIMiddleware
@@ -282,7 +282,7 @@ class FabAuthManager(BaseAuthManager[User]):
 
     def deserialize_user(self, token: dict[str, Any]) -> User:
         with create_session() as session:
-            return session.get(User, int(token["sub"]))
+            return session.scalars(select(User).where(User.id == int(token["sub"]))).one()
 
     def serialize_user(self, user: User) -> dict[str, Any]:
         return {"sub": str(user.id)}
@@ -292,7 +292,7 @@ class FabAuthManager(BaseAuthManager[User]):
         user = self.get_user()
         return (
             self.appbuilder
-            and self.appbuilder.get_app.config.get("AUTH_ROLE_PUBLIC", None)
+            and self.appbuilder.app.config.get("AUTH_ROLE_PUBLIC", None)
             or (not user.is_anonymous and user.is_active)
         )
 
@@ -470,14 +470,18 @@ class FabAuthManager(BaseAuthManager[User]):
             return {dag.dag_id for dag in session.execute(select(DagModel.dag_id))}
         if isinstance(user, AnonymousUser):
             return set()
-        user_query = session.scalar(
-            select(User)
-            .options(
-                joinedload(User.roles)
-                .subqueryload(Role.permissions)
-                .options(joinedload(Permission.action), joinedload(Permission.resource))
+        user_query = (
+            session.scalars(
+                select(User)
+                .options(
+                    joinedload(User.roles)
+                    .subqueryload(Role.permissions)
+                    .options(joinedload(Permission.action), joinedload(Permission.resource))
+                )
+                .where(User.id == user.id)
             )
-            .where(User.id == user.id)
+            .unique()
+            .one()
         )
         roles = user_query.roles
 
@@ -547,7 +551,7 @@ class FabAuthManager(BaseAuthManager[User]):
         if not self.appbuilder:
             raise AirflowException("AppBuilder is not initialized.")
 
-        sm_from_config = self.appbuilder.get_app.config.get("SECURITY_MANAGER_CLASS")
+        sm_from_config = current_app.config.get("SECURITY_MANAGER_CLASS")
         if sm_from_config:
             if not issubclass(sm_from_config, FabAirflowSecurityManagerOverride):
                 raise AirflowConfigException(

airflow/providers/fab/auth_manager/models/__init__.py

@@ -23,9 +23,8 @@ import datetime
 # 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 flask_appbuilder import Model
 from sqlalchemy import (
     Boolean,
     Column,
@@ -33,7 +32,7 @@ from sqlalchemy import (
     ForeignKey,
     Index,
     Integer,
-    MetaData,
+    Sequence,
     String,
     Table,
     UniqueConstraint,
@@ -41,11 +40,19 @@ from sqlalchemy import (
     func,
     select,
 )
-from sqlalchemy.orm import backref, declared_attr, registry, relationship
+from sqlalchemy.orm import Mapped, backref, declared_attr, relationship
 
-from airflow import __version__ as airflow_version
 from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
-from airflow.models.base import _get_schema, naming_convention
+
+try:
+    from sqlalchemy.orm import mapped_column
+except ImportError:
+    # fallback for SQLAlchemy < 2.0
+    def mapped_column(*args, **kwargs):
+        from sqlalchemy import Column
+
+        return Column(*args, **kwargs)
+
 
 if TYPE_CHECKING:
     try:
@@ -57,25 +64,88 @@ if TYPE_CHECKING:
 Compatibility note: The models in this file are duplicated from Flask AppBuilder.
 """
 
-metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention)
-mapper_registry = registry(metadata=metadata)
+assoc_group_role = Table(
+    "ab_group_role",
+    Model.metadata,
+    Column(
+        "id",
+        Integer,
+        Sequence("ab_group_role_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    ),
+    Column("group_id", Integer, ForeignKey("ab_group.id", ondelete="CASCADE")),
+    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
+    UniqueConstraint("group_id", "role_id"),
+    Index("idx_group_id", "group_id"),
+    Index("idx_group_role_id", "role_id"),
+)
 
-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
+assoc_permission_role = Table(
+    "ab_permission_view_role",
+    Model.metadata,
+    Column(
+        "id",
+        Integer,
+        Sequence(
+            "ab_permission_view_role_id_seq",
+            start=1,
+            increment=1,
+            minvalue=1,
+            cycle=False,
+        ),
+        primary_key=True,
+    ),
+    Column(
+        "permission_view_id",
+        Integer,
+        ForeignKey("ab_permission_view.id", ondelete="CASCADE"),
+    ),
+    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
+    UniqueConstraint("permission_view_id", "role_id"),
+    Index("idx_permission_view_id", "permission_view_id"),
+    Index("idx_role_id", "role_id"),
+)
+
+assoc_user_role = Table(
+    "ab_user_role",
+    Model.metadata,
+    Column(
+        "id",
+        Integer,
+        Sequence("ab_user_role_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    ),
+    Column("user_id", Integer, ForeignKey("ab_user.id", ondelete="CASCADE")),
+    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
+    UniqueConstraint("user_id", "role_id"),
+)
 
-    Model.metadata = Base.metadata
+assoc_user_group = Table(
+    "ab_user_group",
+    Model.metadata,
+    Column(
+        "id",
+        Integer,
+        Sequence("ab_user_group_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    ),
+    Column("user_id", Integer, ForeignKey("ab_user.id", ondelete="CASCADE")),
+    Column("group_id", Integer, ForeignKey("ab_group.id", ondelete="CASCADE")),
+    UniqueConstraint("user_id", "group_id"),
+)
 
 
 class Action(Model):
     """Represents permission actions such as `can_read`."""
 
     __tablename__ = "ab_permission"
-    id = Column(Integer, primary_key=True)
-    name = Column(String(100), unique=True, nullable=False)
+
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_permission_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    name: Mapped[str] = mapped_column(String(100), unique=True, nullable=False)
 
     def __repr__(self):
         return self.name
@@ -85,8 +155,13 @@ class Resource(Model):
     """Represents permission object such as `User` or `Dag`."""
 
     __tablename__ = "ab_view_menu"
-    id = Column(Integer, primary_key=True)
-    name = Column(String(250), unique=True, nullable=False)
+
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_view_menu_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    name: Mapped[str] = mapped_column(String(250), unique=True, nullable=False)
 
     def __eq__(self, other):
         return (isinstance(other, self.__class__)) and (self.name == other.name)
@@ -98,50 +173,18 @@ class Resource(Model):
         return self.name
 
 
-assoc_permission_role = Table(
-    "ab_permission_view_role",
-    Model.metadata,
-    Column("id", Integer, primary_key=True),
-    Column(
-        "permission_view_id",
-        Integer,
-        ForeignKey("ab_permission_view.id", ondelete="CASCADE"),
-    ),
-    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
-    UniqueConstraint("permission_view_id", "role_id"),
-)
-
-assoc_user_group = Table(
-    "ab_user_group",
-    Model.metadata,
-    Column("id", Integer, primary_key=True),
-    Column("user_id", Integer, ForeignKey("ab_user.id", ondelete="CASCADE")),
-    Column("group_id", Integer, ForeignKey("ab_group.id", ondelete="CASCADE")),
-    UniqueConstraint("user_id", "group_id"),
-    Index("idx_user_id", "user_id"),
-    Index("idx_user_group_id", "group_id"),
-)
-
-assoc_group_role = Table(
-    "ab_group_role",
-    Model.metadata,
-    Column("id", Integer, primary_key=True),
-    Column("group_id", Integer, ForeignKey("ab_group.id", ondelete="CASCADE")),
-    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
-    UniqueConstraint("group_id", "role_id"),
-    Index("idx_group_id", "group_id"),
-    Index("idx_group_role_id", "role_id"),
-)
-
-
 class Role(Model):
     """Represents a user role to which permissions can be assigned."""
 
     __tablename__ = "ab_role"
 
-    id = Column(Integer, primary_key=True)
-    name = Column(String(64), unique=True, nullable=False)
-    permissions = relationship(
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_role_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    name: Mapped[str] = mapped_column(String(64), unique=True, nullable=False)
+    permissions: Mapped[list[Permission]] = relationship(
         "Permission",
         secondary=assoc_permission_role,
         backref="role",
@@ -153,93 +196,100 @@ class Role(Model):
         return self.name
 
 
-class Group(Model):
-    """Represents a user group."""
-
-    __tablename__ = "ab_group"
-
-    id = Column(Integer, primary_key=True)
-    name = Column(String(100), unique=True, nullable=False)
-    label = Column(String(150))
-    description = Column(String(512))
-    users = relationship("User", secondary=assoc_user_group, backref="groups", passive_deletes=True)
-    roles = relationship("Role", secondary=assoc_group_role, backref="groups", passive_deletes=True)
-
-    def __repr__(self):
-        return self.name
-
-
 class Permission(Model):
     """Permission pair comprised of an Action + Resource combo."""
 
     __tablename__ = "ab_permission_view"
     __table_args__ = (UniqueConstraint("permission_id", "view_menu_id"),)
-    id = Column(Integer, primary_key=True)
-    action_id = Column("permission_id", Integer, ForeignKey("ab_permission.id"))
-    action = relationship(
-        "Action",
-        uselist=False,
-        lazy="joined",
-    )
-    resource_id = Column("view_menu_id", Integer, ForeignKey("ab_view_menu.id"))
-    resource = relationship(
-        "Resource",
-        uselist=False,
-        lazy="joined",
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_permission_view_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
     )
+    action_id: Mapped[int] = mapped_column("permission_id", Integer, ForeignKey("ab_permission.id"))
+    action: Mapped[Action] = relationship("Action", lazy="joined", uselist=False)
+    resource_id: Mapped[int] = mapped_column("view_menu_id", Integer, ForeignKey("ab_view_menu.id"))
+    resource: Mapped[Resource] = relationship("Resource", lazy="joined", uselist=False)
 
     def __repr__(self):
-        return str(self.action).replace("_", " ") + " on " + str(self.resource)
+        return str(self.action).replace("_", " ") + f" on {str(self.resource)}"
 
 
-assoc_user_role = Table(
-    "ab_user_role",
-    Model.metadata,
-    Column("id", Integer, primary_key=True),
-    Column("user_id", Integer, ForeignKey("ab_user.id", ondelete="CASCADE")),
-    Column("role_id", Integer, ForeignKey("ab_role.id", ondelete="CASCADE")),
-    UniqueConstraint("user_id", "role_id"),
-)
+class Group(Model):
+    """Represents an Airflow user group."""
+
+    __tablename__ = "ab_group"
+
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_group_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    name: Mapped[str] = Column(String(100), unique=True, nullable=False)
+    label: Mapped[str] = Column(String(150))
+    description: Mapped[str] = Column(String(512))
+    users: Mapped[list[User]] = relationship(
+        "User", secondary=assoc_user_group, backref="groups", passive_deletes=True
+    )
+    roles: Mapped[list[Role]] = relationship(
+        "Role", secondary=assoc_group_role, backref="groups", passive_deletes=True
+    )
+
+    def __repr__(self):
+        return self.name
 
 
 class User(Model, BaseUser):
     """Represents an Airflow user which has roles assigned to it."""
 
     __tablename__ = "ab_user"
-    id = Column(Integer, primary_key=True)
-    first_name = Column(String(256), nullable=False)
-    last_name = Column(String(256), nullable=False)
-    username = Column(
+
+    id: Mapped[int] = mapped_column(
+        Integer,
+        Sequence("ab_user_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    first_name: Mapped[str] = mapped_column(String(64), nullable=False)
+    last_name: Mapped[str] = mapped_column(String(64), nullable=False)
+    username: Mapped[str] = mapped_column(
         String(512).with_variant(String(512, collation="NOCASE"), "sqlite"), unique=True, nullable=False
     )
-    password = Column(String(256))
-    active = Column(Boolean, default=True)
-    email = Column(String(512), unique=True, nullable=False)
-    last_login = Column(DateTime)
-    login_count = Column(Integer)
-    fail_login_count = Column(Integer)
-    roles = relationship(
-        "Role", secondary=assoc_user_role, backref="user", lazy="selectin", passive_deletes=True
+    password: Mapped[str | None] = mapped_column(String(256))
+    active: Mapped[bool | None] = mapped_column(Boolean, default=True)
+    email: Mapped[str] = mapped_column(String(320), unique=True, nullable=False)
+    last_login: Mapped[datetime.datetime | None] = mapped_column(DateTime, nullable=True)
+    login_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    fail_login_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    roles: Mapped[list[Role]] = relationship(
+        "Role",
+        secondary=assoc_user_role,
+        backref="user",
+        lazy="selectin",
+        passive_deletes=True,
+    )
+    created_on: Mapped[datetime.datetime | None] = mapped_column(
+        DateTime, default=lambda: datetime.datetime.now(), nullable=True
+    )
+    changed_on: Mapped[datetime.datetime | None] = mapped_column(
+        DateTime, default=lambda: datetime.datetime.now(), nullable=True
     )
-    created_on = Column(DateTime, default=datetime.datetime.now, nullable=True)
-    changed_on = Column(DateTime, default=datetime.datetime.now, nullable=True)
 
     @declared_attr
-    def created_by_fk(self):
+    def created_by_fk(self) -> Column:
         return Column(Integer, ForeignKey("ab_user.id"), default=self.get_user_id, nullable=True)
 
     @declared_attr
-    def changed_by_fk(self):
+    def changed_by_fk(self) -> Column:
         return Column(Integer, ForeignKey("ab_user.id"), default=self.get_user_id, nullable=True)
 
-    created_by = relationship(
+    created_by: Mapped[User] = relationship(
         "User",
         backref=backref("created", uselist=True),
         remote_side=[id],
         primaryjoin="User.created_by_fk == User.id",
         uselist=False,
     )
-    changed_by = relationship(
+    changed_by: Mapped[User] = relationship(
         "User",
         backref=backref("changed", uselist=True),
         remote_side=[id],
@@ -274,7 +324,7 @@ class User(Model, BaseUser):
             if current_app:
                 sm = current_app.appbuilder.sm
                 self._perms: set[tuple[str, str]] = set(
-                    sm.get_session.execute(
+                    sm.session.execute(
                         select(sm.action_model.name, sm.resource_model.name)
                         .join(sm.permission_model.action)
                         .join(sm.permission_model.resource)
@@ -307,16 +357,23 @@ class RegisterUser(Model):
     """Represents a user registration."""
 
     __tablename__ = "ab_register_user"
-    id = Column(Integer, primary_key=True)
-    first_name = Column(String(256), nullable=False)
-    last_name = Column(String(256), nullable=False)
-    username = Column(
+
+    id = mapped_column(
+        Integer,
+        Sequence("ab_register_user_id_seq", start=1, increment=1, minvalue=1, cycle=False),
+        primary_key=True,
+    )
+    first_name: Mapped[str] = mapped_column(String(64), nullable=False)
+    last_name: Mapped[str] = mapped_column(String(64), nullable=False)
+    username: Mapped[str] = mapped_column(
         String(512).with_variant(String(512, collation="NOCASE"), "sqlite"), unique=True, nullable=False
     )
-    password = Column(String(256))
-    email = Column(String(512), nullable=False)
-    registration_date = Column(DateTime, default=datetime.datetime.now, nullable=True)
-    registration_hash = Column(String(256))
+    password: Mapped[str | None] = mapped_column(String(256))
+    email: Mapped[str] = mapped_column(String(320), unique=True, nullable=False)
+    registration_date: Mapped[datetime.datetime | None] = mapped_column(
+        DateTime, default=lambda: datetime.datetime.now(), nullable=True
+    )
+    registration_hash: Mapped[str | None] = mapped_column(String(256))
 
 
 @event.listens_for(User.__table__, "before_create")

airflow/providers/fab/auth_manager/models/db.py

@@ -18,9 +18,10 @@ from __future__ import annotations
 
 from pathlib import Path
 
+from flask_appbuilder import Model
+
 from airflow import settings
 from airflow.exceptions import AirflowException
-from airflow.providers.fab.auth_manager.models import metadata
 from airflow.utils.db import _offline_migration, print_happy_cat
 from airflow.utils.db_manager import BaseDBManager
 
@@ -41,14 +42,14 @@ def _get_flask_db(sql_database_uri):
     flask_app.config["SQLALCHEMY_DATABASE_URI"] = sql_database_uri
     flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
     db = SQLAlchemy(flask_app)
-    AirflowDatabaseSessionInterface(app=flask_app, db=db, table="session", key_prefix="")
-    return db
+    AirflowDatabaseSessionInterface(app=flask_app, client=db, table="session", key_prefix="")
+    return db, flask_app
 
 
 class FABDBManager(BaseDBManager):
     """Manages FAB database."""
 
-    metadata = metadata
+    metadata = Model.metadata
     version_table_name = "alembic_version_fab"
     migration_dir = (PACKAGE_DIR / "migrations").as_posix()
     alembic_file = (PACKAGE_DIR / "alembic.ini").as_posix()
@@ -57,7 +58,9 @@ class FABDBManager(BaseDBManager):
 
     def create_db_from_orm(self):
         super().create_db_from_orm()
-        _get_flask_db(settings.SQL_ALCHEMY_CONN).create_all()
+        db, flask_app = _get_flask_db(settings.SQL_ALCHEMY_CONN)
+        with flask_app.app_context():
+            db.create_all()
 
     def reset_to_2_x(self):
         self.create_db_from_orm()
@@ -126,4 +129,6 @@ class FABDBManager(BaseDBManager):
 
     def drop_tables(self, connection):
         super().drop_tables(connection)
-        _get_flask_db(settings.SQL_ALCHEMY_CONN).drop_all()
+        db, flask_app = _get_flask_db(settings.SQL_ALCHEMY_CONN)
+        with flask_app.app_context():
+            db.drop_all()