Playing around with migrations to update roles/permissions in Airflow

airflow__migratons__operations.py

# airflow/migratons/operatons.py
# To enable we need to import this module from within airflow/migrations/env.py

from alembic.operations import MigrateOperation, Operations

from sqlalchemy import String, Integer, Sequence
from sqlalchemy.sql import column, table

from flask_appbuilder.security.sqla import models


class BaseOperation(MigrateOperation):
    """
    A base migration class that inserts/deletes to Tables of (id int, name string)

    Requires that the class has a ``table`` property, and that the instance has
    a .name attribute
    """
    @classmethod
    def invoke_op(cls, operations, *args, **kwargs):
        op = cls(*args, **kwargs)
        operations.invoke(op)


def basic_add_by_name(operations, operation):
    operations.execute(operation.table.insert().values(name=operation.name))


def basic_delete_by_name(operations, operation):
    tbl = operation.table

    operations.execute(tbl.delete().where(tbl.columns['name'] == operation.permission))


class PermissionOperation(MigrateOperation):
    table = table(
        'ab_permission',
        column('id', Integer, Sequence('ab_permission_id_seq')),
        column('name', String),
    )


class PermissionViewOperation(MigrateOperation):
    table = table(
        'ab_permission_view',
        column('id', Integer),
        column('permission_id', Integer),
        column('view_id', Integer),
    )


@Operations.register_operation("add_permission", "invoke_op")
class AddPermission(PermissionOperation, BaseOperation):
    def __init__(self, name):
        self.name = name


Operations.implementation_for(AddPermission)(basic_add_by_name)


@Operations.register_operation("remove_permission", "invoke_op")
class RemovePermission(PermissionOperation, BaseOperation):
    def __init__(self, name):
        self.name = name


Operations.implementation_for(RemovePermission)(basic_delete_by_name)


class ViewMenuOperation(MigrateOperation):
    table = table(
        'ab_view_menu',
        column('id', Integer),
        column('name', String),
    )


@Operations.register_operation("add_view_menu", "invoke_op")
class AddViewMenu(ViewMenuOperation, BaseOperation):
    def __init__(self, name):
        self.name = name


Operations.implementation_for(AddViewMenu)(basic_add_by_name)


@Operations.register_operation("remove_view_menu", "invoke_op")
class RemoveViewMenu(PermissionOperation, BaseOperation):
    def __init__(self, name):
        self.name = name


Operations.implementation_for(RemoveViewMenu)(basic_delete_by_name)


@Operations.register_operation("add_permission_view", "invoke_op")
class AddPermissionView(PermissionViewOperation, BaseOperation):
    permission_table = PermissionOperation.table
    view_table = ViewMenuOperation.table

    def __init__(self, permission, view):
        self.permission = permission
        self.view = view


@Operations.implementation_for(AddPermissionView)
def basic_add_by_name(operations, operation):
    __import__('ipdb').set_trace()
    pt = operation.permission_table
    pt.select([pt.id]).where()
    operations.execute(operation.table.insert().values(name=operation.name))

test_migraton.py

#
# 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.

"""Update DAG roles for FAB

Revision ID: a83bc0fd663d
Revises: 939bb1e647c8
Create Date: 2019-03-05 21:54:59.710545

"""

from alembic import op

import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'a83bc0fd663d'
down_revision = '939bb1e647c8'
branch_labels = None
depends_on = None


def upgrade():
    op.add_permission("x")
    op.add_view_menu("y")
    op.add_permission_view("x", "y")
    pass


def downgrade():
    op.remove_permission("x")
    pass