-
Notifications
You must be signed in to change notification settings - Fork 14.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: add UUID column to ImportMixin #11098
Merged
Merged
Changes from all commits
Commits
Show all changes
18 commits
Select commit
Hold shift + click to select a range
da65260
Add UUID column to ImportMixin
betodealmeida 17a3971
Fix default value
betodealmeida f117619
Fix lint
betodealmeida c927029
Fix order of downgrade
betodealmeida 8fa4db7
Add logging when downgrade fails
betodealmeida fc5a8cb
Migrate position_json to contain UUIDs, and add schedule tables
betodealmeida 232b171
Save UUID when adding charts to dashboard
betodealmeida 1d4554b
Fix heads
betodealmeida 6bcfdcc
Rename migration file
betodealmeida bcbbb0b
Fix dashboard serialization
betodealmeida 239ba38
Fix migration script with Postgres
betodealmeida 0a4ea63
Fix unique contraint name
betodealmeida a8a55d6
Handle UUID when exporting dashboard
betodealmeida f171518
Fix Dataset PUT
betodealmeida 12d56ad
Add UUID JSON serialization
betodealmeida 5ec2537
Fix tests
betodealmeida aad4f87
Simplify logic
betodealmeida 921960a
Try binary=True
betodealmeida File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
154 changes: 154 additions & 0 deletions
154
superset/migrations/versions/b56500de1855_add_uuid_column_to_import_mixin.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,154 @@ | ||
# 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. | ||
"""add_uuid_column_to_import_mixin | ||
|
||
Revision ID: b56500de1855 | ||
Revises: 18532d70ab98 | ||
Create Date: 2020-09-28 17:57:23.128142 | ||
|
||
""" | ||
import json | ||
import logging | ||
import uuid | ||
|
||
import sqlalchemy as sa | ||
from alembic import op | ||
from sqlalchemy.ext.declarative import declarative_base | ||
from sqlalchemy_utils import UUIDType | ||
|
||
from superset import db | ||
from superset.utils import core as utils | ||
|
||
# revision identifiers, used by Alembic. | ||
revision = "b56500de1855" | ||
down_revision = "18532d70ab98" | ||
|
||
|
||
Base = declarative_base() | ||
|
||
|
||
class ImportMixin: | ||
id = sa.Column(sa.Integer, primary_key=True) | ||
uuid = sa.Column(UUIDType(binary=True), primary_key=False, default=uuid.uuid4) | ||
|
||
|
||
table_names = [ | ||
# Core models | ||
"dbs", | ||
"dashboards", | ||
"slices", | ||
# SQLAlchemy connectors | ||
"tables", | ||
"table_columns", | ||
"sql_metrics", | ||
# Druid connector | ||
"clusters", | ||
"datasources", | ||
"columns", | ||
"metrics", | ||
# Dashboard email schedules | ||
"dashboard_email_schedules", | ||
"slice_email_schedules", | ||
] | ||
models = { | ||
table_name: type(table_name, (Base, ImportMixin), {"__tablename__": table_name}) | ||
for table_name in table_names | ||
} | ||
|
||
models["dashboards"].position_json = sa.Column(utils.MediumText()) | ||
|
||
|
||
def add_uuids(objects, session, batch_size=100): | ||
uuid_map = {} | ||
count = len(objects) | ||
for i, object_ in enumerate(objects): | ||
object_.uuid = uuid.uuid4() | ||
uuid_map[object_.id] = object_.uuid | ||
session.merge(object_) | ||
if (i + 1) % batch_size == 0: | ||
session.commit() | ||
print(f"uuid assigned to {i + 1} out of {count}") | ||
|
||
session.commit() | ||
print(f"Done! Assigned {count} uuids") | ||
|
||
return uuid_map | ||
|
||
|
||
def update_position_json(dashboard, session, uuid_map): | ||
layout = json.loads(dashboard.position_json or "{}") | ||
for object_ in layout.values(): | ||
if ( | ||
isinstance(object_, dict) | ||
and object_["type"] == "CHART" | ||
and object_["meta"]["chartId"] | ||
): | ||
chart_id = object_["meta"]["chartId"] | ||
if chart_id in uuid_map: | ||
object_["meta"]["uuid"] = str(uuid_map[chart_id]) | ||
elif object_["meta"].get("uuid"): | ||
del object_["meta"]["uuid"] | ||
|
||
dashboard.position_json = json.dumps(layout, indent=4) | ||
session.merge(dashboard) | ||
session.commit() | ||
|
||
|
||
def upgrade(): | ||
bind = op.get_bind() | ||
session = db.Session(bind=bind) | ||
|
||
uuid_maps = {} | ||
for table_name, model in models.items(): | ||
with op.batch_alter_table(table_name) as batch_op: | ||
batch_op.add_column( | ||
sa.Column( | ||
"uuid", | ||
UUIDType(binary=True), | ||
primary_key=False, | ||
default=uuid.uuid4, | ||
) | ||
) | ||
|
||
# populate column | ||
objects = session.query(model).all() | ||
uuid_maps[table_name] = add_uuids(objects, session) | ||
|
||
# add uniqueness constraint | ||
with op.batch_alter_table(table_name) as batch_op: | ||
batch_op.create_unique_constraint(f"uq_{table_name}_uuid", ["uuid"]) | ||
|
||
# add UUID to Dashboard.position_json | ||
Dashboard = models["dashboards"] | ||
for dashboard in session.query(Dashboard).all(): | ||
update_position_json(dashboard, session, uuid_maps["slices"]) | ||
|
||
|
||
def downgrade(): | ||
bind = op.get_bind() | ||
session = db.Session(bind=bind) | ||
|
||
# remove uuid from position_json | ||
Dashboard = models["dashboards"] | ||
for dashboard in session.query(Dashboard).all(): | ||
update_position_json(dashboard, session, {}) | ||
|
||
# remove uuid column | ||
for table_name, model in models.items(): | ||
with op.batch_alter_table(model) as batch_op: | ||
batch_op.drop_constraint(f"uq_{table_name}_uuid") | ||
batch_op.drop_column("uuid") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@betodealmeida this migration is very slow, it is worth to mention in the changelog e.g. for our staging env it took ~30 min and often it means extra downtime for the service
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point, @bkyryliuk! I'll add it today.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@bkyryliuk @betodealmeida I've managed to rewrite the uuid generation with native SQL queries and sped up the migration process by more than 100x. The whole migration job can now complete in under 5 minutes for our Superset db of more than 200k
slices
and 1 milliontable_columns
. Do you mind taking a look and maybe testing it on your Superset deployments as well?