From ed0a61d72a2b4ad9607c5d3681086110fdf9442d Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Wed, 16 Mar 2022 23:18:32 -0300 Subject: [PATCH 1/4] fix: update sql lab models constraints to support db deletion --- ..._lab_models_database_constraint_updates.py | 111 ++++++++++++++++++ superset/models/sql_lab.py | 10 +- 2 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py diff --git a/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py new file mode 100644 index 0000000000000..d5c2da66253f0 --- /dev/null +++ b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py @@ -0,0 +1,111 @@ +# 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. +"""sql_lab_models_database_constraint_updates + +Revision ID: 8b841273bec3 +Revises: 6766938c6065 +Create Date: 2022-03-16 21:07:48.768425 + +""" + +# revision identifiers, used by Alembic. +revision = "8b841273bec3" +down_revision = "6766938c6065" + +import sqlalchemy as sa +from alembic import op + +from superset.utils.core import generic_find_fk_constraint_name + + +def upgrade(): + bind = op.get_bind() + insp = sa.engine.reflection.Inspector.from_engine(bind) + + with op.batch_alter_table("tab_state") as batch_op: + batch_op.drop_constraint( + generic_find_fk_constraint_name("tab_state", {"id"}, "dbs", insp), + type_="foreignkey", + ) + + batch_op.drop_constraint( + generic_find_fk_constraint_name("tab_state", {"client_id"}, "query", insp), + type_="foreignkey", + ) + + batch_op.create_foreign_key( + "tab_state_database_id_fkey", + "dbs", + ["database_id"], + ["id"], + ondelete="CASCADE", + ) + + batch_op.create_foreign_key( + "tab_state_latest_query_id_fkey", + "query", + ["latest_query_id"], + ["client_id"], + ondelete="SET NULL", + ) + + with op.batch_alter_table("table_schema") as batch_op: + batch_op.drop_constraint( + generic_find_fk_constraint_name("table_schema", {"id"}, "dbs", insp), + type_="foreignkey", + ) + batch_op.create_foreign_key( + "table_schema_database_id_fkey", + "dbs", + ["database_id"], + ["id"], + ondelete="CASCADE", + ) + + +def downgrade(): + bind = op.get_bind() + insp = sa.engine.reflection.Inspector.from_engine(bind) + + with op.batch_alter_table("tab_state") as batch_op: + batch_op.drop_constraint( + generic_find_fk_constraint_name("tab_state", {"id"}, "dbs", insp), + type_="foreignkey", + ) + batch_op.drop_constraint( + generic_find_fk_constraint_name("tab_state", {"client_id"}, "query", insp), + type_="foreignkey", + ) + + batch_op.create_foreign_key( + "tab_state_database_id_fkey", "dbs", ["database_id"], ["id"] + ) + batch_op.create_foreign_key( + "tab_state_latest_query_id_fkey", + "query", + ["latest_query_id"], + ["client_id"], + ) + + with op.batch_alter_table("table_schema") as batch_op: + batch_op.drop_constraint( + generic_find_fk_constraint_name("table_schema", {"id"}, "dbs", insp), + type_="foreignkey", + ) + batch_op.create_foreign_key( + "table_schema_database_id_fkey", "dbs", ["database_id"], ["id"] + ) diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py index 6a3b4ad8bfd7c..3c117c629d7df 100644 --- a/superset/models/sql_lab.py +++ b/superset/models/sql_lab.py @@ -265,7 +265,7 @@ class TabState(Model, AuditMixinNullable, ExtraJSONMixin): active = Column(Boolean, default=False) # selected DB and schema - database_id = Column(Integer, ForeignKey("dbs.id")) + database_id = Column(Integer, ForeignKey("dbs.id", ondelete="CASCADE")) database = relationship("Database", foreign_keys=[database_id]) schema = Column(String(256)) @@ -282,7 +282,9 @@ class TabState(Model, AuditMixinNullable, ExtraJSONMixin): query_limit = Column(Integer) # latest query that was run - latest_query_id = Column(Integer, ForeignKey("query.client_id")) + latest_query_id = Column( + Integer, ForeignKey("query.client_id", ondelete="SET NULL") + ) latest_query = relationship("Query") # other properties @@ -322,7 +324,9 @@ class TableSchema(Model, AuditMixinNullable, ExtraJSONMixin): id = Column(Integer, primary_key=True, autoincrement=True) tab_state_id = Column(Integer, ForeignKey("tab_state.id", ondelete="CASCADE")) - database_id = Column(Integer, ForeignKey("dbs.id"), nullable=False) + database_id = Column( + Integer, ForeignKey("dbs.id", ondelete="CASCADE"), nullable=False + ) database = relationship("Database", foreign_keys=[database_id]) schema = Column(String(256)) table = Column(String(256)) From 8e91afa9c926616f8d3982d899329bd76394bb04 Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Wed, 30 Mar 2022 17:42:09 -0300 Subject: [PATCH 2/4] update with master --- .../58df9d617f14_add_on_saved_query_delete_tab_state_.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py b/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py index 57e13cf1488e6..6e608793dd8da 100644 --- a/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py +++ b/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py @@ -17,14 +17,14 @@ """add_on_saved_query_delete_tab_state_null_constraint" Revision ID: 58df9d617f14 -Revises: 6766938c6065 +Revises: 8b841273bec3 Create Date: 2022-03-16 23:24:40.278937 """ # revision identifiers, used by Alembic. revision = "58df9d617f14" -down_revision = "6766938c6065" +down_revision = "8b841273bec3" import sqlalchemy as sa from alembic import op From 20f2906c2c62886fa053fdf2966dd3b69c2ea9c6 Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Fri, 1 Apr 2022 21:03:39 -0300 Subject: [PATCH 3/4] solve db migration conflict --- .../58df9d617f14_add_on_saved_query_delete_tab_state_.py | 4 ++-- ...8b841273bec3_sql_lab_models_database_constraint_updates.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py b/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py index 6e608793dd8da..57e13cf1488e6 100644 --- a/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py +++ b/superset/migrations/versions/58df9d617f14_add_on_saved_query_delete_tab_state_.py @@ -17,14 +17,14 @@ """add_on_saved_query_delete_tab_state_null_constraint" Revision ID: 58df9d617f14 -Revises: 8b841273bec3 +Revises: 6766938c6065 Create Date: 2022-03-16 23:24:40.278937 """ # revision identifiers, used by Alembic. revision = "58df9d617f14" -down_revision = "8b841273bec3" +down_revision = "6766938c6065" import sqlalchemy as sa from alembic import op diff --git a/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py index d5c2da66253f0..4e92438696d15 100644 --- a/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py +++ b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py @@ -17,14 +17,14 @@ """sql_lab_models_database_constraint_updates Revision ID: 8b841273bec3 -Revises: 6766938c6065 +Revises: 2ed890b36b94 Create Date: 2022-03-16 21:07:48.768425 """ # revision identifiers, used by Alembic. revision = "8b841273bec3" -down_revision = "6766938c6065" +down_revision = "2ed890b36b94" import sqlalchemy as sa from alembic import op From 7dc0fe048320109c0759b77c55968b9c649a57ea Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Mon, 4 Apr 2022 18:50:26 -0300 Subject: [PATCH 4/4] fix tests --- ..._lab_models_database_constraint_updates.py | 65 +++++++++++++------ 1 file changed, 46 insertions(+), 19 deletions(-) diff --git a/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py index 4e92438696d15..a497cf80f4bc0 100644 --- a/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py +++ b/superset/migrations/versions/8b841273bec3_sql_lab_models_database_constraint_updates.py @@ -37,15 +37,23 @@ def upgrade(): insp = sa.engine.reflection.Inspector.from_engine(bind) with op.batch_alter_table("tab_state") as batch_op: - batch_op.drop_constraint( - generic_find_fk_constraint_name("tab_state", {"id"}, "dbs", insp), - type_="foreignkey", + table_schema_id_constraint = generic_find_fk_constraint_name( + "tab_state", {"id"}, "dbs", insp ) - - batch_op.drop_constraint( - generic_find_fk_constraint_name("tab_state", {"client_id"}, "query", insp), - type_="foreignkey", + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) + + table_schema_id_constraint = generic_find_fk_constraint_name( + "tab_state", {"client_id"}, "query", insp ) + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) batch_op.create_foreign_key( "tab_state_database_id_fkey", @@ -64,10 +72,15 @@ def upgrade(): ) with op.batch_alter_table("table_schema") as batch_op: - batch_op.drop_constraint( - generic_find_fk_constraint_name("table_schema", {"id"}, "dbs", insp), - type_="foreignkey", + table_schema_id_constraint = generic_find_fk_constraint_name( + "table_schema", {"id"}, "dbs", insp ) + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) + batch_op.create_foreign_key( "table_schema_database_id_fkey", "dbs", @@ -82,14 +95,23 @@ def downgrade(): insp = sa.engine.reflection.Inspector.from_engine(bind) with op.batch_alter_table("tab_state") as batch_op: - batch_op.drop_constraint( - generic_find_fk_constraint_name("tab_state", {"id"}, "dbs", insp), - type_="foreignkey", + table_schema_id_constraint = generic_find_fk_constraint_name( + "tab_state", {"id"}, "dbs", insp ) - batch_op.drop_constraint( - generic_find_fk_constraint_name("tab_state", {"client_id"}, "query", insp), - type_="foreignkey", + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) + + table_schema_id_constraint = generic_find_fk_constraint_name( + "tab_state", {"client_id"}, "query", insp ) + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) batch_op.create_foreign_key( "tab_state_database_id_fkey", "dbs", ["database_id"], ["id"] @@ -102,10 +124,15 @@ def downgrade(): ) with op.batch_alter_table("table_schema") as batch_op: - batch_op.drop_constraint( - generic_find_fk_constraint_name("table_schema", {"id"}, "dbs", insp), - type_="foreignkey", + table_schema_id_constraint = generic_find_fk_constraint_name( + "table_schema", {"id"}, "dbs", insp ) + if table_schema_id_constraint: + batch_op.drop_constraint( + table_schema_id_constraint, + type_="foreignkey", + ) + batch_op.create_foreign_key( "table_schema_database_id_fkey", "dbs", ["database_id"], ["id"] )