fix(migration): handle permalink edge cases correctly (#23980)

This commit is contained in:
Ville Brofeldt
2023-05-09 15:28:24 +03:00
committed by GitHub
parent d96b72d46f
commit 7a4117097a
2 changed files with 20 additions and 3 deletions

View File

@@ -45,7 +45,7 @@ class CreateExplorePermalinkCommand(BaseExplorePermalinkCommand):
value = {
"chartId": self.chart_id,
"datasourceId": datasource_id,
"datasourceType": datasource_type,
"datasourceType": datasource_type.value,
"datasource": self.datasource,
"state": self.state,
}

View File

@@ -45,7 +45,10 @@ RESOURCES_TO_MIGRATE = ("app", "dashboard_permalink", "explore_permalink")
class RestrictedUnpickler(pickle.Unpickler):
def find_class(self, module, name):
raise pickle.UnpicklingError(f"Unpickling of {module}.{name} is forbidden")
if not (module == "superset.utils.core" and name == "DatasourceType"):
raise pickle.UnpicklingError(f"Unpickling of {module}.{name} is forbidden")
return super().find_class(module, name)
class KeyValueEntry(Base):
@@ -58,14 +61,28 @@ class KeyValueEntry(Base):
def upgrade():
bind = op.get_bind()
session: Session = db.Session(bind=bind)
truncated_count = 0
for entry in paginated_update(
session.query(KeyValueEntry).filter(
KeyValueEntry.resource.in_(RESOURCES_TO_MIGRATE)
)
):
value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
try:
value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
except pickle.UnpicklingError as ex:
if str(ex) == "pickle data was truncated":
# make truncated values that were created prior to #20385 an empty
# dict so that downgrading will work properly.
truncated_count += 1
value = {}
else:
raise
entry.value = bytes(json.dumps(value), encoding="utf-8")
if truncated_count:
print(f"Replaced {truncated_count} corrupted values with an empty value")
def downgrade():
bind = op.get_bind()