Skip to content

Commit

Permalink
change constraint alters to non blocking
Browse files Browse the repository at this point in the history
  • Loading branch information
eschutho committed Sep 2, 2023
1 parent bbfaeb0 commit 424ed12
Show file tree
Hide file tree
Showing 2 changed files with 174 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,12 @@
revision = "8ace289026f3"
down_revision = "2e826adca42c"

from alembic import op
from sqlalchemy import create_engine
from sqlalchemy.engine.reflection import Inspector

from superset.migrations.shared.constraints import ForeignKey, redefine
from superset.utils.core import generic_find_fk_constraint_name

foreign_keys = [
ForeignKey(
Expand All @@ -47,9 +51,90 @@

def upgrade():
for foreign_key in foreign_keys:
redefine(foreign_key, on_delete="CASCADE")
bind = op.get_bind()
insp = Inspector.from_engine(bind)
db_dialect = bind.engine.url.drivername

if db_dialect == "postgresql":
if constraint := generic_find_fk_constraint_name(
table=foreign_key.table,
columns=set(foreign_key.remote_cols),
referenced=foreign_key.referent_table,
insp=insp,
):
# rename the constraint so that we can drop it later- non-blocking
# this only works on postgres
op.execute(
f"ALTER TABLE {foreign_key.table} RENAME CONSTRAINT {constraint} TO {foreign_key.constraint_name}_old"
)

# create the new constraint- uses a non-blocking ALTER TABLE Not Valid property
op.create_foreign_key(
constraint_name=foreign_key.constraint_name,
source_table=foreign_key.table,
referent_table=foreign_key.referent_table,
local_cols=foreign_key.local_cols,
remote_cols=foreign_key.remote_cols,
ondelete="CASCADE",
postgresql_not_valid=True,
)

# validate the constraint- non-blocking
op.execute(
f"ALTER TABLE {foreign_key.table} VALIDATE CONSTRAINT {foreign_key.constraint_name}"
)

# drop the old constraint which isn't needed anymore- non-blocking
if constraint:
op.drop_constraint(
f"{foreign_key.constraint_name}_old",
foreign_key.table,
type_="foreignkey",
)
else:
redefine(foreign_key, on_delete="CASCADE")


def downgrade():
for foreign_key in foreign_keys:
redefine(foreign_key)
bind = op.get_bind()
insp = Inspector.from_engine(bind)
db_dialect = bind.engine.url.drivername

if db_dialect == "postgresql":
if constraint := generic_find_fk_constraint_name(
table=foreign_key.table,
columns=set(foreign_key.remote_cols),
referenced=foreign_key.referent_table,
insp=insp,
):
# rename the constraint so that we can drop it later- non-blocking
# this only works on postgres
op.execute(
f"ALTER TABLE {foreign_key.table} RENAME CONSTRAINT {constraint} TO {foreign_key.constraint_name}_old"
)

# create the new constraint- uses a non-blocking ALTER TABLE Not Valid property
op.create_foreign_key(
constraint_name=foreign_key.constraint_name,
source_table=foreign_key.table,
referent_table=foreign_key.referent_table,
local_cols=foreign_key.local_cols,
remote_cols=foreign_key.remote_cols,
postgresql_not_valid=True,
)

# validate the constraint- non-blocking
op.execute(
f"ALTER TABLE {foreign_key.table} VALIDATE CONSTRAINT {foreign_key.constraint_name}"
)

# drop the old constraint which isn't needed anymore- non-blocking
if constraint:
op.drop_constraint(
f"{foreign_key.constraint_name}_old",
foreign_key.table,
type_="foreignkey",
)
else:
redefine(foreign_key)
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,11 @@
revision = "4448fa6deeb1"
down_revision = "8ace289026f3"

from alembic import op
from sqlalchemy.engine.reflection import Inspector

from superset.migrations.shared.constraints import ForeignKey, redefine
from superset.utils.core import generic_find_fk_constraint_name

foreign_keys = [
ForeignKey(
Expand All @@ -40,9 +44,90 @@

def upgrade():
for foreign_key in foreign_keys:
redefine(foreign_key, on_delete="CASCADE")
bind = op.get_bind()
insp = Inspector.from_engine(bind)
db_dialect = bind.engine.url.drivername

if db_dialect == "postgresql":
if constraint := generic_find_fk_constraint_name(
table=foreign_key.table,
columns=set(foreign_key.remote_cols),
referenced=foreign_key.referent_table,
insp=insp,
):
# rename the constraint so that we can drop it later- non-blocking
# this only works on postgres
op.execute(
f"ALTER TABLE {foreign_key.table} RENAME CONSTRAINT {constraint} TO {foreign_key.constraint_name}_old"
)

# create the new constraint- uses a non-blocking ALTER TABLE Not Valid property
op.create_foreign_key(
constraint_name=foreign_key.constraint_name,
source_table=foreign_key.table,
referent_table=foreign_key.referent_table,
local_cols=foreign_key.local_cols,
remote_cols=foreign_key.remote_cols,
ondelete="CASCADE",
postgresql_not_valid=True,
)

# validate the constraint- non-blocking
op.execute(
f"ALTER TABLE {foreign_key.table} VALIDATE CONSTRAINT {foreign_key.constraint_name}"
)

# drop the old constraint which isn't needed anymore- non-blocking
if constraint:
op.drop_constraint(
f"{foreign_key.constraint_name}_old",
foreign_key.table,
type_="foreignkey",
)
else:
redefine(foreign_key, on_delete="CASCADE")


def downgrade():
for foreign_key in foreign_keys:
redefine(foreign_key)
bind = op.get_bind()
insp = Inspector.from_engine(bind)
db_dialect = bind.engine.url.drivername

if db_dialect == "postgresql":
if constraint := generic_find_fk_constraint_name(
table=foreign_key.table,
columns=set(foreign_key.remote_cols),
referenced=foreign_key.referent_table,
insp=insp,
):
# rename the constraint so that we can drop it later- non-blocking
# this only works on postgres
op.execute(
f"ALTER TABLE {foreign_key.table} RENAME CONSTRAINT {constraint} TO {foreign_key.constraint_name}_old"
)

# create the new constraint- uses a non-blocking ALTER TABLE Not Valid property
op.create_foreign_key(
constraint_name=foreign_key.constraint_name,
source_table=foreign_key.table,
referent_table=foreign_key.referent_table,
local_cols=foreign_key.local_cols,
remote_cols=foreign_key.remote_cols,
postgresql_not_valid=True,
)

# validate the constraint- non-blocking
op.execute(
f"ALTER TABLE {foreign_key.table} VALIDATE CONSTRAINT {foreign_key.constraint_name}"
)

# drop the old constraint which isn't needed anymore- non-blocking
if constraint:
op.drop_constraint(
f"{foreign_key.constraint_name}_old",
foreign_key.table,
type_="foreignkey",
)
else:
redefine(foreign_key)

0 comments on commit 424ed12

Please sign in to comment.