airbnb/caravel

# 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.
"""Delete None permissions

Revision ID: e786798587de
Revises: 6f139c533bea
Create Date: 2022-05-18 16:07:47.648514

"""

# revision identifiers, used by Alembic.
revision = "e786798587de"
down_revision = "6f139c533bea"

from alembic import op  # noqa: E402
from sqlalchemy import (  # noqa: E402
    Column,
    ForeignKey,
    Integer,
    Sequence,
    String,
    Table,
    UniqueConstraint,
)
from sqlalchemy.ext.declarative import declarative_base  # noqa: E402
from sqlalchemy.orm import relationship, Session  # noqa: E402

Base = declarative_base()


class Permission(Base):
    __tablename__ = "ab_permission"
    id = Column(Integer, Sequence("ab_permission_id_seq"), primary_key=True)
    name = Column(String(100), unique=True, nullable=False)

    def __repr__(self):
        return self.name


class ViewMenu(Base):
    __tablename__ = "ab_view_menu"
    id = Column(Integer, Sequence("ab_view_menu_id_seq"), primary_key=True)
    name = Column(String(250), unique=True, nullable=False)

    def __repr__(self) -> str:
        return self.name


assoc_permissionview_role = Table(
    "ab_permission_view_role",
    Base.metadata,
    Column("id", Integer, Sequence("ab_permission_view_role_id_seq"), primary_key=True),
    Column("permission_view_id", Integer, ForeignKey("ab_permission_view.id")),
    Column("role_id", Integer, ForeignKey("ab_role.id")),
    UniqueConstraint("permission_view_id", "role_id"),
)


class Role(Base):
    __tablename__ = "ab_role"

    id = Column(Integer, Sequence("ab_role_id_seq"), primary_key=True)
    name = Column(String(64), unique=True, nullable=False)
    permissions = relationship(
        "PermissionView", secondary=assoc_permissionview_role, backref="role"
    )

    def __repr__(self) -> str:
        return f"{self.name}"


class PermissionView(Base):
    __tablename__ = "ab_permission_view"
    __table_args__ = (UniqueConstraint("permission_id", "view_menu_id"),)
    id = Column(Integer, Sequence("ab_permission_view_id_seq"), primary_key=True)
    permission_id = Column(Integer, ForeignKey("ab_permission.id"))
    permission = relationship("Permission")
    view_menu_id = Column(Integer, ForeignKey("ab_view_menu.id"))
    view_menu = relationship("ViewMenu")

    def __repr__(self) -> str:
        return f"{self.permission.name} on {self.view_menu.name}"


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    bind = op.get_bind()
    session = Session(bind=bind)

    pvms = (
        session.query(PermissionView)
        .join(ViewMenu)
        .join(Permission)
        .filter(
            Permission.name.in_(("datasource_access", "schema_access")),
            ViewMenu.name.like("[None].%"),
        )
        .all()
    )

    roles = (
        session.query(Role)
        .outerjoin(Role.permissions)
        .join(ViewMenu)
        .join(Permission)
        .filter(
            Permission.name.in_(("datasource_access", "schema_access")),
            ViewMenu.name.like("[None].%"),
        )
        .all()
    )

    for pvm in pvms:
        for role in roles:
            if pvm in role.permissions:
                print(
                    f"Going to delete a data access permission [{pvm}] on Role [{role}]"
                )
                role.permissions.remove(pvm)
        print(f"Going to delete a data access permission [{pvm}]")
        session.delete(pvm)
        session.delete(pvm.view_menu)
    session.commit()
    session.close()
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    ...
    # ### end Alembic commands ###