[superset] 02/02: working migration and need to fill in owners values

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hugh pushed a commit to branch add-owners-new-dataset
in repository https://gitbox.apache.org/repos/asf/superset.git

commit c1a5e5adb8e771d0481db4b5a4ea7dc0d602c449
Author: hughhhh <hu...@gmail.com>
AuthorDate: Thu Mar 31 23:06:36 2022 -0400

    working migration and need to fill in owners values
---
 .../89f17d951737_add_owners_dataset_model.py       | 67 ++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/superset/migrations/versions/89f17d951737_add_owners_dataset_model.py b/superset/migrations/versions/89f17d951737_add_owners_dataset_model.py
index 5b83609..fe2f38e 100644
--- a/superset/migrations/versions/89f17d951737_add_owners_dataset_model.py
+++ b/superset/migrations/versions/89f17d951737_add_owners_dataset_model.py
@@ -29,11 +29,63 @@ down_revision = "2ed890b36b94"
 import sqlalchemy as sa
 from alembic import op
 from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+
+from superset import db, security_manager
+
+Base = declarative_base()
+
+
+class User(Base):
+    """Declarative class to do query in upgrade"""
+
+    __tablename__ = "ab_user"
+    id = sa.Column(sa.Integer, primary_key=True)
+
+
+class SqlaTable(Base):
+    __tablename__ = "tables"
+
+    id = sa.Column(sa.Integer, primary_key=True)
+    owners = relationship(
+        "User",
+        secondary=sa.Table(
+            "sqlatable_user",
+            Base.metadata,
+            sa.Column("id", sa.Integer, primary_key=True),
+            sa.Column("user_id", sa.Integer, sa.ForeignKey("ab_user.id")),
+            sa.Column("table_id", sa.Integer, sa.ForeignKey("tables.id")),
+        ),
+        backref="tables",
+    )
+
+
+class Dataset(Base):
+    __tablename__ = "sl_datasets"
+
+    id = sa.Column(sa.Integer, primary_key=True)
+    sqlatable_id = sa.Column(sa.Integer, nullable=True, unique=True)
+    owners = relationship(
+        "User",
+        secondary=sa.Table(
+            "sl_dataset_users",
+            Base.metadata,
+            sa.Column("id", sa.Integer, primary_key=True),
+            sa.Column("user_id", sa.ForeignKey("ab_user.id")),
+            sa.Column("dataset_id", sa.ForeignKey("sl_datasets.id")),
+        ),
+        backref="sl_datasets",
+    )
 
 
 def upgrade():
+    bind = op.get_bind()
+    session = db.Session(bind=bind)
+
     op.create_table(
         "sl_dataset_users",
+        sa.Column("id", sa.Integer(), primary_key=True),
         sa.Column("user_id", sa.Integer(), nullable=True),
         sa.Column("dataset_id", sa.Integer(), nullable=True),
         sa.ForeignKeyConstraint(
@@ -46,6 +98,21 @@ def upgrade():
         ),
     )
 
+    for sqlatable in session.query(SqlaTable).all():
+        try:
+            ds = (
+                session.query(Dataset)
+                .filter(Dataset.sqlatable_id == sqlatable.id)
+                .one()
+            )
+        except sa.orm.exc.NoResultFound:
+            continue
+
+        ds.owners = sqlatable.owners
+
+    session.commit()
+    session.close()
+
 
 def downgrade():
     op.drop_table("sl_dataset_users")