fix: upgrade_catalog_perms and downgrade_catalog_perms implementation

[michael-s-molina]

SUMMARY

This PR fixes the implementation of upgrade_catalog_perms and downgrade_catalog_perms methods to handle tables with millions of rows. To achieve this, I changed the algorithm to execute the migrations using a batched approach invoking the update/delete commands directly using a subquery instead of changing the records one by one. It also commits the transaction on every batch to avoid keeping a long standing transaction which leads to memory and timeout issues. Finally, I added detailed logging to help system administrators to track the migration progress and prepare for downtime.

Fixes #29801

TESTING INSTRUCTIONS

1 - Run any migration that uses the changed commands
2 - Check that both upgrades and downgrades work

ADDITIONAL INFORMATION

• Has associated issue:
• Required feature flags:
• Changes UI
• Includes DB Migration (follow approval process in SIP-59)
  • Migration is atomic, supports rollback & is backwards-compatible
  • Confirm DB migration upgrade and downgrade tested
  • Runtime estimates and downtime expectations provided
• Introduces new feature or API
• Removes existing feature or API

[codecov]

Codecov Report

Attention: Patch coverage is 87.83784% with 9 lines in your changes missing coverage. Please review.

Project coverage is 83.69%. Comparing base (76d897e) to head (adc0996).
Report is 1094 commits behind head on master.

Files with missing lines	Patch %	Lines
superset/migrations/shared/catalogs.py	87.83%	9 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #29860       +/-   ##
===========================================
+ Coverage   60.48%   83.69%   +23.20%     
===========================================
  Files        1931      527     -1404     
  Lines       76236    38088    -38148     
  Branches     8568        0     -8568     
===========================================
- Hits        46114    31878    -14236     
+ Misses      28017     6210    -21807     
+ Partials     2105        0     -2105     

Flag	Coverage Δ
hive	48.97% <0.00%> (-0.19%)	⬇️
javascript	?
mysql	76.71% <0.00%> (?)
postgres	76.78% <0.00%> (?)
presto	53.51% <0.00%> (-0.29%)	⬇️
python	83.69% <87.83%> (+20.20%)	⬆️
sqlite	76.27% <0.00%> (?)
unit	60.37% <87.83%> (+2.75%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[eschutho]

I had also suggested making transactions smaller on migrations, but the limitation that I've been told is that if the second batch fails, we won't be able to roll back the first batch. Do you have an idea how we can do this so that it can roll back gracefully?

[michael-s-molina]

In this case, rolling back means setting the column to None, which can be done independently of previous failures.

[michael-s-molina]

@eschutho Answering your question more generically, there are some possible strategies to revert batch migrations such as logging changes to be reverted, making migrations reversible by storing the original value in a dedicated column (we do that for chart migrations) or using database checkpoints.

[eschutho]

Gotcha, would the other "possible strategies" involve another custom script to get back to the pre-migration state? I'm just concerned that for most people just running an upgrade that fails on a batch other than the first batch would get them in a state where they could have some columns with a catalog value, and some without. @betodealmeida may be able to answer if this would be problematic, knowing how the codebase uses these values.

[michael-s-molina]

@eschutho I was able to remove the commit statement after verifying that we could hold/rollback a transaction with an appropriate block size.

[eschutho]

ok great. 🙏

[michael-s-molina]

I'll also need to fix this part which generates a query for each dataset, which in Airbnb's case is more than 25,000 queries.

# update `schema_perm` and `catalog_perm` for tables and charts
for table in session.query(SqlaTable).filter_by(
    database_id=database.id,
    catalog=None,
):
    schema_perm = security_manager.get_schema_perm(
        database.database_name,
        default_catalog,
        table.schema,
    )

    table.catalog = default_catalog
    table.catalog_perm = catalog_perm
    table.schema_perm = schema_perm

    for chart in session.query(Slice).filter_by(
        datasource_id=table.id,
        datasource_type="table",
    ):
        chart.catalog_perm = catalog_perm
        chart.schema_perm = schema_perm

[betodealmeida]

This is great, thanks @michael-s-molina! I was worried about performance when I introduced the feature, but it's hard to test things at the AirBnB scale.

[betodealmeida]

Oops, approved too early.

[betodealmeida]

I'll also need to fix this part which generates a query for each dataset, which in Airbnb's case is more than 25,000 queries.

# update `schema_perm` and `catalog_perm` for tables and charts
for table in session.query(SqlaTable).filter_by(
    database_id=database.id,
    catalog=None,
):
    schema_perm = security_manager.get_schema_perm(
        database.database_name,
        default_catalog,
        table.schema,
    )

    table.catalog = default_catalog
    table.catalog_perm = catalog_perm
    table.schema_perm = schema_perm

    for chart in session.query(Slice).filter_by(
        datasource_id=table.id,
        datasource_type="table",
    ):
        chart.catalog_perm = catalog_perm
        chart.schema_perm = schema_perm

You might be able to rewrite this as a single query for some dialects:

UPDATE tables
SET
  catalog=${default_catalog},
  catalog_perm=${catalog_perm},
  schema_perm=REGEXP_REPLACE(schema_perm, '\[([^\]]+)\]\.\[([^\]]+)\]', '[\1].[${default_catalog}].[\2]', 'g')
WHERE
  database_id=${database_id} AND
  catalog IS NULL
RETURNING id;

For charts you can then use the returned IDs from the query above to build a similar UPDATE query.

[michael-s-molina]

You might be able to rewrite this as a single query for some dialects:

@betodealmeida Given that the number of rows for tables and charts will hardly pass 1 million, I opted for preserving a lot of the previous logic and eliminating the part where we queried the slices table for each dataset. Now we query for all slices of type table for each database and I use a map to get the appropriate permissions. This will significantly reduce the number of queries to the number of databases.

[sadpandajoe]

@supersetbot label 4.1

[michael-s-molina]

@betodealmeida I found another problem with how we were dealing with catalog and schema permissions that come from the security manager. Looking at the interface definition, they could be None but we were not handling these cases which resulted in errors when inserting to ab_view_menu because of NULL names. I fixed the problem here but I don't have much context about these permissions so please take a look.

[sadpandajoe]

@betodealmeida @eschutho any other concerns about this pr? I think @michael-s-molina has addressed both your comments.

[betodealmeida]

@betodealmeida I found another problem with how we were dealing with catalog and schema permissions that come from the security manager. Looking at the interface definition, they could be None but we were not handling these cases which resulted in errors when inserting to ab_view_menu because of NULL names. I fixed the problem here but I don't have much context about these permissions so please take a look.

Ah, good point! This looks correct.

[michael-s-molina]

@betodealmeida Can you approve or remove the request for changes?

[eschutho]

lgtm

[betodealmeida]

Thank you!