apache-airflow-providers-fab 1.5.3rc1__py3-none-any.whl → 2.0.0b1__py3-none-any.whl

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

@@ -31,6 +31,20 @@ _REVISION_HEADS_MAP: dict[str, str] = {
 }
 
 
+def _get_flask_db(sql_database_uri):
+    from flask import Flask
+    from flask_sqlalchemy import SQLAlchemy
+
+    from airflow.providers.fab.www.session import AirflowDatabaseSessionInterface
+
+    flask_app = Flask(__name__)
+    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
+
+
 class FABDBManager(BaseDBManager):
     """Manages FAB database."""
 
@@ -40,6 +54,10 @@ class FABDBManager(BaseDBManager):
     alembic_file = (PACKAGE_DIR / "alembic.ini").as_posix()
     supports_table_dropping = True
 
+    def _create_db_from_orm(self):
+        super()._create_db_from_orm()
+        _get_flask_db(settings.SQL_ALCHEMY_CONN).create_all()
+
     def upgradedb(self, to_revision=None, from_revision=None, show_sql_only=False):
         """Upgrade the database."""
         if from_revision and not show_sql_only:
@@ -68,11 +86,6 @@ class FABDBManager(BaseDBManager):
             _offline_migration(command.upgrade, config, f"{from_revision}:{to_revision}")
             return  # only running sql; our job is done
 
-        if not self.get_current_revision():
-            # New DB; initialize and exit
-            self.initdb()
-            return
-
         command.upgrade(config, revision=to_revision or "heads")
 
     def downgrade(self, to_revision, from_revision=None, show_sql_only=False):
@@ -104,3 +117,7 @@ class FABDBManager(BaseDBManager):
         else:
             self.log.info("Applying FAB downgrade migrations.")
             command.downgrade(config, revision=to_revision, sql=show_sql_only)
+
+    def drop_tables(self, connection):
+        super().drop_tables(connection)
+        _get_flask_db(settings.SQL_ALCHEMY_CONN).drop_all()

airflow/providers/fab/auth_manager/schemas/user_schema.py

@@ -21,9 +21,9 @@ from typing import NamedTuple
 from marshmallow import Schema, fields
 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
 
-from airflow.api_connexion.parameters import validate_istimezone
 from airflow.providers.fab.auth_manager.models import User
 from airflow.providers.fab.auth_manager.schemas.role_and_permission_schema import RoleSchema
+from airflow.providers.fab.www.api_connexion.parameters import validate_istimezone
 
 
 class UserCollectionItemSchema(SQLAlchemySchema):